AMX-native bf16 packed operands: pack once at prepare time, not per tile

[AdaWorldAPI]

Follow-up to #5's B0-B3 gap diagnosis (see the diagnostic comment/bench on that PR), which proved the B1-to-B3 gap was entirely add_mat_mul_bf16 re-converting and re-VNNI-packing both operands on every 16x16 tile call, not TDPBF16PS itself. This PR builds tract's packing layer for that fix: two new MMMInputFormat implementations (NdarrayAmxBf16A, NdarrayAmxBf16B in linalg/src/x86_64/ndarray_amx_native_pack.rs) that convert to bf16 (and, for B, VNNI-pack via ndarray::simd::PackedBf16B) once at prepare_one time, plus a new kernel (ndarray_avx512_bf16_native_mmm_f32_16x16, packing index 1, in ndarray_bf16_native_gemm.rs) whose AddMatMul step calls bf16_tile_gemm_16x16_packed directly on the already-packed panels — zero allocation, zero conversion, zero VNNI packing inside the tile call. Note: this PR is stacked on #5 (it merges that branch's diagnostic work as its base) since it depends on ndarray_bf16_gemm.rs and the shared bench harness.

This branch is additive-only, zero effect on automatic dispatch. Same non-inventory::submit! registration trick as #5's kernel — reachable only by direct construction (benches/tests), invisible to MmmDispatch::native(), verified by a dedicated test (amx_native_candidate_is_not_reachable_through_automatic_dispatch). No core::ops::einsum / model-optimizer wiring in this PR.

W0 survey findings

