Skip to content

Update mlir-aie to v1.4.0 and migrate the operator library to its APIs - #145

Open
atassis wants to merge 1 commit into
amd:develfrom
atassis:migrate/mlir-aie-v1.4.0
Open

Update mlir-aie to v1.4.0 and migrate the operator library to its APIs#145
atassis wants to merge 1 commit into
amd:develfrom
atassis:migrate/mlir-aie-v1.4.0

Conversation

@atassis

@atassis atassis commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

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.txt alone. I hit this running the
operator 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.

before after
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) call fn directly
rt.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 of
the handle, so it is bound once on .prod(tile=...) / .cons(tile=...) instead
of 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-elf and
--no-compile-host. iron/common/compilation/base.py still 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.

before after
--aie-generate-xclbin --get-xclbin
--aie-generate-npu-insts --get-npu-insts
--generate-full-elf --get-full-elf
--no-compile-host dropped
--no-compile dropped

The last two are deletions rather than renames. Requesting only --get-npu-insts
and not --get-xclbin is precisely what --no-compile used to mean, so output
selection by omission replaces it.

Two places that needed more than a rename

gemm records tensor access patterns from inside the sequence body. The body
now runs at resolve_program() rather than at Runtime() construction, so with a
literal translation generate_taps=True would have returned three empty
sequences, quietly and with no error. The program is now resolved before the taps
are read.

mem_copy guarded rt.start on not bypass. Since workers moved to the
Program constructor, that becomes workers=None if bypass else my_workers;
translating the guard literally would have started workers in the bypass path.

Testing

pytest iron/operators on Strix (npu2), full suite including the extensive
set, against a local build of mlir-aie at the v1.4.0 API.

2810 passed, 5 skipped in 1021.64s

All 21 operator suites green. That includes the configurations that exercise the
parts of this change that needed judgement rather than a rename: mem_copy with
bypass=True (the conditional-workers case), mha at 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 operators
taking the full-ELF path.

Three branches no test reaches, checked by generating their MLIR directly:

  • gemm with generate_taps=True returns 8/8/4 non-empty access patterns. Under
    a literal translation this returns three empty sequences with no error, which
    is what motivated the reorder above.
  • gemm with separate_c_tiles=True, used by iron/applications/llama_3.2_1b
    but by no test, still builds.
  • strided_copy has no test.py at all. Both paths build, and the
    offset-parameter path still emits aiex.sync_scratchpad_parameters_from_host,
    the offset_parameter BD attributes and the aiex.scratchpad_parameter
    declaration.

Still unverified: softmax with vector_size_parameter set, and the
llama_3.2_1b end-to-end chain. Both are reachable only from that application,
whose tests skip without IRON_EXAMPLE_WEIGHTS_DIR, so I could not run them.

black and reuse lint clean.

Not included

The nightly update-mlir-aie.yml run that opened #144 reported failure even
though the PR was created: after pushing the branch it calls workflow_dispatch
on 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.

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 andrej left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants