diff --git a/src/hpc/bf16_tile_gemm.rs b/src/hpc/bf16_tile_gemm.rs index 3a2d2c8f..22e02a9e 100644 --- a/src/hpc/bf16_tile_gemm.rs +++ b/src/hpc/bf16_tile_gemm.rs @@ -53,7 +53,7 @@ pub fn bf16_tile_gemm_16x16(a_bf16: &[u16], b_bf16: &[u16], c: &mut [f32], k: us // `bf16_tile_gemm_16x16_packed` to hoist this out of hot loops). let mut b_vnni = vec![0u16; k * 16]; vnni_pack_bf16(b_bf16, &mut b_vnni, k, 16); - dispatch_vnni(a_bf16, &b_vnni, c, k); + dispatch_vnni(a_bf16, k, &b_vnni, c, k); } else { // Pure-FMA tier consumes row-major B directly — no pack needed. fallback_path(a_bf16, b_bf16, c, k); @@ -204,21 +204,74 @@ pub fn bf16_tile_gemm_16x16_packed(a_bf16: &[u16], b: &PackedBf16B, c: &mut [f32 let k = b.k; assert_eq!(a_bf16.len(), 16 * k); assert_eq!(c.len(), 16 * 16); - dispatch_vnni(a_bf16, &b.data, c, k); + dispatch_vnni(a_bf16, k, &b.data, c, k); } -/// Shared tier dispatch over a VNNI-packed B. -fn dispatch_vnni(a_bf16: &[u16], b_vnni: &[u16], c: &mut [f32], k: usize) { +/// `C[16, 16] += A[16, K] × B[K, 16]` with B pre-packed and A read in place +/// from a strided buffer: row `i` of A is `a_bf16[i·a_row_stride ..][..K]`. +/// +/// This is the entry for an A operand that already lives in memory with a +/// row pitch other than K, for example 16 equally spaced bf16 lanes of one +/// slab (row = lane, column = position within the lane). The tile tiers load +/// it where it is: AMX `TILELOADD` takes the row stride as an operand, and +/// the AVX-512 tier reads each row at its offset. Nothing is copied on those +/// tiers. The decode+FMA tier gathers the 16 rows into a dense 16×K buffer +/// first (it allocates), because the polyfill kernel consumes dense rows. +/// +/// `a_row_stride` must be at least K, and `a_bf16` must hold +/// `15·a_row_stride + K` elements; both are checked. With `a_row_stride == K` +/// this is exactly [`bf16_tile_gemm_16x16_packed`]. +/// +/// # Examples +/// ``` +/// use ndarray::simd::{bf16_tile_gemm_16x16_packed_strided, PackedBf16B}; +/// let k = 32; +/// let pitch = 100; // 16 lanes, each 100 bf16 long, in one slab +/// let mut slab = vec![0u16; 16 * pitch]; +/// slab[3 * pitch + 7] = 0x3F80; // lane 3, position 7 = 1.0 +/// let mut b = PackedBf16B::zeroed(k); +/// b.data_mut()[PackedBf16B::vnni_index(7, 2)] = 0x3F80; // B[7][2] = 1.0 +/// let mut c = vec![0.0f32; 256]; +/// bf16_tile_gemm_16x16_packed_strided(&slab, pitch, &b, &mut c); +/// assert_eq!(c[3 * 16 + 2], 1.0); // C[3][2] = A[3][7]·B[7][2] +/// ``` +pub fn bf16_tile_gemm_16x16_packed_strided(a_bf16: &[u16], a_row_stride: usize, b: &PackedBf16B, c: &mut [f32]) { + let k = b.k; + assert!(a_row_stride >= k, "A row stride {a_row_stride} is shorter than K = {k}"); + let needed = a_row_stride + .checked_mul(15) + .and_then(|n| n.checked_add(k)) + .expect("15·a_row_stride + K overflows usize"); + assert!(a_bf16.len() >= needed, "A holds {} bf16, needs {needed}", a_bf16.len()); + assert_eq!(c.len(), 16 * 16); + dispatch_vnni(a_bf16, a_row_stride, &b.data, c, k); +} + +/// The 16 rows of a strided A, copied into a dense 16×K buffer. +fn dense_rows(a_bf16: &[u16], a_row_stride: usize, k: usize) -> Vec { + let mut dense = vec![0u16; 16 * k]; + for (i, row) in dense.chunks_exact_mut(k).enumerate() { + row.copy_from_slice(&a_bf16[i * a_row_stride..i * a_row_stride + k]); + } + dense +} + +/// Shared tier dispatch over a VNNI-packed B. Row `i` of A starts at +/// `a_bf16[i * a_row_stride]`; callers have checked the bounds. +fn dispatch_vnni(a_bf16: &[u16], a_row_stride: usize, b_vnni: &[u16], c: &mut [f32], k: usize) { if amx_available() { // SAFETY: amx_available() confirmed CPUID + XCR0 + arch_prctl grant. unsafe { - amx_path(a_bf16, b_vnni, c, k); + amx_path(a_bf16, a_row_stride, b_vnni, c, k); } } else if is_x86_feature_detected!("avx512bf16") { // SAFETY: runtime detection confirmed avx512f+avx512bf16. unsafe { - avx512bf16_path(a_bf16, b_vnni, c, k); + avx512bf16_path(a_bf16, a_row_stride, b_vnni, c, k); } + } else if a_row_stride != k { + // The polyfill consumes dense rows: gather the 16 strided rows. + dispatch_vnni(&dense_rows(a_bf16, a_row_stride, k), k, b_vnni, c, k); } else { // Slow tier from a packed buffer: unpack back to row-major and take // the decode+FMA polyfill. Allocates — acceptable off the tile tiers @@ -252,20 +305,21 @@ fn dispatch_vnni(a_bf16: &[u16], b_vnni: &[u16], c: &mut [f32], k: usize) { /// # Safety /// Caller must have verified `is_x86_feature_detected!("avx512bf16")`. #[target_feature(enable = "avx512f,avx512bf16")] -unsafe fn avx512bf16_path(a_bf16: &[u16], b_vnni: &[u16], c: &mut [f32], k: usize) { +unsafe fn avx512bf16_path(a_bf16: &[u16], a_row_stride: usize, b_vnni: &[u16], c: &mut [f32], k: usize) { use core::arch::x86_64::{ __m512bh, __m512i, _mm512_dpbf16_ps, _mm512_loadu_ps, _mm512_loadu_si512, _mm512_set1_epi32, _mm512_storeu_ps, }; debug_assert_eq!(k % 2, 0); let pairs = k / 2; for i in 0..16 { - // SAFETY (whole block): a_bf16 is 16×k and c is 16×16 (asserted by - // the public callers), b_vnni is k×16 pairs; every pointer below + // SAFETY (whole block): row i of A is a_bf16[i·a_row_stride..][..k] + // and c is 16×16 (both asserted by the public callers), b_vnni is + // k×16 pairs; every pointer below // stays inside those bounds. read_unaligned handles the u16→u32 // pair load without alignment requirements. unsafe { let mut acc = _mm512_loadu_ps(c.as_ptr().add(i * 16)); - let a_row = a_bf16.as_ptr().add(i * k); + let a_row = a_bf16.as_ptr().add(i * a_row_stride); for p in 0..pairs { // Broadcast the A pair (a[2p], a[2p+1]) to all 16 lanes; lane j // then accumulates A[i][2p]·B[2p][j] + A[i][2p+1]·B[2p+1][j]. @@ -295,7 +349,7 @@ unsafe fn avx512bf16_path(a_bf16: &[u16], b_vnni: &[u16], c: &mut [f32], k: usiz /// # Safety /// Caller must have verified `amx_available() == true`. #[inline] -unsafe fn amx_path(a_bf16: &[u16], b_vnni: &[u16], c: &mut [f32], k: usize) { +unsafe fn amx_path(a_bf16: &[u16], a_row_stride: usize, b_vnni: &[u16], c: &mut [f32], k: usize) { // Tile config: shapes at K_bytes=64 match BF16 K=32 case let cfg = TileConfig::for_dpbusd(64); tile_loadconfig(&cfg); @@ -305,7 +359,7 @@ unsafe fn amx_path(a_bf16: &[u16], b_vnni: &[u16], c: &mut [f32], k: usize) { // Accumulate over K/32 tile blocks let k_blocks = k / 32; - let a_stride = (k * 2) as usize; // full A row stride in bytes (bf16 = 2B) + let a_stride = a_row_stride * 2; // A row stride in bytes (bf16 = 2B); TILELOADD takes it as an operand let b_stride = 64usize; // VNNI row stride in bytes // Operand placement (verified empirically — see `tile_dpbusd` doc): the @@ -510,6 +564,97 @@ mod tests { assert_eq!(c, c_ref, "GEMM over LE-roundtripped buffer not exact"); } + /// Scatter a dense 16×k A into a slab where row i starts at i·pitch, + /// with a distinct sentinel in every gap so a wrong offset shows up. + fn strided_slab(a: &[u16], k: usize, pitch: usize) -> Vec { + let mut slab = vec![bf16_int(99); 15 * pitch + k]; + for i in 0..16 { + slab[i * pitch..i * pitch + k].copy_from_slice(&a[i * k..(i + 1) * k]); + } + slab + } + + #[test] + fn strided_entry_reads_rows_in_place_exactly() { + let k = 64; + let (a, b, c_ref) = integer_case(k); + let packed = PackedBf16B::pack(&b, k); + for pitch in [k, k + 1, k + 32, 4 * k + 7] { + let slab = strided_slab(&a, k, pitch); + let mut c = vec![0.0f32; 256]; + bf16_tile_gemm_16x16_packed_strided(&slab, pitch, &packed, &mut c); + assert_eq!(c, c_ref, "pitch {pitch}, tier [{}]", bf16_tile_gemm_tier()); + } + } + + #[test] + fn strided_rows_are_exact_on_every_tier_this_host_has() { + // The dispatch above takes one tier per host; drive the others here. + let k = 64; + let pitch = k + 32; + let (a, b, c_ref) = integer_case(k); + let packed = PackedBf16B::pack(&b, k); + let slab = strided_slab(&a, k, pitch); + + let mut c = vec![0.0f32; 256]; + let mut b_rm = vec![0u16; k * 16]; + for row in 0..k { + for col in 0..16 { + b_rm[row * 16 + col] = packed.data()[PackedBf16B::vnni_index(row, col)]; + } + } + fallback_path(&dense_rows(&slab, pitch, k), &b_rm, &mut c, k); + assert_eq!(c, c_ref, "gather tier not exact"); + + if is_x86_feature_detected!("avx512bf16") { + let mut c = vec![0.0f32; 256]; + // SAFETY: detection checked above; slab holds 15·pitch + k. + unsafe { avx512bf16_path(&slab, pitch, packed.data(), &mut c, k) }; + assert_eq!(c, c_ref, "VDPBF16PS strided not exact"); + } else { + eprintln!("SKIP (honestly): no avx512bf16 — strided VDPBF16PS not exercised"); + } + if amx_available() { + let mut c = vec![0.0f32; 256]; + // SAFETY: amx_available() confirmed; slab holds 15·pitch + k. + unsafe { amx_path(&slab, pitch, packed.data(), &mut c, k) }; + assert_eq!(c, c_ref, "TDPBF16PS strided not exact"); + } else { + eprintln!("SKIP (honestly): no AMX — strided TDPBF16PS not exercised"); + } + } + + #[test] + fn strided_entry_does_not_read_the_gap() { + // Reading the slab densely (stride k) instead of at its pitch picks + // up the sentinel gap and must give a different result. + let k = 32; + let (a, b, c_ref) = integer_case(k); + let packed = PackedBf16B::pack(&b, k); + let slab = strided_slab(&a, k, k + 32); + let mut dense_misread = vec![0.0f32; 256]; + bf16_tile_gemm_16x16_packed(&slab[..16 * k], &packed, &mut dense_misread); + assert_ne!(dense_misread, c_ref, "fixture cannot tell the pitch from K"); + } + + #[test] + #[should_panic(expected = "shorter than K")] + fn strided_entry_rejects_a_stride_below_k() { + let packed = PackedBf16B::zeroed(32); + let a = vec![0u16; 16 * 32]; + let mut c = vec![0.0f32; 256]; + bf16_tile_gemm_16x16_packed_strided(&a, 31, &packed, &mut c); + } + + #[test] + #[should_panic(expected = "needs")] + fn strided_entry_rejects_a_short_buffer() { + let packed = PackedBf16B::zeroed(32); + let a = vec![0u16; 15 * 40 + 31]; + let mut c = vec![0.0f32; 256]; + bf16_tile_gemm_16x16_packed_strided(&a, 40, &packed, &mut c); + } + #[test] fn packed_entry_accumulates() { // C += A·B semantics: pre-existing C values must survive. @@ -535,7 +680,7 @@ mod tests { let packed = PackedBf16B::pack(&b, k); let mut c = vec![0.0f32; 256]; // SAFETY: detection checked above. - unsafe { avx512bf16_path(&a, packed.data(), &mut c, k) }; + unsafe { avx512bf16_path(&a, k, packed.data(), &mut c, k) }; assert_eq!(c, c_ref, "VDPBF16PS not exact on integer operands"); // Float parity vs the decode+FMA fallback: same products (exact), @@ -555,7 +700,7 @@ mod tests { let packed_f = PackedBf16B::pack(&b_bf, k); let mut c_v = vec![0.0f32; 256]; // SAFETY: detection checked above. - unsafe { avx512bf16_path(&a_bf, packed_f.data(), &mut c_v, k) }; + unsafe { avx512bf16_path(&a_bf, k, packed_f.data(), &mut c_v, k) }; let mut c_fb = vec![0.0f32; 256]; fallback_path(&a_bf, &b_bf, &mut c_fb, k); for i in 0..256 { diff --git a/src/simd.rs b/src/simd.rs index 8c6d9628..28e7cb1c 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -873,7 +873,8 @@ pub use crate::simd_runtime::matmul_i8_to_i32; // run reports. #[cfg(all(feature = "std", target_arch = "x86_64"))] pub use crate::hpc::bf16_tile_gemm::{ - bf16_tile_gemm_16x16 as bf16_tile_gemm_16x16_amx, bf16_tile_gemm_16x16_packed, bf16_tile_gemm_tier, PackedBf16B, + bf16_tile_gemm_16x16 as bf16_tile_gemm_16x16_amx, bf16_tile_gemm_16x16_packed, bf16_tile_gemm_16x16_packed_strided, + bf16_tile_gemm_tier, PackedBf16B, }; // CPU-generation detection (cached): SPR / EMR / GNR / Sierra Forest. Lets a // consumer report which silicon a run landed on and distinguish "no AMX