Follow-up: upstream the last in-tree AIE kernels and remove aie_kernels/

[efurst]

Summary

PR #203 moved every C++ compute kernel out of IRON's tree and made each operator
source its .cc/.h from the installed mlir_aie package (context.kernels_dir).
Two files could not make the move and remain in-tree, for two unrelated reasons:

• aie_kernels/generic/mm_fused.cc — a signature the port doesn't have yet
  (see Exception 1).
• aie_kernels/aie2/mm.cc — a rounding-mode regression in the port
  (see Exception 2).

This issue tracks upstreaming both to mlir-aie and deleting the in-tree copies,
which completes the migration and lets us remove the aie_kernels/ directory
from IRON entirely. The two are independent and can be filed / fixed separately.

---

Exception 1: mm_fused.cc runtime-epilogue signature

Why it was left behind

Two PRs crossed within ~10 minutes on 2026-09-16:

• mlir-aie #3740 (the upstream half of the migration) froze the kernel set at
  merge commit 20a9c2f, carrying the old compile-time mm_fused.cc.
• IRON flm.GEMM: take M, K, N and the activation as runtime parameters #200 (flm.GEMM runtime M/K/N + activation) then merged a new runtime
  mm_fused.cc into IRON's tree — after the upstream cut, so it was never
  included upstream.

The pinned wheel (mlir_aie==1.4.4.dev4+g20a9c2f) therefore ships the stale
kernel. Sourcing mm_fused.cc from the package would compile #200's design
against the wrong signature, so #203 keeps the file in-tree as a deliberate,
documented exception.

The divergence

mm_fused_epilogue_chunk gained runtime parameters in #200.

Wheel / upstream (20a9c2f) — compile-time mode + clamp:

void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer,
                             int32_t half);
// mode/clamp selected by -DMM_FUSED_EPILOGUE_MODE / -DMM_FUSED_CLAMP

IRON in-tree (current) — runtime mode + clamp:

void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer,
                             int32_t half, int32_t mode,
                             int32_t clamp_min_bits, int32_t clamp_max_bits);
// which modes are compiled in is chosen by -DMM_FUSED_EPILOGUE_MODE_MASK

Scope check: only mm_fused.cc diverged in behavior. Its #include companions
(mm_fused_mmul.h, activations.h) are code-identical between IRON and the
wheel. aie_kernel_utils.h and the per-arch zero.cc differ textually, but the
wheel's versions are a compatible superset (already validated — see Testing), so
they do NOT need porting. The port is a single file.

Work required

