From 0bdebc0d2b2bd6478f874dab7a64770f6ba2a3e4 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Fri, 31 Jul 2026 00:05:33 +0300 Subject: [PATCH] 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