Update mlir-aie to v1.4.0 and migrate the operator library to its APIs - #145
Open
atassis wants to merge 1 commit into
Open
Update mlir-aie to v1.4.0 and migrate the operator library to its APIs#145atassis wants to merge 1 commit into
atassis wants to merge 1 commit into
Conversation
v1.4.0 carries two breaking changes that the pinned v1.3.5.dev20 predates, so the version bump alone does not build. mlir-aie #3387 reworked the IRON Runtime into a callback body. Runtime() plus 'with rt.sequence(...)' is gone; the constructor now takes (seq_fn, fn_args), fill/drain moved onto the ObjectFifo handle, workers moved to Program(workers=), task_group()/finish_task_group() became TaskGroup()/tg.finish(), set_barrier() became barrier.set(), inline_ops() became a direct call, enable_trace() moved to Program and sync_parameters() to module scope. The per-transfer tile= argument is now bound on .prod()/.cons(), since placement is a property of the handle. mlir-aie #3364 unified aiecc's output selection under --get-<name>, removing --aie-generate-xclbin, --aie-generate-npu-insts and --no-compile-host. Asking only for the insts is what --no-compile used to mean, so that flag goes away rather than being renamed. Two spots needed more than a rename. gemm returns tensor access patterns that are recorded while the sequence body runs, and the body now runs at resolve_program() rather than at construction, so the program is resolved before the taps are read. mem_copy guarded rt.start on 'not bypass', which becomes a conditional workers= argument. Verified on Strix (npu2): the axpy suite passes 20/20 on device.
andrej
approved these changes
Jul 31, 2026
andrej
left a comment
Collaborator
There was a problem hiding this comment.
Much appreciated -- thanks for picking up the numerous upstream API changes!
Side note, I have activated the merge queue so I will be able to enqueue multiple PRs at once and have them rebase and test automatically, so hopefully I can merge a swath of your PRs today. Thanks for your patience so far.
3 tasks
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
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.
This updates mlir-aie from v1.3.5.dev20+g167f34d to v1.4.0 and migrates the
operator library to the APIs that release ships, so the bump builds.
Supersedes #144, which changes
requirements.txtalone. I hit this running theoperator library against a v1.4.0-era mlir-aie on Strix, so the changes here are
what it took to get back to green rather than a reading of the changelog.
Why the bump alone does not build
Two breaking changes landed in mlir-aie after the pinned v1.3.5.dev20 (which is
mlir-aie 167f34d, 8 July) and before v1.4.0.
#3387 reworked the IRON Runtime into a callback body.
Runtime()+with rt.sequence(*tys) as (...)Runtime(seq_fn, fn_args)rt.fill(h, src, ...)/rt.drain(h, dst, ...)h.fill(src, ...)/h.drain(dst, ...)rt.start(*workers)Program(dev, rt, workers=workers)rt.task_group()/rt.finish_task_group(tg)TaskGroup()/tg.finish()rt.set_barrier(b, 1)b.set(1)rt.inline_ops(fn, args)fndirectlyrt.enable_trace(...)Program.enable_trace(...)rt.sync_parameters()iron.sync_parameters()task_group=group=The per-transfer
tile=argument also goes away: placement is now a property ofthe handle, so it is bound once on
.prod(tile=...)/.cons(tile=...)insteadof being repeated on every
fill/drain.#3364 unified aiecc's output selection under
--get-<name>, removing--aie-generate-xclbin,--aie-generate-npu-insts,--generate-full-elfand--no-compile-host.iron/common/compilation/base.pystill emitted all of them,so nothing compiled at all. The xclbin path and the full-ELF path both died on
the first argument aiecc did not recognise.
--aie-generate-xclbin--get-xclbin--aie-generate-npu-insts--get-npu-insts--generate-full-elf--get-full-elf--no-compile-host--no-compileThe last two are deletions rather than renames. Requesting only
--get-npu-instsand not
--get-xclbinis precisely what--no-compileused to mean, so outputselection by omission replaces it.
Two places that needed more than a rename
gemmrecords tensor access patterns from inside the sequence body. The bodynow runs at
resolve_program()rather than atRuntime()construction, so with aliteral translation
generate_taps=Truewould have returned three emptysequences, quietly and with no error. The program is now resolved before the taps
are read.
mem_copyguardedrt.startonnot bypass. Since workers moved to theProgramconstructor, that becomesworkers=None if bypass else my_workers;translating the guard literally would have started workers in the bypass path.
Testing
pytest iron/operatorson Strix (npu2), full suite including theextensiveset, against a local build of mlir-aie at the v1.4.0 API.
All 21 operator suites green. That includes the configurations that exercise the
parts of this change that needed judgement rather than a rename:
mem_copywithbypass=True(the conditional-workers case),mhaat both 4 and 8 pipelines(the 4-pipeline arm is the one where the second Q/O handle pair is absent) and
its GQA case, and
swiglu_decode/swiglu_prefill, which are the only operatorstaking the full-ELF path.
Three branches no test reaches, checked by generating their MLIR directly:
gemmwithgenerate_taps=Truereturns 8/8/4 non-empty access patterns. Undera literal translation this returns three empty sequences with no error, which
is what motivated the reorder above.
gemmwithseparate_c_tiles=True, used byiron/applications/llama_3.2_1bbut by no test, still builds.
strided_copyhas notest.pyat all. Both paths build, and theoffset-parameter path still emits
aiex.sync_scratchpad_parameters_from_host,the
offset_parameterBD attributes and theaiex.scratchpad_parameterdeclaration.
Still unverified:
softmaxwithvector_size_parameterset, and thellama_3.2_1bend-to-end chain. Both are reachable only from that application,whose tests skip without
IRON_EXAMPLE_WEIGHTS_DIR, so I could not run them.blackandreuse lintclean.Not included
The nightly
update-mlir-aie.ymlrun that opened #144 reported failure eventhough the PR was created: after pushing the branch it calls
workflow_dispatchon a workflow that has no such trigger, gets HTTP 422, and exits 1. That is why
#144 carries no checks. It is unrelated to this migration, and would fail on any
bump, so it is worth a separate fix.