1. Upstream to mlir-aie: replace aie_kernels/generic/mm_fused.cc with the
   runtime-epilogue version currently in this repo. (Source of truth: this
   branch's aie_kernels/generic/mm_fused.cc.)
2. Publish a wheel containing that change (nightly buildRyzenWheels.yml, or
   cut a release), then bump requirements.txt to it.
3. Delete the in-tree copy and simplify iron/operators/flm/gemm/op.py
get_kernel_artifacts():
   • drop in_tree_generic = self.context.base_dir / "aie_kernels" / "generic"
   • source mm_fused.cc from generic (the kernels_dir path) like its companions
   • drop the extra -I{generic} include added only so the in-tree source could
     find its package-hosted headers (the arch -I for zero.cc stays)
   • remove the "last kernel IRON keeps in-tree" comment

Once both exceptions below are resolved, the aie_kernels/ directory is empty
and can be removed from IRON entirely.

Acceptance criteria

• iron/operators/flm/gemm/op.py sources mm_fused.cc from context.kernels_dir.
• flm.GEMM test suite passes on NPU against the new wheel (see Testing).

Testing

Validated for #203 on NPU Strix Halo (aie2p): full iron/operators/flm/gemm/test.py
passes (115 cases, including #200's test_one_xclbin_serves_every_shape and
test_one_xclbin_serves_every_clamp_bound). Re-run the same suite after the port
to confirm the package-sourced kernel is equivalent.

source /opt/xilinx/xrt/setup.sh
python -m pytest iron/operators/flm/gemm/test.py -q

---

Exception 2: aie2/mm.cc dropped conv_even rounding

Summary

The mlir-aie port of aie2/mm.cc omits the conv_even rounding-mode set that
the previous IRON in-tree aie2/mm.cc applied and that the port's own
aie2p/mm.cc still applies. Without it, the AIE core runs in its power-up
rounding_mode::floor (round toward negative infinity), so every bf16 store in
the matmul is biased low and the bias accumulates over the K reduction instead of
cancelling.

On NPU1 / Phoenix (aie2) this pushed test_swiglu_prefill[...hidden_dim_2048 ...prio_accuracy_False] past its output tolerance (the chained
GEMM→SiLU→mul→GEMM path amplifies the bias). NPU2 / Strix (aie2p) is
unaffected because aie2p/mm.cc keeps the rounding set. The same test passes on
devel (which still had the in-tree aie2/mm.cc with the rounding), confirming
the regression came from sourcing the port's kernel.

As an interim, #203 keeps a patched aie2/mm.cc in IRON's tree with the
rounding restored; aie2p sources from the package unchanged.

The divergence

aie2p/mm.cc selects conv_even for the duration of the bf16 matmul and restores
the caller's mode on return, gated on AIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16:

// aie2p/mm.cc — present
constexpr bool emulated_bf16 = std::is_same_v<T_in, bfloat16>;
aie::rounding_mode saved_rounding = aie::rounding_mode::floor;
if constexpr (emulated_bf16)
  saved_rounding = aie::swap_rounding(aie::rounding_mode::conv_even);
// ... matmul ...
if constexpr (emulated_bf16)
  aie::set_rounding(saved_rounding);

aie2/mm.cc in the wheel has no set_rounding/swap_rounding call at all
(grep returns zero hits), so -DROUND_CONV_EVEN — which the GEMM operator passes
by default — has nothing to act on.

IRON's interim patch adds the equivalent set/restore to the two bf16 dispatch
functions (matmul_vectorized_4x8x4_bf16_bf16 and _bf16_f32), gated on
ROUND_CONV_EVEN.

Scope check: only aie2/mm.cc needs the fix. aie2p/mm.cc is already correct.
The in-tree file's #includes (zero.cc, ../aie_kernel_utils.h) are unchanged
and resolve from the package via -I.

Work required

1. Upstream to mlir-aie: add the conv_even set/restore to aie2/mm.cc,
   matching aie2p/mm.cc. Prefer the aie2p structure (gate on
   AIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16); IRON's interim uses the
   ROUND_CONV_EVEN flag it already passes, but the aie2p idiom is the cleaner
   upstream form. (Source of truth for the intended behavior: this branch's
   aie_kernels/aie2/mm.cc.)
2. Publish a wheel with the fix, then bump requirements.txt to it.
3. Delete the in-tree copy and simplify iron/operators/gemm/op.py
get_kernel_artifacts():
   • drop the if kernel_dir == "aie2": branch that sources mm.cc from
     base_dir / "aie_kernels" and appends -I{kernels_dir}/aie2
   • source mm.cc from self.context.kernels_dir / kernel_dir for all arches
   • remove the INTERIM comment

Acceptance criteria

• iron/operators/gemm/op.py sources mm.cc from context.kernels_dir for all arches.
• test_swiglu_prefill passes on NPU1 / Phoenix against the new wheel.

Testing

NOTE: the regression is aie2-only. It cannot be reproduced on a Strix
(aie2p) box — get_kernel_dir() returns aie2p there and the aie2p kernel was
never broken. Verify on Phoenix (npu1) hardware or CI.

source /opt/xilinx/xrt/setup.sh
python -m pytest "iron/operators/swiglu_prefill/test.py::test_swiglu_prefill" -q

---

References

• IRON aie_kernels: source all C++ kernels from mlir-aie, remove in-tree copies #203 (this migration):
  • f4e6894 (migration), 81f49fe (pin bump)
  • mm_fused.cc exception: iron/operators/flm/gemm/op.py::get_kernel_artifacts
  • aie2/mm.cc exception: iron/operators/gemm/op.py::get_kernel_artifacts,
    kernel patch in aie_kernels/aie2/mm.cc
• IRON flm.GEMM: take M, K, N and the activation as runtime parameters #200 (introduced the runtime mm_fused.cc)
• mlir-aie #3740 (upstream kernel set, merge 20a9c2f)