xdna-gemm

xdna-gemm is a small Rust crate that captures the host-side shape planning, layout packing, and CPU-verifiable GEMM semantics described in Unlocking the AMD Neural Processing Unit for ML Training on the Client Using Bare-Metal-Programming Tools.

The paper implements a bare-metal AMD XDNA/NPU GEMM kernel for GPT-2 fine-tuning using the IRON / MLIR-AIE tool flow. This crate does not include AMD's XRT, IRON, generated .xclbin, or AI Engine C++ kernels. Instead, it provides:

• GPT-2 small GEMM problem-size constants from the paper.
• The article's default tile shape: m = 64, k = 64, n = 32.
• A planner that pads dimensions to the 4-by-4 compute-core partition used by the design.
• L3 shim transfer schedules for A, B, and C tiles, following the paper's data-movement formulas.
• Packing helpers for row-major A, column-major B, row-major C, and VMAC sub-tile layouts (4x8, 8x4, and 4x4).
• A portable CPU backend that emulates the NPU kernel's numerical contract: bfloat16 inputs with float32 accumulation and float32 outputs.

What this crate is good for

Use it to validate host-side integration logic before wiring in real XRT calls, to generate per-problem manifests for IRON build scripts, or to test numerical divergence from bfloat16 input quantization.

What it deliberately leaves out

The actual AMD XDNA execution path requires generated IRON/MLIR-AIE artifacts, XRT device handles, shared buffers, and command-processor instruction streams. Those parts are vendor-toolchain-specific and are represented here by manifest and schedule data rather than by a fake hardware backend.

Quick start

use xdna_gemm::{CpuBackend, MatrixLayout, Problem, TileShape};

let problem = Problem::new(2, 3, 2); // A: 2x3, B: 3x2, C: 2x2
let a = [1.0, 2.0, 3.0,
         4.0, 5.0, 6.0]; // row-major
let b = [7.0, 8.0,
         9.0, 10.0,
         11.0, 12.0]; // row-major
let mut c = [0.0_f32; 4];

let backend = CpuBackend::new(TileShape::ARTICLE);
backend.gemm(
    problem,
    &a,
    MatrixLayout::RowMajor,
    &b,
    MatrixLayout::RowMajor,
    &mut c,
).unwrap();

assert_eq!(c.len(), 4);

The CPU backend internally packs A and B to bfloat16, so results are not bit-identical to full-float32 BLAS GEMM. They match the crate's simulated NPU contract.

Generate the GPT-2 manifest

cargo run --example gpt2_manifest

This prints JSON containing the 12 GPT-2 small GEMM sizes listed in the paper's Figure 6, padded dimensions, per-core runtime parameters, and tile counts.