Skip to content

Commit e4e8266

Browse files
committed
examples: sub-pel interpolation = a BF16 16×16 tile op — measured (HEVC 8-tap)
Next probe from the HEVC amortization map: HEVC fractional-pel motion compensation applies a separable 8-tap FIR (half-pel luma taps [-1,4,-11,40,40,-11,4,-1]/64). An FIR is a linear operator, so filtering a 16×16 block is a matrix product — one bf16_tile_gemm_16x16 per pass, and the full separable H+V is Hᵀ·(X·H): the SAME two-tile sandwich shape as the splat covariance projection (Mᵀ·Σ·M, #233) and the codec transform (M·X, #232). Measured on real AMX (this host: AMX TDPBF16PS, amx_available = true), vs a direct 8-tap FIR f64 reference (same 16×16 band operator, reflected edges, so only BF16 rounding differs): (1) horizontal 8-tap out = X·H frobenius_rel_err = 0.157% (one tile op) (2) separable H+V out = Hᵀ·(X·H) frobenius_rel_err = 0.215% (two tile ops) (3) throughput: 1.22 M sub-pel tile ops/s Both asserts pass (rel err < BF16 tolerance). Fractional-pel motion is therefore "just another tile op", closing the MC story past integer-pel. R-5 dispatch noted (per-block 8-tap butterfly below the batch crossover, batched tile GEMM above). x86-gated via the same `mod amx` / non-x86 fallback-main pattern as gridlake_field_tile; canonical ndarray::simd imports; fmt + clippy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MLBnPuScZy6w9di2QEjsXM
1 parent 1ea77e8 commit e4e8266

2 files changed

Lines changed: 235 additions & 0 deletions

File tree

Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,10 @@ required-features = ["codec"]
9595
name = "motion_transform_noop"
9696
required-features = ["codec"]
9797

98+
[[example]]
99+
name = "subpel_tap_tile"
100+
required-features = ["std"]
101+
98102
[[example]]
99103
name = "entropy_ladder_probe"
100104
required-features = ["std"]

examples/subpel_tap_tile.rs

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
//! Sub-pel interpolation = a BF16 16×16 tile op — measured (HEVC 8-tap).
2+
//!
3+
//! Next probe from the HEVC amortization map (`.claude/knowledge/hevc-amortization-map.md`):
4+
//! HEVC fractional-pel motion compensation applies a **separable 8-tap FIR** to a
5+
//! block (half-pel luma taps `[-1, 4, -11, 40, 40, -11, 4, -1] / 64`). An FIR is a
6+
//! **linear operator**, so filtering a 16×16 block is a matrix product — one
7+
//! `bf16_tile_gemm_16x16` per pass, and the full separable H+V interpolation is
8+
//! `H_v · X · H_h`: **the same two-tile "sandwich" shape as the splat covariance
9+
//! projection** (`Mᵀ·Σ·M`, #233) and the codec transform (`M·X`, #232). This
10+
//! probe measures that the tile-GEMM form matches the direct 8-tap FIR to BF16
11+
//! precision, on the shipped kernel, and reports the tier.
12+
//!
13+
//! Fairness / honesty:
14+
//! • Both the tile path and the reference use the SAME 16×16 band operator `H`
15+
//! (8 taps per output column, reflected at the block edge), so this is an
16+
//! equal-result comparison — only BF16 rounding differs.
17+
//! • R-5 dispatch still applies: below the per-arch batch crossover (SPR 64,
18+
//! ICX 32, …) a per-block 8-tap butterfly wins; the tile GEMM wins in the
19+
//! batched regime. This probe shows the *equivalence*, not that the tile path
20+
//! is always the right dispatch for one block.
21+
//!
22+
//! The x86-only AMX/tile path is wrapped in a `cfg(target_arch = "x86_64")`
23+
//! module with a non-x86 fallback `main` (see `gridlake_field_tile`).
24+
//!
25+
//! Run: `cargo run --release --example subpel_tap_tile --features std`
26+
27+
#[cfg(target_arch = "x86_64")]
28+
mod amx {
29+
use ndarray::hpc::quantized::{f32_to_bf16_rounded, BF16};
30+
use ndarray::simd::{amx_available, bf16_tile_gemm_16x16_amx as bf16_tile_gemm_16x16, bf16_tile_gemm_tier};
31+
32+
const M: usize = 16; // block edge
33+
const KPAD: usize = 32; // K padded to a multiple of 32 for the tile GEMM
34+
35+
/// HEVC half-pel luma 8-tap (sum = 64).
36+
const TAPS: [f32; 8] = [-1.0, 4.0, -11.0, 40.0, 40.0, -11.0, 4.0, -1.0];
37+
38+
fn mix(mut z: u64) -> u64 {
39+
z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
40+
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
41+
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
42+
z ^ (z >> 31)
43+
}
44+
45+
fn to_bf16(src: &[f32]) -> Vec<u16> {
46+
let mut b = vec![BF16(0); src.len()];
47+
f32_to_bf16_rounded(src, &mut b);
48+
b.iter().map(|x| x.0).collect()
49+
}
50+
51+
/// Reflect an index into `[0, M)` (block-edge boundary handling).
52+
fn reflect(i: i32) -> usize {
53+
let mut i = i;
54+
while !(0..M as i32).contains(&i) {
55+
if i < 0 {
56+
i = -i - 1;
57+
} else if i >= M as i32 {
58+
i = 2 * M as i32 - 1 - i;
59+
}
60+
}
61+
i as usize
62+
}
63+
64+
/// The 16×16 band operator `H`: `out[·][j] = Σ_t TAPS[t]/64 · in[·][reflect(j-3+t)]`,
65+
/// i.e. `H[k][j]` accumulates the tap weight of input column `k` for output
66+
/// column `j`. `out = X · H`. This IS the sub-pel FIR as a fixed matrix.
67+
fn band_operator() -> [f32; M * M] {
68+
let mut h = [0.0f32; M * M];
69+
for j in 0..M {
70+
for (t, &w) in TAPS.iter().enumerate() {
71+
let k = reflect(j as i32 - 3 + t as i32);
72+
h[k * M + j] += w / 64.0;
73+
}
74+
}
75+
h
76+
}
77+
78+
/// `C[16×16] = A[16×16] · B[16×16]` via the shipped BF16 tile op (K padded 16→32).
79+
fn tile_matmul(a: &[f32; M * M], b: &[f32; M * M]) -> [f32; M * M] {
80+
let mut a_pad = [0.0f32; M * KPAD];
81+
let mut b_pad = [0.0f32; KPAD * M];
82+
for i in 0..M {
83+
for k in 0..M {
84+
a_pad[i * KPAD + k] = a[i * M + k];
85+
b_pad[k * M + i] = b[k * M + i];
86+
}
87+
}
88+
let (ab, bb) = (to_bf16(&a_pad), to_bf16(&b_pad));
89+
let mut c = vec![0.0f32; M * M];
90+
bf16_tile_gemm_16x16(&ab, &bb, &mut c, KPAD);
91+
let mut out = [0.0f32; M * M];
92+
out.copy_from_slice(&c);
93+
out
94+
}
95+
96+
fn direct_matmul(a: &[f32; M * M], b: &[f32; M * M]) -> [f64; M * M] {
97+
let mut c = [0.0f64; M * M];
98+
for i in 0..M {
99+
for j in 0..M {
100+
let mut s = 0.0f64;
101+
for k in 0..M {
102+
s += a[i * M + k] as f64 * b[k * M + j] as f64;
103+
}
104+
c[i * M + j] = s;
105+
}
106+
}
107+
c
108+
}
109+
110+
fn transpose(a: &[f32; M * M]) -> [f32; M * M] {
111+
let mut t = [0.0f32; M * M];
112+
for i in 0..M {
113+
for j in 0..M {
114+
t[j * M + i] = a[i * M + j];
115+
}
116+
}
117+
t
118+
}
119+
120+
/// Direct separable H+V 8-tap FIR reference (`f64`), reflected edges — the
121+
/// ground truth the tile-GEMM sandwich must match.
122+
fn direct_hv(x: &[f32; M * M]) -> [f64; M * M] {
123+
// horizontal
124+
let mut h = [0.0f64; M * M];
125+
for i in 0..M {
126+
for j in 0..M {
127+
let mut s = 0.0f64;
128+
for (t, &w) in TAPS.iter().enumerate() {
129+
s += (w as f64 / 64.0) * x[i * M + reflect(j as i32 - 3 + t as i32)] as f64;
130+
}
131+
h[i * M + j] = s;
132+
}
133+
}
134+
// vertical
135+
let mut out = [0.0f64; M * M];
136+
for i in 0..M {
137+
for j in 0..M {
138+
let mut s = 0.0f64;
139+
for (t, &w) in TAPS.iter().enumerate() {
140+
s += (w as f64 / 64.0) * h[reflect(i as i32 - 3 + t as i32) * M + j];
141+
}
142+
out[i * M + j] = s;
143+
}
144+
}
145+
out
146+
}
147+
148+
fn frob_rel(tile: &[f32; M * M], direct: &[f64; M * M]) -> (f64, f64) {
149+
let (mut max_abs, mut num, mut den) = (0.0f64, 0.0f64, 0.0f64);
150+
for i in 0..M * M {
151+
let e = tile[i] as f64 - direct[i];
152+
max_abs = max_abs.max(e.abs());
153+
num += e * e;
154+
den += direct[i] * direct[i];
155+
}
156+
(max_abs, (num / den.max(1e-12)).sqrt())
157+
}
158+
159+
pub fn run() {
160+
println!("Sub-pel interpolation = BF16 16×16 tile op — measured (HEVC 8-tap half-pel)");
161+
println!(
162+
" shipped kernel tier on THIS host: {} (amx_available = {})\n",
163+
bf16_tile_gemm_tier(),
164+
amx_available()
165+
);
166+
167+
// 7-bit luma block (fits u8 and i8; positive operand for the tile GEMM)
168+
let mut x = [0.0f32; M * M];
169+
for r in 0..M {
170+
for c in 0..M {
171+
x[r * M + c] = ((mix((r as u64) << 8 | c as u64) & 0x7F) as f32) * 0.7 + 20.0 * (r as f32 * 0.3).sin();
172+
}
173+
}
174+
let h = band_operator();
175+
let ht = transpose(&h);
176+
177+
// (1) horizontal half-pel = X · H (one tile op)
178+
let horiz_tile = tile_matmul(&x, &h);
179+
let horiz_direct = direct_matmul(&x, &h);
180+
let (ha, hr) = frob_rel(&horiz_tile, &horiz_direct);
181+
println!(" (1) horizontal 8-tap out = X·H (one tile op)");
182+
println!(" max_abs_err={ha:.4} frobenius_rel_err={:.3}%", hr * 100.0);
183+
184+
// (2) full separable H+V = Hᵀ · (X · H) (two tile ops — the sandwich shape)
185+
let hv_tile = tile_matmul(&ht, &horiz_tile);
186+
let hv_direct = direct_hv(&x);
187+
let (va, vr) = frob_rel(&hv_tile, &hv_direct);
188+
println!(" (2) separable H+V out = Hᵀ·(X·H) (two tile ops = Mᵀ·Σ·M sandwich)");
189+
println!(" max_abs_err={va:.4} frobenius_rel_err={:.3}%", vr * 100.0);
190+
191+
// (3) throughput of the sub-pel tile op
192+
let iters = 200_000usize;
193+
let t0 = std::time::Instant::now();
194+
let mut acc = 0.0f32;
195+
for it in 0..iters {
196+
let mut xv = x;
197+
xv[0] += (it & 0x7) as f32 * 0.01;
198+
let c = tile_matmul(&xv, &h);
199+
acc += c[0] + c[M * M - 1];
200+
}
201+
let dt = t0.elapsed().as_secs_f64();
202+
println!(
203+
" (3) throughput: {iters} sub-pel tile ops in {:.1} ms → {:.2} M/s (checksum {acc:.1})",
204+
dt * 1000.0,
205+
iters as f64 / dt / 1e6
206+
);
207+
208+
assert!(hr < 0.05, "horizontal FIR tile-GEMM rel err too high for BF16 class");
209+
assert!(vr < 0.05, "separable H+V tile-GEMM rel err too high for BF16 class");
210+
211+
println!(
212+
"\n MEASURED: HEVC 8-tap sub-pel interpolation IS a tile GEMM — X·H (one op)\n\
213+
\x20 and the separable H+V is Hᵀ·(X·H), the SAME two-tile sandwich shape as the\n\
214+
\x20 splat covariance projection (#233) and the codec transform (#232). Matches\n\
215+
\x20 the direct 8-tap FIR to BF16 precision. Fractional-pel motion is therefore\n\
216+
\x20 'just another tile op', closing the MC story past integer-pel. (Dispatch\n\
217+
\x20 per R-5: per-block 8-tap butterfly below the batch crossover, batched tile\n\
218+
\x20 GEMM above it.)"
219+
);
220+
}
221+
}
222+
223+
#[cfg(target_arch = "x86_64")]
224+
fn main() {
225+
amx::run();
226+
}
227+
228+
#[cfg(not(target_arch = "x86_64"))]
229+
fn main() {
230+
eprintln!("subpel_tap_tile requires x86_64 (AMX TDPBF16PS / AVX-512 VDPBF16PS BF16 tile GEMM).");
231+
}

0 commit comments

Comments
 (0)