aie_kernels: source all C++ kernels from mlir-aie, remove in-tree copies - #203
Merged
Merged
Conversation
Contributor
CI Test Results31360ab (2026_09_18_15_03_07) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Smalliron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Phoenix - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Phoenix - ExamplesIRONTested on Trend tables omitted, the comment hit GitHub's size limit. Full report in the workflow run. |
Collaborator
|
@efurst If you update to latest nightly, hopefully we can get this merged in soon! |
IRON no longer hosts C++ compute kernels. Every operator now sources its kernel .cc from the installed mlir_aie package via AIEContext.kernels_dir (root_path()/include/aie_kernels); the in-tree aie_kernels/ directory is removed. IRON_AIE_KERNELS_DIR overrides that path to a local mlir-aie checkout for kernel development. DEPENDS ON mlir-aie: several kernels needed additions upstream before IRON could source them (runtime-size eltwise/activation entry points, runtime-epsilon rms_norm, device-generic rope, the fused-GEMM cluster, aie2 mask_bf16 and layer_norm). This change builds against a wheel that contains those additions; until it ships, the default path resolves to a wheel that lacks the symbols, so set IRON_AIE_KERNELS_DIR at a local mlir-aie checkout. The requirements.txt pin should be bumped to the post-merge wheel before this lands. Co-Authored-By: Claude <noreply@anthropic.com>
mlir-aie #3740 merged (commit 20a9c2f) and published as mlir_aie 1.4.4.dev4+g20a9c2f. Point IRON at that build so the kernels it now sources from the mlir_aie package resolve to the new symbols. Co-Authored-By: Claude <noreply@anthropic.com>
efurst
force-pushed
the
aie-kernels-dedup
branch
from
September 17, 2026 17:20
969e0e9 to
81f49fe
Compare
efurst
marked this pull request as ready for review
September 17, 2026 17:21
Contributor
Author
|
@hunhoffe this should be ready for review and merging pending checks passing. After this closes, I can also create an issue describing the contact points for follow on FLM clean up work! |
andrej
approved these changes
Sep 17, 2026
Collaborator
There was a problem hiding this comment.
Thanks for this.
The Phoenix failure is on test_swiglu_prefill:
----------------------------- Captured stdout call -----------------------------
Latency (us): 15634.28
Effective Bandwidth: 0.1341 GB/s
Mismatch in output[15]: expected 235520.000000, got -5603328.000000
Mismatch in output[18]: expected 18743296.000000, got 15400960.000000
Mismatch in output[24]: expected 2850816.000000, got -1769472.000000
Mismatch in output[25]: expected -15204352.000000, got -18874368.000000
Mismatch in output[31]: expected 13565952.000000, got 10813440.000000
Mismatch in output[36]: expected 5603328.000000, got 3817472.000000
Mismatch in output[39]: expected 5767168.000000, got 3883008.000000
Mismatch in output[47]: expected 970752.000000, got -84992.000000
Mismatch in output[51]: expected 12517376.000000, got 10158080.000000
Mismatch in output[54]: expected -8716288.000000, got -12124160.000000
output: 78791 errors (15.03%) exceeds allowed rate of 5.00% (26214 errors)
I wasn't able to quickly identify the root cause. This PR does not touch the swiglu_prefill operator directly (note the stream_swiglu_prefill is a separate, different operator that seems to pass), so the error must be a regression in the upstream kernels that swiglu prefill uses (deleted in this PR) or something else.
The mlir-aie port's aie2/mm.cc dropped the conv_even rounding that aie2p/mm.cc still applies, so bf16 GEMM on aie2 (npu1/Phoenix) fell back to the core's power-up floor mode -- biasing every bf16 store low across the K reduction and tipping swiglu_prefill past its output tolerance. Source a patched aie2/mm.cc from the tree for aie2 only (aie2p sources from the package unchanged) until the rounding fix is upstreamed. Its zero.cc / aie_kernel_utils.h includes resolve from the unchanged package copies via -I. Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
Author
looks like the kernel port dropped a conv_even rounding mode on only the aie2 version of the mm kernel in mlir-aie. I'll note this as a follow on issue and retain the aie2 mm.cc file in IRON for now until the mlir-aie fix lands. |
The in-tree kernel is subject to IRON's clang-format check (aie_kernels/ is not excluded); the vendored copy was not formatted to IRON's config. Formatting only -- the conv_even rounding patch is unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
github-merge-queue
Bot
removed this pull request from the merge queue due to a conflict with the base branch
Sep 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
IRON no longer hosts C++ compute kernels. Every operator now sources its kernel
.ccfrom the installedmlir_aiepackage viaAIEContext.kernels_dir(
root_path()/include/aie_kernels), and the in-treeaie_kernels/directory isremoved. This is the IRON half of the kernel consolidation started by
mlir-aie#3599.
Added
AIEContext.kernels_dir— resolves kernel sources to the installedmlir_aiepackage (aie.utils.config.root_path()/include/aie_kernels).IRON_AIE_KERNELS_DIRenv override onkernels_dir, to point IRON at a localmlir-aie checkout for kernel development. Happy to remove this if it isn't wanted.
Changed
get_kernel_artifacts()(and the stream-op / flm-gemmkernel sourcing) now reads from
context.kernels_dirinstead of the in-treeaie_kernels/. Where mlir-aie renamed or re-homed a symbol, the operator'sKernel(...)binding was updated to match (e.g.convert_copy_f32_to_bf16→cast_f32_bf16_row,expand_int4_to_bfloat16→expand_uint4_to_bfloat16,the
_size/_epsruntime variants, andrope's two-symbol scheme).Removed
aie_kernels/directory (41 files: all.cc/.hkernels across
generic/,aie2/,aie2p/, plus its README and thekernels_from_mlir_aiemigration flag once no operator needed it).PR Merge Checklist
develcommit and pointing todevel.