Auto-generate matmul tiling schedules in the driver

[erwei-xilinx]

Summary

Makes f32/bf16/i8 GEMMs run on the NPU with zero configuration. Previously each matmul needed a hand-written (or manually generated) transform-dialect tiling script passed via AIR_TRANSFORM_TILING_SCRIPT; the built-in default is only a trivial non-vectorized 32×32 tiling.

When no user script is supplied and the lowered TritonShared IR is a single plain linalg.matmul, the driver now derives tiling parameters directly from the IR (operand element types + tile shapes) and calls the existing matmul transform generator.

• Move the generator into the backend (amd_triton_npu/backend/matmul_transform.py) so it ships in the wheel. No logic change; it already produced every committed matmul_* script.
• _detect_matmul / _matmul_transform_params in driver.py, wired into _get_transform_ir_string and the launch() path.
  • dtype → schedule: f32 ⇒ accum=f32, contract_in=bf16 + auto-enable bf16 emulation (no native f32 MAC on AIE); bf16 ⇒ accum=f32; i8 ⇒ accum=i32, contract_in=i16.
  • pack sizes + herd caps per NPU generation (npu2 (8,8,8)/4×4, npu1 (4,4,8)/4×2).
  • L2 K-tile capped at 2*pack_k — larger tiles miscompute padded boundary tiles under air-split-launch-for-padding.
• Priority preserved: a user-supplied AIR_TRANSFORM_TILING_SCRIPT still wins; non-matmul IR (elementwise, softmax, matvec, …) is untouched (_detect_matmul returns None).
• setup.py: ship config.py and matmul_transform.py in the pre-installed-triton copy path.

Validation (NPU2/Strix, scriptless)

Correctness passes for f32 padded (500×500×1024), bf16, and i8. Same-driver A/B (committed script vs auto-gen) shows parity in compile time and per-dispatch latency, including a K=2048 stress where the L2 K-tile choice matters most.

dtype	committed script	scriptless auto-gen
bf16	12.56 ms	12.85 ms
i8	12.39 ms	12.19 ms
f32 (padded)	18.80 ms	20.53 ms

Validation (NPU1/Phoenix, scriptless)

Tested on NPU Phoenix (XRT 2.23.0, fw 1.5.5.391), detect_npu_version() == "npu1", xclbin output. Run against this branch's pinned toolchain (mlir-air d3c5f87 / mlir-aie v1.4.0).

dtype	scriptless result on npu1
bf16	passes — 27 kernels auto-generated, pack_sizes=(4,4,8)
i8	fails to compile — see below, needs a fix in this PR + an upstream fix
f32 (padded)	fails in aircc/aiecc, pre-existing — the committed hand-written transform_aie2.mlir fails at the identical stage, so auto-gen is not the cause

bf16 A/B, fresh TRITON_CACHE_DIR each, full 27-shape sweep: committed script 392 s vs scriptless 352 s. No compile-time regression. (Wall-clock over the sweep; per-dispatch latency not separately measured on npu1.)

npu1 i8 needs a pack-size change in this PR

pack_sizes is ordered (M, N, K) while the hardware MAC shape is (M, K, N). The two coincide only when the shape is cubic, which is why npu2 (8,8,8) hides the distinction. On npu1 the current (4,4,8) means (M,K,N) = (4,8,4), and:

error: failed to legalize operation 'vector.contract' that was explicitly marked illegal:
  vector<1x1x4x8xi16>, vector<1x1x8x4xi16> into vector<1x1x4x4xi32>

IsValidAIE2MatMulShapeAndType in mlir-aie's AIEVecTypeConstraints.td has exactly one i8×i8→i32 entry for AIE2: 4x8 × 8x8 → 4x8, i.e. (M,K,N) = (4,8,8). It has exactly one bf16 entry, (M,K,N) = (4,8,4) — which is what (4,4,8) produces, so bf16 is correct as written. Pack sizes on AIE2 are dtype-dependent, not just generation-dependent:

