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/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/__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 8b06de537..c6f83f65e 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, compile_mlir_module # Global Functions # ########################################################################## @@ -487,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) @@ -514,35 +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')}", - "--no-compile-host", - ] - if self.use_chess: - compile_cmd += [ - "--xchesscc", - "--xbridge", - ] - else: - compile_cmd += [ - "--no-xchesscc", - "--no-xbridge", - "--peano", - str(self.peano_dir), - ] - compile_cmd += [ "--expand-load-pdis", - "--generate-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 @@ -569,58 +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')}", - "--no-compile-host", - ] - if self.use_chess: - compile_cmd += [ - "--xchesscc", - "--xbridge", - ] - else: - compile_cmd += [ - "--no-xchesscc", - "--no-xbridge", - "--peano", - str(self.peano_dir), - ] - compile_cmd += [ - "--dynamic-objFifos", - ] + 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 + [ - "--aie-generate-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 - if not do_compile_xclbin: - compile_cmd += ["--no-compile"] - compile_cmd += first_insts_bin.extra_flags + [ - "--aie-generate-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]: @@ -713,7 +727,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 = [] @@ -733,41 +746,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: 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/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) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index c4034cad4..0b1ac7115 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 @@ -288,9 +287,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 @@ -572,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 ( @@ -606,28 +607,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") @@ -698,13 +692,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 57ab48992..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,108 +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=None): - """ - 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. When given, - 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 - # 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) - - @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" - if self._parent is not None: - 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). - """ - 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) - - @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, - ) 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/_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/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/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/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/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/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/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..4df9bdf9b 100755 --- a/requirements.txt +++ b/requirements.txt @@ -13,8 +13,8 @@ --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 -llvm-aie==21.0.0.2026062301+cb664e8c +mlir_aie==1.4.2.dev12+ga5b4788 +llvm-aie==22.0.0.2026081401+bc4c83bc black reuse