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 d7f397a54..5ecc8746a 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,41 +537,32 @@ def compile(self, graph): commands = [] for artifact in worklist: - compile_cmd = [ - str(self.aiecc_path), - "-v", + mlir_source = artifact.mlir_input + work_dir = _aiecc_work_dir(mlir_source.filename) + options = [ f"-j{os.environ.get('AIECC_JOBS', '1')}", - ] - if self.use_chess: - compile_cmd += [ - "--xchesscc", - "--xbridge", - ] - else: - compile_cmd += [ - "--no-xchesscc", - "--no-xbridge", - "--peano", - str(self.peano_dir), - ] - compile_cmd += [ "--expand-load-pdis", - "--get-full-elf", - "--full-elf-name", - os.path.abspath(artifact.filename), - # Unrequested, params.txt is never written and - # SequenceFullELFCallable.params silently returns None. It has - # no name flag of its own, so --output-dir is what puts it in - # the project directory that property reads. "--get-scratchpad-parameters", - "--output-dir", - os.path.abspath(artifact.mlir_input.filename) + ".prj", - *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 @@ -575,57 +589,53 @@ def compile(self, graph): commands = [] # Now we know for each mlir source if we need to generate an xclbin, an insts.bin or both for it for mlir_source in mlir_sources: - compile_cmd = [ - str(self.aiecc_path), - "-v", - f"-j{os.environ.get('AIECC_JOBS', '1')}", - ] - if self.use_chess: - compile_cmd += [ - "--xchesscc", - "--xbridge", - ] - else: - compile_cmd += [ - "--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 + [ - "--get-xclbin", - "--xclbin-name=" + os.path.abspath(first_xclbin.filename), - "--xclbin-kernel-name=" + first_xclbin.kernel_name, + xclbin_path = os.path.abspath(first_xclbin.filename) + options += first_xclbin.extra_flags + [ + f"--xclbin-kernel-name={first_xclbin.kernel_name}", ] if first_xclbin.xclbin_input is not None: - compile_cmd += [ + options.append( "--xclbin-input=" + os.path.abspath(first_xclbin.xclbin_input.filename) - ] + ) if do_compile_insts_bin: first_insts_bin = mlir_sources_to_insts[mlir_source][ 0 ] # TODO: this does not handle the case of multiple insts.bins with different flags from the same MLIR - # Outputs are selected by --get-; asking only for the insts is what - # "--no-compile" used to mean, so there is nothing to opt out of here. - compile_cmd += first_insts_bin.extra_flags + [ - "--get-npu-insts", - "--npu-insts-name=" + os.path.abspath(first_insts_bin.filename), - ] - compile_cmd += [os.path.abspath(mlir_source.filename)] + insts_path = os.path.abspath(first_insts_bin.filename) + options += first_insts_bin.extra_flags - commands.append( - ShellCompilationCommand(compile_cmd, cwd=str(self.build_dir)) - ) + work_dir = _aiecc_work_dir(mlir_source.filename) + + def _compile( + mlir_source=mlir_source, + xclbin_path=xclbin_path, + insts_path=insts_path, + options=options, + work_dir=work_dir, + ): + work_dir.mkdir(parents=True, exist_ok=True) + _link_build_outputs_into(work_dir, Path(mlir_source.filename).parent) + compile_mlir_module( + Path(mlir_source.filename).read_text(), + insts_path=insts_path, + xclbin_path=xclbin_path, + work_dir=str(work_dir), + options=options, + use_chess=self.use_chess, + verbose=True, + ) + + commands.append(PythonCallbackCompilationCommand(_compile)) # There may be multiple targets that require an xclbin/insts.bin from the same MLIR with different names; copy them for sources_to in [mlir_sources_to_xclbins, mlir_sources_to_insts]: @@ -718,7 +728,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 = [] @@ -738,41 +747,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. + 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 c504718e7..0d33b482c 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,7 @@ 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. self.extra_flags = extra_flags or [] self.share_designs = share_designs self._dispatch = dispatch @@ -571,16 +568,18 @@ def __init__(self, op, device_name="main", sequence_name="sequence"): 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: the file is still written, but holds a count of - zero and there is no ctrl scratchpad buffer object to bind to. + The ``params.txt`` describing the runtime parameters is requested from + aiecc via ``--get-scratchpad-parameters``; it is a graph output, so it + lands in aiecc's ``--output-dir``, which compile_mlir_module() points at + the work dir (see ``_aiecc_work_dir``) for the fused MLIR source. + Returns ``None`` if the sequence declared no runtime parameters: the + file is still written, but holds a count of zero and there is no ctrl + scratchpad buffer object to bind to. """ 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) / "params.txt" if not params_path.exists(): return None if params_path.read_text().split("\n", 1)[0].strip() == "0": @@ -609,28 +608,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") @@ -701,13 +693,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..99ad7327b 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: @@ -25,7 +23,7 @@ def float_to_name(v: float) -> str: """Convert a float to a filesystem-safe string for use in operator names. Uses repr() for the shortest exact round-trip representation, then sanitizes - characters that are problematic in filenames or shell scripts: + characters that are problematic in filenames or shell scripts, for instance: '.' -> 'p' (decimal point) '-' -> 'n' (negative sign / negative exponent) '+' -> '' (positive exponent, redundant) @@ -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/gemm/design.py b/iron/operators/gemm/design.py index 169d848ae..c8a4f6e7c 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -757,8 +757,6 @@ def sequence(A, B, C, A_prods, B_prods, C_conses): 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) diff --git a/iron/operators/mem_copy/design.py b/iron/operators/mem_copy/design.py index 788345dbc..8eadffba0 100644 --- a/iron/operators/mem_copy/design.py +++ b/iron/operators/mem_copy/design.py @@ -401,8 +401,7 @@ def sequence(a_in, b_out, of_ins_prods, of_outs_conss): ], ) # Place components (assign them resources on the device) and generate an MLIR module - # bypass means the DMAs run without any compute worker, as `rt.start` was - # previously guarded by the same condition. + # bypass means the DMAs run without any compute worker prog = Program(dev, rt, workers=None if bypass else my_workers) if not bypass: maybe_enable_trace(prog, trace_size, my_workers) 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 61397fbd2..f17d26fad 100644 --- a/iron/operators/mha/design.py +++ b/iron/operators/mha/design.py @@ -776,8 +776,6 @@ def legalize_tas(tas: TensorAccessSequence): # print_tap_seq_info(O_tiles, "O") # Runtime operations to move data to/from the AIE-array - # 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)) @@ -786,7 +784,6 @@ def legalize_tas(tas: TensorAccessSequence): 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): mha_rtps_list[j][i][0] = num_q_block_per_pipeline diff --git a/iron/operators/mha/op.py b/iron/operators/mha/op.py index 80ee0fb32..5d9e2364b 100644 --- a/iron/operators/mha/op.py +++ b/iron/operators/mha/op.py @@ -108,9 +108,6 @@ def get_kernel_artifacts(self): ), ] - def get_artifacts(self): - return super().get_artifacts(dynamic_obj_fifos=True) - def get_arg_spec(self): seq_padding = self._calculate_seq_padding(self.seq_len, self.num_of_pipelines) buffer_size = self.num_heads * self.d * seq_padding diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 68de6187a..4a4098f57 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -167,7 +167,6 @@ def __init__( ], input_args=inputs, output_args=outputs, - extra_flags=["--dynamic-objFifos"], share_designs=share_designs, context=context, ) diff --git a/requirements.txt b/requirements.txt index fa84c59f4..f626f2f26 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,12 +9,12 @@ # CUDA build served from PyPI. We therefore also pin torch to the "+cpu" local # version below, which is only available from the PyTorch CPU index. --index-url https://download.pytorch.org/whl/cpu ---find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.0 +--find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4 --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.0 -llvm-aie==21.0.0.2026062301+cb664e8c +mlir_aie==1.4.2.dev16+g7e00b57 +llvm-aie==22.0.0.2026082001+84660bc3 black reuse