xilinx/mlir-aie/aie-kernel-opt

AIE kernel optimization

This skill makes one compiled kernel faster. It assumes the kernel is already correct on hardware (aie-hw-bringup) and worth the work (aie-dataflow-opt ranks kernels by ablation).

Evidence rule. The static report tells you what to try and predicts the result. Only two arms measured back to back on the NPU, as traced cycles per call of the kernel itself, show a speedup. Only a raw output-word diff shows bit-exactness. Without a device, report a static change as "candidate, HW unconfirmed", never as faster: an II33 → 18 change measured 594 → 594 (it edited a symbol no factory calls), an II37 → 31 change measured slower, and rolling int8 mm's K loop, a uniform static win, measured +3..+28% on half the shapes (51d0a8e000a).

Setup

source /opt/xilinx/xrt/setup.sh            # XRT first, or pyxrt is missing
source <venv>/bin/activate
source utils/env_setup.sh install          # from the repository root (docs/Building.md)
export MLIR_AIE_KERNEL_SOURCES=$PWD        # else "this tree" is build/include's stale copy
W=$(mktemp -d); K=<factory>; CASE=<case name from kernel_cases.py>; CPUS=<fixed host CPU list>

Remarks takes --target aie2p for npu2 (Strix, Krackan) or --target aie2 for npu1 (Phoenix); xrt-smi examine shows which device you have. What an intrinsic lowers to is in third_party/aie_api/include/aie_api/detail/<arch>/.

Your own kernel or design

The commands below are written for a library kernel (an aie.iron.kernels factory with cases in test/python/npu/kernel_cases.py). The workflow is the same for a kernel in your own design; only the drivers change:

Every rule and lever holds either way. The MobileNet, ResNet and llama numbers below are examples of a consumer design exercising a kernel, not requirements; the short hashes cite the mlir-aie commits that measured them (git fetch origin pull/3807/head brings them in after the squash merge).

Workflow

1. Resolve what runs

grep -n "def $K\b" -A40 python/iron/kernels/*.py      # source, -D flags, extern "C" symbol, trace=
grep -n "Case(\"$K" test/python/npu/kernel_cases.py     # cases, calls=, kwargs
python -m aie.utils.compile.remarks --target aie2p --cases test/python/npu/kernel_cases.py \
  --only "^$K" --out $W/rows.json --meta $W/meta.json

• Remarks prints [OK] <build>: <symbol> from <source> per build. Follow

that symbol to the function it calls and edit only that. A .cc often holds siblings no factory selects (in-place, _scalar, another dtype).

• --cases compiles exactly the builds the cases run, named by case

(4c8c4d91ff4); it first prints cases: N builds for M cases ..., and a case missing from M is a finding. Without it remarks builds only factory defaults and kernels.DTYPES, which missed every stride-2 and sized -D build of the conv kernels. For a shape no case has, use --build FACTORY:KEY=VALUE,... (repeatable; 3dfa82d81af).

• Read the contract's trace=. Only Trace.whole_call() gets a cycles row.

Trace.none(reason) and Trace.partial(reason) leave you with wall clock.

• Grep the vector path's guard. Every condition on a runtime scalar

(not width or alignment) needs a case on each side, at the values the consumer design passes. MobileNet passed skip_scale 0, the guard was > 0, no case used 0, and three layers ran scalar at 579k-927k cycles a call, 16.2 of the network's 16.3 ms. Relaxing the guard took the network to 0.77 ms (9e649eda9d7).

• Time the case at the design's scalars. A case at other widths can

take another path. xy_pool read 820 at outC 120, but MobileNet passes 960 padded to 1280, and a scalar pad loop ran at 3066 a call until it stored vectors (875, 8fafbef845b). bn2's dw row (56 wide) ran the generic path because the 8-pixel path stopped at 32 pixels: 4302 → 840 (eb89013e5d8). A gap that is the same on every row is per-call work, not a DMA stall.

2. Gate, and prove the gate

pytest test/python/test_kernel_contracts.py -k "$K" -q           # host only
pytest test/python/npu/test_kernels_e2e.py -m extensive -k "$CASE" --seeds 3 \
  --report-error $W/err.json > $W/gate.log 2>&1

• -m extensive runs every case, every data case (random plus the

contract's edge cases) and every seed. -k is a substring match (-k add also runs mul_add); check the selection with --collect-only -q. Keep the whole log: an unreproducible failure needs its words.

• --report-error records per case max_ulp, not_correctly_rounded,

the worst index and the reference precision judged against (8f749a06d52); a "correctly rounded" claim must quote a float64 entry. For a unary bf16 kernel run the whole domain: aie.utils.accuracy.all_bf16(nan=False) against round_to(ref64, bfloat16).

• Mutation-prove it. Break the kernel on purpose (drop a term, skip the

tail, swap two loads), watch the gate fail, revert. A surviving mutant is a missing case or dead work: decide which. Report it with the reason; don't swap in a louder one. Mutate data, not addresses, on aligned paths (a 64 B-aligned load_v ignores a 32 B offset), and count only random data runs (zeros pass most mutants).

• Derive every tolerance from the arithmetic, over the longest case

(one kernel's error was 0.07 at 4 calls and 0.25 at 256).

• The reference must not replay the kernel. Write a single-pass fp32

(or float64) reference, and judge against the mathematics (see the exp2 and tanh traps below). Define the semantics at every scalar the kernel accepts; a scalar path that is undefined there (a shift by -1) is not a reference.

• One case per parameter and per path. A check case for each factory

parameter away from its default (conv2dk14 second tile group: 404/1024 wrong, fae15d5b998), each fallback dtype (cascade_mm bf16 fallback: max

versioned loop actually runs, not only its fallback.

• Misaligned views. A design that packs several weights into one

buffer hands them out at offsets like 16 mod 64; kernel tests bind each argument alone and aligned. Test it with Case(..., arg_byte_offsets=((idx, 16),)) (c7d229b9316), and run it on the base kernels first so it fails. Declare alignments in the contract for what the fast path loads (d8e877d26fc).

• Keep outputs poisoned (kd.upload(..., poison=True); the core's output

tile is poisoned too, 36bd3b637a7) and guarded (kd.design(..., guard=True)), so an unwritten output or a write past a tile's end fails. Gate every entry point alone, and the composition once.

3. Build the base arm

An arm is a directory holding aie_kernels/ and aie_runtime_lib/. Remarks and the JIT compile from MLIR_AIE_KERNEL_SOURCES=<arm> when it is set, and both take --baseline-sources <arm> for the comparison.

BASE=$W/base; mkdir -p $BASE
git archive HEAD aie_kernels aie_runtime_lib | tar -x -C $BASE
# Shared tree, or your file includes a sibling (mha.cc includes mm, softmax, zero):
cp -r aie_kernels aie_runtime_lib $BASE/ && git show <rev>:<file> > $BASE/<file>
diff -rq aie_kernels $BASE/aie_kernels     # must list only your files

The factories, cases and -D flags come from the installed Python in both arms, so snapshot Python-side parameters (stack sizes, geometry tables) you change along with the arm. Archive right before the A/B: on a shared branch an older archive lacks others' later commits and credits their gains to you.

4. Read the static report

python -m aie.utils.compile.remarks --target aie2p --cases test/python/npu/kernel_cases.py \
  --only "^$CASE" --out $W/rows.json --meta $W/meta.json --keep $W/objs --baseline-sources $BASE