npu	dtype	required pack_sizes	this PR
npu2	all	(8,8,8)	(8,8,8) ✅
npu1	bf16 / f32	(4,4,8)	(4,4,8) ✅
npu1	i8	(4,8,8)	(4,4,8) ❌

mlir-air's programming_examples/matrix_multiplication/i8 agrees: run.py sets mmul_mkn = [4,8,8] for aie2 and [8,8,8] for aie2p.

Separately, contract_input_type="i16" is inert for i8 and can be dropped — getSourceOfWideningOp in mlir-aie peels the extension off the contract operands either way. I measured identical results for None / "i8" / "i16".

Dependency: npu1 i8 also needs Xilinx/mlir-aie#3489

Correcting the pack size makes npu1 i8 compile, but results are still wrong — an upstream AIE2 defect, not something this PR can fix. The lowered MAC gets configuration word 0x108 (signX=0, signY=1), so the LHS is zero-extended while the RHS is sign-extended, computing zext8(A) * sext8(B):

A	B	expected	actual
-1	1	-1	255
-7	5	-35	1245
1	-3	-3	-3 ✅

Non-negative A is always correct, so this is invisible to positive-only smoke tests. I confirmed layout, K-reduction and herd distribution are all fine in isolation (single-element placement exact, K-sweep exact to K=256, A=ones,B=ones exact) — only signed LHS values are wrong. Fix filed as Xilinx/mlir-aie#3489.

Suggested sequencing: land the (4,8,8) pack fix here, and until #3489 is in the pinned mlir-aie, gate i8 off on npu1 with an explicit error rather than emitting a kernel that silently returns wrong values for negative inputs. bf16 on npu1 is unaffected and ready.

Notes

• npu1 i8 is blocked on Xilinx/mlir-aie#3489 plus the (4,8,8) pack fix above; bf16 on npu1 is validated and unblocked.
• Stacked on Cache NPU device probe to avoid per-dispatch xrt-smi subprocess #86 (fix-xrt-smi-memoization) — the scriptless path reaches detect_npu_version(), and that PR removes the per-dispatch xrt-smi subprocess it would otherwise trigger. Retarget this PR to main once Cache NPU device probe to avoid per-dispatch xrt-smi subprocess #86 merges.

Test plan

• py_compile + black --check on changed files
• Scriptless f32 / bf16 / i8 correctness on NPU2 hardware
• Compile-time + latency parity vs committed scripts (incl. K=2048)
• Scriptless bf16 correctness + compile-time A/B on NPU1 hardware
• npu1 i8 / f32 triage (i8 = pack size + upstream mlir-aie#3489; f32 = pre-existing aircc failure, reproduces with the committed script)
• npu1 i8 end-to-end once mlir-aie#3489 lands and the pack size is corrected
• scripts/run_tests.py on NPU hardware (examples still ship their scripts, so behavior is unchanged there)

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR adds a “scriptless” path for NPU matmul kernels by detecting plain linalg.matmul in lowered TritonShared IR and auto-generating an MLIR Transform dialect tiling script in the driver, while still respecting user-provided AIR_TRANSFORM_TILING_SCRIPT overrides.

Changes:

• Add matmul detection + parameter derivation in driver.py, and wire auto-generated transform scripts into the AIR lowering/launch path.
• Move the matmul transform generator into the backend (matmul_transform.py) so it ships with the wheel and can be invoked programmatically.
• Update setup.py to copy additional backend Python files into Triton’s installed backend directory.

Reviewed changes

Copilot reviewed 2 out of 3 changed files in this pull request and generated 2 comments.

File	Description
setup.py	Ensures the backend copy step includes config.py and the new matmul transform generator.
amd_triton_npu/backend/matmul_transform.py	Adds a parameterized generator that emits the transform-dialect script used for matmul tiling/vectorization.
amd_triton_npu/backend/driver.py	Detects eligible matmul IR, derives schedule params, enables bf16 emulation for f32 matmul, and routes the generated script into AIR compilation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.