tract already has exactly the extension point this task needed: MMMInputFormat/MMMInputValue (linalg/src/frame/mmm/input_store.rs) plus the packing[$pnum] = $pid => |k| k.with_packing(...) clause on MMMKernel!/MMMRustKernel!. A production kernel already uses this pattern for real AMX bf16 packing — avx512amx_mmm_f32_16x16 (linalg/src/x86_64/mmm.rs, gated #[cfg(tract_amx_bf16)]) with PackedAmxBf16A/PackedBf16K2 (linalg/src/x86_64/amx_bf16.rs) — and its panel geometry (r=16, K padded to a multiple of 32, row-major bf16 for A, K=2-inner VNNI blocks for B) turns out to be bit-identical to ndarray::simd::PackedBf16B's internal layout, which made mirroring that existing format straightforward. Const vs Value operand provenance (A activation / B weight) is visible in core/src/ops/einsum, confirming the packing lifetimes named in the brief are plausible, without needing to wire into it.

Real complication: bf16_tile_gemm_16x16_packed takes &PackedBf16B, a struct whose fields (Vec<u16>, usize) are private — there is no public constructor that borrows bytes zero-copy (from_le_bytes allocates and copies). A byte-blob MMMInputValue (the pattern every existing packed format uses) would therefore force a per-tile-call allocation+copy to reconstitute a PackedBf16B, defeating the point. The fix: NdarrayAmxBf16BValue stores owned Vec<PackedBf16B> panels directly (not raw bytes in a Blob), and panel_bytes() returns a pointer to the PackedBf16B value itself (cast to *const u8), which AddMatMul reinterprets back (&*(pb as *const PackedBf16B)) — zero-copy, zero-alloc, at the cost of a custom (non-EagerPackedInput) MMMInputValue impl with hand-written Clone/Debug/Hash/Eq (via PackedBf16B::as_le_bytes/data()/k(), since the type itself derives none of them).

Benchmark — the corrected four-rung methodology

The initial brief's P0-under-prepare_one-outside-the-loop and B2's "convert A once per invocation" lifetime were not the same measurement; per the coordinator's correction this PR reports two P cases at matched lifetimes against their respective baselines, both through the real MatMatMulKer::run path (linalg/benches/amx_bf16_gap_decomposition.rs):

• P0 — A and B both prepare_one'd outside the timed loop (matches B1's lifetime: pure tract packed-execution tax over the raw AMX ceiling).
• P1 — B prepare_one'd once outside the loop and reused (persistent weight); A re-prepare_one'd fresh once per whole matrix inside each timed iteration (matches B2's runtime-activation lifetime) — the realistic inference-shaped path.

RUSTFLAGS="-C target-cpu=native", CARGO_PROFILE_BENCH_DEBUG=0, this Sapphire Rapids Xeon, cargo bench -p tract-linalg --bench amx_bf16_gap_decomposition -- --sample-size 20 (median of 20 samples):

	512³	1024³
B1 — raw AMX ceiling (pre-packed)	845.4 µs / 317.5 Gelem/s	6.766 ms / 317.4 Gelem/s
P0 — tract packed execution	1241.5 µs / 216.2 Gelem/s	11.887 ms / 180.7 Gelem/s
P0/B1	1.47×	1.76×
B2 — raw runtime-A oracle	899.25 µs / 298.5 Gelem/s	6.865 ms / 312.8 Gelem/s
P1 — tract inference-like path	1356.7 µs / 197.9 Gelem/s	12.187 ms / 176.2 Gelem/s
P1/B2	1.51×	1.78×

(B0/B3, unchanged from #5, still show the original 22-114 ms b3 kernel and the 2.7/22.4 ms asm baseline for cross-reference.)

Ruling: FAIL against the 1.0-1.2 acceptance band — a real, honestly-reported tract-abstraction tax

P0/B1 and P1/B2 both land at ~1.5-1.8×, not the 1.0-1.2× that would mean "dense BF16 execution is essentially solved." The packing layer itself (prepare_one converting/VNNI-packing once, AddMatMul doing zero conversion) does what it was built to do — the correctness suite (17 auto-generated mmm_packed_packed_tests for packing 1, plus a dedicated tolerance test) all pass, and P0/P1 sit at roughly a third of B3's 512³ time and a tenth of its 1024³ time. But there is a real, unexplained ~50-80% tax between the raw hand-rolled tiling loop (B1/B2) and going through MatMatMulKer::run's generic executor (P0/P1) even with every operand pre-packed. Candidate sources, not yet profiled: panel_bytes() is a vtable call (dyn MMMInputValue) invoked per output tile inside the generic panel walk, versus B1/B2 indexing flat pre-built Rust arrays directly; MatMatMulKer's fused-op interpreter (FusedKerSpec linked-list walk, one match per op per tile) versus B1/B2's tight for loop over bf16_tile_gemm_16x16_packed calls; and general mmm.run() m/n blocking/stride bookkeeping. This PR does not chase that down further, per scope.

Follow-up work named, not attempted here

• VCVTNE2PS2BF16/VCVTNEPS2BF16 conversion optimization — this PR keeps using f32_to_bf16_batch_rne throughout.
• K-register/opmask/ternlog scheduling and the sparsity crossover probe.
• (New, surfaced by this PR's own P0/P1 finding) Profiling the ~1.5-1.8× tract-executor tax identified above — vtable dispatch cost of panel_bytes(), FusedKerSpec interpreter overhead, or mmm.run() blocking overhead.

🍍

Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht

---

Generated by Claude Code

[cursor]

Bugbot couldn't run - usage limit reached

Bugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit.

A user or team admin can review and increase usage limits in the Cursor dashboard.

(requestId: serverGenReqId_01060589-cf0c-43ec-9eed-283d1216fa51)

[AdaWorldAPI]

Fixed the aarch64-apple-darwin build failure (3116f22) — same class of bug fixed earlier on #5: ndarray_amx_native_pack.rs/ndarray_bf16_native_gemm.rs use ndarray types throughout (not just in one or two functions), so rather than per-item #[cfg] stubs (the pattern the other pilot kernels use), both files get a whole-module #[cfg(target_arch = "x86_64")] gate, with the mmm.rs registration for this kernel gated to match.

Verified: cargo check -p tract-linalg clean on x86_64; cross-checked cargo check --target aarch64-apple-darwin --features foreign-inventory — it now gets past the Rust-level compile entirely (the only remaining failure is this sandbox's known lack of a real macOS cc toolchain for the unrelated .S asm build scripts, not an ndarray import error). Full cargo test -p tract-linalg --lib is 4463/4463 passing, cargo fmt --all -- --check clean, cargo clippy -p tract-linalg --all-targets -- -D warnings shows only the same 2 pre-existing unrelated errors.

---

Generated by Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3116f22201

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Obtain approval before adding unsafe pointer packing

This introduces raw-pointer reads for the new packing path, and the B/kernel side additionally round-trips a PackedBf16B through *const u8; the repository explicitly prohibits adding unsafe without permission. Replace these operations with safe tensor/packing APIs or obtain explicit approval before merging.

AGENTS.md reference: AGENTS.md:L279-L282

Useful? React with 👍 / 👎.

[AdaWorldAPI]

Fixed the banner-comment finding in the sibling thread (2aae5c2).

On this one: I'm not removing the unsafe here. The repo's own sibling AMX kernel, amx_bf16.rs's PackedAmxBf16A/PackedBf16K2 (the production kernel this PR's packing formats mirror, per the panel-geometry note in the PR description), does the identical class of raw-pointer packing unconditionally (Blob::new_for_size_and_align + raw writes in its prepare_one_view, no extra annotation). CLAUDE.md's rule is "no new unsafe outside linalg kernels without explicit permission" — this file lives under linalg/src/x86_64/, is doing the same class of tile-packing work as its unsafe-using neighbor, and the *const u8 round-trip for PackedBf16B is the documented zero-copy design point of this PR (see the "Real complication" section of the PR description) — a safe API doesn't exist for it, that's the actual finding from W0's survey.

So this reads as in-scope, precedented unsafe in a linalg kernel rather than a new-kind-of-unsafe needing separate approval. Happy to be overruled by a human maintainer if the convention is read more narrowly than that.

---

Generated by Claude Code

[chatgpt-codex-connector]

Remove the section-banner comments

This newly added decorative section banner directly violates the repository convention against section banners. Rely on the following item documentation or split the A/B implementations structurally rather than retaining banner comments.

AGENTS.md reference: AGENTS.md:L268-L273

Useful? React with 👍 / 👎.

[AdaWorldAPI]

The report check failure (on since-superseded commits 3116f22/e1ac408) isn't this PR's: it's the Bench comment workflow (triggered via workflow_run after Bench completes) failing to download a tract-cli-x86_64-unknown-linux-gnu artifact — a workflow-orchestration issue, not a code defect. Confirmed not diff-specific: the same Bench comment workflow is currently queued against main's own head (d8a60c1, the #5 merge commit) with no PR involved. No fix exists in this PR's scope to port; this is an artifact-passing gap in .github/workflows/bench-comment.yml/bench.yml itself. Current head 2aae5c2's own Bench/Rust crates runs are in progress — watching for whether report reproduces on this head before doing anything further.

---

Generated by Claude Code