Update mlir-aie to v1.4.0 and migrate the operator library to its APIs - #145
Update mlir-aie to v1.4.0 and migrate the operator library to its APIs#145atassis wants to merge 1 commit into
Conversation
andrej
left a comment
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.
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.
4017de7 to
0bdebc0
Compare
|
Rebased. The tracing changes conflicted but it was an easy fix. As anticipated, @asyms the Stream codegen will need some changes to work with MLIR-AIE 1.4.0. If they're not too involved, I think it would be easiest if you can push those changes to this branch as well (or open a PR with this branch as the base, since you likely won't have push rights to this branch). There are also a couple tests that fail for me locally: Edit: Turns out some of these failures are just because I ran this on a machine with the wrong XRT/driver version. |
CI Test Resultsb154dba (2026_08_03_17_37_33) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Smalliron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/relu
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Trends: IRON Trendsiron/applications/llama_3.2_1btest_llama_3_2_1b[llama_3.2_1b_prompt_1024_tokens_1]
test_llama_3_2_1b[llama_3.2_1b_prompt_1024_tokens_40]
test_llama_3_2_1b[llama_3.2_1b_prompt_13_tokens_1]
test_llama_3_2_1b[llama_3.2_1b_prompt_13_tokens_40]
Phoenix - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/relu
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Trends: IRON Trendsiron/operators/axpytest_axpy[input_length_2048-num_aie_columns_1-tile_size_2048-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_2-tile_size_1024-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_4-tile_size_512-scalar_factor_3.0]
iron/operators/dequanttest_dequant[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-group_size_32]
test_dequant[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-group_size_32]
test_dequant[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-group_size_32]
test_dequant[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-group_size_32]
test_dequant[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-group_size_32]
test_dequant[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-group_size_32]
iron/operators/elementwise_addtest_elementwise_add[input_length_2048-num_aie_columns_1-tile_size_2048]
test_elementwise_add[input_length_2048-num_aie_columns_2-tile_size_1024]
test_elementwise_add[input_length_2048-num_aie_columns_4-tile_size_512]
iron/operators/elementwise_multest_elementwise_mul[input_length_2048-num_aie_columns_1-tile_size_2048]
test_elementwise_mul[input_length_2048-num_aie_columns_2-tile_size_1024]
test_elementwise_mul[input_length_2048-num_aie_columns_4-tile_size_512]
iron/operators/gelutest_gelu[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048]
test_gelu[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024]
test_gelu[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024]
test_gelu[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512]
test_gelu[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512]
test_gelu[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256]
iron/operators/gemmtest_gemm[M_192-K_384-N_64-num_aie_columns_4-b_col_maj_False-c_col_maj_False-m_48-k_96-n_16-trace_size_0-partition_N_1]
test_gemm[M_192-K_384-N_64-num_aie_columns_4-b_col_maj_True-c_col_maj_True-m_48-k_96-n_16-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_1-b_col_maj_False-c_col_maj_False-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_2-b_col_maj_True-c_col_maj_False-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_384-K_1536-N_1792-num_aie_columns_4-b_col_maj_True-c_col_maj_False-m_32-k_48-n_64-trace_size_0-partition_N_1]
test_gemm[M_64-K_512-N_256-num_aie_columns_4-b_col_maj_True-c_col_maj_False-m_16-k_64-n_64-trace_size_0-partition_N_4]
iron/operators/gemvtest_gemv[M_128-K_128-num_aie_columns_1-tile_size_input_32-tile_size_output_128]
test_gemv[M_2048-K_8192-num_aie_columns_1-tile_size_input_1-tile_size_output_2048]
test_gemv[M_2048-K_8192-num_aie_columns_2-tile_size_input_1-tile_size_output_1024]
test_gemv[M_2048-K_8192-num_aie_columns_4-tile_size_input_1-tile_size_output_512]
test_gemv[M_8192-K_2048-num_aie_columns_1-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_2-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_4-tile_size_input_4-tile_size_output_1024]
test_gemv_batched[M_1024-K_1024-num_aie_columns_1-tile_size_input_1-tile_size_output_64-num_batches_2]
test_gemv_batched[M_1026-K_64-num_aie_columns_1-tile_size_input_1-tile_size_output_2-num_batches_2]
test_gemv_batched[M_256-K_128-num_aie_columns_1-tile_size_input_1-tile_size_output_256-num_batches_4]
test_gemv_batched[M_64-K_1536-num_aie_columns_1-tile_size_input_1-tile_size_output_64-num_batches_8]
test_gemv_gelu[M_128-K_128-num_aie_columns_1-tile_size_input_32-tile_size_output_128]No metrics available. test_gemv_gelu[M_2048-K_8192-num_aie_columns_1-tile_size_input_1-tile_size_output_2048]No metrics available. test_gemv_gelu[M_8192-K_2048-num_aie_columns_1-tile_size_input_4-tile_size_output_1024]No metrics available. iron/operators/layer_normtest_layer_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048]
test_layer_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024]
test_layer_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024]
test_layer_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512]
test_layer_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512]
test_layer_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256]
iron/operators/leaky_relutest_leaky_relu[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-alpha_0.01]
test_leaky_relu[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-alpha_0.1]
test_leaky_relu[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-alpha_0.25]
test_leaky_relu[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-alpha_0.01]
test_leaky_relu[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-alpha_0.01]
test_leaky_relu[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-alpha_0.01]
test_leaky_relu[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-alpha_0.01]
test_leaky_relu[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-alpha_0.01]
iron/operators/mem_copytest_mem_copy[input_length_2048-num_cores_1-num_channels_1-bypass_False-tile_size_2048]
test_mem_copy[input_length_2048-num_cores_2-num_channels_1-bypass_False-tile_size_1024]
test_mem_copy[input_length_2048-num_cores_2-num_channels_2-bypass_False-tile_size_1024]
test_mem_copy[input_length_2048-num_cores_4-num_channels_1-bypass_False-tile_size_512]
test_mem_copy[input_length_2048-num_cores_4-num_channels_2-bypass_False-tile_size_512]
test_mem_copy[input_length_2048-num_cores_8-num_channels_2-bypass_False-tile_size_256]
iron/operators/rms_normtest_rms_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-weighted_False]
iron/operators/ropetest_rope[rows_32-cols_512-angle_rows_32-aie_columns_1-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_2-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_4-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_1-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_2-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_4-method_type_0]
iron/operators/softmaxtest_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_1024]
test_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_2048]
test_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_512]
iron/operators/swiglu_decodetest_swiglu_decode[embedding_dim_1024-hidden_dim_3584]
test_swiglu_decode[embedding_dim_2048-hidden_dim_2048]
iron/operators/swiglu_prefilltest_swiglu_prefill[seq_len_256-embedding_dim_2048-hidden_dim_2048-prio_accuracy_False]
iron/operators/transposetest_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8-num_batches_1]
test_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8-num_batches_2]
test_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8]
test_transpose[M_2048-N_64-aie_columns_1-channels_2-m_64-n_64-s_8-num_batches_1]
test_transpose[M_2048-N_64-aie_columns_1-channels_2-m_64-n_64-s_8]
Phoenix - ExamplesIRONTested on Trends: IRON Trends |
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.