feat(hpc): TD-T2 — AMX TDPBUSD tile kernel + matmul_i8_to_i32 wiring

Mirror of the BF16 AMX work (TD-T1 / TD-T1b in PR #182) for the
integer operand family. Builds the missing int8 tile kernel from
scratch (the BF16 equivalent shipped in PR #104; the int8 one had
never been built despite the primitives existing in simd_amx since
day one) and wires matmul_i8_to_i32's AMX arm through it.

New module `hpc::int8_tile_gemm`:

  * `int8_tile_gemm_16x16(a_u8, b_i8, c, k)` — public tile kernel,
    K must be multiple of 64. Mirror shape of
    `bf16_tile_gemm_16x16` but for the `u8 × i8 → i32` operand
    family that TDPBUSD natively supports. **One TDPBUSD = 16 384
    multiply-accumulates per instruction** (16×16 output tile × 64
    K-elements per A row × 4 K-elements per inner-product). That's
    256× the VPDPBUSD-zmm throughput per instruction.
  * Internal `amx_path()` uses the existing primitives in
    `amx_matmul`: TileConfig::for_dpbusd(64) → tile_loadconfig →
    tile_zero → K/64 iterations of (tile_load A, tile_load B,
    tile_dpbusd) → tile_store → tile_release.
  * `fallback_path()` for non-AMX hosts: scalar u8 × i8 → i32
    triple-loop reference.

New primitive `amx_matmul::vnni_pack_i8(src, dst, k, n)`:

  * Packs K × N row-major i8 into K/4 outer rows × (N*4) VNNI quad
    layout required by TDPBUSD tile 2.
  * `dst[kb*N*4 + j*4 + p] = src[(4*kb + p) * N + j]`
  * Sibling of `vnni_pack_bf16` (which uses K/2 × (N*2) pair layout
    for TDPBF16PS — both kernels reach the same 64-byte tile row
    width via element-width × pack-factor symmetry: BF16 is 2B × 2,
    INT8 is 1B × 4).

Wiring `matmul_i8_to_i32`'s AMX arm (was placebo):

Pre-commit the AMX branch shifted i8 → u8 then called the SCALAR
`int8_gemm_i32` reference and subtracted the bias — TDPBUSD itself
was never reached even on real AMX silicon. Now:

  1. Shift A: i8 → u8 via (+128).
  2. Tile-loop over M/16 i_tile × N/16 j_tile blocks, calling
     int8_tile_gemm_16x16 per (i_tile, j_tile). B sub-block
     extracted into K × 16 scratch once per j_tile, reused across
     i_tile iterations.
  3. Subtract bias: c[i, j] -= 128 × colsum(B[:, j]).

The shape requirement is m%16 == 0 && n%16 == 0 && k%64 == 0;
misaligned shapes fall back to the scalar reference. Phase-4 work
will land mixed AMX-tile + per-axis scalar tail handling for
arbitrary shapes (same shape of Phase-4 work TD-T1 deferred).

Verification:
  * Default v3 build: 2092 lib tests pass (was 2087 — adds 5 new
    tests: 4 in int8_tile_gemm + the existing matmul_i8_to_i32 test
    now exercises the actual TDPBUSD path because this host has
    amx_int8 + amx_tile in /proc/cpuinfo; the test continues to
    pass with bit-identical results to the scalar reference).
  * `vnni_pack_i8_roundtrip` test verifies the pack layout matches
    the spec exactly for an 8 × 4 sample.
  * `fallback_matches_scalar_reference_k64` test verifies the
    non-AMX path produces the same i32 output as a hand-written
    reference for a 64-K, pseudo-random u8/i8 matrix pair.
  * `public_api_diagonal_k128` test asserts a structured pattern
    (A = identity-like, B = constant 2) gives the expected
    accumulation through the full dispatch chain.
  * `cargo clippy --lib -D warnings` clean.
  * `cargo fmt --all --check` clean.

Dropped: `int8_gemm_i32` import in `amx_matmul.rs` since the AMX
arm no longer falls back to it (the scalar else-branch uses an
inline triple-loop directly).

After this commit, the per-CPU dispatch table from PR #180 has the
AMX tier wired for BOTH operand families on Sapphire Rapids+:

  BF16 GEMM:  SPR+ → TDPBF16PS  (TD-T1 / TD-T1b in PR #182)
  INT8 GEMM:  SPR+ → TDPBUSD    (this commit)

Out of scope (separate PRs):
  * VPDPBUSD-zmm arm of matmul_i8_to_i32 for Cooper Lake / Cascade
    Lake / Zen 4+ (avx512vnni without AMX). The kernel function
    `vnni_dot_u8_i8` and `vnni_matvec` exist in simd_amx.rs; just
    need to assemble them into a m×n×k GEMM and wire as the
    middle dispatch tier (analogous to the VDPBF16PS arm in PR
    #182's bf16_gemm_dispatch).
  * AMX tile path for `simd_int_ops::gemm_u8_i8` (the slice-level
    surface from PR #182) — it's u8 × i8 natively so no sign-shift
    needed, simpler to wire than matmul_i8_to_i32.

https://claude.ai/code/session_01HbqooFZHAjaUtFEzhA1R2u

Author: claude

src/hpc/amx_matmul.rs

@@ -193,6 +193,32 @@ pub fn vnni_pack_bf16(src: &[u16], dst: &mut [u16], k: usize, n: usize) {
     }
 }
 
+/// Pack B[K, N] i8 row-major into K/4 × (N*4) VNNI quads for `TDPBUSD`.
+///
+/// Output layout required by `TDPBUSD` tile 2 (16 rows × 64 bytes):
+///   dst[kb*N*4 + j*4 + p] = src[(4*kb + p) * N + j]
+///
+/// For N=16 (AMX tile width), each output "row" holds 16 i8 quads = 64
+/// bytes (matches the 64-byte tile row width). K must be a multiple of
+/// 4. The same layout is used for `u8` operands (just bit-cast through
+/// — VNNI doesn't care about sign at the packing layer; sign
+/// interpretation happens inside TDPBUSD which treats A as u8 and B
+/// as i8 for the multiply).
+#[inline]
+pub fn vnni_pack_i8(src: &[i8], dst: &mut [i8], k: usize, n: usize) {
+    debug_assert_eq!(src.len(), k * n);
+    debug_assert_eq!(dst.len(), k * n);
+    debug_assert_eq!(k % 4, 0, "K must be multiple of 4 for VNNI INT8 quads");
+    for kb in 0..(k / 4) {
+        let dst_row = kb * n * 4;
+        for j in 0..n {
+            for p in 0..4 {
+                dst[dst_row + j * 4 + p] = src[(4 * kb + p) * n + j];
+            }
+        }
+    }
+}
+
 // ═══════════════════════════════════════════════════════════════════════════
 // Public ndarray-typed matmul API (sprint A4 / Burn parity item 6)
 // ═══════════════════════════════════════════════════════════════════════════
@@ -207,7 +233,7 @@ pub fn vnni_pack_bf16(src: &[u16], dst: &mut [u16], k: usize, n: usize) {
 // strided (e.g. `view.slice(s![.., ..;2])`). Strided inputs are repacked
 // into contiguous staging buffers before the kernel runs.
 
-use crate::hpc::quantized::{bf16_gemm_f32, int8_gemm_i32, BF16};
+use crate::hpc::quantized::{bf16_gemm_f32, BF16};
 use crate::{ArrayView2, ArrayViewMut2};
 
 /// Errors returned by the public AMX matmul API.
@@ -537,14 +563,17 @@ pub fn matmul_f32(
 
 /// Matrix multiply i8 × i8 → i32: `out = lhs · rhs`.
 ///
-/// On AMX hosts uses `TDPBUSD` (256 MACs/instr); otherwise falls back to
-/// the scalar `int8_gemm_i32`.
+/// On AMX hosts with 16/16/64-aligned shapes uses `TDPBUSD` via the
+/// 16×16 tile kernel in [`crate::hpc::int8_tile_gemm::int8_tile_gemm_16x16`]
+/// — 16 384 MACs per instruction. Mis-aligned shapes (or non-AMX hosts)
+/// fall back to the scalar i8×i8 → i32 reference.
 ///
-/// Note: `TDPBUSD` natively expects unsigned-by-signed (u8 × i8). For the
-/// signed-by-signed surface required here, the LHS is shifted into the
-/// unsigned domain and the bias subtracted from the accumulator (only on
-/// the AMX path; the scalar path operates directly in i8). The public
-/// result is identical.
+/// Note: `TDPBUSD` natively expects unsigned-by-signed (u8 × i8). For
+/// the signed-by-signed surface required here, the LHS is shifted into
+/// the unsigned domain (i8 + 128 → u8) and the bias `128 · sum(B[:, j]
+/// over k)` is subtracted from the accumulator. The public result is
+/// bit-identical to the scalar reference because all arithmetic stays
+/// in i32 (no float rounding).
 ///
 /// `out` must be row-contiguous; inputs may be strided.
 pub fn matmul_i8_to_i32(
@@ -556,13 +585,39 @@ pub fn matmul_i8_to_i32(
     let b_i8 = pack_contig(&rhs);
     let mut c = vec![0i32; m * n];
 
-    if amx_available() {
-        // AMX TDPBUSD path: shift LHS i8 → u8 via (+128) and subtract the
-        // bias 128·sum(B[:, j] over k) afterwards. This keeps numerics exact.
+    if amx_available() && m % 16 == 0 && n % 16 == 0 && k % 64 == 0 {
+        // AMX TDPBUSD path: shift LHS i8 → u8 via (+128), tile-GEMM into
+        // i32, subtract bias 128·colsum(B). The tile kernel zeroes its
+        // internal accumulator (TILEZERO + TDPBUSD accumulate); we need
+        // fresh per-tile output here so we tile manually over M/N and
+        // call int8_tile_gemm_16x16 per (i, j) block.
         let a_u8: Vec<u8> = a_i8.iter().map(|&v| (v as i32 + 128) as u8).collect();
 
-        // Compute C' = A_u8 · B_i8 in i32, then subtract 128 · colsum(B).
-        int8_gemm_i32(&a_u8, &b_i8, &mut c, m, n, k);
+        // B sub-block extraction per j-tile (B is row-major K × N; the
+        // tile kernel wants K × 16 contiguous). Reused across i-tiles.
+        let mut b_tile = vec![0i8; k * 16];
+        let mut tile_c = vec![0i32; 256];
+
+        for j_tile in (0..n).step_by(16) {
+            // Pack B[0..k, j_tile..j_tile+16] into 16-wide K-rows.
+            for kk in 0..k {
+                let row = kk * n + j_tile;
+                b_tile[kk * 16..(kk + 1) * 16]
+                    .copy_from_slice(unsafe { core::slice::from_raw_parts(b_i8.as_ptr().add(row), 16) });
+            }
+            for i_tile in (0..m).step_by(16) {
+                let a_tile = &a_u8[i_tile * k..(i_tile + 16) * k];
+                tile_c.fill(0);
+                crate::hpc::int8_tile_gemm::int8_tile_gemm_16x16(a_tile, &b_tile, &mut tile_c, k);
+                // Write tile_c (16 × 16) into c at (i_tile, j_tile).
+                for ii in 0..16 {
+                    let dst_off = (i_tile + ii) * n + j_tile;
+                    c[dst_off..dst_off + 16].copy_from_slice(&tile_c[ii * 16..(ii + 1) * 16]);
+                }
+            }
+        }
+
+        // Subtract bias: c[i, j] -= 128 · colsum(B[:, j]).
         let mut colsum = vec![0i32; n];
         for p in 0..k {
             for j in 0..n {
@@ -575,7 +630,8 @@ pub fn matmul_i8_to_i32(
             }
         }
     } else {
-        // Scalar i8×i8 → i32 reference.
+        // Scalar i8×i8 → i32 reference — used for non-AMX hosts and for
+        // shapes that don't fit the 16/16/64 tile alignment.
         for i in 0..m {
             for p in 0..k {
                 let av = a_i8[i * k + p] as i32;

src/hpc/int8_tile_gemm.rs

@@ -0,0 +1,214 @@
+//! INT8 tile GEMM polyfill — AMX (TDPBUSD) tile kernel.
+//!
+//! Mirror of `hpc::bf16_tile_gemm` for the `u8 × i8 → i32` shape, the
+//! native TDPBUSD operand type. One TDPBUSD: 16×16 output tile × 64
+//! K-elements per A row × 4 K-elements per inner product = **16 384
+//! multiply-accumulates per instruction**. That's 256× the VPDPBUSD
+//! zmm throughput per instruction (which does 16 × 4 = 64 MACs).
+//!
+//! Public surface:
+//!   * [`int8_tile_gemm_16x16`] — the 16×16 tile kernel; M=16, N=16,
+//!     K a multiple of 64. AMX path requires runtime feature
+//!     detection (`amx_available()`); falls back to a scalar reference
+//!     when AMX isn't OS-enabled.
+//!
+//! Caller responsibility:
+//!   * B comes in row-major K × 16 i8; the kernel pre-packs it into
+//!     VNNI quad layout via [`super::amx_matmul::vnni_pack_i8`].
+//!   * A is row-major 16 × K u8 (TDPBUSD's unsigned operand).
+//!   * C accumulates into the caller's i32 buffer (16 × 16 = 256 i32).
+//!
+//! Same shape as `bf16_tile_gemm::bf16_tile_gemm_16x16`. The two kernels
+//! together cover the SPR/GNR AMX dispatch tier for both `BF16 × BF16
+//! → f32` and `u8 × i8 → i32` — the two operand families that AMX
+//! supports natively.
+
+use crate::hpc::amx_matmul::{
+    amx_available, tile_dpbusd, tile_load, tile_loadconfig, tile_release, tile_store, tile_zero, vnni_pack_i8,
+    TileConfig,
+};
+
+// ═════════════════════════════════════════════════════════════════════
+// Public API — safe dispatching wrapper
+// ═════════════════════════════════════════════════════════════════════
+
+/// Compute C[16, 16] += A[16, K] × B[K, 16] where A is u8 row-major,
+/// B is i8 row-major, C is i32 row-major. K must be a multiple of 64.
+///
+/// Tier dispatch (runtime):
+///   AMX available   → TDPBUSD tile GEMM  (16×16 × K/64 tile iterations,
+///                                          16 384 MACs per instruction)
+///   AMX unavailable → scalar u8 × i8 → i32 reference
+///
+/// Output behavior: this function **accumulates** into `c` (does NOT
+/// zero it first). Callers wanting fresh `C = A·B` semantics should
+/// zero `c` before calling, the same convention `bf16_tile_gemm_16x16`
+/// uses.
+pub fn int8_tile_gemm_16x16(a_u8: &[u8], b_i8: &[i8], c: &mut [i32], k: usize) {
+    assert_eq!(k % 64, 0, "K must be multiple of 64 for TDPBUSD tiles");
+    assert_eq!(a_u8.len(), 16 * k);
+    assert_eq!(b_i8.len(), k * 16);
+    assert_eq!(c.len(), 16 * 16);
+
+    if amx_available() {
+        // AMX path: pack B into VNNI quad layout, call tile GEMM.
+        let mut b_vnni = vec![0i8; k * 16];
+        vnni_pack_i8(b_i8, &mut b_vnni, k, 16);
+        // SAFETY: amx_available() just confirmed CPUID + XCR0 + prctl.
+        unsafe {
+            amx_path(a_u8, &b_vnni, c, k);
+        }
+    } else {
+        fallback_path(a_u8, b_i8, c, k);
+    }
+}
+
+// ═════════════════════════════════════════════════════════════════════
+// AMX path (TDPBUSD)
+// ═════════════════════════════════════════════════════════════════════
+
+/// AMX tile GEMM. B must be pre-VNNI-packed (see `vnni_pack_i8`).
+/// # Safety
+/// Caller must have verified `amx_available() == true`.
+#[inline]
+unsafe fn amx_path(a_u8: &[u8], b_vnni: &[i8], c: &mut [i32], k: usize) {
+    // Tile config: 16×64-byte tiles, identical shape to the BF16 tile
+    // (BF16 is 32 elements × 2 bytes per row, INT8 is 64 elements × 1
+    // byte — same 64-byte row width either way).
+    let cfg = TileConfig::for_dpbusd(64);
+    tile_loadconfig(&cfg);
+    tile_zero(0);
+
+    // Accumulate over K/64 tile blocks. Each TDPBUSD consumes 64
+    // K-elements per A row × 4 K-elements per inner-product = 256 MACs
+    // per output cell × 16 × 16 = 16 384 MACs per instruction.
+    let k_blocks = k / 64;
+    let a_stride = k; // bytes per A row (u8 = 1 byte each)
+    let b_stride = 64usize; // VNNI: 16 columns × 4 bytes per row
+
+    for kb in 0..k_blocks {
+        let a_ptr = a_u8.as_ptr().add(kb * 64);
+        // B sits in VNNI layout: K/4 outer rows × 64 bytes. Each
+        // 64-K-element block spans 16 outer rows × 64 bytes = 1024
+        // bytes.
+        let b_ptr = b_vnni.as_ptr().add(kb * 16 * 64) as *const u8;
+        tile_load(1, a_ptr, a_stride);
+        tile_load(2, b_ptr, b_stride);
+        tile_dpbusd();
+    }
+
+    tile_store(0, c.as_mut_ptr() as *mut u8, 64);
+    tile_release();
+}
+
+// ═════════════════════════════════════════════════════════════════════
+// Scalar fallback (i32 reference)
+// ═════════════════════════════════════════════════════════════════════
+
+/// Direct scalar u8 × i8 → i32 reference. Accumulates into `c`.
+fn fallback_path(a_u8: &[u8], b_i8: &[i8], c: &mut [i32], k: usize) {
+    for i in 0..16 {
+        for kk in 0..k {
+            let a_val = a_u8[i * k + kk] as i32;
+            for j in 0..16 {
+                c[i * 16 + j] += a_val * b_i8[kk * 16 + j] as i32;
+            }
+        }
+    }
+}
+
+// ═════════════════════════════════════════════════════════════════════
+// Tests
+// ═════════════════════════════════════════════════════════════════════
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Reference: scalar u8 × i8 → i32 (matches `fallback_path`).
+    fn ref_gemm(a: &[u8], b: &[i8], c: &mut [i32], k: usize) {
+        for i in 0..16 {
+            for j in 0..16 {
+                let mut s = 0i32;
+                for kk in 0..k {
+                    s += a[i * k + kk] as i32 * b[kk * 16 + j] as i32;
+                }
+                c[i * 16 + j] = s;
+            }
+        }
+    }
+
+    #[test]
+    fn fallback_matches_scalar_reference_k64() {
+        let k = 64;
+        // Deterministic pseudo-random inputs covering the u8 / i8 ranges.
+        let a: Vec<u8> = (0..16 * k).map(|i| ((i * 7 + 3) % 256) as u8).collect();
+        let b: Vec<i8> = (0..k * 16)
+            .map(|i| (((i * 11 + 5) % 256) as u8 as i8))
+            .collect();
+
+        let mut c_ref = vec![0i32; 256];
+        ref_gemm(&a, &b, &mut c_ref, k);
+
+        let mut c_fb = vec![0i32; 256];
+        fallback_path(&a, &b, &mut c_fb, k);
+
+        for i in 0..256 {
+            assert_eq!(c_fb[i], c_ref[i], "fallback mismatch at {}", i);
+        }
+    }
+
+    #[test]
+    fn public_api_runs_on_any_hardware_k64() {
+        let k = 64;
+        let a = vec![0u8; 16 * k];
+        let b = vec![0i8; k * 16];
+        let mut c = vec![0i32; 256];
+        int8_tile_gemm_16x16(&a, &b, &mut c, k);
+        for v in c.iter() {
+            assert_eq!(*v, 0, "zero × zero must be 0");
+        }
+    }
+
+    #[test]
+    fn public_api_diagonal_k128() {
+        // A = identity-like (only A[i, i] = 1, but we need 16 × 128), so
+        // pick A[i, i*8..i*8+8] = 1 (8 ones per i-row). B = constant 2.
+        // Expected: C[i, j] = sum_{kk in i*8..i*8+8}(1 × 2) = 16.
+        let k = 128;
+        let mut a = vec![0u8; 16 * k];
+        for i in 0..16 {
+            for off in 0..8 {
+                a[i * k + i * 8 + off] = 1;
+            }
+        }
+        let b = vec![2i8; k * 16];
+        let mut c = vec![0i32; 256];
+        int8_tile_gemm_16x16(&a, &b, &mut c, k);
+        for i in 0..16 {
+            for j in 0..16 {
+                assert_eq!(c[i * 16 + j], 16, "diagonal accumulator at ({}, {})", i, j);
+            }
+        }
+    }
+
+    #[test]
+    fn vnni_pack_i8_roundtrip() {
+        // Pack then verify the VNNI layout matches the spec:
+        // dst[kb*N*4 + j*4 + p] = src[(4*kb + p) * N + j]
+        let k = 8usize;
+        let n = 4usize;
+        let src: Vec<i8> = (0..(k * n) as i8).collect();
+        let mut dst = vec![0i8; k * n];
+        vnni_pack_i8(&src, &mut dst, k, n);
+        for kb in 0..(k / 4) {
+            for j in 0..n {
+                for p in 0..4 {
+                    let dst_idx = kb * n * 4 + j * 4 + p;
+                    let expected = src[(4 * kb + p) * n + j];
+                    assert_eq!(dst[dst_idx], expected, "vnni quad mismatch at kb={} j={} p={}", kb, j, p);
+                }
+            }
+        }
+    }
+}

src/hpc/mod.rs

@@ -66,6 +66,10 @@ pub mod heel_f64x8;
 pub mod amx_matmul;
 #[cfg(target_arch = "x86_64")]
 pub mod bf16_tile_gemm;
+/// INT8 (`u8 × i8 → i32`) tile GEMM via AMX `TDPBUSD` — mirror of
+/// `bf16_tile_gemm` for the integer operand family.
+#[cfg(target_arch = "x86_64")]
+pub mod int8_tile_gemm;
 #[allow(missing_docs)]
 pub mod bf16_truth;
 #[allow(missing_docs)]