From 0bdebc0d2b2bd6478f874dab7a64770f6ba2a3e4 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Fri, 31 Jul 2026 00:05:33 +0300 Subject: [PATCH 1/6] Update mlir-aie to v1.4.0 and migrate the operator library to its APIs 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-, 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. --- iron/common/compilation/base.py | 12 +-- iron/operators/_trace.py | 10 +- iron/operators/axpy/design.py | 43 ++++---- iron/operators/binary_elementwise_design.py | 43 ++++---- iron/operators/channeled_unary_design.py | 36 ++++--- iron/operators/dequant/design.py | 38 ++++--- iron/operators/gemm/design.py | 85 ++++++++-------- iron/operators/gemv/design.py | 36 ++++--- iron/operators/leaky_relu/design.py | 34 ++++--- iron/operators/mem_copy/design.py | 101 ++++++++++--------- iron/operators/mha/design.py | 105 +++++++++----------- iron/operators/repeat/design.py | 22 ++-- iron/operators/rms_norm/design.py | 31 +++--- iron/operators/rms_norm/design_weighted.py | 38 ++++--- iron/operators/rope/design.py | 43 ++++---- iron/operators/softmax/design.py | 47 +++++---- iron/operators/strided_copy/design.py | 39 +++++--- iron/operators/transpose/design.py | 31 +++--- requirements.txt | 4 +- 19 files changed, 439 insertions(+), 359 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 8b06de537..6affb7ddb 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -518,7 +518,6 @@ def compile(self, graph): str(self.aiecc_path), "-v", f"-j{os.environ.get('AIECC_JOBS', '1')}", - "--no-compile-host", ] if self.use_chess: compile_cmd += [ @@ -534,7 +533,7 @@ def compile(self, graph): ] compile_cmd += [ "--expand-load-pdis", - "--generate-full-elf", + "--get-full-elf", "--full-elf-name", os.path.abspath(artifact.filename), *artifact.extra_flags, @@ -573,7 +572,6 @@ def compile(self, graph): str(self.aiecc_path), "-v", f"-j{os.environ.get('AIECC_JOBS', '1')}", - "--no-compile-host", ] if self.use_chess: compile_cmd += [ @@ -597,7 +595,7 @@ def compile(self, graph): 0 ] # TODO: this does not handle the case of multiple xclbins with different kernel names or flags from the same MLIR compile_cmd += first_xclbin.extra_flags + [ - "--aie-generate-xclbin", + "--get-xclbin", "--xclbin-name=" + os.path.abspath(first_xclbin.filename), "--xclbin-kernel-name=" + first_xclbin.kernel_name, ] @@ -610,10 +608,10 @@ def compile(self, graph): first_insts_bin = mlir_sources_to_insts[mlir_source][ 0 ] # TODO: this does not handle the case of multiple insts.bins with different flags from the same MLIR - if not do_compile_xclbin: - compile_cmd += ["--no-compile"] + # Outputs are selected by --get-; asking only for the insts is what + # "--no-compile" used to mean, so there is nothing to opt out of here. compile_cmd += first_insts_bin.extra_flags + [ - "--aie-generate-npu-insts", + "--get-npu-insts", "--npu-insts-name=" + os.path.abspath(first_insts_bin.filename), ] compile_cmd += [os.path.abspath(mlir_source.filename)] diff --git a/iron/operators/_trace.py b/iron/operators/_trace.py index 617584845..4dcb8943d 100644 --- a/iron/operators/_trace.py +++ b/iron/operators/_trace.py @@ -39,13 +39,11 @@ def _default_coretile_events(): ] -def maybe_enable_trace(rt, trace_size, workers, coretile_events=None): - """Configure per-op hardware trace on ``rt`` if tracing is requested. - - Call inside the ``rt.sequence(...)`` block, before ``rt.start(...)``. +def maybe_enable_trace(prog, trace_size, workers, coretile_events=None): + """Configure per-op hardware trace if tracing is requested. Args: - rt: the ``Runtime`` being built. + prog: the ``Program`` being built. trace_size: the design's ``trace_size`` argument (may be None/0). workers: the design's workers; the first ``IRON_TRACE_NTILES`` are traced. coretile_events: override the default core-tile event set. @@ -61,7 +59,7 @@ def maybe_enable_trace(rt, trace_size, workers, coretile_events=None): # meaningless (a negative slice index would silently drop the LAST worker). ntiles = max(0, int(os.environ.get("IRON_TRACE_NTILES", "1"))) - rt.enable_trace( + prog.enable_trace( ts, workers=list(workers)[:ntiles], coretile_events=( diff --git a/iron/operators/axpy/design.py b/iron/operators/axpy/design.py index 12685bd63..e9421c8ae 100644 --- a/iron/operators/axpy/design.py +++ b/iron/operators/axpy/design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ from iron.operators._trace import maybe_enable_trace @@ -84,38 +84,45 @@ def core_body(of_in1, of_in2, of_out, axpy): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C): - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) - + def sequence(A, B, C, in1_prods, in2_prods, out_conses): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): - rt.fill( - of_in1s[i].prod(), + in1_prods[i].fill( A, taps[i], - task_group=tg, + group=tg, ) - rt.fill( - of_in2s[i].prod(), + in2_prods[i].fill( B, taps[i], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): - rt.drain( - of_outs[i].cons(), + out_conses[i].drain( C, taps[i], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + + rt = Runtime( + sequence, + [ + tensor_ty, + tensor_ty, + tensor_ty, + [of_in1s[i].prod() for i in range(num_columns)], + [of_in2s[i].prod() for i in range(num_columns)], + [of_outs[i].cons() for i in range(num_columns)], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/binary_elementwise_design.py b/iron/operators/binary_elementwise_design.py index 5b75f8152..fea333f40 100644 --- a/iron/operators/binary_elementwise_design.py +++ b/iron/operators/binary_elementwise_design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ from iron.operators._trace import maybe_enable_trace @@ -83,37 +83,44 @@ def core_body(of_in1, of_in2, of_out, eltwise_fn): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C): - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) - - tg = rt.task_group() + def sequence(A, B, C, in1_prods, in2_prods, out_conses): + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): - rt.fill( - of_in1s[i].prod(), + in1_prods[i].fill( A, taps[i], - task_group=tg, + group=tg, ) - rt.fill( - of_in2s[i].prod(), + in2_prods[i].fill( B, taps[i], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): - rt.drain( - of_outs[i].cons(), + out_conses[i].drain( C, taps[i], wait=True, - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + + rt = Runtime( + sequence, + [ + tensor_ty, + tensor_ty, + tensor_ty, + [of_in1s[i].prod() for i in range(num_columns)], + [of_in2s[i].prod() for i in range(num_columns)], + [of_outs[i].cons() for i in range(num_columns)], + ], + ) # Place program components and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/channeled_unary_design.py b/iron/operators/channeled_unary_design.py index ab8e0fcb5..7cff67c60 100644 --- a/iron/operators/channeled_unary_design.py +++ b/iron/operators/channeled_unary_design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ from iron.operators._trace import maybe_enable_trace @@ -97,33 +97,39 @@ def core_fn(of_in, of_out, kernel_line): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(transfer_type, transfer_type) as (a_in, b_out): - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) - - tg = rt.task_group() + def sequence(a_in, b_out, in_prods, out_conses): + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.fill( - of_ins[i * num_channels + j].prod(), + in_prods[i * num_channels + j].fill( a_in, taps[i * num_channels + j], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.drain( - of_outs[i * num_channels + j].cons(), + out_conses[i * num_channels + j].drain( b_out, taps[i * num_channels + j], wait=True, - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + + rt = Runtime( + sequence, + [ + transfer_type, + transfer_type, + [of.prod() for of in of_ins], + [of.cons() for of in of_outs], + ], + ) # Place components and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/dequant/design.py b/iron/operators/dequant/design.py index e613e08fd..ad5cdb59d 100644 --- a/iron/operators/dequant/design.py +++ b/iron/operators/dequant/design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ @@ -118,35 +118,41 @@ def core_body(of_in1, of_out, dequant_kernel): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(in_tensor_ty, out_tensor_ty) as (A, C): - if enable_trace: - rt.enable_trace(trace_size) - rt.start(*my_workers) + def sequence(A, C, of_in1s_prods, of_outs_conss): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.fill( - of_in1s[i * num_channels + j].prod(), + of_in1s_prods[i * num_channels + j].fill( A, taps_in[i * num_channels + j], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.drain( - of_outs[i * num_channels + j].cons(), + of_outs_conss[i * num_channels + j].drain( C, taps_out[i * num_channels + j], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) - + tg.finish() + + rt = Runtime( + sequence, + [ + in_tensor_ty, + out_tensor_ty, + [of.prod() for of in of_in1s], + [of.cons() for of in of_outs], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + if enable_trace: + prog.enable_trace(trace_size) + return prog.resolve_program() diff --git a/iron/operators/gemm/design.py b/iron/operators/gemm/design.py index bfdb84426..169d848ae 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -14,6 +14,7 @@ Program, Buffer, Runtime, + TaskGroup, Worker, WorkerRuntimeBarrier, str_to_dtype, @@ -551,28 +552,21 @@ def core_fn( ) # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(A_ty, B_ty, C_ty) as (A, B, C): - maybe_enable_trace(rt, trace_size, workers) - rt.start(*workers) - + def sequence(A, B, C, A_prods, B_prods, C_conses): # Set runtime parameters - def set_rtps(*args): - for row, rtps_row in enumerate(args): - for col, rtp_row_col in enumerate(rtps_row): - rtp_row_col[0] = K_div_k - rtp_row_col[1] = n_c_row_tiles_per_core * n_c_col_tiles_per_core - - rt.inline_ops(set_rtps, rtps) + for rtps_row in rtps: + for rtp_row_col in rtps_row: + rtp_row_col[0] = K_div_k + rtp_row_col[1] = n_c_row_tiles_per_core * n_c_col_tiles_per_core # Set the barriers to 1 to allow the worker to read the # runtime parameters and start the computation for row in range(n_aie_rows): for col in range(n_aie_cols): - rt.set_barrier(workerBarriers[row][col], 1) + workerBarriers[row][col].set(1) # Task groups will be used to determine when to sync/await/free DMA runtime ops - tg = rt.task_group() + tg = TaskGroup() for tb in range(ceildiv(n_c_row_tiles_per_core, tb_max_n_rows)): for pingpong in [0, 1]: row_base = tb * tb_max_n_rows + pingpong * tb_max_n_rows // 2 @@ -630,13 +624,11 @@ def set_rtps(*args): # This line does not change MLIR output at all - it's just for recording data movement C_taps.append(C_tile) - rt.drain( - C_l2l3_fifos[col].cons(), + C_conses[col].drain( C, tap=C_tile, wait=True, - task_group=tg, - tile=Tile(col, 0), + group=tg, ) for tile_row in range(current_tb_n_rows): @@ -685,13 +677,11 @@ def set_rtps(*args): sizes=C_sizes, strides=C_strides, ) - rt.drain( - C_l2l3_fifos[col].cons(), + C_conses[col].drain( C, tap=C_tile, wait=True, - task_group=tg, - tile=Tile(col, 0), + group=tg, ) # This line does not change MLIR output at all - it's just for recording data movement C_taps.append(C_tile) @@ -720,14 +710,10 @@ def set_rtps(*args): # always equal to n_aie_rows since we have n_aie_rows row tiles for matrix A if col < n_aie_rows: - rt.fill( - A_l3l2_fifos[col].prod(), + A_prods[col].fill( A, tap=A_tiles[tile_offset], - task_group=tg, - tile=Tile( - 2 * col if n_aie_cols == 8 else col, 0 - ), # alternate columns in full 4x8 NPU2 case + group=tg, ) # Use the calculated sizes/strides/offsets to record the data movement # caused by the above call to npu_dma_memcpy_nd. @@ -751,21 +737,45 @@ def set_rtps(*args): # |0011 0011 | # |0011 0011 | # ---------------- - rt.fill( - B_l3l2_fifos[col].prod(), + B_prods[col].fill( B, tap=B_tiles[col], - task_group=tg, - tile=Tile(col, 0), + group=tg, ) # These lines do not change MLIR output at all - they are just for recording data movement A_taps.append(A_tiles[tile_offset]) B_taps.append(B_tiles[col]) if tb > 0 or (tb == 0 and pingpong > 0): - rt.finish_task_group(tg) - tg = rt.task_group() - rt.finish_task_group(tg) + tg.finish() + tg = TaskGroup() + tg.finish() + + rt = Runtime( + sequence, + [ + A_ty, + B_ty, + C_ty, + # The shim tile that used to be named per-transfer is now a property + # of the handle, so it is bound here instead. + [ + f.prod(tile=Tile(2 * c if n_aie_cols == 8 else c, 0)) + for c, f in enumerate(A_l3l2_fifos) + ], + [f.prod(tile=Tile(c, 0)) for c, f in enumerate(B_l3l2_fifos)], + [f.cons(tile=Tile(c, 0)) for c, f in enumerate(C_l2l3_fifos)], + ], + ) + + # Create the program from the device type and runtime + my_program = Program(dev_ty, rt, workers=workers) + maybe_enable_trace(my_program, trace_size, workers) + + # Place components (assign them resources on the device) and generate an MLIR module. + # This is what runs the sequence body, so it must happen before the taps it + # records are read. + module = my_program.resolve_program() if generate_taps: # If generate taps is true, return a representation of tensor access patterns @@ -776,11 +786,6 @@ def set_rtps(*args): TensorAccessSequence.from_taps(C_taps), ) - # Create the program from the device type and runtime - my_program = Program(dev_ty, rt) - - # Place components (assign them resources on the device) and generate an MLIR module - module = my_program.resolve_program() return module diff --git a/iron/operators/gemv/design.py b/iron/operators/gemv/design.py index f08b44f58..28d8b42c4 100644 --- a/iron/operators/gemv/design.py +++ b/iron/operators/gemv/design.py @@ -8,7 +8,7 @@ from aie.dialects.aie import T from aie.helpers.dialects.scf import _for as range_ from aie.helpers.taplib import TensorAccessPattern -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker """ Matrix-vector design @@ -250,33 +250,41 @@ def coalesced_tap(L3_ty, col_off, split, bstride): for col in range(cols) ] - rt = Runtime() - with rt.sequence(L3_A_ty, L3_B_ty, L3_C_ty) as (A, B, C): - rt.start(*workers) - tg_b = rt.task_group() + def sequence(A, B, C, B_L3L1_fifos_prods, A_L3L1_fifos_prods, C_L1L3_fifos_conss): + tg_b = TaskGroup() for col in range(cols): # Simple linear transfer of B, includes all batches in sequence - rt.fill(B_L3L1_fifos[col].prod(), B, B_tap, task_group=tg_b) + B_L3L1_fifos_prods[col].fill(B, B_tap, group=tg_b) # Coalesced: one iterated BD per column covers all batches (num_waits==1, a # single drain wait for the whole column). Fallback (incl. num_batches==1): the # stock per-batch unroll (num_waits==num_batches, one wait per batch). The fills # and drains are otherwise identical; only the TAP and the wait count differ. num_waits = 1 if coalesce else num_batches for w in range(num_waits): - tg_ac = rt.task_group() + tg_ac = TaskGroup() for col in range(cols): a_tap = A_taps_coalesced[col] if coalesce else A_taps[col][w] - rt.fill(A_L3L1_fifos[col].prod(), A, a_tap, task_group=tg_ac) + A_L3L1_fifos_prods[col].fill(A, a_tap, group=tg_ac) for col in range(cols): c_tap = C_taps_coalesced[col] if coalesce else C_taps[col][w] - rt.drain( - C_L1L3_fifos[col].cons(), + C_L1L3_fifos_conss[col].drain( C, c_tap, - task_group=tg_ac, + group=tg_ac, wait=True, ) - rt.finish_task_group(tg_ac) - rt.finish_task_group(tg_b) + tg_ac.finish() + tg_b.finish() - return Program(dev, rt).resolve_program() + rt = Runtime( + sequence, + [ + L3_A_ty, + L3_B_ty, + L3_C_ty, + [of.prod() for of in B_L3L1_fifos], + [of.prod() for of in A_L3L1_fifos], + [of.cons() for of in C_L1L3_fifos], + ], + ) + return Program(dev, rt, workers=workers).resolve_program() diff --git a/iron/operators/leaky_relu/design.py b/iron/operators/leaky_relu/design.py index b3ea1fb4b..408a311ba 100644 --- a/iron/operators/leaky_relu/design.py +++ b/iron/operators/leaky_relu/design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ from iron.operators._trace import maybe_enable_trace @@ -93,34 +93,40 @@ def core_fn(of_in, of_out, leaky_relu_line): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(transfer_type, transfer_type) as (a_in, b_out): - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) + def sequence(a_in, b_out, of_ins_prods, of_outs_conss): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.fill( - of_ins[i * num_channels + j].prod(), + of_ins_prods[i * num_channels + j].fill( a_in, taps[i * num_channels + j], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.drain( - of_outs[i * num_channels + j].cons(), + of_outs_conss[i * num_channels + j].drain( b_out, taps[i * num_channels + j], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + rt = Runtime( + sequence, + [ + transfer_type, + transfer_type, + [of.prod() for of in of_ins], + [of.cons() for of in of_outs], + ], + ) # Place components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/mem_copy/design.py b/iron/operators/mem_copy/design.py index 1eb3685eb..788345dbc 100644 --- a/iron/operators/mem_copy/design.py +++ b/iron/operators/mem_copy/design.py @@ -9,6 +9,7 @@ import math from aie.iron import ( + TaskGroup, Kernel, ObjectFifo, Program, @@ -242,13 +243,7 @@ def core_fn(of_in, of_out, mem_copy_line): # -------------------------------------------------------------------------- # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(transfer_type, transfer_type) as (a_in, b_out): - # Start the workers if not bypass - if not bypass: - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) - + def sequence(a_in, b_out, of_ins_prods, of_outs_conss): # Calculate how much of workload can be partitioned evenly and what's remaining minimum_work_size = ( line_size * num_cores @@ -263,20 +258,19 @@ def core_fn(of_in, of_out, mem_copy_line): size, num_cores, line_size, whole_partition_size ) - tg_out = rt.task_group() # Use taskgroup for parallel drain tasks + tg_out = TaskGroup() # Use taskgroup for parallel drain tasks # Fill the input objectFIFOs with data for i in range(num_cores): - rt.fill(of_ins[i].prod(), a_in, taps[i], task_group=tg_out) + of_ins_prods[i].fill(a_in, taps[i], group=tg_out) # Drain the output objectFIFOs with data for i in range(num_cores): - rt.drain( - of_outs[i].cons(), + of_outs_conss[i].drain( b_out, taps[i], wait=True, # wait for the transfer to complete and data to be available - task_group=tg_out, + group=tg_out, ) - rt.finish_task_group(tg_out) + tg_out.finish() # Runtime for the part of the workload partially partitionable to the cores utilized if partial_work_size > 0: @@ -310,46 +304,42 @@ def core_fn(of_in, of_out, mem_copy_line): and partial_config.partial_tap is not None ): # Fill the last objfifo with padding+real data - tg_out = rt.task_group() + tg_out = TaskGroup() tg_count = 0 for padding_tap_repeat, padding_tap in zip( partial_config.padding_tap_repeats, partial_config.padding_taps ): for _ in range(padding_tap_repeat): if tg_count % TASK_GROUP_SIZE == 0: - rt.fill( - of_ins[objfifo_idx].prod(), + of_ins_prods[objfifo_idx].fill( a_in, padding_tap, wait=True, - task_group=tg_out, + group=tg_out, ) - rt.finish_task_group(tg_out) - tg_out = rt.task_group() + tg_out.finish() + tg_out = TaskGroup() else: - rt.fill( - of_ins[objfifo_idx].prod(), + of_ins_prods[objfifo_idx].fill( a_in, padding_tap, - task_group=tg_out, + group=tg_out, ) tg_count += 1 if tg_count % TASK_GROUP_SIZE == 0: - rt.fill( - of_ins[objfifo_idx].prod(), + of_ins_prods[objfifo_idx].fill( a_in, partial_config.partial_tap, wait=True, - task_group=tg_out, + group=tg_out, ) - rt.finish_task_group(tg_out) - tg_out = rt.task_group() + tg_out.finish() + tg_out = TaskGroup() else: - rt.fill( - of_ins[objfifo_idx].prod(), + of_ins_prods[objfifo_idx].fill( a_in, partial_config.partial_tap, - task_group=tg_out, + group=tg_out, ) tg_count += 1 # Drain the last objfifo with padding+real data @@ -358,53 +348,62 @@ def core_fn(of_in, of_out, mem_copy_line): ): for _ in range(padding_tap_repeat): if tg_count % TASK_GROUP_SIZE == 0: - rt.drain( - of_outs[objfifo_idx].cons(), + of_outs_conss[objfifo_idx].drain( b_out, padding_tap, wait=True, - task_group=tg_out, + group=tg_out, ) - rt.finish_task_group(tg_out) - tg_out = rt.task_group() + tg_out.finish() + tg_out = TaskGroup() else: - rt.drain( - of_outs[objfifo_idx].cons(), + of_outs_conss[objfifo_idx].drain( b_out, padding_tap, - task_group=tg_out, + group=tg_out, ) tg_count += 1 - rt.drain( - of_outs[objfifo_idx].cons(), + of_outs_conss[objfifo_idx].drain( b_out, partial_config.partial_tap, wait=True, - task_group=tg_out, + group=tg_out, ) - rt.finish_task_group(tg_out) + tg_out.finish() objfifo_idx += 1 else: - tg_out = rt.task_group() # Use taskgroup for parallel drain tasks + tg_out = TaskGroup() # Use taskgroup for parallel drain tasks for j in range(partial_config.num_cores_with_full_tiles): # Fill the input objectFIFOs with valid data - rt.fill( - of_ins[objfifo_idx + j].prod(), + of_ins_prods[objfifo_idx + j].fill( a_in, partial_config.full_taps[j], - task_group=tg_out, + group=tg_out, ) for j in range(partial_config.num_cores_with_full_tiles): # Drain the output objectFIFOs with valid data - rt.drain( - of_outs[objfifo_idx + j].cons(), + of_outs_conss[objfifo_idx + j].drain( b_out, partial_config.full_taps[j], wait=True, - task_group=tg_out, + group=tg_out, ) - rt.finish_task_group(tg_out) + tg_out.finish() objfifo_idx += partial_config.num_cores_with_full_tiles + rt = Runtime( + sequence, + [ + transfer_type, + transfer_type, + [of.prod() for of in of_ins], + [of.cons() for of in of_outs], + ], + ) # Place components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + # bypass means the DMAs run without any compute worker, as `rt.start` was + # previously guarded by the same condition. + prog = Program(dev, rt, workers=None if bypass else my_workers) + if not bypass: + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/mha/design.py b/iron/operators/mha/design.py index d5dac245a..61397fbd2 100644 --- a/iron/operators/mha/design.py +++ b/iron/operators/mha/design.py @@ -15,6 +15,7 @@ ObjectFifo, Program, Runtime, + TaskGroup, Worker, Buffer, WorkerRuntimeBarrier, @@ -775,31 +776,27 @@ def legalize_tas(tas: TensorAccessSequence): # print_tap_seq_info(O_tiles, "O") # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(Q_ty, KV_ty, KV_ty, Q_ty) as (Q, K, V, O): - - def set_mha_rtps(): - for j in range(3): - for i in range(number_of_pipelines): - mha_rtps_list[j][i][0] = num_q_block_per_pipeline - mha_rtps_list[j][i][1] = num_kv_blocks - mha_rtps_list[j][i][2] = S_q_eff - mha_rtps_list[j][i][3] = S_kv_eff - - rt.inline_ops(set_mha_rtps, ()) - + # The shim tile that used to be named per-transfer is now a property of the + # handle, so the handles are bound up front and passed into the sequence. + inQ_h = inQ.prod(tile=Tile(col=4, row=0)) + inQ2_h = inQ2.prod(tile=Tile(col=4, row=0)) if number_of_pipelines > 6 else None + inK_h = inK.prod(tile=Tile(col=5, row=0)) + inV_h = inV.prod(tile=Tile(col=6, row=0)) + memO_h = memO.cons(tile=Tile(col=7, row=0)) + memO2_h = memO2.cons(tile=Tile(col=7, row=0)) if number_of_pipelines > 6 else None + + def sequence(Q, K, V, O, inQ_h, inQ2_h, inK_h, inV_h, memO_h, memO2_h): + # The body is eager now, so the RTP writes are a plain loop (was inline_ops). for j in range(3): for i in range(number_of_pipelines): - rt.set_barrier(worker_barrier_list[j][i], 1) + mha_rtps_list[j][i][0] = num_q_block_per_pipeline + mha_rtps_list[j][i][1] = num_kv_blocks + mha_rtps_list[j][i][2] = S_q_eff + mha_rtps_list[j][i][3] = S_kv_eff - maybe_enable_trace( - rt, trace_size, matmul_workers + softmax_workers + matmul_pv_workers - ) - - for i in range(number_of_pipelines): - rt.start(matmul_workers[i]) - rt.start(softmax_workers[i]) - rt.start(matmul_pv_workers[i]) + for j in range(3): + for i in range(number_of_pipelines): + worker_barrier_list[j][i].set(1) for head_idx in range(heads): @@ -808,67 +805,54 @@ def set_mha_rtps(): for q_block_idx in range(num_q_block_per_pipeline): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() if number_of_pipelines > 6: - rt.fill( - inQ.prod(), + inQ_h.fill( Q, tap=Q_tiles[ 2 * head_idx * num_q_block_per_pipeline + q_block_idx * 2 ], - tile=Tile(col=4, row=0), - task_group=tg, + group=tg, ) - rt.fill( - inQ2.prod(), + inQ2_h.fill( Q, tap=Q_tiles[ 2 * head_idx * num_q_block_per_pipeline + q_block_idx * 2 + 1 ], - tile=Tile(col=4, row=0), - task_group=tg, + group=tg, ) else: - rt.fill( - inQ.prod(), + inQ_h.fill( Q, tap=Q_tiles[head_idx * num_q_block_per_pipeline + q_block_idx], - tile=Tile(col=4, row=0), - task_group=tg, + group=tg, ) # Thow on bd containing the full K and V in the object fifo, then does it transfer cunks of inKV size at the time? - rt.fill( - inK.prod(), + inK_h.fill( K, tap=K_tiles[kv_head_idx], - tile=Tile(col=5, row=0), - task_group=tg, + group=tg, ) - rt.fill( - inV.prod(), + inV_h.fill( V, tap=V_tiles[kv_head_idx], - tile=Tile(col=6, row=0), - task_group=tg, + group=tg, ) if number_of_pipelines > 6: - rt.drain( - memO.cons(), + memO_h.drain( O, tap=O_tiles[ 2 * head_idx * num_q_block_per_pipeline + q_block_idx * 2 ], wait=True, - tile=Tile(col=7, row=0), - task_group=tg, + group=tg, ) - rt.drain( - memO2.cons(), + memO2_h.drain( O, tap=O_tiles[ 2 * head_idx * num_q_block_per_pipeline @@ -876,24 +860,31 @@ def set_mha_rtps(): + 1 ], wait=True, - tile=Tile(col=7, row=0), - task_group=tg, + group=tg, ) else: - rt.drain( - memO.cons(), + memO_h.drain( O, tap=O_tiles[head_idx * num_q_block_per_pipeline + q_block_idx], wait=True, - tile=Tile(col=7, row=0), - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + + rt = Runtime( + sequence, + [Q_ty, KV_ty, KV_ty, Q_ty, inQ_h, inQ2_h, inK_h, inV_h, memO_h, memO2_h], + ) # Create the program from the device type and runtime dev_ty = NPU2() - my_program = Program(dev_ty, rt) + my_program = Program( + dev_ty, rt, workers=matmul_workers + softmax_workers + matmul_pv_workers + ) + maybe_enable_trace( + my_program, trace_size, matmul_workers + softmax_workers + matmul_pv_workers + ) # Place components (assign them resources on the device) and generate an MLIR module module = my_program.resolve_program() diff --git a/iron/operators/repeat/design.py b/iron/operators/repeat/design.py index 4e6f2ac17..4058394e6 100644 --- a/iron/operators/repeat/design.py +++ b/iron/operators/repeat/design.py @@ -8,7 +8,7 @@ import numpy as np from aie.dialects.aiex import TensorAccessPattern -from aie.iron import ObjectFifo, Program, Runtime +from aie.iron import ObjectFifo, Program, Runtime, TaskGroup def repeat(dev, dtype, rows, cols, repeat, transfer_size=None): @@ -61,11 +61,19 @@ def repeat(dev, dtype, rows, cols, repeat, transfer_size=None): fifo_in = ObjectFifo(transfer_ty, name="fifo_in", depth=2) fifo_out = fifo_in.cons().forward(name="fifo_out", depth=2) - rt = Runtime() - with rt.sequence(inp_ty, out_ty) as (inp, out): - tg = rt.task_group() - rt.fill(fifo_in.prod(), inp, input_tap, task_group=tg) - rt.drain(fifo_out.cons(), out, output_tap, task_group=tg, wait=True) - rt.finish_task_group(tg) + def sequence(inp, out, fifo_in_prod, fifo_out_cons): + tg = TaskGroup() + fifo_in_prod.fill(inp, input_tap, group=tg) + fifo_out_cons.drain(out, output_tap, group=tg, wait=True) + tg.finish() + rt = Runtime( + sequence, + [ + inp_ty, + out_ty, + fifo_in.prod(), + fifo_out.cons(), + ], + ) return Program(dev, rt).resolve_program() diff --git a/iron/operators/rms_norm/design.py b/iron/operators/rms_norm/design.py index d2183f624..1e5f7159e 100644 --- a/iron/operators/rms_norm/design.py +++ b/iron/operators/rms_norm/design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.iron.device import NPU1, NPU2 from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ @@ -94,33 +94,38 @@ def core_body(of_in1, of_out, rms_norm_kernel): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, tensor_ty) as (A, C): - rt.start(*my_workers) + def sequence(A, C, of_in1s_prods, of_outs_conss): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.fill( - of_in1s[i * num_channels + j].prod(), + of_in1s_prods[i * num_channels + j].fill( A, taps[i * num_channels + j], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.drain( - of_outs[i * num_channels + j].cons(), + of_outs_conss[i * num_channels + j].drain( C, taps[i * num_channels + j], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + rt = Runtime( + sequence, + [ + tensor_ty, + tensor_ty, + [of.prod() for of in of_in1s], + [of.cons() for of in of_outs], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + return Program(dev, rt, workers=my_workers).resolve_program() diff --git a/iron/operators/rms_norm/design_weighted.py b/iron/operators/rms_norm/design_weighted.py index e5333feb5..f7c55bd68 100644 --- a/iron/operators/rms_norm/design_weighted.py +++ b/iron/operators/rms_norm/design_weighted.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.iron.device import NPU1, NPU2 from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ @@ -140,42 +140,48 @@ def core_body_mul(of_in1, of_in2, of_out2, eltwise_mul): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, weights_ty, tensor_ty) as (A, B, C): - rt.start(*my_workers) + def sequence(A, B, C, of_in1s_prods, of_in2s_prods, of_out2s_conss): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): for j in range(num_channels): idx = i * num_channels + j - rt.fill( - of_in1s[idx].prod(), + of_in1s_prods[idx].fill( A, taps[idx], - task_group=tg, + group=tg, ) # Fill weights (one per channel) for j in range(num_channels): - rt.fill( - of_in2s[j].prod(), + of_in2s_prods[j].fill( B, - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_columns): for j in range(num_channels): idx = i * num_channels + j - rt.drain( - of_out2s[idx].cons(), + of_out2s_conss[idx].drain( C, taps[idx], wait=True, - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + rt = Runtime( + sequence, + [ + tensor_ty, + weights_ty, + tensor_ty, + [of.prod() for of in of_in1s], + [of.prod() for of in of_in2s], + [of.cons() for of in of_out2s], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + return Program(dev, rt, workers=my_workers).resolve_program() diff --git a/iron/operators/rope/design.py b/iron/operators/rope/design.py index 5d8e2ccff..a1105be0c 100644 --- a/iron/operators/rope/design.py +++ b/iron/operators/rope/design.py @@ -17,7 +17,7 @@ import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.iron.device import NPU1, NPU2 from aie.helpers.taplib.tap import TensorAccessPattern from aie.helpers.dialects.scf import _for as range_ @@ -128,38 +128,45 @@ def core_body(of_in, of_lut, of_out, rope_kernel): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, angle_ty, tensor_ty) as (A, B, C): - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) + def sequence(A, B, C, of_in_prods, of_lut_prods, of_out_conss): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_aie_columns): - rt.fill( - of_in[i].prod(), + of_in_prods[i].fill( A, tensor_taps[i], - task_group=tg, + group=tg, ) - rt.fill( - of_lut[i].prod(), + of_lut_prods[i].fill( B, angle_taps[i], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_aie_columns): - rt.drain( - of_out[i].cons(), + of_out_conss[i].drain( C, tensor_taps[i], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) - + tg.finish() + + rt = Runtime( + sequence, + [ + tensor_ty, + angle_ty, + tensor_ty, + [of.prod() for of in of_in], + [of.prod() for of in of_lut], + [of.cons() for of in of_out], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/softmax/design.py b/iron/operators/softmax/design.py index 1aca73796..e798956da 100644 --- a/iron/operators/softmax/design.py +++ b/iron/operators/softmax/design.py @@ -10,9 +10,11 @@ ScratchpadParameter, Program, Runtime, + TaskGroup, Worker, Buffer, WorkerRuntimeBarrier, + sync_parameters, ) from aie.iron.device import NPU1, NPU2 from aie.helpers.taplib.tap import TensorAccessPattern @@ -156,51 +158,54 @@ def worker_args(i, j): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, tensor_ty) as (A, C): - maybe_enable_trace(rt, trace_size, my_workers) - rt.start(*my_workers) - + def sequence(A, C, in1_prods, out_conses): if use_scratchpad: # The host writes vector_size into the scratchpad via # ParameterScratchpad before each dispatch; sync delivers it to the # per-core parameter buffer. - rt.sync_parameters() + sync_parameters() else: # Set the static (compile-time) run-time parameter controlling how # many elements each core processes. - def set_rtps(*args): - for rtp in args: - rtp[0] = rtp_vector_size - - rt.inline_ops(set_rtps, rtps) + for rtp in rtps: + rtp[0] = rtp_vector_size for i in range(num_aie_columns * num_channels): - rt.set_barrier(barriers[i], 1) + barriers[i].set(1) # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_aie_columns): for j in range(num_channels): - rt.fill( - of_in1s[i * num_channels + j].prod(), + in1_prods[i * num_channels + j].fill( A, taps[i * num_channels + j], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs with data for i in range(num_aie_columns): for j in range(num_channels): - rt.drain( - of_outs[i * num_channels + j].cons(), + out_conses[i * num_channels + j].drain( C, taps[i * num_channels + j], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + + rt = Runtime( + sequence, + [ + tensor_ty, + tensor_ty, + [of.prod() for of in of_in1s], + [of.cons() for of in of_outs], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + prog = Program(dev, rt, workers=my_workers) + maybe_enable_trace(prog, trace_size, my_workers) + return prog.resolve_program() diff --git a/iron/operators/strided_copy/design.py b/iron/operators/strided_copy/design.py index e6ef483f1..8d6813d85 100644 --- a/iron/operators/strided_copy/design.py +++ b/iron/operators/strided_copy/design.py @@ -11,7 +11,14 @@ import numpy as np from aie.dialects.aiex import TensorAccessPattern -from aie.iron import ObjectFifo, ScratchpadParameter, Program, Runtime +from aie.iron import ( + ObjectFifo, + Program, + Runtime, + ScratchpadParameter, + TaskGroup, + sync_parameters, +) def strided_copy( @@ -130,27 +137,33 @@ def strided_copy( for c in range(num_aie_channels) ] - rt = Runtime() - with rt.sequence(inp_ty, out_ty) as (inp, out): + def sequence(inp, out, fifos_in_prods, fifos_out_conss): if in_offset_param is not None or out_offset_param is not None: - rt.sync_parameters() - tg = rt.task_group() + sync_parameters() + tg = TaskGroup() for c in range(num_aie_channels): - rt.fill( - fifos_in[c].prod(), + fifos_in_prods[c].fill( inp, input_taps[c], - task_group=tg, + group=tg, offset_parameter=in_offset_param, ) - rt.drain( - fifos_out[c].cons(), + fifos_out_conss[c].drain( out, output_taps[c], - task_group=tg, + group=tg, wait=True, offset_parameter=out_offset_param, ) - rt.finish_task_group(tg) - + tg.finish() + + rt = Runtime( + sequence, + [ + inp_ty, + out_ty, + [of.prod() for of in fifos_in], + [of.cons() for of in fifos_out], + ], + ) return Program(dev, rt).resolve_program() diff --git a/iron/operators/transpose/design.py b/iron/operators/transpose/design.py index afbc7a21a..bb0c3348f 100644 --- a/iron/operators/transpose/design.py +++ b/iron/operators/transpose/design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ @@ -152,36 +152,41 @@ def core_body(of_in1, of_out, transpose_kernel): ] # Runtime operations to move data to/from the AIE-array - rt = Runtime() - with rt.sequence(tensor_ty, tensor_ty) as (A, C): - rt.start(*my_workers) + def sequence(A, C, of_in1s_L3L2_prods, of_outs_conss): # One task group per batch (each a parallel fill+drain over all columns/channels), so the # num_batches contiguous matrices stream through the same FIFOs in sequence. for batch in range(num_batches): # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. - tg = rt.task_group() + tg = TaskGroup() # Fill the input objectFIFOs with data for i in range(num_columns): for j in range(num_channels): - rt.fill( - of_in1s_L3L2[i * num_channels + j].prod(), + of_in1s_L3L2_prods[i * num_channels + j].fill( A, taps_in_L3L2[i * num_channels + j][batch], - task_group=tg, + group=tg, ) # Drain the output objectFIFOs of data for i in range(num_columns): for j in range(num_channels): - rt.drain( - of_outs[i * num_channels + j].cons(), + of_outs_conss[i * num_channels + j].drain( C, taps_out_L1L3[i * num_channels + j][batch], wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + group=tg, ) - rt.finish_task_group(tg) + tg.finish() + rt = Runtime( + sequence, + [ + tensor_ty, + tensor_ty, + [of.prod() for of in of_in1s_L3L2], + [of.cons() for of in of_outs], + ], + ) # Place program components (assign them resources on the device) and generate an MLIR module - return Program(dev, rt).resolve_program() + return Program(dev, rt, workers=my_workers).resolve_program() diff --git a/requirements.txt b/requirements.txt index 9636453e4..fa84c59f4 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,11 +9,11 @@ # CUDA build served from PyPI. We therefore also pin torch to the "+cpu" local # version below, which is only available from the PyTorch CPU index. --index-url https://download.pytorch.org/whl/cpu ---find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4 +--find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.0 --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.3.5.dev20+g167f34d +mlir_aie==1.4.0 llvm-aie==21.0.0.2026062301+cb664e8c black From 43117d292cd88973ef316afc015a931113d66764 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 14 Aug 2026 09:41:48 -0600 Subject: [PATCH 2/6] Update mlir-aie to latest nightly (1.4.2.dev12) and llvm-aie to 22.0.0.2026081401 Builds on #145's v1.4.0 migration. Fixes further upstream breaking changes: - aiecc flipped its Peano/Chess default; --no-xchesscc/--no-xbridge no longer exist, so drop them (Peano is now the implicit default) - XRTTensor's residency tracking was reworked into a Storage/coherence-map model; XRTSubBuffer now shares its parent's _storage instead of bootstrapping incompatible state - drop --dynamic-objFifos everywhere it was passed (default upstream now) Also fixes iron/operators/__init__.py unconditionally importing the stream-dse-backed operator, which broke its documented self-skip behavior when onnx/stream-dse aren't installed. Verified: 146/146 fast + 2830/2830 extensive operator tests pass on NPU2 hardware (some extensive failures under -n auto were pre-existing xdist build-directory races, confirmed passing when rerun serially). Co-Authored-By: Claude --- iron/common/base.py | 5 +--- iron/common/compilation/base.py | 7 ----- iron/common/sequence.py | 5 ++-- iron/common/utils.py | 32 ++++++++++++---------- iron/operators/__init__.py | 1 - iron/operators/mem_copy/op.py | 3 -- iron/operators/mha/op.py | 3 -- iron/operators/swiglu_prefill_stream/op.py | 1 - requirements.txt | 6 ++-- 9 files changed, 23 insertions(+), 40 deletions(-) diff --git a/iron/common/base.py b/iron/common/base.py index 701e90dfe..e2ab3ceed 100644 --- a/iron/common/base.py +++ b/iron/common/base.py @@ -156,23 +156,20 @@ def get_kernel_artifacts(self) -> list[CompilationArtifact]: pass def get_artifacts( - self, prefix: str = "", dynamic_obj_fifos: bool = False + self, prefix: str = "" ) -> tuple[XclbinArtifact, InstsBinArtifact]: operator_name = prefix + self.name mlir_artifact = self.get_mlir_artifact() kernel_deps = self.get_kernel_artifacts() - extra_flags = ["--dynamic-objFifos"] if dynamic_obj_fifos else [] xclbin_artifact = XclbinArtifact( f"{operator_name}.xclbin", mlir_input=mlir_artifact, dependencies=[mlir_artifact] + kernel_deps, - extra_flags=extra_flags, ) insts_artifact = InstsBinArtifact( f"{operator_name}.bin", mlir_input=mlir_artifact, dependencies=[mlir_artifact], - extra_flags=extra_flags, ) return xclbin_artifact, insts_artifact diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 6affb7ddb..4fa1422b1 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -526,8 +526,6 @@ def compile(self, graph): ] else: compile_cmd += [ - "--no-xchesscc", - "--no-xbridge", "--peano", str(self.peano_dir), ] @@ -580,14 +578,9 @@ def compile(self, graph): ] else: compile_cmd += [ - "--no-xchesscc", - "--no-xbridge", "--peano", str(self.peano_dir), ] - compile_cmd += [ - "--dynamic-objFifos", - ] do_compile_xclbin = mlir_source in mlir_sources_to_xclbins do_compile_insts_bin = mlir_source in mlir_sources_to_insts if do_compile_xclbin: diff --git a/iron/common/sequence.py b/iron/common/sequence.py index c4034cad4..aef197cc2 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -288,9 +288,8 @@ def __init__( self.explicit_buffer_sizes = ( buffer_sizes or {} ) # Optional dict: buffer_name -> size_in_bytes - # Extra aiecc flags forwarded to the full-ELF build (e.g. --dynamic-objFifos - # for placed/routed whole-array designs that would otherwise overflow AIE2p - # program memory). Empty by default, so other sequences are unaffected. + # Extra aiecc flags forwarded to the full-ELF build. Empty by default, so + # other sequences are unaffected. self.extra_flags = extra_flags or [] self.share_designs = share_designs self._dispatch = dispatch diff --git a/iron/common/utils.py b/iron/common/utils.py index 57ab48992..4060147e8 100644 --- a/iron/common/utils.py +++ b/iron/common/utils.py @@ -49,7 +49,7 @@ class XRTSubBuffer(XRTTensor): The parent XRTTensor must remain alive as long as this sub-buffer is in use. """ - def __init__(self, parent_bo, offset_bytes, size_bytes, shape, dtype, parent=None): + def __init__(self, parent_bo, offset_bytes, size_bytes, shape, dtype, parent): """ Args: parent_bo: The parent pyxrt.bo object. @@ -57,18 +57,23 @@ def __init__(self, parent_bo, offset_bytes, size_bytes, shape, dtype, parent=Non size_bytes: Size of this sub-region in bytes. shape: Tuple giving the logical shape of this sub-buffer. dtype: numpy dtype for interpreting the buffer contents. - parent: The parent XRTTensor this sub-buffer views into. When given, + parent: The parent XRTTensor this sub-buffer views into. Its + storage/coherence tracking is shared rather than duplicated, so moving this sub-buffer between devices propagates the resulting device state to the parent (they share the same memory), so a later whole-parent sync stays consistent with the sub-views. """ # Skip XRTTensor.__init__ (which would allocate a new bo); set base attrs directly. - self.device = "npu" self.dtype = np.dtype(dtype) self._parent = parent + self._shape = tuple(shape) + # Share the parent's storage/coherence tracking (mirrors upstream + # NpuTensor._subview) instead of building our own: this is a sub-region + # of memory the parent already owns and syncs, not a fresh allocation. + self._storage = parent._storage + self._offset_bytes = parent.storage_offset + offset_bytes # TODO: replace with XRTTensor.__getitem__ slice support when available upstream self._bo = _pyxrt.bo(parent_bo, size_bytes, offset_bytes) - self._shape = tuple(shape) ptr = self._bo.map() self._data = np.frombuffer(ptr, dtype=self.dtype).reshape(self._shape) @@ -87,8 +92,7 @@ def data(self) -> np.ndarray: # so the op computes on stale init-zeros. A redundant re-read sync is cheap; # a silently-skipped write sync is a correctness bug. self.device = "cpu" - if self._parent is not None: - self._parent.device = "cpu" + self._parent.device = "cpu" return self._data def buffer_object(self): @@ -115,15 +119,13 @@ def to(self, target_device: str): keep a now-stale ``device`` flag. Revisit once XRT's sub-buffer sync semantics are pinned down (or track per-region dirtiness). """ - if self._parent is not None: - # Reflect this sub-view's current residency onto the parent (e.g. - # "cpu" after a torch_view() write) so the parent's own sync fires - # instead of no-opping, then sync the whole parent buffer. - self._parent.device = self.device - result = self._parent.to(target_device) - self.device = self._parent.device - return result - return super().to(target_device) + # Reflect this sub-view's current residency onto the parent (e.g. + # "cpu" after a torch_view() write) so the parent's own sync fires + # instead of no-opping, then sync the whole parent buffer. + self._parent.device = self.device + result = self._parent.to(target_device) + self.device = self._parent.device + return result @classmethod def from_parent(cls, parent, shape, offset_elements, length_elements, dtype): diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 6d62e215b..4a6c56044 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -12,7 +12,6 @@ from .softmax.op import Softmax from .swiglu_decode.op import SwiGLUDecode from .swiglu_prefill.op import SwiGLUPrefill -from .swiglu_prefill_stream.op import SwiGLUPrefillStream from .transpose.op import Transpose from .strided_copy.op import StridedCopy from .repeat.op import Repeat diff --git a/iron/operators/mem_copy/op.py b/iron/operators/mem_copy/op.py index dd1f056a0..37058b61b 100644 --- a/iron/operators/mem_copy/op.py +++ b/iron/operators/mem_copy/op.py @@ -71,9 +71,6 @@ def get_kernel_artifacts(self): ) ] - def get_artifacts(self): - return super().get_artifacts(dynamic_obj_fifos=True) - def get_arg_spec(self): return [ AIERuntimeArgSpec("in", (self.size,)), diff --git a/iron/operators/mha/op.py b/iron/operators/mha/op.py index 80ee0fb32..5d9e2364b 100644 --- a/iron/operators/mha/op.py +++ b/iron/operators/mha/op.py @@ -108,9 +108,6 @@ def get_kernel_artifacts(self): ), ] - def get_artifacts(self): - return super().get_artifacts(dynamic_obj_fifos=True) - def get_arg_spec(self): seq_padding = self._calculate_seq_padding(self.seq_len, self.num_of_pipelines) buffer_size = self.num_heads * self.d * seq_padding diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 68de6187a..4a4098f57 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -167,7 +167,6 @@ def __init__( ], input_args=inputs, output_args=outputs, - extra_flags=["--dynamic-objFifos"], share_designs=share_designs, context=context, ) diff --git a/requirements.txt b/requirements.txt index fa84c59f4..4df9bdf9b 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,12 +9,12 @@ # CUDA build served from PyPI. We therefore also pin torch to the "+cpu" local # version below, which is only available from the PyTorch CPU index. --index-url https://download.pytorch.org/whl/cpu ---find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.0 +--find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4 --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.0 -llvm-aie==21.0.0.2026062301+cb664e8c +mlir_aie==1.4.2.dev12+ga5b4788 +llvm-aie==22.0.0.2026081401+bc4c83bc black reuse From 3cded00613488b0e9bfb1c298bbb597d2291c8e2 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 14 Aug 2026 09:51:27 -0600 Subject: [PATCH 3/6] Replace hand-rolled XRTSubBuffer with upstream NpuTensor.subview() mlir-aie's nightly (1.4.2.dev12) added a first-class subview() API backed by a shared storage/coherence map between a view and its parent. This directly supersedes XRTSubBuffer, which existed only because that capability didn't exist upstream (per its own TODO). Deleting it also fixes a real bug it had: XRTSubBuffer.to() always resynced the *entire* parent buffer, which could clobber a sibling sub-view whose fresh device data hadn't been read back yet. subview()'s shared coherence map tracks dirtiness per byte range instead, so each view's .to() only moves what it needs to. Converts all six call sites (iron/common/sequence.py, and four in llama_npu.py) from XRTSubBuffer(...)/XRTSubBuffer.from_parent(...) to parent.subview(offset_bytes, shape, dtype), then deletes the class. Verified: 146/146 fast + 2830/2830 extensive operator tests pass on NPU2 hardware. llama_npu.py's call sites are syntax/import-checked only; no model weights available locally to run it end-to-end. Co-Authored-By: Claude --- iron/applications/llama_3.2_1b/llama_npu.py | 33 +++--- iron/common/sequence.py | 29 ++---- iron/common/utils.py | 109 -------------------- 3 files changed, 20 insertions(+), 151 deletions(-) diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index 223487e10..99963a1c6 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -24,7 +24,6 @@ sys.path.insert(0, str(repo_root)) from iron.common.context import AIEContext -from iron.common.utils import XRTSubBuffer from iron.common.sequence import OperatorSequence from iron.operators import ( RMSNorm, @@ -672,12 +671,10 @@ def __init__(self, prompt_len, emb_dim, hidden_dim, n_heads, n_kv_groups, head_d (n_heads * prompt_len, head_dim), dtype=ml_dtypes.bfloat16 ) self.attn_scores_queries_per_head = [ - XRTSubBuffer.from_parent( - self.attn_scores_queries_all, + self.attn_scores_queries_all.subview( + h * prompt_len * head_dim * np.dtype(ml_dtypes.bfloat16).itemsize, (prompt_len, head_dim), - offset_elements=h * prompt_len * head_dim, - length_elements=prompt_len * head_dim, - dtype=ml_dtypes.bfloat16, + ml_dtypes.bfloat16, ) for h in range(n_heads) ] @@ -686,12 +683,10 @@ def __init__(self, prompt_len, emb_dim, hidden_dim, n_heads, n_kv_groups, head_d (n_kv_groups * head_dim, prompt_len), dtype=ml_dtypes.bfloat16 ) self.attn_scores_keys_per_kv_group = [ - XRTSubBuffer.from_parent( - self.attn_scores_keys_all, + self.attn_scores_keys_all.subview( + g * head_dim * prompt_len * np.dtype(ml_dtypes.bfloat16).itemsize, (head_dim, prompt_len), - offset_elements=g * head_dim * prompt_len, - length_elements=head_dim * prompt_len, - dtype=ml_dtypes.bfloat16, + ml_dtypes.bfloat16, ) for g in range(n_kv_groups) ] @@ -700,12 +695,10 @@ def __init__(self, prompt_len, emb_dim, hidden_dim, n_heads, n_kv_groups, head_d (n_heads * prompt_len, prompt_len), dtype=ml_dtypes.bfloat16 ) self.attn_scores_per_head = [ - XRTSubBuffer.from_parent( - self.attn_scores, + self.attn_scores.subview( + h * prompt_len * prompt_len * np.dtype(ml_dtypes.bfloat16).itemsize, (prompt_len, prompt_len), - offset_elements=h * prompt_len * prompt_len, - length_elements=prompt_len * prompt_len, - dtype=ml_dtypes.bfloat16, + ml_dtypes.bfloat16, ) for h in range(n_heads) ] @@ -840,15 +833,13 @@ def __init__(self, config, prompt_len, aie_ops): config.padded_vocab_size // config.vocab_partitions ) self.prefill.logits_parts = [ - XRTSubBuffer.from_parent( - self.prefill.logits, + self.prefill.logits.subview( + i * logits_part_len * np.dtype(ml_dtypes.bfloat16).itemsize, ( prompt_len, config.padded_vocab_size // config.vocab_partitions, ), - offset_elements=i * logits_part_len, - length_elements=logits_part_len, - dtype=ml_dtypes.bfloat16, + ml_dtypes.bfloat16, ) for i in range(config.vocab_partitions) ] diff --git a/iron/common/sequence.py b/iron/common/sequence.py index aef197cc2..6c9f4b1e1 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -11,7 +11,6 @@ import torch from . import compilation as comp from .base import AIEOperatorBase, MLIROperator -from .utils import XRTSubBuffer import aie.utils as aie_utils from aie.iron.device import NPU2 from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor @@ -605,28 +604,21 @@ def get_buffer(self, buffer_name): "output": self.output_buffer, "scratch": self.scratch_buffer, }[buf_type] - sub = XRTSubBuffer( - parent_bo=parent.buffer_object(), - offset_bytes=offset, - size_bytes=length, - shape=(length // BF16.itemsize,), - dtype=ml_dtypes.bfloat16, - parent=parent, - ) + sub = parent.subview(offset, (length // BF16.itemsize,), ml_dtypes.bfloat16) self._buffer_cache[buffer_name] = sub return sub def _sync_inputs(self): - # Sub-views handed out by get_buffer() mark this parent host-dirty on .data - # access (XRTSubBuffer.data), so `to("npu")` here actually fires the host->device - # sync for the freshly written inputs. + # Sub-views handed out by get_buffer() share the parent's coherence map, so + # a write through one (e.g. torch_view()) marks its byte range host-dirty + # there too, and `to("npu")` here syncs every dirty range in one pass. self.input_buffer.to("npu") def _sync_outputs(self): # _run just rewrote the output arena on the device, so the device holds the # authoritative copy. Force the device->host sync: assert device residency first - # so `to("cpu")` fires even if a prior read of get_buffer(...).data marked the - # buffer "cpu" (otherwise a looped dispatch would read stale output). + # so `to("cpu")` fires even if a prior read of get_buffer(...) marked some + # range "cpu" (otherwise a looped dispatch would read stale output). self.output_buffer.device = "npu" self.output_buffer.to("cpu") @@ -697,13 +689,8 @@ def _make_buffer(self, n_elements): return XRTTensor((n_elements,), dtype=ml_dtypes.bfloat16) def _make_subbuffer(self, parent, offset_bytes, size_bytes): - return XRTSubBuffer( - parent_bo=parent.buffer_object(), - offset_bytes=offset_bytes, - size_bytes=size_bytes, - shape=(size_bytes // BF16.itemsize,), - dtype=ml_dtypes.bfloat16, - parent=parent, + return parent.subview( + offset_bytes, (size_bytes // BF16.itemsize,), ml_dtypes.bfloat16 ) def _allocate_buffers(self): diff --git a/iron/common/utils.py b/iron/common/utils.py index 4060147e8..d0d9b6ba3 100644 --- a/iron/common/utils.py +++ b/iron/common/utils.py @@ -1,9 +1,7 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import numpy as np from aie.dialects.aie import get_target_model, WireBundle -from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor, xrt as _pyxrt def get_shim_dma_limit(dev) -> int: @@ -37,110 +35,3 @@ def float_to_name(v: float) -> str: 1e-10 -> '1en10' """ return repr(v).replace(".", "p").replace("-", "n").replace("+", "") - - -class XRTSubBuffer(XRTTensor): - """ - A view into a sub-region of an XRTTensor's underlying pyxrt.bo buffer. - - Inherits from XRTTensor so that isinstance checks in the runtime pass. - Bypasses XRTTensor.__init__ to avoid allocating a new buffer object. - - The parent XRTTensor must remain alive as long as this sub-buffer is in use. - """ - - def __init__(self, parent_bo, offset_bytes, size_bytes, shape, dtype, parent): - """ - Args: - parent_bo: The parent pyxrt.bo object. - offset_bytes: Byte offset into the parent buffer. - size_bytes: Size of this sub-region in bytes. - shape: Tuple giving the logical shape of this sub-buffer. - dtype: numpy dtype for interpreting the buffer contents. - parent: The parent XRTTensor this sub-buffer views into. Its - storage/coherence tracking is shared rather than duplicated, so - moving this sub-buffer between devices propagates the resulting - device state to the parent (they share the same memory), so a - later whole-parent sync stays consistent with the sub-views. - """ - # Skip XRTTensor.__init__ (which would allocate a new bo); set base attrs directly. - self.dtype = np.dtype(dtype) - self._parent = parent - self._shape = tuple(shape) - # Share the parent's storage/coherence tracking (mirrors upstream - # NpuTensor._subview) instead of building our own: this is a sub-region - # of memory the parent already owns and syncs, not a fresh allocation. - self._storage = parent._storage - self._offset_bytes = parent.storage_offset + offset_bytes - # TODO: replace with XRTTensor.__getitem__ slice support when available upstream - self._bo = _pyxrt.bo(parent_bo, size_bytes, offset_bytes) - ptr = self._bo.map() - self._data = np.frombuffer(ptr, dtype=self.dtype).reshape(self._shape) - - @property - def shape(self) -> tuple[int, ...]: - return self._shape - - @property - def data(self) -> np.ndarray: - # `.data` is the write handle for this sub-view. Callers get it to write fresh - # host data (inputs, resident weights), but numpy gives us no write hook, so we - # conservatively mark this sub-view AND its parent host-dirty ("cpu") on any - # access. That makes a subsequent parent `.to("npu")` actually fire the - # host->device sync -- otherwise the residency guard no-ops (device already - # "npu" from allocation) and the freshly written bytes never reach the device, - # so the op computes on stale init-zeros. A redundant re-read sync is cheap; - # a silently-skipped write sync is a correctness bug. - self.device = "cpu" - self._parent.device = "cpu" - return self._data - - def buffer_object(self): - """Return the underlying pyxrt.bo (required by NPUKernel).""" - return self._bo - - def to(self, target_device: str): - """Move this sub-buffer to ``target_device`` by syncing the whole parent. - - The sub-buffer and its parent alias the same underlying memory. Rather - than syncing only this sub-region's bo (whose effect on the parent is - unclear), the parent's current residency is set to this sub-view's - residency and the *entire parent buffer* is synced. This makes the - behaviour explicit and consistent with a caller that writes a sub-view - and then pushes it to the device. - - FIXME: This assumes a sub-buffer sync means a whole-parent sync, which - is ambiguous in XRT: it is unclear whether ``bo.sync()`` on a sub-buffer - transfers only its slice or the whole parent. Because we sync the whole - parent here, moving one sub-buffer to a device can clobber sibling - sub-buffers that view the same parent (e.g. a host->device sync will - overwrite the device side of a sibling whose fresh device data has not - been synced back to the host yet). Those siblings are not notified and - keep a now-stale ``device`` flag. Revisit once XRT's sub-buffer sync - semantics are pinned down (or track per-region dirtiness). - """ - # Reflect this sub-view's current residency onto the parent (e.g. - # "cpu" after a torch_view() write) so the parent's own sync fires - # instead of no-opping, then sync the whole parent buffer. - self._parent.device = self.device - result = self._parent.to(target_device) - self.device = self._parent.device - return result - - @classmethod - def from_parent(cls, parent, shape, offset_elements, length_elements, dtype): - """Create an XRTSubBuffer into a sub-region of a parent XRTTensor. - - Accepts element-count offsets/lengths and converts to bytes internally. - XRTTensor has no built-in slice API; use this until mlir-aie gains - XRTTensor.__getitem__ slice support. - """ - itemsize = np.dtype(dtype).itemsize - return cls( - parent_bo=parent.buffer_object(), - offset_bytes=offset_elements * itemsize, - size_bytes=length_elements * itemsize, - shape=shape, - dtype=dtype, - parent=parent, - ) From 5ab6c2662fb734c5e26082640e3f93796daa5a9b Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 14 Aug 2026 09:59:11 -0600 Subject: [PATCH 4/6] device_utils: delegate arch string mapping to upstream resolve_target_arch() get_kernel_dir() reimplemented the device-arch-to-kernel-dir-string mapping that mlir-aie's aie.utils.compile.utils.resolve_target_arch() now provides. Keep our own no-arg auto-detect default (resolve_target_arch(None) returns "aie2" unconditionally rather than checking the current device, which would silently regress every no-arg call site on NPU2 hardware), but delegate the actual device->arch-string logic once a concrete device is known. Verified: 2830/2830 operator tests pass on NPU2 hardware. Co-Authored-By: Claude --- iron/common/device_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/iron/common/device_utils.py b/iron/common/device_utils.py index 549ab7803..2705ad20f 100644 --- a/iron/common/device_utils.py +++ b/iron/common/device_utils.py @@ -2,11 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 import aie.utils as aie_utils -from aie.iron.device import NPU2 +from aie.utils.compile.utils import resolve_target_arch def get_kernel_dir(dev=None) -> str: """Returns 'aie2p' for NPU2 (Strix, Krackan), 'aie2' for NPU1 (Phoenix).""" if dev is None: dev = aie_utils.get_current_device() - return "aie2p" if isinstance(dev, NPU2) else "aie2" + return resolve_target_arch(dev) From 19d8950dcd37e26d07cf3d7185ad1d9096dda1e1 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 14 Aug 2026 10:41:41 -0600 Subject: [PATCH 5/6] KernelCompilationRule: delegate kernel .o compilation to upstream Replaces the hand-built clang++/xchesscc invocation in KernelCompilationRule.compile() with aie.utils.compile.utils's compile_cxx_core_function(), wrapped in PythonCallbackCompilationCommand (which exists precisely for this: an arbitrary Python callable deferred into the same plan/execute contract as ShellCompilationCommand). This looked architecturally blocked at first (IRON's compile() methods return commands for a separate execute() step, while upstream's function runs synchronously), but IRON's execute() is already fully sequential in Python -- the only real parallelism is aiecc's own internal -j flag inside a single invocation, which is unaffected here since this rule only compiles individual kernel objects, not the aiecc/xclbin step. Nothing is lost by switching to a blocking call under a deferred wrapper. peano_dir/mlir_aie_dir were already always sourced from aie.utils.config.peano_install_dir()/root_path() (iron/common/context.py), the same config module the upstream function reads internally, so there's no override capability to lose. xchesscc_wrapper resolves via `shutil.which` upstream instead of an absolute path IRON constructed itself, but the wheel install puts it on PATH (ironenv/bin/xchesscc_wrapper), so this is a non-issue in the standard install. -Wno-missing-template-arg-list-after-template-kw is kept as an IRON-side extra flag (Peano-only, matching prior behavior) since it's specific to our kernel sources and upstream's default flag list doesn't include it. Symbol renaming/prefixing for operator fusion (_rename_symbols/ _prefix_symbols) stays IRON-side; upstream has no equivalent multi-symbol bulk-prefix operation. Verified: 146/146 fast + 2830/2830 extensive operator tests pass on NPU2 hardware, including a from-scratch rebuild (rm -rf build/) to exercise every kernel compile through the new path. The use_chess=True path is unexercised by any test (also true before this change). Co-Authored-By: Claude --- iron/common/compilation/base.py | 57 +++++++++++++-------------------- 1 file changed, 22 insertions(+), 35 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 4fa1422b1..e07681e03 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -49,6 +49,7 @@ import sys from iron.common.device_utils import get_kernel_dir +from aie.utils.compile.utils import compile_cxx_core_function # Global Functions # ########################################################################## @@ -704,7 +705,6 @@ def matches(self, artifacts): return any(artifacts.get_worklist(KernelObjectArtifact)) def compile(self, artifacts): - include_path = Path(self.mlir_aie_dir) / "include" worklist = artifacts.get_worklist(KernelObjectArtifact) commands = [] @@ -724,41 +724,28 @@ def compile(self, artifacts): "Expected KernelObject dependency to be a C source file" ) - if self.use_chess: - wrapper_path = Path(self.mlir_aie_dir) / "bin" / "xchesscc_wrapper" - cmd = ( - [ - str(wrapper_path), - kernel_dir, # e.g. "aie2" or "aie2p" - f"-I{str(include_path)}", - f"-I{str(runtime_lib_include_path)}", - ] - + artifact.extra_flags - + ["-c", source_file.filename, "-o", artifact.filename] - ) - else: - clang_path = Path(self.peano_dir) / "bin" / "clang++" - target = f"{kernel_dir}-none-unknown-elf" - cmd = ( - [ - str(clang_path), - "-O2", - "-std=c++20", - f"--target={target}", - "-D__AIE_API_AIE_ADF_HPP__", - "-Wno-parentheses", - "-Wno-attributes", - "-Wno-macro-redefined", - "-Wno-empty-body", - "-Wno-missing-template-arg-list-after-template-kw", - f"-I{str(include_path)}", - f"-I{str(runtime_lib_include_path)}", - ] - + artifact.extra_flags - + ["-c", source_file.filename, "-o", artifact.filename] - ) + # -Wno-missing-template-arg-list-after-template-kw only applies to + # the Peano (clang) path: xchesscc's own front end doesn't + # recognize it, and upstream's chess branch never carried it. + compile_args = list(artifact.extra_flags) + if not self.use_chess: + compile_args = [ + "-Wno-missing-template-arg-list-after-template-kw" + ] + compile_args - commands.append(ShellCompilationCommand(cmd)) + commands.append( + PythonCallbackCompilationCommand( + partial( + compile_cxx_core_function, + source_path=source_file.filename, + target_arch=kernel_dir, + output_path=artifact.filename, + include_dirs=[str(runtime_lib_include_path)], + compile_args=compile_args, + use_chess=self.use_chess, + ) + ) + ) if artifact.rename_symbols: commands.extend(self._rename_symbols(artifact)) if artifact.prefix_symbols: From 4f50d264bfcc0c88ff108d6dfd5b1d66cb11ce6c Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 14 Aug 2026 11:27:40 -0600 Subject: [PATCH 6/6] Aiecc*CompilationRule: delegate to upstream compile_mlir_module() Replaces the hand-built aiecc invocation in AieccXclbinInstsCompilationRule and AieccFullElfCompilationRule with aie.utils.compile.utils's compile_mlir_module(), wrapped in PythonCallbackCompilationCommand (same pattern as the prior KernelCompilationRule change). Two real gaps had to be worked around, not just papered over: - compile_mlir_module() always names its own copy of the MLIR source "aie.mlir" inside a work_dir it controls, rather than accepting an existing file path. IRON's build previously kept every artifact's .mlir file (and aiecc's ".prj" companion) at a distinctive flat filename in one shared build/ directory. Introduced _aiecc_work_dir(), giving each MLIR source its own ".mlir.d/" subdirectory, and updated the one place that depended on the old ".mlir.prj" convention (OperatorSequence's params.txt lookup in sequence.py) to derive the new path the same way. - aiecc resolves an MLIR module's relative kernel-object references (e.g. link_with = "axpy.o") against that work_dir, not against the flat build_dir where KernelCompilationRule/ArchiveCompilationRule actually produced them. Added _link_build_outputs_into() to symlink the build directory's existing outputs into each artifact's work_dir before compiling, rather than reworking where kernel objects get built. Known, accepted regression: compile_mlir_module() resolves aiecc via aie.utils.compile.utils's own config.aiecc_path(), which does not honor the AIECC_PATH env var IRON previously supported for pointing a build at a locally-built aiecc without reinstalling the wheel. No test exercises this, but it's a real dev-workflow feature being dropped, not just an implementation detail. peano_dir/mlir_aie_dir/build_dir are no longer threaded into AieccCompilationRule's constructor (dead after this change), simplifying the context.py call sites accordingly. Verified: 146/146 fast + 2830/2830 extensive operator tests pass on NPU2 hardware, from a clean build/ (rm -rf) to exercise every aiecc invocation through the new path, and again from a warm cache. Co-Authored-By: Claude --- iron/common/compilation/__init__.py | 1 + iron/common/compilation/base.py | 160 ++++++++++++++++------------ iron/common/context.py | 8 +- iron/common/sequence.py | 11 +- 4 files changed, 101 insertions(+), 79 deletions(-) diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index 046eb4e3e..de748ffb1 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -3,6 +3,7 @@ from .base import ( DesignGenerator, + _aiecc_work_dir, plan, execute, compile, diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index e07681e03..c6f83f65e 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -49,7 +49,7 @@ import sys from iron.common.device_utils import get_kernel_dir -from aie.utils.compile.utils import compile_cxx_core_function +from aie.utils.compile.utils import compile_cxx_core_function, compile_mlir_module # Global Functions # ########################################################################## @@ -488,20 +488,42 @@ def generate_mlir(output_artifact, generator): f.write(mlir_code) +def _aiecc_work_dir(mlir_filename: str) -> Path: + """Directory aiecc writes its own 'aie.mlir' copy and '.prj' project directory + into for the given MLIR source artifact's filename. + + compile_mlir_module() always names its copy of the source "aie.mlir" inside + the work_dir it's given, rather than reusing the artifact's own filename, so + each MLIR source needs its own work_dir to avoid colliding with every other + artifact's aiecc output in the flat build directory. Callers that need to + find aiecc's project directory afterward (e.g. for a runtime-parameters + scratchpad) should derive it from this same function rather than + re-deriving the convention. + """ + p = Path(mlir_filename) + return p.parent / (p.name + ".d") + + +def _link_build_outputs_into(work_dir: Path, build_dir: Path) -> None: + """Symlink every file already built in build_dir into work_dir. + + aiecc resolves an MLIR module's relative kernel-object references (e.g. + ``link_with = "axpy.o"``, produced by KernelCompilationRule / + ArchiveCompilationRule into the flat build_dir) against work_dir, since + that's where compile_mlir_module() writes its own copy of the MLIR + source. Symlinking makes those lookups succeed without copying kernel + objects into every artifact's own work_dir. + """ + for entry in build_dir.iterdir(): + if entry.is_dir(): + continue + link = work_dir / entry.name + if not link.exists(): + link.symlink_to(entry.resolve()) + + class AieccCompilationRule(CompilationRule): - def __init__( - self, build_dir, peano_dir, mlir_aie_dir, use_chess=False, *args, **kwargs - ): - self.build_dir = build_dir - # AIECC_PATH lets a build point at a locally-built aiecc (e.g. a compiler under - # development) without replacing the installed one. Default = the installed aiecc. - _aiecc_override = os.environ.get("AIECC_PATH") - self.aiecc_path = ( - Path(_aiecc_override) - if _aiecc_override - else Path(mlir_aie_dir) / "bin" / "aiecc" - ) - self.peano_dir = peano_dir + def __init__(self, use_chess=False, *args, **kwargs): self.use_chess = use_chess super().__init__(*args, **kwargs) @@ -515,32 +537,31 @@ def compile(self, graph): commands = [] for artifact in worklist: - compile_cmd = [ - str(self.aiecc_path), - "-v", + mlir_source = artifact.mlir_input + work_dir = _aiecc_work_dir(mlir_source.filename) + options = [ f"-j{os.environ.get('AIECC_JOBS', '1')}", - ] - if self.use_chess: - compile_cmd += [ - "--xchesscc", - "--xbridge", - ] - else: - compile_cmd += [ - "--peano", - str(self.peano_dir), - ] - compile_cmd += [ "--expand-load-pdis", - "--get-full-elf", - "--full-elf-name", - os.path.abspath(artifact.filename), - *artifact.extra_flags, - os.path.abspath(artifact.mlir_input.filename), - ] - commands.append( - ShellCompilationCommand(compile_cmd, cwd=str(self.build_dir)) - ) + ] + artifact.extra_flags + + def _compile( + artifact=artifact, + mlir_source=mlir_source, + work_dir=work_dir, + options=options, + ): + work_dir.mkdir(parents=True, exist_ok=True) + _link_build_outputs_into(work_dir, Path(mlir_source.filename).parent) + compile_mlir_module( + Path(mlir_source.filename).read_text(), + full_elf_path=os.path.abspath(artifact.filename), + work_dir=str(work_dir), + options=options, + use_chess=self.use_chess, + verbose=True, + ) + + commands.append(PythonCallbackCompilationCommand(_compile)) artifact.available = True return commands @@ -567,52 +588,53 @@ def compile(self, graph): commands = [] # Now we know for each mlir source if we need to generate an xclbin, an insts.bin or both for it for mlir_source in mlir_sources: - compile_cmd = [ - str(self.aiecc_path), - "-v", - f"-j{os.environ.get('AIECC_JOBS', '1')}", - ] - if self.use_chess: - compile_cmd += [ - "--xchesscc", - "--xbridge", - ] - else: - compile_cmd += [ - "--peano", - str(self.peano_dir), - ] + options = [f"-j{os.environ.get('AIECC_JOBS', '1')}"] + xclbin_path = None + insts_path = None do_compile_xclbin = mlir_source in mlir_sources_to_xclbins do_compile_insts_bin = mlir_source in mlir_sources_to_insts if do_compile_xclbin: first_xclbin = mlir_sources_to_xclbins[mlir_source][ 0 ] # TODO: this does not handle the case of multiple xclbins with different kernel names or flags from the same MLIR - compile_cmd += first_xclbin.extra_flags + [ - "--get-xclbin", - "--xclbin-name=" + os.path.abspath(first_xclbin.filename), - "--xclbin-kernel-name=" + first_xclbin.kernel_name, + xclbin_path = os.path.abspath(first_xclbin.filename) + options += first_xclbin.extra_flags + [ + f"--xclbin-kernel-name={first_xclbin.kernel_name}", ] if first_xclbin.xclbin_input is not None: - compile_cmd += [ + options.append( "--xclbin-input=" + os.path.abspath(first_xclbin.xclbin_input.filename) - ] + ) if do_compile_insts_bin: first_insts_bin = mlir_sources_to_insts[mlir_source][ 0 ] # TODO: this does not handle the case of multiple insts.bins with different flags from the same MLIR - # Outputs are selected by --get-; asking only for the insts is what - # "--no-compile" used to mean, so there is nothing to opt out of here. - compile_cmd += first_insts_bin.extra_flags + [ - "--get-npu-insts", - "--npu-insts-name=" + os.path.abspath(first_insts_bin.filename), - ] - compile_cmd += [os.path.abspath(mlir_source.filename)] + insts_path = os.path.abspath(first_insts_bin.filename) + options += first_insts_bin.extra_flags - commands.append( - ShellCompilationCommand(compile_cmd, cwd=str(self.build_dir)) - ) + work_dir = _aiecc_work_dir(mlir_source.filename) + + def _compile( + mlir_source=mlir_source, + xclbin_path=xclbin_path, + insts_path=insts_path, + options=options, + work_dir=work_dir, + ): + work_dir.mkdir(parents=True, exist_ok=True) + _link_build_outputs_into(work_dir, Path(mlir_source.filename).parent) + compile_mlir_module( + Path(mlir_source.filename).read_text(), + insts_path=insts_path, + xclbin_path=xclbin_path, + work_dir=str(work_dir), + options=options, + use_chess=self.use_chess, + verbose=True, + ) + + commands.append(PythonCallbackCompilationCommand(_compile)) # There may be multiple targets that require an xclbin/insts.bin from the same MLIR with different names; copy them for sources_to in [mlir_sources_to_xclbins, mlir_sources_to_insts]: diff --git a/iron/common/context.py b/iron/common/context.py index 0dc1eee48..a7a5136ca 100644 --- a/iron/common/context.py +++ b/iron/common/context.py @@ -55,10 +55,6 @@ def compilation_rules(self): comp.GenerateMLIRFromPythonCompilationRule(), comp.KernelCompilationRule(peano_dir, mlir_aie_dir, use_chess=use_chess), comp.ArchiveCompilationRule(peano_dir, mlir_aie_dir), - comp.AieccXclbinInstsCompilationRule( - self.build_dir, peano_dir, mlir_aie_dir, use_chess=use_chess - ), - comp.AieccFullElfCompilationRule( - self.build_dir, peano_dir, mlir_aie_dir, use_chess=use_chess - ), + comp.AieccXclbinInstsCompilationRule(use_chess=use_chess), + comp.AieccFullElfCompilationRule(use_chess=use_chess), ] diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 6c9f4b1e1..0b1ac7115 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -570,14 +570,17 @@ def params(self): """Lazy ParameterScratchpad bound to this ELF's ctrl scratchpad BO. The ``params.txt`` describing the runtime parameters is written by - ``aie-lower-parameters`` into the ``.prj`` project directory next - to the fused MLIR source. Returns ``None`` if the sequence declared no - runtime parameters (in which case the file is not written). + ``aie-lower-parameters`` into the aiecc project directory (see + ``_aiecc_work_dir``) for the fused MLIR source. Returns ``None`` if the + sequence declared no runtime parameters (in which case the file is not + written). """ if self._params is not None: return self._params mlir_filename = self.op.artifacts[0].mlir_input.filename - params_path = Path(mlir_filename + ".prj") / "params.txt" + params_path = ( + comp._aiecc_work_dir(mlir_filename) / "aie.mlir.prj" / "params.txt" + ) if not params_path.exists(): return None from aie.utils.hostruntime.xrtruntime.parameter_scratchpad import (