From aecc68ace2371788aa06cd2301d03470838ca9a1 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:13:26 -0600 Subject: [PATCH 01/22] Build kernels from mlir-aie kernel factories; port transpose and gemv Operators can now hand the compilation system an aie.iron.kernels factory's ExternalFunction (KernelObjectArtifact.from_extern) instead of a hand-built source + flags recipe. mlir-aie compiles, prefixes and stamps the object; its file name and symbols carry a digest of the recipe, so a changed recipe never reuses a stale object, and fused sequences share identical recipes instead of re-prefixing them per operator. - A generated MLIR module is only fresh if it links every kernel its design was given, since a new recipe means a new object name. - Operators pin the probed NPU as the selected device: the factories pick their architecture from the selected device only and otherwise fall back to aie2. - kernels_dir resolves through mlir-aie's config (MLIR_AIE_KERNEL_SOURCES); IRON_AIE_KERNELS_DIR is gone. - transpose uses datamovement.transpose; gemv uses linalg.mv, with the gelu epilogue bound from the gelu factory's object. Co-Authored-By: Claude --- iron/common/base.py | 2 + iron/common/compilation/base.py | 90 +++++++++++++++++++++++++++++- iron/common/context.py | 11 ++-- iron/common/device_utils.py | 11 ++++ iron/common/sequence.py | 19 ++++++- iron/operators/gemv/design.py | 38 +++---------- iron/operators/gemv/op.py | 74 ++++++++++++------------ iron/operators/transpose/design.py | 13 +---- iron/operators/transpose/op.py | 21 ++----- 9 files changed, 175 insertions(+), 104 deletions(-) diff --git a/iron/common/base.py b/iron/common/base.py index e2ab3ceed0..1e3dd68d77 100644 --- a/iron/common/base.py +++ b/iron/common/base.py @@ -17,6 +17,7 @@ from . import compilation as comp from .context import AIEContext +from .device_utils import pin_current_device from .utils import float_to_name from .compilation import ( CompilationArtifact, @@ -35,6 +36,7 @@ class AIEOperatorBase(ABC): def __init__(self, context: AIEContext | None = None) -> None: self.artifacts = comp.CompilationArtifactGraph() + pin_current_device() if context is None: context = self.get_default_context() self.context = context diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 22e7ef1bb9..372b1928dc 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -50,7 +50,13 @@ import sys from iron.common.device_utils import get_kernel_dir -from aie.utils.compile.utils import compile_cxx_core_function, compile_mlir_module +from aie.iron.kernel import ExternalFunction, Kernel +from aie.utils.compile.utils import ( + _has_current_symbol_prefix_stamp, + compile_cxx_core_function, + compile_external_kernel, + compile_mlir_module, +) # Global Functions # ########################################################################## @@ -73,6 +79,20 @@ def __call__(self) -> str: spec.loader.exec_module(module) return str(getattr(module, self.fn_name)(*self.args, **self.kwargs)) + def kernels(self) -> list[Kernel]: + """The kernels passed to the design, which its module must declare.""" + values = [*self.args, *self.kwargs.values()] + found = [] + while values: + value = values.pop() + if isinstance(value, Kernel): + found.append(value) + elif isinstance(value, (list, tuple)): + values.extend(value) + elif isinstance(value, dict): + values.extend(value.values()) + return found + def plan( rules: Sequence[CompilationRule], @@ -386,11 +406,41 @@ def __init__( extra_flags: list[str] | None = None, rename_symbols: dict[str, str] | None = None, prefix_symbols: str | None = None, + extern: ExternalFunction | None = None, ) -> None: super().__init__(filename, dependencies) self.extra_flags = extra_flags if extra_flags is not None else [] self.rename_symbols = rename_symbols if rename_symbols is not None else {} self.prefix_symbols = prefix_symbols + # The mlir-aie kernel factory recipe this object is built from, if any. + # Such an object is compiled by mlir-aie itself, and its file name and + # symbols already carry a digest of the recipe. + self.extern = extern + + @classmethod + def from_extern(cls, fn: ExternalFunction) -> KernelObjectArtifact: + """The object an ``aie.iron.kernels`` factory's ExternalFunction links. + + Every symbol bound from the same object (``fn.object_file.bind(...)``) + is served by this one artifact. + """ + if fn.source_file is None: + raise ValueError(f"{fn.name}: only file-backed kernels are supported") + return cls( + fn.object_file_name, + dependencies=[SourceArtifact(fn.source_file)], + extern=fn, + ) + + def is_available_in_filesystem(self) -> bool: + if not super().is_available_in_filesystem(): + return False + # A prefixed object whose prefix pass never completed exports the + # unprefixed symbols; mlir-aie stamps the object once the pass is done. + prefix = self.extern._symbol_prefix if self.extern is not None else None + return prefix is None or _has_current_symbol_prefix_stamp( + self.filename, f"{prefix}_" + ) class KernelArchiveArtifact(CompilationArtifact): @@ -408,6 +458,18 @@ def __init__( self.generator = generator super().__init__(filename, dependencies=[SourceArtifact(generator.source_path)]) + def is_available_in_filesystem(self) -> bool: + if not super().is_available_in_filesystem(): + return False + # A factory kernel's object is named after its recipe, so a changed + # recipe leaves a module that is newer than its design yet links an + # object that is no longer built -- or worse, an old one still on disk. + text = Path(self.filename).read_text() + return all( + f'link_with = "{kernel.object_file_name}"' in text + for kernel in self.generator.kernels() + ) + def _sha256_of(path: Path) -> str: with open(path, "rb") as f: @@ -845,7 +907,20 @@ def compile(self, artifacts): Path(self.mlir_aie_dir) / "aie_runtime_lib" / kernel_dir.upper() ) + compiled_externs = set() for artifact in worklist: + if artifact.extern is not None: + # Operators sharing a recipe share its object, so a fused + # sequence may list the same one several times. + if artifact.filename not in compiled_externs: + compiled_externs.add(artifact.filename) + commands.append( + PythonCallbackCompilationCommand( + partial(self._compile_extern, artifact, kernel_dir) + ) + ) + artifact.available = True + continue if len(artifact.dependencies) < 1: raise RuntimeError( "Expected at least one dependency (the C source code) for KernelObjectArtifact" @@ -886,6 +961,19 @@ def compile(self, artifacts): return commands + def _compile_extern(self, artifact, kernel_dir): + fn = artifact.extern + if fn.use_chess != self.use_chess: + raise RuntimeError( + f"{fn.name} is a {'Chess' if fn.use_chess else 'Peano'} kernel, " + f"but this context compiles with {'Chess' if self.use_chess else 'Peano'}" + ) + # The artifact is only on the worklist if its object is missing, older + # than its source, or half-prefixed. mlir-aie reuses any object already + # at the output path, so remove it to make mlir-aie rebuild it. + Path(artifact.filename).unlink(missing_ok=True) + compile_external_kernel(fn, str(Path(artifact.filename).parent), kernel_dir) + def _find_tool(self, name): return _find_tool(name, self.peano_dir, self.mlir_aie_dir) diff --git a/iron/common/context.py b/iron/common/context.py index 6979f18388..57eb823399 100644 --- a/iron/common/context.py +++ b/iron/common/context.py @@ -27,16 +27,13 @@ class AIEContext: @property def kernels_dir(self) -> Path: - """C++ kernel sources bundled with the installed mlir-aie package. + """C++ kernel sources the mlir-aie kernel factories build from. - IRON_AIE_KERNELS_DIR overrides this to point at a local mlir-aie + MLIR_AIE_KERNEL_SOURCES overrides this to point at a local mlir-aie checkout for kernel development. """ - # Lazy: root_path() needs the package importable at call time. - override = os.environ.get("IRON_AIE_KERNELS_DIR") - if override: - return Path(override) - return Path(aie.utils.config.root_path()) / "include" / "aie_kernels" + # Lazy: the config needs the package importable at call time. + return Path(aie.utils.config.aie_kernels_dir()) def __post_init__(self) -> None: """Normalize build_dir to a Path object.""" diff --git a/iron/common/device_utils.py b/iron/common/device_utils.py index 2705ad20f3..b7722eccde 100644 --- a/iron/common/device_utils.py +++ b/iron/common/device_utils.py @@ -10,3 +10,14 @@ def get_kernel_dir(dev=None) -> str: if dev is None: dev = aie_utils.get_current_device() return resolve_target_arch(dev) + + +def pin_current_device() -> None: + """Bind the probed NPU as the explicitly selected device. + + The mlir-aie kernel factories choose their sources by architecture from the + explicitly selected device only; with none selected they fall back to aie2, + which on an NPU2 machine silently builds aie2 kernels. + """ + if aie_utils.get_current_device(probe_runtime=False) is None: + aie_utils.set_current_device(aie_utils.get_current_device()) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 66c51f84ef..bc41c33888 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -103,6 +103,18 @@ def _trace_tag(seq): return f"_traced{seq.trace_size}" if seq.trace_size else "" +def _hand_built_kernels(op, objs=None): + """``op``'s kernel artifacts that IRON builds itself, rather than an + mlir-aie kernel factory, and so must be prefixed apart within a sequence.""" + if objs is None: + objs = op.get_kernel_artifacts() + return [ + obj + for obj in objs + if not (isinstance(obj, comp.KernelObjectArtifact) and obj.extern is not None) + ] + + class FusedDispatch(SequenceDispatch): """Single-ELF dispatch (NPU2 only): all operators fused into one ELF.""" @@ -141,7 +153,7 @@ def build_fused_mlir(self, seq): for idx, op in enumerate(designs): mlir_artifact = op.get_mlir_artifact() - if len(op.get_kernel_artifacts()) > 0: + if _hand_built_kernels(op): mlir_artifact.generator.kwargs["func_prefix"] = f"op{idx}_" op_name = f"op{idx}_{op.__class__.__name__}" design_names.append(op_name) @@ -161,11 +173,12 @@ def build_fused_mlir(self, seq): ) def _collect_kernel_artifacts(self, seq): - """Kernel artifacts from all child operators, prefixed per operator index.""" + """Kernel artifacts from all child operators, hand-built ones prefixed per + operator index. Factory-built objects are already unique per recipe.""" kernel_artifacts = [] for idx, op in enumerate(seq.unique_designs()[0]): objs = op.get_kernel_artifacts() - for obj in objs: + for obj in _hand_built_kernels(op, objs): obj.filename = f"op{idx}_{obj.filename}" obj.prefix_symbols = f"op{idx}_" kernel_artifacts.extend(objs) diff --git a/iron/operators/gemv/design.py b/iron/operators/gemv/design.py index 5fffe70d30..82281bba8c 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, TaskGroup, Worker +from aie.iron import ObjectFifo, Program, Runtime, TaskGroup, Worker """ Matrix-vector design @@ -33,10 +33,10 @@ def my_matvec( m_input, m_output=None, num_batches=1, - kernel_object="mv.o", - func_prefix="", verbose=False, - epilogue="none", + *, + matvec_fn, + epilogue_fn=None, ): if m_output is None: m_output = m_input @@ -56,11 +56,8 @@ def my_matvec( assert m_input <= M // cols, "m_input must be less than or equal to M/cols" assert (M // cols) % m_input == 0, "m_input must evenly divide M/cols" - vectorized = True dtype_in = np.dtype[bfloat16] - dtype_in_str = "bf16" dtype_out = np.dtype[bfloat16] - dtype_out_str = "bf16" assert M % cols == 0 @@ -80,26 +77,9 @@ def my_matvec( L3_B_ty = np.ndarray[(num_batches * K,), dtype_in] L3_C_ty = np.ndarray[(num_batches * M,), dtype_out] - func_type = "vectorized" if vectorized else "scalar" - matvec = Kernel( - f"{func_prefix}matvec_{func_type}_{dtype_in_str}_{dtype_out_str}", - f"{func_prefix}{kernel_object}", - [np.int32, np.int32, L1_A_ty, L1_B_ty, L1_C_ty], - ) - # Optional fused activation over the full m_output C-tile, applied once per tile in core_body - # (after the matvec inner-loop has filled all rows) rather than per matvec call, whose m_input - # tile can be smaller than the 16-wide activation vector. - assert epilogue in ("none", "gelu") - gelu_kernel = None - if epilogue == "gelu": - assert ( - m_output % 16 == 0 - ), f"gelu epilogue needs m_output % 16 == 0 (got {m_output})" - gelu_kernel = Kernel( - f"{func_prefix}gelu_tile_bf16", - f"{func_prefix}{kernel_object}", - [np.int32, L1_C_ty], - ) + # epilogue_fn: optional fused activation over the full m_output C-tile, applied once per + # tile in core_body (after the matvec inner-loop has filled all rows) rather than per + # matvec call, whose m_input tile can be smaller than the 16-wide activation vector. A_L3L1_fifos = [ ObjectFifo(L1_A_ty, name=f"A_L3L1_{i}", depth=2) for i in range(cols) @@ -137,9 +117,9 @@ def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, gelu_kernel=None): A_L3L1_fifos[i].cons(), B_L3L1_fifos[i].cons(), C_L1L3_fifos[i].prod(), - matvec, + matvec_fn, ] - + ([gelu_kernel] if epilogue == "gelu" else []), + + ([epilogue_fn] if epilogue_fn is not None else []), ) for i in range(cols) ] diff --git a/iron/operators/gemv/op.py b/iron/operators/gemv/op.py index c929980569..c9aea02459 100644 --- a/iron/operators/gemv/op.py +++ b/iron/operators/gemv/op.py @@ -4,16 +4,18 @@ from dataclasses import dataclass, field from typing import ClassVar, Dict +import numpy as np +from ml_dtypes import bfloat16 + from iron.common import ( MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - KernelArchiveArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) import aie.utils as aie_utils +from aie.iron.kernels import activation, linalg from iron.common.device_utils import get_kernel_dir @@ -84,16 +86,35 @@ def name(self) -> str: return base return f"{base}_epi{self.epilogue}" - @property - def _kernel_link_file(self): - # With the gelu epilogue the core also links the gelu kernel, so the object becomes an - # archive of (matvec, gelu); the plain matvec stays a single object. - if self.epilogue == "gelu": - return f"gemv_{self.K}k_{self.kernel_vector_size}vs_gelu_kernels.a" - return f"gemv_{self.K}k_{self.kernel_vector_size}vs.o" + def _matvec(self): + return linalg.mv( + self.tile_size_input, + self.K, + bfloat16, + bfloat16, + vec_size=self.kernel_vector_size, + output_rows=self.tile_size_output, + use_chess=self.context.compiler == "chess", + ) + + def _gelu(self): + # The epilogue is gelu.cc's in-place gelu_tile_bf16, which only aie2p's + # gelu.cc exports; it rides in the object the gelu factory builds. + if get_kernel_dir() != "aie2p": + raise NotImplementedError( + "gemv gelu epilogue is only available on NPU2 (aie2p); " + f"current kernel dir is {get_kernel_dir()!r}" + ) + return activation.gelu() def get_mlir_artifact(self): mlir_verbose = getattr(self.context, "mlir_verbose", False) + epilogue_fn = None + if self.epilogue == "gelu": + epilogue_fn = self._gelu().object_file.bind( + "gelu_tile_bf16", + [np.int32, np.ndarray[(self.tile_size_output,), np.dtype[bfloat16]]], + ) return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", @@ -111,42 +132,17 @@ def get_mlir_artifact(self): ), { "verbose": mlir_verbose, - "kernel_object": self._kernel_link_file, - "epilogue": self.epilogue, + "matvec_fn": self._matvec(), + "epilogue_fn": epilogue_fn, }, ), ) def get_kernel_artifacts(self): - matvec_obj = KernelObjectArtifact( - f"gemv_{self.K}k_{self.kernel_vector_size}vs.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / "generic" / "mv.cc") - ], - extra_flags=[ - f"-DDIM_K={self.K}", - f"-DVEC_SIZE={self.kernel_vector_size}", - ], - ) + fns = [self._matvec()] if self.epilogue == "gelu": - # The gelu kernel lives in aie2p/gelu.cc, so the fused epilogue is NPU2-only. - if get_kernel_dir() != "aie2p": - raise NotImplementedError( - "gemv gelu epilogue is only available on NPU2 (aie2p); " - f"current kernel dir is {get_kernel_dir()!r}" - ) - gelu_obj = KernelObjectArtifact( - "gelu.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / "aie2p" / "gelu.cc") - ], - ) - return [ - KernelArchiveArtifact( - self._kernel_link_file, dependencies=[matvec_obj, gelu_obj] - ) - ] - return [matvec_obj] + fns.append(self._gelu()) + return [KernelObjectArtifact.from_extern(fn) for fn in fns] def get_arg_spec(self): batch_dim = (self.num_batches,) if self.num_batches > 1 else () diff --git a/iron/operators/transpose/design.py b/iron/operators/transpose/design.py index bb0c3348fe..5002a00deb 100644 --- a/iron/operators/transpose/design.py +++ b/iron/operators/transpose/design.py @@ -4,13 +4,13 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker +from aie.iron import ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ def shuffle_transpose( - dev, M, N, num_columns, num_channels, m, n, s, num_batches=1, func_prefix="" + dev, M, N, num_columns, num_channels, m, n, s, num_batches=1, *, transpose_fn ): num_elements = M * N per_tile_elements = m * n @@ -116,13 +116,6 @@ def shuffle_transpose( for j in range(num_channels) ] - # AIE Core Function declaration - transpose_kernel = Kernel( - f"{func_prefix}transpose_{s}x{s}", - f"{func_prefix}transpose_{m}x{n}.o", - [tile_ty, tile_ty], - ) - # Define a task that will run on a compute tile def core_body(of_in1, of_out, transpose_kernel): # Process num_batches contiguous matrices through the same FIFOs: num_batches x the per-matrix @@ -144,7 +137,7 @@ def core_body(of_in1, of_out, transpose_kernel): [ of_in1s_L2L1[i * num_channels + j].cons(), of_outs[i * num_channels + j].prod(), - transpose_kernel, + transpose_fn, ], ) for i in range(num_columns) diff --git a/iron/operators/transpose/op.py b/iron/operators/transpose/op.py index 0e304fcb7c..56d4e2fa38 100644 --- a/iron/operators/transpose/op.py +++ b/iron/operators/transpose/op.py @@ -5,11 +5,11 @@ from typing import ClassVar, Dict import aie.utils as aie_utils +from aie.iron.kernels import datamovement from iron.common import ( MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) @@ -96,24 +96,15 @@ def get_mlir_artifact(self): self.s, self.num_batches, ), + {"transpose_fn": self._kernel()}, ), ) + def _kernel(self): + return datamovement.transpose(self.m, self.n, self.s) + def get_kernel_artifacts(self): - return [ - KernelObjectArtifact( - f"transpose_{self.m}x{self.n}.o", - dependencies=[ - SourceArtifact( - self.context.kernels_dir / "generic" / "transpose.cc" - ) - ], - extra_flags=[ - f"-DDIM_m={self.m}", - f"-DDIM_n={self.n}", - ], - ), - ] + return [KernelObjectArtifact.from_extern(self._kernel())] def get_arg_spec(self): batch_dim = (self.num_batches,) if self.num_batches > 1 else () From 29fd2a70e0b7ff907a54ec0b8ce24dfaada8f45f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:24:47 -0600 Subject: [PATCH 02/22] Port the Llama decode ops and elementwise bases to kernel factories rms_norm (plain and weighted), rope, softmax, and every ChanneledUnary / BinaryElementwise operator (silu, gelu, relu, sigmoid, tanh, leaky_relu, layer_norm, elementwise_add/mul, axpy) now take their kernels from aie.iron.kernels. Each op builds its ExternalFunctions once in _kernel(s) and hands the same objects to the design and to KernelObjectArtifact.from_extern, so the design no longer re-declares the symbol, object name, or func_prefix by hand. softmax binds mask_bf16 from the softmax object with object_file.bind. Co-Authored-By: Claude --- iron/common/operator_bases.py | 96 ++++++------------- iron/operators/axpy/design.py | 8 +- iron/operators/axpy/op.py | 20 +--- iron/operators/binary_elementwise_design.py | 13 +-- iron/operators/channeled_unary_design.py | 15 +-- iron/operators/elementwise_add/op.py | 9 +- iron/operators/elementwise_mul/op.py | 9 +- iron/operators/gelu/op.py | 8 +- iron/operators/layer_norm/op.py | 6 +- iron/operators/leaky_relu/design.py | 11 +-- iron/operators/leaky_relu/op.py | 8 +- iron/operators/relu/op.py | 7 +- iron/operators/rms_norm/design.py | 9 +- iron/operators/rms_norm/design_weighted.py | 18 +--- iron/operators/rms_norm/op.py | 34 +++---- iron/operators/rope/design.py | 20 +--- iron/operators/rope/op.py | 17 ++-- iron/operators/sigmoid/op.py | 8 +- iron/operators/silu/op.py | 8 +- iron/operators/softmax/design.py | 18 +--- iron/operators/softmax/op.py | 42 +++----- iron/operators/tanh/op.py | 8 +- .../kernel_object_arch_isolation.py | 2 +- 23 files changed, 136 insertions(+), 258 deletions(-) diff --git a/iron/common/operator_bases.py b/iron/common/operator_bases.py index 8342db30cc..f7e4da43aa 100644 --- a/iron/common/operator_bases.py +++ b/iron/common/operator_bases.py @@ -8,17 +8,16 @@ from typing import Any, ClassVar import aie.utils as aie_utils +from aie.iron.kernel import ExternalFunction from .base import MLIROperator, AIERuntimeArgSpec from .context import AIEContext from .compilation import ( - KernelArchiveArtifact, KernelObjectArtifact, SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) -from .device_utils import get_kernel_dir from .utils import get_shim_dma_limit @@ -43,19 +42,15 @@ def lut_based_ops_artifacts(kernel_dir: str) -> list[KernelObjectArtifact]: class ChanneledUnaryOperator(MLIROperator): """Base class for channeled unary AIE operators (single input, single output). - Assumes a single kernel source file and a standard design.py callback - with args [device, size, num_aie_columns, num_channels, tile_size, trace_size]. + Assumes a single kernel and a standard design.py callback with args + [device, size, num_aie_columns, num_channels, tile_size, trace_size]. - Subclasses must define ClassVar attributes: - kernel_name: name of the kernel object file (e.g. "gelu" → gelu.o / gelu.cc) - callback_fn: design.py callback function name (e.g. "my_gelu") - needs_lut_ops: set True for operators that require lut_based_ops.o on aie2 + Subclasses must implement _kernel(), returning the mlir-aie kernel factory's + ExternalFunction for one line of _line_size elements. Customization points: - For operators with extra parameters (e.g. alpha, trace_size), add dataclass fields and override _mlir_callback_args(). - - For operators requiring multiple kernels, extra compile flags, or - external source files, override get_kernel_artifacts() directly. - For non-standard arg specs, override get_arg_spec() directly. - If none of these fit, subclass MLIROperator instead. """ @@ -66,10 +61,7 @@ class ChanneledUnaryOperator(MLIROperator): tile_size: int context: AIEContext | None = field(default=None, repr=False) - kernel_name: ClassVar[str] - kernel_fn_name: ClassVar[str] callback_fn: ClassVar[str] - needs_lut_ops: ClassVar[bool] = False tile_cap: ClassVar[int] = 4096 def __post_init__(self) -> None: @@ -111,23 +103,16 @@ def _mlir_callback_args(self) -> list[Any]: ] @property - def _kernel_link_file(self) -> str: - """The file name that the MLIR Kernel declaration should link_with. + def _line_size(self) -> int: + """Elements each core processes per kernel call.""" + return min(self.tile_size, self.tile_cap) - When auxiliary objects are required (e.g. lut_based_ops.o on aie2), - all objects are bundled into an archive and the archive name is - returned so that aiecc links the entire archive. - """ - if self.needs_lut_ops and get_kernel_dir() == "aie2": - return f"{self.name}_kernels.a" - return f"{self.kernel_name}.o" + def _kernel(self) -> ExternalFunction: + """The kernel each core runs over one line of _line_size elements.""" + raise NotImplementedError def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: - callback_args = self._mlir_callback_args() + [ - self.kernel_fn_name, - self._kernel_link_file, - self.tile_cap, - ] + callback_args = self._mlir_callback_args() + [self._kernel(), self.tile_cap] return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", DesignGenerator( @@ -137,43 +122,23 @@ def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: ), ) - def get_kernel_artifacts(self) -> list: - dev = aie_utils.get_current_device() - kernel_dir = get_kernel_dir(dev) - kernel_obj = KernelObjectArtifact( - f"{self.kernel_name}.o", - dependencies=[ - SourceArtifact( - self.context.kernels_dir / kernel_dir / f"{self.kernel_name}.cc" - ) - ], - ) - if self.needs_lut_ops and kernel_dir == "aie2": - lut_objs = lut_based_ops_artifacts(kernel_dir) - return [ - KernelArchiveArtifact( - f"{self.name}_kernels.a", - dependencies=[kernel_obj] + lut_objs, - ) - ] - return [kernel_obj] + def get_kernel_artifacts(self) -> list[KernelObjectArtifact]: + return [KernelObjectArtifact.from_extern(self._kernel())] @dataclass class BinaryElementwiseOperator(MLIROperator): """Base class for binary element-wise AIE operators (two inputs, one output). - Assumes a single kernel source file and a standard design.py callback - with args [device, size, num_aie_columns, tile_size, trace_size]. + Assumes a single kernel and a standard design.py callback with args + [device, size, num_aie_columns, tile_size, trace_size]. Unlike ChanneledUnaryOperator, binary operators have no explicit num_channels parameter — each core uses 2 DMA channels (one per input), so the ShimDMA limit is enforced as num_aie_columns * 2 <= 16. - Subclasses must define ClassVar attributes: - kernel_name: name of the kernel object file (e.g. "add" → add.o / add.cc) - kernel_subdir: subdirectory under aie_kernels/ (e.g. "generic") - callback_fn: design.py callback function name (e.g. "my_eltwise_add") + Subclasses must implement _kernel(), returning the mlir-aie kernel factory's + ExternalFunction for one tile of _tile_elements elements. """ size: int @@ -181,9 +146,6 @@ class BinaryElementwiseOperator(MLIROperator): num_aie_columns: int = 8 context: AIEContext | None = field(default=None, repr=False) - kernel_name: ClassVar[str] - kernel_fn_name: ClassVar[str] - kernel_subdir: ClassVar[str] callback_fn: ClassVar[str] # Override parent's "c" alias with "col" so binary-elementwise operator names # are unambiguous when num_aie_columns and num_channels both appear in the @@ -231,11 +193,17 @@ def _mlir_callback_args(self) -> list[Any]: 0, ] + @property + def _tile_elements(self) -> int: + """Elements each core processes per kernel call.""" + return min(self.tile_size, 4096) + + def _kernel(self) -> ExternalFunction: + """The kernel each core runs over one tile of _tile_elements elements.""" + raise NotImplementedError + def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: - callback_args = self._mlir_callback_args() + [ - self.kernel_fn_name, - f"{self.kernel_name}.o", - ] + callback_args = self._mlir_callback_args() + [self._kernel()] return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", DesignGenerator( @@ -246,10 +214,4 @@ def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: ) def get_kernel_artifacts(self) -> list[KernelObjectArtifact]: - source = self.context.kernels_dir / get_kernel_dir() / f"{self.kernel_name}.cc" - return [ - KernelObjectArtifact( - f"{self.kernel_name}.o", - dependencies=[SourceArtifact(source)], - ), - ] + return [KernelObjectArtifact.from_extern(self._kernel())] diff --git a/iron/operators/axpy/design.py b/iron/operators/axpy/design.py index e9421c8aeb..a10737ba7d 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, TaskGroup, Worker +from aie.iron import 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 @@ -17,6 +17,7 @@ def my_axpy( tile_size, trace_size, scalar_factor, + axpy_bf16_vector, ): factor = scalar_factor per_tile_elements = 4096 if tile_size > 4096 else tile_size @@ -38,11 +39,6 @@ def my_axpy( of_in2s = [ObjectFifo(tile_ty, name=f"in2_{i}") for i in range(num_columns)] of_outs = [ObjectFifo(tile_ty, name=f"out_{i}") for i in range(num_columns)] - # AIE Core Function declaration - axpy_bf16_vector = Kernel( - "saxpy", "axpy.o", [tile_ty, tile_ty, np.float32, tile_ty, np.int32] - ) - # Define a task that will run on a compute tile def core_body(of_in1, of_in2, of_out, axpy): # Number of sub-vector "tile" iterations diff --git a/iron/operators/axpy/op.py b/iron/operators/axpy/op.py index 6c03dd9148..87ec4d2823 100644 --- a/iron/operators/axpy/op.py +++ b/iron/operators/axpy/op.py @@ -4,10 +4,10 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import datamovement + from iron.common import ( BinaryElementwiseOperator, - KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) @@ -19,23 +19,13 @@ class AXPY(BinaryElementwiseOperator): scalar_factor: float = 3.0 - kernel_name: ClassVar[str] = "axpy" - kernel_fn_name: ClassVar[str] = "saxpy" callback_fn: ClassVar[str] = "my_axpy" - def get_kernel_artifacts(self) -> list[KernelObjectArtifact]: - # axpy.cc lives under aie_kernels/generic/ (not device-specific) - return [ - KernelObjectArtifact( - "axpy.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / "generic" / "axpy.cc") - ], - ) - ] + def _kernel(self): + return datamovement.axpy(self._tile_elements) def _mlir_callback_args(self): - return super()._mlir_callback_args() + [self.scalar_factor] + return super()._mlir_callback_args() + [self.scalar_factor, self._kernel()] def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: return PythonGeneratedMLIRArtifact( diff --git a/iron/operators/binary_elementwise_design.py b/iron/operators/binary_elementwise_design.py index fea333f404..a56bad10bc 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, TaskGroup, Worker +from aie.iron import 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 @@ -16,9 +16,7 @@ def binary_elementwise_design( num_columns, tile_size, trace_size, - kernel_fn_name, - kernel_obj_file, - func_prefix="", + eltwise_kernel, ): per_tile_elements = 4096 if tile_size > 4096 else tile_size n = per_tile_elements * num_columns @@ -39,13 +37,6 @@ def binary_elementwise_design( of_in2s = [ObjectFifo(tile_ty, name=f"in2_{i}") for i in range(num_columns)] of_outs = [ObjectFifo(tile_ty, name=f"out_{i}") for i in range(num_columns)] - # AIE Core Function declaration - eltwise_kernel = Kernel( - f"{func_prefix}{kernel_fn_name}", - f"{func_prefix}{kernel_obj_file}", - [tile_ty, tile_ty, tile_ty, np.int32], - ) - # Define a task that will run on a compute tile def core_body(of_in1, of_in2, of_out, eltwise_fn): for _ in range_(N_div_n): diff --git a/iron/operators/channeled_unary_design.py b/iron/operators/channeled_unary_design.py index 7cff67c609..fd2c540a73 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, TaskGroup, Worker +from aie.iron import 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 @@ -17,10 +17,8 @@ def channeled_unary_design( num_channels, tile_size, trace_size, - kernel_fn_name, - kernel_obj_file, + kernel_fn, tile_cap=4096, - func_prefix="", ): xfr_dtype = bfloat16 line_size = tile_cap if tile_size > tile_cap else tile_size @@ -54,13 +52,6 @@ def channeled_unary_design( for j in range(num_channels) ] - # External, binary kernel definition - kernel_fcn = Kernel( - f"{func_prefix}{kernel_fn_name}", - f"{func_prefix}{kernel_obj_file}", - [line_type, line_type, np.int32], - ) - # Task for the core to perform def core_fn(of_in, of_out, kernel_line): for _ in range_(N_div_n): @@ -77,7 +68,7 @@ def core_fn(of_in, of_out, kernel_line): [ of_ins[i * num_channels + j].cons(), of_outs[i * num_channels + j].prod(), - kernel_fcn, + kernel_fn, ], ) for i in range(num_columns) diff --git a/iron/operators/elementwise_add/op.py b/iron/operators/elementwise_add/op.py index d129233bde..5ced5b1eaa 100644 --- a/iron/operators/elementwise_add/op.py +++ b/iron/operators/elementwise_add/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import eltwise + from iron.common import BinaryElementwiseOperator @@ -11,11 +13,10 @@ class ElementwiseAdd(BinaryElementwiseOperator): """AIE-accelerated element-wise addition""" - kernel_name: ClassVar[str] = "add" - kernel_fn_name: ClassVar[str] = "eltwise_add_bf16_vector_size" - kernel_subdir: ClassVar[str] = "generic" callback_fn: ClassVar[str] = "my_eltwise_add" - kernels_from_mlir_aie: ClassVar[bool] = True + + def _kernel(self): + return eltwise.add_sized(self._tile_elements) def reference(self, a, b): from iron.operators.elementwise_add.reference import reference diff --git a/iron/operators/elementwise_mul/op.py b/iron/operators/elementwise_mul/op.py index cc7cc7761e..8a7f4748f1 100644 --- a/iron/operators/elementwise_mul/op.py +++ b/iron/operators/elementwise_mul/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import eltwise + from iron.common import BinaryElementwiseOperator @@ -11,11 +13,10 @@ class ElementwiseMul(BinaryElementwiseOperator): """AIE-accelerated element-wise multiplication""" - kernel_name: ClassVar[str] = "mul" - kernel_fn_name: ClassVar[str] = "eltwise_mul_bf16_vector_size" - kernel_subdir: ClassVar[str] = "generic" callback_fn: ClassVar[str] = "my_eltwise_mul" - kernels_from_mlir_aie: ClassVar[bool] = True + + def _kernel(self): + return eltwise.mul_sized(self._tile_elements) def reference(self, a, b): from iron.operators.elementwise_mul.reference import reference diff --git a/iron/operators/gelu/op.py b/iron/operators/gelu/op.py index c67c036ea9..1644f56dfb 100644 --- a/iron/operators/gelu/op.py +++ b/iron/operators/gelu/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import activation + from iron.common import ChanneledUnaryOperator @@ -11,8 +13,8 @@ class GELU(ChanneledUnaryOperator): """AIE-accelerated GELU activation function""" - kernel_name: ClassVar[str] = "gelu" - kernel_fn_name: ClassVar[str] = "gelu_bf16_size" - needs_lut_ops: ClassVar[bool] = True callback_fn: ClassVar[str] = "my_gelu" tile_cap: ClassVar[int] = 8192 + + def _kernel(self): + return activation.gelu_sized(self._line_size) diff --git a/iron/operators/layer_norm/op.py b/iron/operators/layer_norm/op.py index 2a55054fc6..0d398d395d 100644 --- a/iron/operators/layer_norm/op.py +++ b/iron/operators/layer_norm/op.py @@ -5,6 +5,7 @@ from typing import ClassVar import aie.utils as aie_utils +from aie.iron.kernels import norm from iron.common import ChanneledUnaryOperator @@ -14,8 +15,6 @@ class LayerNorm(ChanneledUnaryOperator): trace_size: InitVar[int] = 0 - kernel_name: ClassVar[str] = "layer_norm" - kernel_fn_name: ClassVar[str] = "layer_norm" callback_fn: ClassVar[str] = "my_layer_norm" tile_cap: ClassVar[int] = 8192 @@ -23,6 +22,9 @@ def __post_init__(self, trace_size): self.trace_size = trace_size super().__post_init__() + def _kernel(self): + return norm.layer_norm(self._line_size) + def _mlir_callback_args(self): return [ aie_utils.get_current_device(), diff --git a/iron/operators/leaky_relu/design.py b/iron/operators/leaky_relu/design.py index 408a311ba3..c4a094b3b8 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, TaskGroup, Worker +from aie.iron import 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 @@ -18,6 +18,7 @@ def my_leaky_relu( tile_size, trace_size, alpha, + leaky_relu_fcn, ): xfr_dtype = bfloat16 # Cap to 4096 bfloat16 elements (8 KB) to fit AIE core local memory @@ -45,14 +46,6 @@ def my_leaky_relu( for j in range(num_channels) ] - # External, binary kernel definition - # Leaky RELU kernel takes: input, output, input_size, alpha - leaky_relu_fcn = Kernel( - "leaky_relu_bf16", - "leaky_relu.o", - [line_type, line_type, np.int32, xfr_dtype], - ) - # Task for the core to perform def core_fn(of_in, of_out, leaky_relu_line): for _ in range_(N_div_n): diff --git a/iron/operators/leaky_relu/op.py b/iron/operators/leaky_relu/op.py index cfd13dfb7c..00d0b18458 100644 --- a/iron/operators/leaky_relu/op.py +++ b/iron/operators/leaky_relu/op.py @@ -5,6 +5,7 @@ from typing import ClassVar, Dict import aie.utils as aie_utils +from aie.iron.kernels import activation from iron.common import ( ChanneledUnaryOperator, PythonGeneratedMLIRArtifact, @@ -18,8 +19,6 @@ class LeakyReLU(ChanneledUnaryOperator): alpha: float = 0.01 - kernel_name: ClassVar[str] = "leaky_relu" - kernel_fn_name: ClassVar[str] = "leaky_relu_bf16" callback_fn: ClassVar[str] = "my_leaky_relu" _name_aliases: ClassVar[Dict[str, str]] = { **ChanneledUnaryOperator._name_aliases, @@ -46,8 +45,11 @@ def __post_init__(self) -> None: ) super().__post_init__() + def _kernel(self): + return activation.leaky_relu(self._line_size) + def _mlir_callback_args(self): - return super()._mlir_callback_args() + [self.alpha] + return super()._mlir_callback_args() + [self.alpha, self._kernel()] def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: return PythonGeneratedMLIRArtifact( diff --git a/iron/operators/relu/op.py b/iron/operators/relu/op.py index 2e070b7b0f..df1e5716fc 100644 --- a/iron/operators/relu/op.py +++ b/iron/operators/relu/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import eltwise + from iron.common import ChanneledUnaryOperator @@ -11,10 +13,11 @@ class ReLU(ChanneledUnaryOperator): """AIE-accelerated ReLU activation function""" - kernel_name: ClassVar[str] = "relu" - kernel_fn_name: ClassVar[str] = "relu_bf16_size" callback_fn: ClassVar[str] = "my_relu" + def _kernel(self): + return eltwise.relu_sized(self._line_size) + def reference(self, x): from iron.operators.relu.reference import reference diff --git a/iron/operators/rms_norm/design.py b/iron/operators/rms_norm/design.py index 2daeea9c6e..3bc0dfd810 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, TaskGroup, Worker +from aie.iron import 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_ @@ -18,6 +18,8 @@ def my_rms_norm( tile_size, trace_size, epsilon=1e-5, + *, + rms_norm_kernel, ): per_tile_elements = 8192 if tile_size > 8192 else tile_size total_cores = num_columns * num_channels @@ -48,11 +50,6 @@ def my_rms_norm( for j in range(num_channels) ] - # AIE Core Function declaration - rms_norm_kernel = Kernel( - "rms_norm_eps", "rms_norm.o", [tile_ty, tile_ty, np.int32, np.float32] - ) - # Define a task that will run on a compute tile def core_body(of_in1, of_out, rms_norm_kernel): # Number of sub-vector "tile" iterations diff --git a/iron/operators/rms_norm/design_weighted.py b/iron/operators/rms_norm/design_weighted.py index 8f82774d6f..b4993e769e 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, TaskGroup, Worker +from aie.iron import 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_ @@ -18,7 +18,9 @@ def my_weighted_rms_norm( weight_length, trace_size, epsilon=1e-5, - func_prefix="", + *, + rms_norm_kernel, + eltwise_mul_kernel, ): per_tile_elements = weight_length total_cores = num_columns * num_channels @@ -60,18 +62,6 @@ def my_weighted_rms_norm( for j in range(num_channels) ] - # AIE Core Function declaration - rms_norm_kernel = Kernel( - f"{func_prefix}rms_norm_eps", - f"{func_prefix}rms_norm.o", - [tile_ty, tile_ty, np.int32, np.float32], - ) - eltwise_mul_kernel = Kernel( - f"{func_prefix}eltwise_mul_bf16_vector_size", - f"{func_prefix}mul.o", - [tile_ty, weights_ty, tile_ty, np.int32], - ) - # Define a task that will run on a compute tile def core_body_norm(of_in1, of_out1, rms_norm): # Number of sub-vector "tile" iterations diff --git a/iron/operators/rms_norm/op.py b/iron/operators/rms_norm/op.py index fcc6a60e7f..468caf7b75 100644 --- a/iron/operators/rms_norm/op.py +++ b/iron/operators/rms_norm/op.py @@ -8,12 +8,11 @@ MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) import aie.utils as aie_utils -from iron.common.device_utils import get_kernel_dir +from aie.iron.kernels import eltwise, norm from iron.common.utils import get_shim_dma_limit @@ -68,6 +67,16 @@ def __post_init__(self): ) MLIROperator.__init__(self, context=self.context) + def _kernels(self): + """The rms_norm kernel, then (if weighted) the weight multiply.""" + # The unweighted design caps a core's tile at 8192 elements; the + # weighted one normalizes whole weight-length rows. + line = self.tile_size if self.weighted else min(self.tile_size, 8192) + kernels = {"rms_norm_kernel": norm.rms_norm_eps(line)} + if self.weighted: + kernels["eltwise_mul_kernel"] = eltwise.mul_sized(line) + return kernels + def get_mlir_artifact(self): if self.weighted: source_path = self.operator_dir / "design_weighted.py" @@ -90,29 +99,12 @@ def get_mlir_artifact(self): 0, # trace_size self.epsilon, ), + self._kernels(), ), ) def get_kernel_artifacts(self): - arch_dir = get_kernel_dir() - artifacts = [ - KernelObjectArtifact( - "rms_norm.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / arch_dir / "rms_norm.cc") - ], - ), - ] - if self.weighted: - artifacts.append( - KernelObjectArtifact( - "mul.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / arch_dir / "mul.cc") - ], - ) - ) - return artifacts + return [KernelObjectArtifact.from_extern(k) for k in self._kernels().values()] def get_arg_spec(self): specs = [AIERuntimeArgSpec("in", (self.size // self.tile_size, self.tile_size))] diff --git a/iron/operators/rope/design.py b/iron/operators/rope/design.py index e9e65dab02..086802034a 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, TaskGroup, Worker +from aie.iron import 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_ @@ -32,18 +32,13 @@ def rope( angle_rows=None, num_aie_columns=1, trace_size=0, - method_type=None, - func_prefix="", + *, + rope_kernel, ): dtype = bfloat16 if angle_rows is None: angle_rows = rows - kernel_object = ( - f"{func_prefix}rope" - + (f"_{method_type}" if method_type is not None else "") - + ".o" - ) assert cols % (16 * 2) == 0 and cols >= ( 16 * 2 @@ -73,15 +68,6 @@ def rope( ObjectFifo(tensor_tile_ty, name=f"out_{i}") for i in range(num_aie_columns) ] - # AIE Core Function declaration. method_type 0 = two-halves (HF), 1 = - # interleaved/Llama (the "rope" symbol). - rope_symbol = "rope_two_halves" if method_type == 0 else "rope" - rope_kernel = Kernel( - f"{func_prefix}{rope_symbol}", - kernel_object, - [tensor_tile_ty, angle_tile_ty, tensor_tile_ty, np.int32], - ) - # Define a task that will run on a compute tile def core_body(of_in, of_lut, of_out, rope_kernel): # Number of sub-vector "tile" iterations diff --git a/iron/operators/rope/op.py b/iron/operators/rope/op.py index 8e084ed265..27bc702b38 100644 --- a/iron/operators/rope/op.py +++ b/iron/operators/rope/op.py @@ -8,11 +8,11 @@ MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) import aie.utils as aie_utils +from aie.iron.kernels import datamovement @dataclass @@ -53,6 +53,10 @@ def __post_init__(self): MLIROperator.__init__(self, context=self.context) + def _kernel(self): + # method_type 0 = two-halves (HF), 1 = interleaved (Llama paper). + return datamovement.rope(self.cols, two_halves=self.method_type == 0) + def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", @@ -66,20 +70,13 @@ def get_mlir_artifact(self): self.angle_rows, self.num_aie_columns, 0, - self.method_type, ), + {"rope_kernel": self._kernel()}, ), ) def get_kernel_artifacts(self): - return [ - KernelObjectArtifact( - f"rope_{self.method_type}.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / "generic" / "rope.cc") - ], - ), - ] + return [KernelObjectArtifact.from_extern(self._kernel())] def get_arg_spec(self): return [ diff --git a/iron/operators/sigmoid/op.py b/iron/operators/sigmoid/op.py index a8daacbe43..ea4005afd6 100644 --- a/iron/operators/sigmoid/op.py +++ b/iron/operators/sigmoid/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import activation + from iron.common import ChanneledUnaryOperator @@ -11,7 +13,7 @@ class Sigmoid(ChanneledUnaryOperator): """AIE-accelerated Sigmoid activation function""" - kernel_name: ClassVar[str] = "sigmoid" - kernel_fn_name: ClassVar[str] = "sigmoid_bf16" - needs_lut_ops: ClassVar[bool] = True callback_fn: ClassVar[str] = "my_sigmoid" + + def _kernel(self): + return activation.sigmoid(self._line_size) diff --git a/iron/operators/silu/op.py b/iron/operators/silu/op.py index 7e3f4a5e0f..4353ac948f 100644 --- a/iron/operators/silu/op.py +++ b/iron/operators/silu/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass, field from typing import ClassVar +from aie.iron.kernels import activation + from iron.common import ChanneledUnaryOperator @@ -13,10 +15,10 @@ class SiLU(ChanneledUnaryOperator): num_channels: int = field(default=1, init=False, repr=False) - kernel_name: ClassVar[str] = "silu" - kernel_fn_name: ClassVar[str] = "silu_bf16_size" callback_fn: ClassVar[str] = "my_silu" - needs_lut_ops: ClassVar[bool] = True + + def _kernel(self): + return activation.silu_sized(self._line_size) def reference(self, x): from iron.operators.silu.reference import reference diff --git a/iron/operators/softmax/design.py b/iron/operators/softmax/design.py index e798956da8..7061b4fe6d 100644 --- a/iron/operators/softmax/design.py +++ b/iron/operators/softmax/design.py @@ -5,7 +5,6 @@ import numpy as np from aie.iron import ( - Kernel, ObjectFifo, ScratchpadParameter, Program, @@ -32,8 +31,9 @@ def softmax( tile_size, rtp_vector_size=None, vector_size_parameter=None, - func_prefix="", - kernel_obj_file="softmax.o", + *, + softmax_kernel, + mask_kernel, ): per_tile_elements = tile_size if rtp_vector_size is None: @@ -64,18 +64,6 @@ def softmax( for j in range(num_channels) ] - # AIE Core Function declaration - softmax_kernel = Kernel( - f"{func_prefix}softmax_bf16", - f"{func_prefix}{kernel_obj_file}", - [tile_ty, tile_ty, np.int32], - ) - mask_kernel = Kernel( - f"{func_prefix}mask_bf16", - f"{func_prefix}{kernel_obj_file}", - [tile_ty, np.int32, np.int32], - ) - # Vector size source: either a scratchpad Parameter (synced from host each # dispatch) or a write-RTP buffer set via rt.inline_ops at compile time. use_scratchpad = vector_size_parameter is not None diff --git a/iron/operators/softmax/op.py b/iron/operators/softmax/op.py index a1aa7994f1..f534b17f29 100644 --- a/iron/operators/softmax/op.py +++ b/iron/operators/softmax/op.py @@ -3,16 +3,16 @@ from dataclasses import dataclass, field +import numpy as np +from ml_dtypes import bfloat16 + import aie.utils as aie_utils +from aie.iron.kernels import activation -from iron.common.device_utils import get_kernel_dir -from iron.common.operator_bases import lut_based_ops_artifacts from iron.common import ( MLIROperator, AIERuntimeArgSpec, - KernelArchiveArtifact, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) @@ -45,14 +45,16 @@ def __post_init__(self): ) MLIROperator.__init__(self, context=self.context) - @property - def _kernel_link_file(self): - kernel_dir = get_kernel_dir() - if kernel_dir == "aie2": - return f"{self.name}_kernels.a" - return "softmax.o" + def _softmax(self): + return activation.softmax(self.cols) def get_mlir_artifact(self): + softmax_fn = self._softmax() + # mask_bf16 is exported by the same softmax.cc translation unit. + mask_fn = softmax_fn.object_file.bind( + "mask_bf16", + [np.ndarray[(self.cols,), np.dtype[bfloat16]], np.int32, np.int32], + ) return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", DesignGenerator( @@ -68,28 +70,14 @@ def get_mlir_artifact(self): "tile_size": self.cols, "rtp_vector_size": self.rtp_vector_size, "vector_size_parameter": self.vector_size_parameter, - "kernel_obj_file": self._kernel_link_file, + "softmax_kernel": softmax_fn, + "mask_kernel": mask_fn, }, ), ) def get_kernel_artifacts(self): - kernel_dir = get_kernel_dir() - softmax_obj = KernelObjectArtifact( - "softmax.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / kernel_dir / "softmax.cc") - ], - ) - lut_objs = lut_based_ops_artifacts(kernel_dir) - if lut_objs: - return [ - KernelArchiveArtifact( - f"{self.name}_kernels.a", - dependencies=[softmax_obj] + lut_objs, - ) - ] - return [softmax_obj] + return [KernelObjectArtifact.from_extern(self._softmax())] def get_arg_spec(self): return [ diff --git a/iron/operators/tanh/op.py b/iron/operators/tanh/op.py index ac25c814df..541303472f 100644 --- a/iron/operators/tanh/op.py +++ b/iron/operators/tanh/op.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import ClassVar +from aie.iron.kernels import activation + from iron.common import ChanneledUnaryOperator @@ -11,7 +13,7 @@ class Tanh(ChanneledUnaryOperator): """AIE-accelerated Tanh activation function""" - kernel_name: ClassVar[str] = "tanh" - kernel_fn_name: ClassVar[str] = "tanh_bf16" - needs_lut_ops: ClassVar[bool] = True callback_fn: ClassVar[str] = "my_tanh" + + def _kernel(self): + return activation.tanh(self._line_size) diff --git a/iron/tests/compilation/kernel_object_arch_isolation.py b/iron/tests/compilation/kernel_object_arch_isolation.py index 4c9f560115..031e6829a4 100644 --- a/iron/tests/compilation/kernel_object_arch_isolation.py +++ b/iron/tests/compilation/kernel_object_arch_isolation.py @@ -45,7 +45,7 @@ def _mul_kernel_object(build_dir, device): def test_two_arches_do_not_resolve_the_same_kernel_object_path(tmp_path): """aie_kernels/generic/mul.cc is one source shared by aie2 and aie2p - (ElementwiseMul.kernel_subdir); its object must not collide in build_dir.""" + (eltwise.mul_sized); its object must not collide in build_dir.""" aie2 = _mul_kernel_object(tmp_path, NPU1()) aie2p = _mul_kernel_object(tmp_path, NPU2()) assert aie2.filename != aie2p.filename From 6268a7485fc7432a9f700470c49a785d1109220a Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:29:48 -0600 Subject: [PATCH 03/22] Port gemm to the linalg.mm / zero / convert_copy kernel factories The accumulator L1 type is now flat (m*n,), as the matmul, zero and convert_copy factories all declare it. ROUND_CONV_EVEN is requested via linalg.mm(round_conv_even=...), which needs the matching mlir-aie change. Co-Authored-By: Claude --- iron/operators/gemm/design.py | 79 +++++++++++------------------ iron/operators/gemm/op.py | 94 ++++++++++++++--------------------- 2 files changed, 65 insertions(+), 108 deletions(-) diff --git a/iron/operators/gemm/design.py b/iron/operators/gemm/design.py index 1a5f50e671..76252d7727 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -9,7 +9,6 @@ import numpy as np from aie.iron import ( - Kernel, ObjectFifo, Program, Buffer, @@ -22,7 +21,8 @@ from aie.iron.device import NPU1Col1, NPU1Col2, NPU1, NPU2, Tile from aie.helpers.taplib import TensorTiler2D, TensorAccessPattern from aie.iron.controlflow import range_ -from iron.common.kernels import zero_object_name +from aie.utils import set_current_device +from aie.iron.kernels import datamovement, linalg, zero from iron.operators._trace import maybe_enable_trace microkernel_mac_dim_map = { @@ -64,12 +64,6 @@ def main(): ) argparser.add_argument("--prio-accuracy", action="store_true", default=False) argparser.add_argument("--separate-c-tiles", type=int, choices=[0, 1], default=0) - argparser.add_argument( - "--archive", - type=str, - default=None, - help="Name of the archive file for the AIE kernels", - ) argparser.add_argument("--dtype_in", type=str, choices=["bf16"], default="bf16") argparser.add_argument( "--dtype_out", @@ -86,6 +80,25 @@ def main(): ) args = argparser.parse_args() + # The kernel factories pick their source and flags by the current device. + set_current_device(NPU1() if args.dev == "npu1" else NPU2()) + dtype_acc = str_to_dtype("f32" if args.prio_accuracy else args.dtype_out) + kernels = { + "matmul_kernel": linalg.mm( + args.m, + args.k, + args.n, + input_dtype=str_to_dtype(args.dtype_in), + output_dtype=dtype_acc, + vectorized=not args.scalar, + b_col_maj=bool(args.b_col_maj), + c_col_maj=bool(args.c_col_maj), + emulate_bf16_mmul_with_bfp16=args.emulate_bf16_mmul_with_bfp16, + ), + "zero_kernel": zero(args.m * args.n, dtype_acc, vectorized=not args.scalar), + } + if args.prio_accuracy: + kernels["convert_copy_kernel"] = datamovement.convert_copy(args.m * args.n) module = my_matmul( args.dev, args.M, @@ -104,7 +117,7 @@ def main(): args.prio_accuracy, args.separate_c_tiles, args.trace_size, - kernel_object=args.archive, + **kernels, ) output_file_path = Path(args.output_file_path) @@ -134,9 +147,10 @@ def my_matmul( prio_accuracy, separate_c_tiles, trace_size, - kernel_object=None, - zero_object=None, - func_prefix="", + *, + matmul_kernel, + zero_kernel, + convert_copy_kernel=None, ): n_aie_rows = 4 @@ -273,54 +287,19 @@ def _hw_stride_ok(stride_elems, itemsize): B_l1_ty = np.ndarray[(k, n), np.dtype[dtype_in]] C_l1_ty = np.ndarray[(m, n), np.dtype[dtype_out]] - # AIE Core Function declarations - scalar_suffix = "_scalar" if use_scalar else "" - gemm_object = ( - f"{func_prefix}{kernel_object}" - if kernel_object - else f"{func_prefix}gemm_{m}x{k}x{n}.o" - ) - # zero.cc is its own translation unit in mlir-aie, exporting a single `zero` - # specialized by -DZERO_TYPE/-DTILE_SIZE, so the zero kernel names a - # different object than the matmuls do. - zero_dtype_str = "f32" if use_larger_internal_buffer else dtype_out_str - zero_object = func_prefix + ( - zero_object or zero_object_name(zero_dtype_str, m * n, use_scalar) - ) - zero_func_name = f"{func_prefix}zero" if use_larger_internal_buffer: # Fix fifo depth for C objfifo to 1 since 1 buffer will be used for accumulation # and another for transfer to L2 fifo_depth_out = 1 # Set the type for accumulation - C_l1_ty_internal = np.ndarray[(m, n), np.dtype[dtype_out_internal]] + # Flat, as the matmul, zero and convert_copy kernels all declare it + C_l1_ty_internal = np.ndarray[(m * n,), np.dtype[dtype_out_internal]] # A kernel to convert from the internal f32 accumulation to bf16 for transfer to L2 is needed - convert_copy_kernel = Kernel( - f"{func_prefix}cast_f32_bf16_row", - f"{func_prefix}cast_f32_bf16.o", - [C_l1_ty_internal, C_l1_ty, np.int32], - ) - # Fix the kernels to use f32 outputs - zero_kernel = Kernel(zero_func_name, zero_object, [C_l1_ty_internal]) - matmul_func_name = f"{func_prefix}matmul{scalar_suffix}_{dtype_in_str}_f32" - matmul_kernel = Kernel( - matmul_func_name, - gemm_object, - [A_l1_ty, B_l1_ty, C_l1_ty_internal], - ) + assert convert_copy_kernel is not None else: # No need to use separate buffers for accumulation and transfer to L2, so # we only need the zero and matmul kernels fifo_depth_out = fifo_depth - zero_kernel = Kernel(zero_func_name, zero_object, [C_l1_ty]) - matmul_func_name = ( - f"{func_prefix}matmul{scalar_suffix}_{dtype_in_str}_{dtype_out_str}" - ) - matmul_kernel = Kernel( - matmul_func_name, - gemm_object, - [A_l1_ty, B_l1_ty, C_l1_ty], - ) # Tile declarations as tile[row][col] tiles = [[(col, row) for col in range(0, n_aie_cols)] for row in range(0, 6)] diff --git a/iron/operators/gemm/op.py b/iron/operators/gemm/op.py index 4290aabd9f..3dbeb757b9 100644 --- a/iron/operators/gemm/op.py +++ b/iron/operators/gemm/op.py @@ -10,13 +10,11 @@ MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) -from iron.common.device_utils import get_kernel_dir -from iron.common.kernels import zero_artifact, zero_object_name from aie.iron import str_to_dtype +from aie.iron.kernels import datamovement, linalg, zero import aie.utils as aie_utils @@ -88,19 +86,42 @@ def __post_init__(self): MLIROperator.__init__(self, context=self.context) - @property - def _kernel_flags_suffix(self): - """Suffix encoding compile-time flags that affect the kernel binary.""" - return f"_{int(self.prio_accuracy)}_{int(self.emulate_bf16_mmul_with_bfp16)}_{int(self.round_conv_even)}" - - @property - def _zero_dtype(self): - """The dtype the zero kernel clears: the accumulator's, not always C's. + def _kernels(self): + """The matmul, the zero that clears its accumulator and, under + prio_accuracy, the f32 -> bf16 copy out of that accumulator. prio_accuracy accumulates in f32 in L1 and converts on the way out, so - the buffer that gets zeroed is f32 even when C is bf16. + the matmul's C, and the buffer that gets zeroed, are f32 even when C + is bf16. """ - return "f32" if self.prio_accuracy else self.dtype_out + use_chess = self.context.compiler == "chess" + dtype_acc = np.float32 if self.prio_accuracy else str_to_dtype(self.dtype_out) + kernels = { + "matmul_kernel": linalg.mm( + self.tile_m, + self.tile_k, + self.tile_n, + input_dtype=str_to_dtype(self.dtype_in), + output_dtype=dtype_acc, + vectorized=not self.use_scalar, + b_col_maj=self.b_col_maj, + c_col_maj=self.c_col_maj, + use_chess=use_chess, + emulate_bf16_mmul_with_bfp16=self.emulate_bf16_mmul_with_bfp16, + round_conv_even=self.round_conv_even, + ), + "zero_kernel": zero( + self.tile_m * self.tile_n, + dtype_acc, + vectorized=not self.use_scalar, + use_chess=use_chess, + ), + } + if self.prio_accuracy: + kernels["convert_copy_kernel"] = datamovement.convert_copy( + self.tile_m * self.tile_n + ) + return kernels def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( @@ -127,56 +148,13 @@ def get_mlir_artifact(self): "prio_accuracy": self.prio_accuracy, "separate_c_tiles": int(self.separate_c_tiles), "trace_size": 0, - "kernel_object": f"gemm_{self.tile_m}x{self.tile_k}x{self.tile_n}_{int(self.b_col_maj)}_{int(self.c_col_maj)}{self._kernel_flags_suffix}.o", - "zero_object": zero_object_name( - self._zero_dtype, self.tile_m * self.tile_n, self.use_scalar - ), + **self._kernels(), }, ), ) def get_kernel_artifacts(self): - kernel_flags = [ - f"-DDIM_M={self.tile_m}", - f"-DDIM_K={self.tile_k}", - f"-DDIM_N={self.tile_n}", - ] - if self.prio_accuracy: - kernel_flags.append("-Dbf16_f32_ONLY") - else: - kernel_flags.append("-Dbf16_bf16_ONLY") - if self.round_conv_even: - kernel_flags.append("-DROUND_CONV_EVEN") - if self.emulate_bf16_mmul_with_bfp16: - kernel_flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16") - if self.b_col_maj: - kernel_flags.append("-DB_COL_MAJ") - if self.c_col_maj: - kernel_flags.append("-DC_COL_MAJ") - - kernel_dir = get_kernel_dir() - mm_source = self.context.kernels_dir / kernel_dir / "mm.cc" - return [ - KernelObjectArtifact( - f"gemm_{self.tile_m}x{self.tile_k}x{self.tile_n}_{int(self.b_col_maj)}_{int(self.c_col_maj)}{self._kernel_flags_suffix}.o", - extra_flags=kernel_flags, - dependencies=[SourceArtifact(mm_source)], - ), - KernelObjectArtifact( - "cast_f32_bf16.o", - [ - SourceArtifact( - self.context.kernels_dir / "aie2p" / "cast_f32_bf16.cc" - ) - ], - ), - zero_artifact( - self.context.kernels_dir, - self._zero_dtype, - self.tile_m * self.tile_n, - self.use_scalar, - ), - ] + return [KernelObjectArtifact.from_extern(k) for k in self._kernels().values()] def get_arg_spec(self): dtype_in = str_to_dtype(self.dtype_in) From 17624f37e4e5dc3280f557f3a9b96f51b4e87aa4 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:31:01 -0600 Subject: [PATCH 04/22] Port mha to the linalg.mha / zero / passthrough kernel factories The QK^T product comes from linalg.mha(b_col_maj=True, emulate_bf16_mmul_with_bfp16=True); partial_softmax, matmul_PV, rescale_O and init_scale_buffer are bound from its object with the design's types, and passThroughLine from a 16-bit passthrough object. Needs the matching mlir-aie linalg.mha b_col_maj/emulate kwargs. Co-Authored-By: Claude --- iron/operators/mha/design.py | 51 +++++++++++++++++-------------- iron/operators/mha/op.py | 59 ++++++++++++------------------------ 2 files changed, 48 insertions(+), 62 deletions(-) diff --git a/iron/operators/mha/design.py b/iron/operators/mha/design.py index 9255df6b48..5528f87cce 100644 --- a/iron/operators/mha/design.py +++ b/iron/operators/mha/design.py @@ -11,7 +11,6 @@ import numpy as np from aie.iron import ( - Kernel, ObjectFifo, Program, Runtime, @@ -24,7 +23,8 @@ from aie.iron.controlflow import range_ from aie.helpers.taplib import TensorTiler2D, TensorAccessSequence, TensorAccessPattern from aie.helpers.dialects.scf import if_, else_ -from iron.common.kernels import zero_object_name +from aie.iron.kernels import eltwise, linalg, zero +from aie.utils import set_current_device from iron.operators._trace import maybe_enable_trace, resolve_trace_size dtype_map = { @@ -82,6 +82,8 @@ def main(): args = argparser.parse_args() dev = NPU2() + # The kernel factories pick their source and flags by the current device. + set_current_device(dev) maybe_module = fused_mha( dev=dev, @@ -96,6 +98,15 @@ def main(): emulate_bf16_mmul_with_bfp16=args.emulate_bf16_mmul_with_bfp16, trace_size=args.trace_size, verbose=args.verbose, + matmul_QK=linalg.mha( + args.B_q, + args.d, + args.B_kv, + b_col_maj=True, + emulate_bf16_mmul_with_bfp16=True, + ), + zero_kernel=zero((args.B_q, args.B_kv), bfloat16), + passthrough_kernel=eltwise.passthrough(4 * args.B_q, np.int16), ) output_file_path = Path(args.output_file_path) @@ -120,6 +131,10 @@ def fused_mha( emulate_bf16_mmul_with_bfp16: bool, trace_size: int = 0, verbose: bool = False, + *, + matmul_QK, + zero_kernel, + passthrough_kernel, ): of_depth = 2 @@ -213,20 +228,20 @@ def fused_mha( s_ty = np.ndarray[(4 * B_q,), np.dtype[dtype]] # AIE kernel declarations - func_type = "" if vectorized else "_scalar" - # mha.cc uses zero.cc's templates internally but no longer re-exports a - # zero_ entry point, so the zero kernel comes from its own object. - zero_kernel = Kernel("zero", zero_object_name(dtype_str, B_q * B_kv), [qk_ty]) - - memcopy_kernel_scale = Kernel( - f"passThroughLine", "mha_passThrough.o", [s_ty, s_ty, np.int32] + # matmul_QK is mha.cc's QK^T product; the rest of mha.cc's toolkit is + # bound from the same object. + mha_object = matmul_QK.object_file + + # passthrough_kernel is the 16-bit passThroughLine; the scale buffers it + # copies are bf16. + memcopy_kernel_scale = passthrough_kernel.object_file.bind( + "passThroughLine", [s_ty, s_ty, np.int32] ) - scale_buffer_init_kernel = Kernel("init_scale_buffer", "mha.o", [s_ty, np.int32]) + scale_buffer_init_kernel = mha_object.bind("init_scale_buffer", [s_ty, np.int32]) - partial_softmax_kernel = Kernel( + partial_softmax_kernel = mha_object.bind( "partial_softmax", - "mha.o", [ qk_ty, qk_ty, @@ -240,15 +255,8 @@ def fused_mha( ], ) - matmul_QK = Kernel( - f"matmul_bf16_bf16_wrapper{func_type}", - "mha.o", - [q_ty, k_ty, qk_ty, np.ndarray[(2,), np.dtype[np.int32]]], - ) - - matmul_PV = Kernel( + matmul_PV = mha_object.bind( "matmul_PV", - "mha.o", [ qk_ty, k_ty, @@ -260,9 +268,8 @@ def fused_mha( ], ) - rescale_O = Kernel( + rescale_O = mha_object.bind( "rescale_O", - "mha.o", [qk_ty, s_ty, np.int32, np.ndarray[(2,), np.dtype[np.int32]]], ) diff --git a/iron/operators/mha/op.py b/iron/operators/mha/op.py index 490e178e0f..5a62c0d793 100644 --- a/iron/operators/mha/op.py +++ b/iron/operators/mha/op.py @@ -10,12 +10,12 @@ MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) -from iron.common.kernels import zero_artifact import aie.utils as aie_utils +from aie.iron.kernels import eltwise, linalg, zero +from ml_dtypes import bfloat16 @dataclass @@ -43,6 +43,21 @@ def __post_init__(self): raise ValueError(f"Only d=64 is supported in this version, got d={self.d}") MLIROperator.__init__(self, context=self.context) + def _kernels(self): + return { + # QK^T; the rest of mha.cc's symbols are bound from its object. + "matmul_QK": linalg.mha( + self.B_q, + self.d, + self.B_kv, + b_col_maj=True, + emulate_bf16_mmul_with_bfp16=True, + ), + "zero_kernel": zero((self.B_q, self.B_kv), bfloat16), + # 16-bit passThroughLine, bound to the bf16 scale buffers. + "passthrough_kernel": eltwise.passthrough(4 * self.B_q, np.int16), + } + def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", @@ -63,49 +78,13 @@ def get_mlir_artifact(self): "emulate_bf16_mmul_with_bfp16": True, "trace_size": 0, "verbose": False, + **self._kernels(), }, ), ) def get_kernel_artifacts(self): - mm_source = str(self.context.kernels_dir / "aie2p" / "mm.cc") - softmax_source = str(self.context.kernels_dir / "aie2p" / "softmax.cc") - mha_source = str(self.context.kernels_dir / "aie2p" / "mha.cc") - passthrough_source = str( - self.context.kernels_dir / "generic" / "passThrough.cc" - ) - - mm_defines_rowmaj = [ - "-Dbf16_bf16_ONLY", - f"-DDIM_M={self.B_q}", - f"-DDIM_K={self.d}", - f"-DDIM_N={self.B_kv}", - "-DROUND_CONV_EVEN", - "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", - ] - mm_defines_colmaj = mm_defines_rowmaj + [ - "-DB_COL_MAJ", - ] - # mha.cc #includes softmax.cc and mm.cc (both col-major and row-major) - # directly, so everything is compiled into a single mha.o translation unit. - return [ - KernelObjectArtifact( - "mha.o", - extra_flags=mm_defines_colmaj, - dependencies=[ - SourceArtifact(mha_source), - SourceArtifact(mm_source), - SourceArtifact(softmax_source), - ], - ), - KernelObjectArtifact( - "mha_passThrough.o", - extra_flags=["-DBIT_WIDTH=16"], - dependencies=[SourceArtifact(passthrough_source)], - ), - # The design zeroes one B_q x B_kv scores tile. - zero_artifact(self.context.kernels_dir, "bf16", self.B_q * self.B_kv), - ] + return [KernelObjectArtifact.from_extern(k) for k in self._kernels().values()] def get_arg_spec(self): seq_padding = self._calculate_seq_padding(self.seq_len, self.num_of_pipelines) From 13979d3f5f20add31df0f5021860b785449e2688 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:34:57 -0600 Subject: [PATCH 05/22] Port dequant and mem_copy to the expand / passthrough kernel factories mem_copy binds passThroughLine from the 16-bit passthrough object with its bf16 line type, and no longer takes a func_prefix: the factory object is already unique per recipe. Co-Authored-By: Claude --- iron/operators/dequant/design.py | 13 +++---------- iron/operators/dequant/op.py | 20 ++++++-------------- iron/operators/mem_copy/design.py | 26 ++++++++++++++++++-------- iron/operators/mem_copy/op.py | 23 +++++++++++------------ 4 files changed, 38 insertions(+), 44 deletions(-) diff --git a/iron/operators/dequant/design.py b/iron/operators/dequant/design.py index 213a2c48b9..1ea027a3bd 100644 --- a/iron/operators/dequant/design.py +++ b/iron/operators/dequant/design.py @@ -4,12 +4,10 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker +from aie.iron import ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ -from iron.common.device_utils import get_kernel_dir - def my_dequant_kernel( dev, @@ -19,6 +17,8 @@ def my_dequant_kernel( trace_size, tile_size, group_size, + *, + dequant_kernel, ): per_tile_elements = ( 16384 if tile_size > 16384 else tile_size @@ -61,13 +61,6 @@ def my_dequant_kernel( for j in range(num_channels) ] - # AIE Core Function declaration - dequant_kernel = Kernel( - "expand_uint4_to_bfloat16", - f"expand_{get_kernel_dir(dev)}_{tile_size}.o", - [in_tile_ty, out_tile_ty], - ) - # Define a task that will run on a compute tile def core_body(of_in1, of_out, dequant_kernel): # Number of sub-vector "tile" iterations diff --git a/iron/operators/dequant/op.py b/iron/operators/dequant/op.py index b919b6f87b..94c0bef221 100644 --- a/iron/operators/dequant/op.py +++ b/iron/operators/dequant/op.py @@ -10,12 +10,11 @@ MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) -from iron.common.device_utils import get_kernel_dir import aie.utils as aie_utils +from aie.iron.kernels import datamovement @dataclass @@ -59,22 +58,15 @@ def get_mlir_artifact(self): self.tile_size, self.group_size, ), + {"dequant_kernel": self._kernel()}, ), ) + def _kernel(self): + return datamovement.expand(self.tile_size, self.group_size) + def get_kernel_artifacts(self): - return [ - KernelObjectArtifact( - f"expand_{get_kernel_dir()}_{self.tile_size}.o", - dependencies=[ - SourceArtifact(self.context.kernels_dir / "generic" / "expand.cc") - ], - extra_flags=[ - f"-DTILE_SIZE={self.tile_size}", - f"-DGROUP_SIZE={self.group_size}", - ], - ) - ] + return [KernelObjectArtifact.from_extern(self._kernel())] def get_arg_spec(self): return [ diff --git a/iron/operators/mem_copy/design.py b/iron/operators/mem_copy/design.py index cd04bd724c..98940eed04 100644 --- a/iron/operators/mem_copy/design.py +++ b/iron/operators/mem_copy/design.py @@ -10,7 +10,6 @@ from aie.iron import ( TaskGroup, - Kernel, ObjectFifo, Program, Runtime, @@ -164,14 +163,27 @@ def create_partial_workload_config( # +def mem_copy_line_size(tile_size): + """Elements per ObjectFifo line, and per passThroughLine call.""" + return 8192 if tile_size > 8192 else tile_size + + def my_mem_copy( - dev, size, num_cores, num_channels, bypass, tile_size, trace_size, func_prefix="" + dev, + size, + num_cores, + num_channels, + bypass, + tile_size, + trace_size, + *, + passthrough_kernel=None, ): # -------------------------------------------------------------------------- # Configuration # -------------------------------------------------------------------------- xfr_dtype = bfloat16 - line_size = 8192 if tile_size > 8192 else tile_size + line_size = mem_copy_line_size(tile_size) fifodepth = 1 if line_size > 4096 else 2 line_type = np.ndarray[(line_size,), np.dtype[xfr_dtype]] transfer_type = np.ndarray[(size,), np.dtype[xfr_dtype]] @@ -199,11 +211,9 @@ def my_mem_copy( # Task core will run # -------------------------------------------------------------------------- - # External, binary kernel definition - mem_copy_fcn = Kernel( - f"{func_prefix}passThroughLine", - f"{func_prefix}mem_copy.o", - [line_type, line_type, np.int32], + # passthrough_kernel is the 16-bit passThroughLine; the lines are bf16. + mem_copy_fcn = passthrough_kernel.object_file.bind( + "passThroughLine", [line_type, line_type, np.int32] ) # Task for the core to perform diff --git a/iron/operators/mem_copy/op.py b/iron/operators/mem_copy/op.py index a4dafc670f..bd4234fc00 100644 --- a/iron/operators/mem_copy/op.py +++ b/iron/operators/mem_copy/op.py @@ -8,11 +8,14 @@ MLIROperator, AIERuntimeArgSpec, KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) import aie.utils as aie_utils +import numpy as np +from aie.iron.kernels import eltwise + +from iron.operators.mem_copy.design import mem_copy_line_size @dataclass @@ -51,23 +54,19 @@ def get_mlir_artifact(self): self.tile_size, 0, ), + {"passthrough_kernel": self._kernel()}, ), ) + def _kernel(self): + if self.bypass: + return None + return eltwise.passthrough(mem_copy_line_size(self.tile_size), np.int16) + def get_kernel_artifacts(self): if self.bypass: return [] - return [ - KernelObjectArtifact( - "mem_copy.o", - extra_flags=["-DBIT_WIDTH=16"], - dependencies=[ - SourceArtifact( - self.context.kernels_dir / "generic" / "passThrough.cc" - ) - ], - ) - ] + return [KernelObjectArtifact.from_extern(self._kernel())] def get_arg_spec(self): return [ From 774e2481f51e05aeeeeff6166b3b5d30b6970a26 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:48:35 -0600 Subject: [PATCH 06/22] Port flm/dequant to quant.q4nx_dequant; retire iron/common/kernels.py flm/dequant binds q4nx_dequant_bfp from the factory's object with its bfp16ebs8 output type. The factory also carries upstream's --aie-pipeliner-max-stagecount=5; the output stays byte-exact. flm/gemm stays hand-built: fused_mm compiles in a single epilogue mode, always rounds to nearest-even and wraps mm_fused.cc in fused_mm_tile.cc, while this operator selects among several modes at runtime from one xclbin. It is now the only user of the aie2 lut_based_ops helper, which moves into it from operator_bases. swiglu_prefill_stream also stays hand-built (its test is skipped at module level), but its silu/mul sources move to generic/, where they now live. Co-Authored-By: Claude --- AGENTS.md | 22 ++++++++--- iron/common/kernels.py | 57 ---------------------------- iron/common/operator_bases.py | 19 ---------- iron/common/stream/ops.py | 8 ++-- iron/operators/flm/dequant/design.py | 19 ++++++---- iron/operators/flm/dequant/op.py | 30 +++++---------- iron/operators/flm/gemm/op.py | 23 ++++++++++- 7 files changed, 63 insertions(+), 115 deletions(-) delete mode 100644 iron/common/kernels.py diff --git a/AGENTS.md b/AGENTS.md index 1483d26ed4..b74b256949 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -131,7 +131,9 @@ reuse lint 2. **AIE Kernels** ([mlir-aie `aie_kernels/`](https://github.com/Xilinx/mlir-aie/tree/main/aie_kernels)) - Architecture-specific C++ compute kernels, sourced from the installed - mlir-aie package (`AIEContext.kernels_dir`), not from this repo: + mlir-aie package, not from this repo. Operators get them from mlir-aie's + kernel factories (`aie.iron.kernels`), each of which returns an + `ExternalFunction` carrying its source, flags, symbol and argument types: - `generic/`: Works on both AIE2 and AIE2P - `aie2/`: AIE2-specific (NPU1) - `aie2p/`: AIE2P-specific (NPU2) @@ -244,16 +246,23 @@ Data movement pattern: L3 → Shim DMA → L2 → L1 (tile local) → Compute 2. Implement `op.py`: - Subclass `MLIROperator` - Implement `get_operator_name()`, `get_mlir_artifact()`, `get_kernel_artifacts()`, `get_arg_spec()` + - Build kernels with the `aie.iron.kernels` factories in one `_kernels()` + helper, pass them to the design as keyword arguments, and return + `[KernelObjectArtifact.from_extern(k) for k in self._kernels().values()]` + from `get_kernel_artifacts()` - Add validation for dimension constraints (assert statements) - Define tile sizes and column counts 3. Implement `design.py`: - - Import from `aie.iron` (Program, Runtime, Worker, ObjectFifo, Kernel) + - Import from `aie.iron` (Program, Runtime, Worker, ObjectFifo) + - Take the kernels as keyword arguments rather than declaring `Kernel(...)`; + bind further symbols of the same object with + `fn.object_file.bind(symbol, arg_types)` - Define function that builds MLIR-AIE design - Use `range_()` for loops (not Python `range`) - Handle device-specific logic (NPU1 vs NPU2) if needed 4. If a new C++ compute kernel is needed, add it to the [mlir-aie kernel library](https://github.com/Xilinx/mlir-aie/tree/main/aie_kernels) - and consume it via `AIEContext.kernels_dir`; IRON no longer hosts kernels + with a factory in `aie.iron.kernels`; IRON no longer hosts kernels - Choose appropriate directory: `generic/`, `aie2/`, or `aie2p/` - Use AIE API for portable vectorization when possible - Add `event0()` and `event1()` for performance profiling @@ -454,9 +463,10 @@ logging.basicConfig(level=logging.DEBUG) **"Kernel not found" or "Symbol not defined"** - Verify the kernel `.cc` exists under the installed mlir-aie package's - `include/aie_kernels//` (`AIEContext.kernels_dir`) -- Check `get_kernel_artifacts()` in `op.py` references correct kernel path -- Ensure kernel function signature matches `Kernel()` declaration in `design.py` + `include/aie_kernels//` (`AIEContext.kernels_dir`, overridden by + `MLIR_AIE_KERNEL_SOURCES`) +- Check `get_kernel_artifacts()` in `op.py` returns every factory the design uses +- Ensure the C signature matches the factory's (or `bind()`'s) argument types **Compilation hangs or fails** diff --git a/iron/common/kernels.py b/iron/common/kernels.py deleted file mode 100644 index 5df4580ea5..0000000000 --- a/iron/common/kernels.py +++ /dev/null @@ -1,57 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Kernel objects several operators compile out of mlir-aie's ``aie_kernels``. - -A kernel shared by more than one operator is declared once here, because the -operator builds the object and its design names that object in ``link_with``: -the two have to agree on a file name, and they live in different files. -""" - -from __future__ import annotations - -from pathlib import Path - -from .compilation import KernelObjectArtifact, SourceArtifact - -# mlir-aie's generic/zero.cc exports one entry point, `zero`, specialized by -# -DZERO_TYPE/-DTILE_SIZE. The C spelling of each dtype IRON zeroes: -ZERO_CTYPES = { - "i8": "int8_t", - "i16": "int16_t", - "i32": "int32_t", - "bf16": "bfloat16", - "f32": "float", -} - - -def zero_object_name(dtype_str: str, tile_size: int, scalar: bool = False) -> str: - """Object file name of the zero kernel specialized this way. - - The specialization is baked in at compile time, so it belongs in the name: - two designs zeroing different tiles need different objects. - """ - return f"zero_{dtype_str}_{tile_size}{'_scalar' if scalar else ''}.o" - - -def zero_artifact( - kernels_dir: Path, dtype_str: str, tile_size: int, scalar: bool = False -) -> KernelObjectArtifact: - """The zero-fill kernel object for a ``tile_size``-element ``dtype_str`` tile. - - mm.cc used to carry `zero_` alongside its matmuls, so a design got the - two from one object. mlir-aie split zero.cc out into its own translation unit - (#3732), which is why this is a separate artifact. - """ - try: - ctype = ZERO_CTYPES[dtype_str] - except KeyError: - raise ValueError(f"zero kernel: unsupported dtype {dtype_str}") from None - flags = [f"-DZERO_TYPE={ctype}", f"-DTILE_SIZE={tile_size}"] - if scalar: - flags.append("-DZERO_SCALAR") - return KernelObjectArtifact( - zero_object_name(dtype_str, tile_size, scalar), - dependencies=[SourceArtifact(kernels_dir / "generic" / "zero.cc")], - extra_flags=flags, - ) diff --git a/iron/common/operator_bases.py b/iron/common/operator_bases.py index f7e4da43aa..7b13f04435 100644 --- a/iron/common/operator_bases.py +++ b/iron/common/operator_bases.py @@ -4,7 +4,6 @@ from __future__ import annotations from dataclasses import dataclass, field -from pathlib import Path from typing import Any, ClassVar import aie.utils as aie_utils @@ -14,30 +13,12 @@ from .context import AIEContext from .compilation import ( KernelObjectArtifact, - SourceArtifact, PythonGeneratedMLIRArtifact, DesignGenerator, ) from .utils import get_shim_dma_limit -def lut_based_ops_artifacts(kernel_dir: str) -> list[KernelObjectArtifact]: - """Return the lut_based_ops kernel artifact for aie2 devices, empty list otherwise.""" - if kernel_dir != "aie2": - return [] - mlir_aie_dir = Path(aie_utils.config.root_path()) - return [ - KernelObjectArtifact( - "lut_based_ops.o", - dependencies=[ - SourceArtifact( - mlir_aie_dir / "aie_runtime_lib" / "AIE2" / "lut_based_ops.cpp" - ) - ], - ) - ] - - @dataclass class ChanneledUnaryOperator(MLIROperator): """Base class for channeled unary AIE operators (single input, single output). diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index 9effcbd880..da847c95bd 100644 --- a/iron/common/stream/ops.py +++ b/iron/common/stream/ops.py @@ -28,7 +28,6 @@ from onnxscript.values import Op, Opset from iron.common.layout import TiledStridedLayout, tiled_2d -from iron.common.kernels import ZERO_CTYPES # Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets. The # operand layouts are the contract the generated DMAs and the compiled kernel @@ -113,7 +112,7 @@ def _gemm_artifacts(kernels_dir, kernel_dir, m: int, k: int, n: int): "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", "-DROUND_CONV_EVEN", # zero.cc's entry point, over the m x n output tile. - f"-DZERO_TYPE={ZERO_CTYPES['bf16']}", + "-DZERO_TYPE=bfloat16", f"-DTILE_SIZE={m * n}", f"-include{zero_source}", ], @@ -159,11 +158,14 @@ def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs): GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts) -SILU = StreamKernel(key="silu", layouts=lambda: elementwise_layouts(2), source="silu") +SILU = StreamKernel( + key="silu", layouts=lambda: elementwise_layouts(2), source="silu", subdir="generic" +) ELTWISE_MUL = StreamKernel( key="eltwise_mul", layouts=lambda: elementwise_layouts(3), source="mul", + subdir="generic", ) Silu = custom_op("Silu") diff --git a/iron/operators/flm/dequant/design.py b/iron/operators/flm/dequant/design.py index d13e0d2e4e..854869c3e5 100644 --- a/iron/operators/flm/dequant/design.py +++ b/iron/operators/flm/dequant/design.py @@ -8,9 +8,7 @@ from aie.dialects._aie_enum_gen import AIEArch from aie.helpers.taplib.tap import TensorAccessPattern from aie.helpers.util import v8bfp16ebs8 -from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker - -from iron.common.device_utils import get_kernel_dir +from aie.iron import ObjectFifo, Program, Runtime, TaskGroup, Worker # flm.GEMM's B tiling, imported rather than restated: this design has to write # the buffer in the order that one reads it, and two copies would drift. @@ -91,8 +89,13 @@ def dequant_bfp( run_out_features=None, run_period_out_features=None, trace_size=0, + *, + dequant_kernel, ): - """K in-features, N out-features. B reaches the GEMM as (K, N).""" + """K in-features, N out-features. B reaches the GEMM as (K, N). + + ``dequant_kernel`` is ``quant.q4nx_dequant`` at this module's geometry. + """ if dev.arch != AIEArch.AIE2p: raise NotImplementedError("bfp16ebs8 exists only on AIE2P") if tile_n != N_TILE: @@ -124,10 +127,10 @@ def dequant_bfp( out_half_ty = np.ndarray[(HALF_BLOCKS,), np.dtype[v8bfp16ebs8]] out_blk_ty = np.ndarray[(CORE_BLOCKS,), np.dtype[v8bfp16ebs8]] - kernel = Kernel( - "q4nx_dequant_bfp", - f"q4nx_dequant_{get_kernel_dir(dev)}.o", - [qw_blk_ty, out_blk_ty], + # The factory declares both operands in bytes; the output FIFO carries + # bfp16ebs8 blocks. + kernel = dequant_kernel.object_file.bind( + "q4nx_dequant_bfp", [qw_blk_ty, out_blk_ty] ) def core_body(qw_in, out_of, k): diff --git a/iron/operators/flm/dequant/op.py b/iron/operators/flm/dequant/op.py index d6046ba693..6a997142c6 100644 --- a/iron/operators/flm/dequant/op.py +++ b/iron/operators/flm/dequant/op.py @@ -7,6 +7,7 @@ import aie.utils as aie_utils from aie.dialects._aie_enum_gen import AIEArch +from aie.iron.kernels import quant from iron.common import ( AIERuntimeArgSpec, @@ -14,10 +15,8 @@ KernelObjectArtifact, MLIROperator, PythonGeneratedMLIRArtifact, - SourceArtifact, ) from iron.common.compilation import InstsBinArtifact, XclbinArtifact -from iron.common.device_utils import get_kernel_dir from iron.operators.flm.dequant.design import ( BFP16_GROUP, @@ -126,6 +125,7 @@ def _mlir_artifact(self, filename, K, N): self.run_out_features, self.run_period_out_features, ), + {"dequant_kernel": self._kernel()}, ), ) @@ -151,28 +151,16 @@ def set_up_artifacts(self) -> None: ) self.add_artifacts([self.xclbin_artifact, self.insts_artifact]) - def get_kernel_artifacts(self): + def _kernel(self): dev = aie_utils.get_current_device() if dev.arch != AIEArch.AIE2p: raise NotImplementedError("bfp16ebs8 exists only on AIE2P") - return [ - KernelObjectArtifact( - f"q4nx_dequant_{get_kernel_dir(dev)}.o", - dependencies=[ - SourceArtifact( - self.context.kernels_dir / "generic" / "q4nx_dequant.cc" - ) - ], - extra_flags=[ - f"-DQ4NX_M_TILE={M_TILE}", - f"-DQ4NX_K_TILE={K_TILE}", - f"-DQ4NX_GROUP={GROUP}", - f"-DQ4NX_CT_K={CT_K}", - f"-DQ4NX_S={S}", - f"-DQ4NX_T={T}", - ], - ) - ] + return quant.q4nx_dequant( + m_tile=M_TILE, k_tile=K_TILE, group=GROUP, ct_k=CT_K, s=S, t=T + ) + + def get_kernel_artifacts(self): + return [KernelObjectArtifact.from_extern(self._kernel())] def get_arg_spec(self): # Both buffers are declared in bytes: a q4nx block interleaves three diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index 2f73ce3885..05f4aa1d22 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field +from pathlib import Path import numpy as np from typing import ClassVar, Dict @@ -19,7 +20,6 @@ from aie.dialects._aie_enum_gen import AIEArch from iron.common.device_utils import get_kernel_dir from iron.common.compilation import InstsBinArtifact, XclbinArtifact -from iron.common.operator_bases import lut_based_ops_artifacts import aie.utils as aie_utils from iron.operators.flm.packing import pack_b, packed_b_size @@ -45,6 +45,23 @@ ) +def lut_based_ops_artifacts(kernel_dir: str) -> list[KernelObjectArtifact]: + """Return the lut_based_ops kernel artifact for aie2 devices, empty list otherwise.""" + if kernel_dir != "aie2": + return [] + mlir_aie_dir = Path(aie_utils.config.root_path()) + return [ + KernelObjectArtifact( + "lut_based_ops.o", + dependencies=[ + SourceArtifact( + mlir_aie_dir / "aie_runtime_lib" / "AIE2" / "lut_based_ops.cpp" + ) + ], + ) + ] + + @dataclass class GEMM(MLIROperator): """AIE-accelerated bf16 GEMM on a 4-row grid, with a fused epilogue. @@ -355,6 +372,10 @@ def set_up_artifacts(self) -> None: self.add_artifacts([self.xclbin_artifact, self.insts_artifact]) def get_kernel_artifacts(self): + # Built by hand rather than from aie.iron.kernels.fused_mm: that + # factory compiles in one epilogue mode (this operator selects among + # several at runtime, from one xclbin), always rounds to nearest-even, + # and wraps mm_fused.cc in fused_mm_tile.cc. kernel_dir = get_kernel_dir() kernels_dir = self.context.kernels_dir generic = kernels_dir / "generic" From 07ebe077fbb3c7cd19272e7d871c6036df658a29 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 24 Sep 2026 21:54:39 -0600 Subject: [PATCH 07/22] Resolve the factory-kernel output dir before handing it to mlir-aie compile_external_kernel() runs the compiler with cwd set to its output directory, so a relative AIEContext build_dir (Llama uses "build_elf") resolved twice and clang could not find the staged kernel source. Co-Authored-By: Claude --- iron/common/compilation/base.py | 5 ++- iron/tests/compilation/relative_build_dir.py | 36 ++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 iron/tests/compilation/relative_build_dir.py diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 372b1928dc..85b03b5c6f 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -972,7 +972,10 @@ def _compile_extern(self, artifact, kernel_dir): # than its source, or half-prefixed. mlir-aie reuses any object already # at the output path, so remove it to make mlir-aie rebuild it. Path(artifact.filename).unlink(missing_ok=True) - compile_external_kernel(fn, str(Path(artifact.filename).parent), kernel_dir) + # mlir-aie compiles with cwd set to the output directory, so a + # relative one (e.g. AIEContext(build_dir="build_elf")) resolves twice. + out_dir = Path(artifact.filename).parent.resolve() + compile_external_kernel(fn, str(out_dir), kernel_dir) def _find_tool(self, name): return _find_tool(name, self.peano_dir, self.mlir_aie_dir) diff --git a/iron/tests/compilation/relative_build_dir.py b/iron/tests/compilation/relative_build_dir.py new file mode 100644 index 0000000000..b5e4c4906f --- /dev/null +++ b/iron/tests/compilation/relative_build_dir.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A factory kernel must compile when AIEContext's build_dir is relative. + +mlir-aie's compile_external_kernel() runs the compiler with cwd set to its +output directory, so a relative output directory is resolved twice and the +compiler looks for the kernel source under build_dir//build_dir/. +Llama's AIEContext(build_dir="build_elf") hit exactly this. The operator +tests never did because their build_dir is absolute. +""" + +from pathlib import Path + +import aie.utils as aie_utils +from aie.iron.device import NPU2 + +from iron.common import AIEContext +from iron.common.compilation import KernelObjectArtifact +from iron.operators.elementwise_mul.op import ElementwiseMul + + +def test_factory_kernel_compiles_with_a_relative_build_dir(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + aie_utils.set_current_device(NPU2()) + ctx = AIEContext(build_dir="build_rel") + op = ElementwiseMul(size=4096, tile_size=4096, num_aie_columns=1, context=ctx) + op.compile() + + objects = [a for a in op.artifacts.bfs() if isinstance(a, KernelObjectArtifact)] + assert objects, "ElementwiseMul produced no KernelObjectArtifact" + for obj in objects: + path = Path(obj.filename).resolve() + assert path.is_relative_to(tmp_path / "build_rel"), path + assert path.is_file(), f"{path} was not built" From d8cad999ff9c8b1b405a76b94a389fef23c39374 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 06:01:04 -0600 Subject: [PATCH 08/22] Llama decode: mask softmax to the current context, not a running sum llama_forward_pass_decode wrote a running sum of every step's context length into softmax_vector_size, so only the first decode step was masked correctly. From the second step on the attention softmax also covered unwritten cache slots, and once the sum passed max_seq_len (about the 7th token for a 293-token prompt) it covered the whole 2048-wide row. Those slots score ~0 and soak up attention weight, which is why generated text degenerated after a few tokens. Write context_len, as the ELF-patching code did before #131. Teacher-forced greedy against an fp32 CPU reference (prompt 1024, 24 tokens): top-1 agreement goes from 5/24 to 24/24, decode KL from 1.0-5.5 to <= 0.013, and the greedy text matches the reference exactly. Co-Authored-By: Claude --- iron/applications/llama_3.2_1b/llama_npu.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index 99963a1c63..62d57ef32c 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -1158,13 +1158,12 @@ def llama_forward_pass_decode(config, state): context_len = state.num_preceding_tokens + 1 cache_offset = state.num_preceding_tokens * config.head_dim - state.softmax_vector_size_cum = ( - getattr(state, "softmax_vector_size_cum", 0) + context_len - ) params = aie_ops.decode.fused.params params.write("cache_offset", np.int32(cache_offset)) - params.write("softmax_vector_size", np.int32(state.softmax_vector_size_cum)) + # Softmax masks every score past the first context_len to -inf; the rest of + # the max_seq_len row is unwritten cache. + params.write("softmax_vector_size", np.int32(context_len)) params.sync() # Prefill RoPE angle look-up tables From 480fa95f03413ca10953359d0ac00d7b31d05834 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 06:28:14 -0600 Subject: [PATCH 09/22] Llama test: check NPU logits against an fp32 CPU reference test.py only checked the exit code and scraped TTFT/TPS, so decode emitted garbage after a few tokens for months without failing (the softmax running-sum mask bug fixed in the previous commit). llama_npu.py --check-accuracy feeds the NPU and llama_cpu (fp32 weights) the reference's greedy token each step and reports KL(fp32 || NPU) of the next-token distribution. test_llama_3_2_1b_accuracy gates on it over a 1024-char prompt and 40 steps: prefill KL 0.074 (limit 0.1) decode KL 0.013 (limit 0.05); 9.2 with the mask bug reintroduced Top-1 agreement is reported but not gated: near-ties flip it (1/40 steps at KL 0.004). Co-Authored-By: Claude --- .../llama_3.2_1b/llama_inference_harness.py | 35 ++++++++++ iron/applications/llama_3.2_1b/llama_npu.py | 21 ++++++ iron/applications/llama_3.2_1b/test.py | 64 ++++++++++++++----- 3 files changed, 105 insertions(+), 15 deletions(-) diff --git a/iron/applications/llama_3.2_1b/llama_inference_harness.py b/iron/applications/llama_3.2_1b/llama_inference_harness.py index 232bdce75e..e27513df9d 100644 --- a/iron/applications/llama_3.2_1b/llama_inference_harness.py +++ b/iron/applications/llama_3.2_1b/llama_inference_harness.py @@ -169,6 +169,35 @@ def generate_token(config, forward_pass, state): return next_token.item(), state +def check_accuracy( + config, state, forward_pass, ref_config, ref_state, ref_forward_pass, num_tokens +): + """Teacher-forced comparison of forward_pass's logits against a reference. + + Both models are fed the reference's greedy token at every step, so a + divergence at step N is the candidate's own error at step N rather than the + consequence of an earlier different choice. Step 0 is prefill. + + Returns one (kl, top1) pair per step: KL(reference || candidate) of the + next-token distributions, and whether both rank the same token first. + """ + ref_state.token_ids = state.token_ids + results = [] + for step in range(num_tokens): + logits, state = forward_pass(config, state) + ref_logits, ref_state = ref_forward_pass(ref_config, ref_state) + cand = torch.log_softmax(logits[0, -1].float(), dim=0) + ref = torch.log_softmax(ref_logits[0, -1].float(), dim=0) + kl = torch.sum(ref.exp() * (ref - cand)).item() + next_token = int(ref.argmax()) + top1 = int(cand.argmax()) == next_token + results.append((kl, top1)) + print(f"step {step:3d} KL {kl:.5f} top-1 {'match' if top1 else 'MISMATCH'}") + state.token_ids = torch.tensor([[next_token]], dtype=torch.long) + ref_state.token_ids = state.token_ids + return results + + def parse_args(): parser = argparse.ArgumentParser(description="LLaMA 3.2 1B Inference Harness") parser.add_argument( @@ -189,6 +218,12 @@ def parse_args(): default=40, help="Number of tokens to generate (default: 40)", ) + parser.add_argument( + "--check-accuracy", + action="store_true", + help="Instead of sampling, compare each step's logits against an fp32 CPU " + "reference, feeding both the reference's greedy token", + ) return parser.parse_args() diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index 62d57ef32c..d171a22ed0 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -11,12 +11,14 @@ # [ ] Patching of operators (instantiating new xrt::elf for each token) is slow; find quicker way of patching instruction sequence in-memory # [ ] Spatial fusion of operators +import copy import torch import math from pathlib import Path import sys import numpy as np import ml_dtypes +import llama_cpu import llama_inference_harness as harness import logging @@ -1235,6 +1237,25 @@ def main(): aie_ops = AIELlamaOperators(config, max_seq_len) aie_buffers = AIELlamaBuffers(config, max_seq_len, aie_ops) + if args.check_accuracy: + ref_config = copy.copy(config) + ref_config.weights = {k: v.float() for k, v in config.weights.items()} + results = harness.check_accuracy( + config, + state, + llama_forward_pass, + ref_config, + harness.LlamaModelState(ref_config), + llama_cpu.llama_forward_pass, + args.num_tokens, + ) + kl = [k for k, _ in results] + print(f"[Accuracy] Prefill KL: {kl[0]:.6f}") + if len(kl) > 1: + print(f"[Accuracy] Decode max KL: {max(kl[1:]):.6f}") + print(f"[Accuracy] Top-1 mismatches: {sum(not t for _, t in results)}") + return + print(prompt, end="", flush=True) harness.generate( config, state, llama_forward_pass, use_kv_cache=True, num_tokens=args.num_tokens diff --git a/iron/applications/llama_3.2_1b/test.py b/iron/applications/llama_3.2_1b/test.py index add64d399c..b9e22f0343 100644 --- a/iron/applications/llama_3.2_1b/test.py +++ b/iron/applications/llama_3.2_1b/test.py @@ -5,6 +5,7 @@ import subprocess import pytest import os +import re import sys from pathlib import Path @@ -27,14 +28,39 @@ def generate_test_params(): params, names = generate_test_params() - -@pytest.mark.skipif( +requires_weights = pytest.mark.skipif( not ( (weights_dir / "llama3.2-1b" / "model.safetensors").exists() and (weights_dir / "llama3.2-1b" / "tokenizer.model").exists() ), reason="llama3.2-1b weights not found", ) + + +def run_llama_npu(prompt_len, num_tokens, *extra_args): + command = [ + sys.executable, + str(test_dir / "llama_npu.py"), + str(weights_dir / "llama3.2-1b" / "model.safetensors"), + str(weights_dir / "llama3.2-1b" / "tokenizer.model"), + "--num-tokens", + str(num_tokens), + "--prompt-len", + str(prompt_len), + *extra_args, + ] + result = subprocess.run(command, cwd=test_dir, capture_output=True, text=True) + + print(result.stdout) + print(result.stderr) + + assert ( + result.returncode == 0 + ), f"Command failed with return code {result.returncode}\nStderr: {result.stderr}" + return result + + +@requires_weights @pytest.mark.supported_devices("npu2") @pytest.mark.metrics( TTFT=r"\[Prefill\]\s*Time to first token:\s*(?P[\d\.e\+-]+) s", @@ -42,19 +68,27 @@ def generate_test_params(): ) @pytest.mark.parametrize("prompt_len,num_tokens", params, ids=names) def test_llama_3_2_1b(prompt_len, num_tokens): - command = f"{sys.executable} {test_dir}/llama_npu.py {weights_dir}/llama3.2-1b/model.safetensors {weights_dir}/llama3.2-1b/tokenizer.model --num-tokens {num_tokens} --prompt-len {prompt_len}" + run_llama_npu(prompt_len, num_tokens) - result = subprocess.run( - command, - cwd=test_dir, - shell=True, - capture_output=True, - text=True, - ) - print(result.stdout) - print(result.stderr) +# KL(fp32 CPU || NPU) of the next-token distribution, teacher-forced over 40 +# steps. The NPU measures 0.074 on prefill and at most 0.013 on decode. Decode +# attention over unmasked KV-cache slots measured 9.2. +MAX_PREFILL_KL = 0.1 +MAX_DECODE_KL = 0.05 - assert ( - result.returncode == 0 - ), f"Command failed with return code {result.returncode}\nStderr: {result.stderr}" + +@requires_weights +@pytest.mark.supported_devices("npu2") +@pytest.mark.metrics( + PrefillKL=r"\[Accuracy\] Prefill KL:\s*(?P[\d\.e\+-]+)", + DecodeMaxKL=r"\[Accuracy\] Decode max KL:\s*(?P[\d\.e\+-]+)", + Top1Mismatches=r"\[Accuracy\] Top-1 mismatches:\s*(?P\d+)", +) +def test_llama_3_2_1b_accuracy(): + result = run_llama_npu(1024, 40, "--check-accuracy") + + prefill_kl = float(re.search(r"Prefill KL:\s*(\S+)", result.stdout).group(1)) + decode_kl = float(re.search(r"Decode max KL:\s*(\S+)", result.stdout).group(1)) + assert prefill_kl <= MAX_PREFILL_KL, f"prefill KL {prefill_kl} > {MAX_PREFILL_KL}" + assert decode_kl <= MAX_DECODE_KL, f"decode KL {decode_kl} > {MAX_DECODE_KL}" From e4fed44646f33295a7569cf9fe887a78cad4289a Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 07:14:26 -0600 Subject: [PATCH 10/22] Llama: flush the prefill KV hand-off to the device, not from it After prefill, llama_forward_pass writes every layer's K/V cache into the fused decode operator's scratch buffer through torch_view(), then called scratch_buffer.to("cpu"). That syncs the other direction, and nothing else flushes scratch (the callable only syncs its input buffer), so the host's dirty cache lines reached DRAM whenever the CPU happened to evict them. Decode then either read stale KV rows, or had rows it had written overwritten one 64-byte line at a time by a late eviction. This is the run-to-run nondeterminism in Llama's output: teacher-forced over 40 steps, 12 of 103 runs had different logits (always starting at decode step 1 or 2; prefill was always identical). With the flush, 0 of 103 differ, and all match the previous majority result. Latency is unchanged (paired TTFT 0.993, decode 1.003 over 8 interleaved rounds). Co-Authored-By: Claude --- iron/applications/llama_3.2_1b/llama_npu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index d171a22ed0..58da94f914 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -1213,7 +1213,7 @@ def llama_forward_pass(config, state): aie_ops.decode.fused.get_buffer(f"values_cache_{layer_idx}").torch_view()[ : ] = (aie_buffers.values_cache[layer_idx].to_torch().flatten()) - aie_ops.decode.fused.scratch_buffer.to("cpu") + aie_ops.decode.fused.scratch_buffer.to("npu") return ret else: ret = llama_forward_pass_decode(config, state) From 263b83b3959b3a67b06fbedfbb9901c6193b917f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 08:15:03 -0600 Subject: [PATCH 11/22] Llama test: guard run-to-run determinism of NPU logits Add --check-determinism ROUNDS to llama_npu.py. It prefills two prompts with different text, alternating, and greedily decodes a few steps from a fresh state each round. It then counts the runs whose logits differ bitwise from the first run of the same prompt. Alternating the prompts matters: when a host write never reaches the device, the NPU reads the other prompt's data instead of a leftover identical copy. That turns the 12% race fixed in the previous commit into a failure on every run: 38/38 in each of three trials with the flush reverted. test_llama_3_2_1b_determinism runs 5 rounds of 4 tokens (about 30 s) and requires zero differing runs. Co-Authored-By: Claude --- .../llama_3.2_1b/llama_inference_harness.py | 37 +++++++++++++++++++ iron/applications/llama_3.2_1b/llama_npu.py | 13 +++++++ iron/applications/llama_3.2_1b/test.py | 16 ++++++++ 3 files changed, 66 insertions(+) diff --git a/iron/applications/llama_3.2_1b/llama_inference_harness.py b/iron/applications/llama_3.2_1b/llama_inference_harness.py index e27513df9d..5e87b1f41d 100644 --- a/iron/applications/llama_3.2_1b/llama_inference_harness.py +++ b/iron/applications/llama_3.2_1b/llama_inference_harness.py @@ -198,6 +198,36 @@ def check_accuracy( return results +def check_determinism(config, prompts, forward_pass, num_tokens, rounds): + """Run each prompt `rounds` times, alternating, and compare logits bitwise. + + Each round prefills from a fresh state and decodes greedily. Alternating + prompts with different text matters: a host write that never reaches the + device then reads the other prompt's data, not a leftover copy of its own. + Returns how many rounds differ from the first round of the same prompt. + """ + first = [None] * len(prompts) + n_differ = 0 + for r in range(rounds * len(prompts)): + p = r % len(prompts) + state = LlamaModelState(config) + state.token_ids = prompts[p] + logits = [] + for _ in range(num_tokens): + out, state = forward_pass(config, state) + logits.append(out[0, -1].clone()) + state.token_ids = out[:, -1:].argmax(dim=-1) + logits = torch.stack(logits).view(torch.int16) + if first[p] is None: + first[p] = logits + continue + steps = (logits != first[p]).any(dim=1).nonzero().flatten().tolist() + if steps: + n_differ += 1 + print(f"round {r} (prompt {p}): logits differ at steps {steps}") + return n_differ + + def parse_args(): parser = argparse.ArgumentParser(description="LLaMA 3.2 1B Inference Harness") parser.add_argument( @@ -224,6 +254,13 @@ def parse_args(): help="Instead of sampling, compare each step's logits against an fp32 CPU " "reference, feeding both the reference's greedy token", ) + parser.add_argument( + "--check-determinism", + type=int, + metavar="ROUNDS", + help="Instead of sampling, run two prompts ROUNDS times each, alternating, " + "and count the runs whose logits differ bitwise from the first run", + ) return parser.parse_args() diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index 58da94f914..184167571a 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -1256,6 +1256,19 @@ def main(): print(f"[Accuracy] Top-1 mismatches: {sum(not t for _, t in results)}") return + if args.check_determinism: + # The second prompt is the same amount of the text that follows. + other = harness.get_prompt(2 * args.prompt_len)[args.prompt_len :] + other_ids = [config.special_tokens["<|begin_of_text|>"]] + other_ids += config.tokenizer.encode(other) + prompts = [state.token_ids, torch.tensor([other_ids], dtype=torch.long)] + n_differ = harness.check_determinism( + config, prompts, llama_forward_pass, args.num_tokens, args.check_determinism + ) + n_compared = len(prompts) * (args.check_determinism - 1) + print(f"[Determinism] Differing runs: {n_differ}/{n_compared}") + return + print(prompt, end="", flush=True) harness.generate( config, state, llama_forward_pass, use_kv_cache=True, num_tokens=args.num_tokens diff --git a/iron/applications/llama_3.2_1b/test.py b/iron/applications/llama_3.2_1b/test.py index b9e22f0343..9545019420 100644 --- a/iron/applications/llama_3.2_1b/test.py +++ b/iron/applications/llama_3.2_1b/test.py @@ -92,3 +92,19 @@ def test_llama_3_2_1b_accuracy(): decode_kl = float(re.search(r"Decode max KL:\s*(\S+)", result.stdout).group(1)) assert prefill_kl <= MAX_PREFILL_KL, f"prefill KL {prefill_kl} > {MAX_PREFILL_KL}" assert decode_kl <= MAX_DECODE_KL, f"decode KL {decode_kl} > {MAX_DECODE_KL}" + + +# Repeated runs must produce bit-identical logits. A prefill KV hand-off that +# was never flushed to the device made 12% of runs diverge. Alternating +# two prompts makes such a missing flush fail every run: 38/38 in each of three +# trials. +@requires_weights +@pytest.mark.supported_devices("npu2") +@pytest.mark.metrics( + DifferingRuns=r"\[Determinism\] Differing runs:\s*(?P\d+)/", +) +def test_llama_3_2_1b_determinism(): + result = run_llama_npu(1024, 4, "--check-determinism", "5") + + differing = re.search(r"Differing runs:\s*(\d+)/(\d+)", result.stdout) + assert int(differing.group(1)) == 0, f"{differing.group(0)} (bitwise logits)" From 3731ae09e8b48186d7c3ec7e9c1cdd798e10b43f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 09:17:44 -0600 Subject: [PATCH 12/22] Flush scratch in the full-ELF sequence callable before each dispatch get_buffer() hands out writable views into the fused ELF's scratch buffer (weights, KV caches), but _sync_inputs only flushed the input buffer. The host runtime's own dispatch flushes every argument, and the separate-xclbin callable goes through it; the full-ELF callable calls run_handle.start() directly and so skipped it for scratch. NPU access to these buffers is not cache-coherent, so an unflushed scratch write was a silent race. It was the Llama run-to-run nondeterminism fixed at the call site in e4fed44. _sync_inputs now flushes scratch too. With nothing dirty that transfers nothing, and it leaves scratch marked device-resident, so reading a scratch view after a run pulls the NPU's writes. Llama's explicit hand-off flush is removed; the determinism test now guards the callable. test_non_input_buffers_sync_without_explicit_flush writes a non-input buffer, with new data each dispatch, in separate and fused modes. It also reads a non-output buffer, with no to() from the caller. Without the flush the fused case failed 5/5 runs. Llama A/B, 8 interleaved rounds, prompt 1024 / 40 tokens: paired decode 1.003, TTFT 1.003, identical text. Co-Authored-By: Claude --- iron/applications/llama_3.2_1b/llama_npu.py | 1 - iron/common/sequence.py | 6 +++ iron/tests/infrastructure/sequence.py | 53 ++++++++++++++++++--- 3 files changed, 53 insertions(+), 7 deletions(-) diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index 184167571a..27fcf06a2c 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -1213,7 +1213,6 @@ def llama_forward_pass(config, state): aie_ops.decode.fused.get_buffer(f"values_cache_{layer_idx}").torch_view()[ : ] = (aie_buffers.values_cache[layer_idx].to_torch().flatten()) - aie_ops.decode.fused.scratch_buffer.to("npu") return ret else: ret = llama_forward_pass_decode(config, state) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index bc41c33888..7c449e79ad 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -689,7 +689,13 @@ def _sync_inputs(self): # 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. + # Scratch is flushed as well: get_buffer() hands out writable views into it + # (weights, KV caches), and this dispatch bypasses the host runtime's own + # per-argument flush. With nothing dirty, `to("npu")` transfers nothing. It + # also leaves all of scratch marked device-resident, so a read of a scratch + # view after the run pulls what the NPU wrote. self.input_buffer.to("npu") + self.scratch_buffer.to("npu") def _sync_outputs(self): # _run just rewrote the output arena on the device, so the device holds the diff --git a/iron/tests/infrastructure/sequence.py b/iron/tests/infrastructure/sequence.py index a1399e3d99..3bdeac3cb8 100644 --- a/iron/tests/infrastructure/sequence.py +++ b/iron/tests/infrastructure/sequence.py @@ -38,10 +38,9 @@ def _set_input(run, name, data): """Write a host tensor into an input buffer and push it to the device. - Mirrors the caller contract for the fused single-ELF callable: after - writing a get_buffer() sub-view via torch_view(), the caller is responsible - for calling .to("npu") so the write reaches the NPU (a no-op sync for the - separate/reference callables, whose __call__ syncs inputs themselves). + The explicit push is redundant, since every callable flushes host writes + at dispatch (see test_non_input_buffers_sync_without_explicit_flush), and + is a no-op sync for the reference callable. """ buf = run.get_buffer(name) buf.torch_view()[: data.numel()] = data.reshape(-1) @@ -57,7 +56,7 @@ def _set_input(run, name, data): _ADD_RELU_COLS = 4 -def _build_add_relu_sequence(context, dispatch, name): +def _build_add_relu_sequence(context, dispatch, name, input_args=("a", "b")): """out = relu(a + b), as a 2-step OperatorSequence.""" add = ElementwiseAdd( size=_ADD_RELU_SIZE, @@ -78,7 +77,7 @@ def _build_add_relu_sequence(context, dispatch, name): (add, "a", "b", "temp"), (relu, "temp", "out"), ], - input_args=["a", "b"], + input_args=list(input_args), output_args=["out"], dispatch=dispatch, context=context, @@ -322,3 +321,45 @@ def test_compare_mode_detects_wrong_reference(reference_is_correct, aie_context) else: with pytest.raises(RuntimeError): run() # compare mode reports the wrong reference by itself + + +# --------------------------------------------------------------------------- +# 5. Buffers that are neither inputs nor outputs (weights, KV caches, +# intermediates) sync like the rest in every NPU dispatch mode. The full-ELF +# callable places them in its scratch buffer, and NPU access to it is not +# cache-coherent: an unflushed host write is a race, not an error, so each +# dispatch below writes different data than the one before. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dispatch", ["separate", "fused"]) +def test_non_input_buffers_sync_without_explicit_flush(dispatch, aie_context): + """Host writes through get_buffer() to a non-input buffer reach the NPU at + the next dispatch, and reads of a non-output buffer after a dispatch see + what the NPU wrote there, with no explicit ``to()`` from the caller.""" + if dispatch == "fused" and not isinstance(aie_utils.get_current_device(), NPU2): + pytest.skip("fused (single-ELF) dispatch requires NPU2") + + # b is not an input, so it is held like a weight (in scratch, when fused). + seq = _build_add_relu_sequence( + aie_context, dispatch, f"infra_add_weight_relu_{dispatch}", input_args=["a"] + ) + seq.compile() + run = seq.get_callable() + + torch.manual_seed(0) + for rep in range(4): + a = torch.rand(_ADD_RELU_SIZE, dtype=torch.bfloat16) * 4 - 2 + b = torch.rand(_ADD_RELU_SIZE, dtype=torch.bfloat16) * 4 - 2 + run.get_buffer("a").torch_view()[:] = a + run.get_buffer("b").torch_view()[:] = b + run() + + temp = run.get_buffer("temp").to_torch()[:_ADD_RELU_SIZE] + out = run.get_buffer("out").to_torch()[:_ADD_RELU_SIZE] + errors = verify_buffer(temp, "temp", a + b, rel_tol=0.04, abs_tol=1e-6) + assert not errors, f"rep {rep}: temp has {len(errors)} mismatches" + errors = verify_buffer( + out, "out", torch.nn.functional.relu(a + b), rel_tol=0.04, abs_tol=1e-6 + ) + assert not errors, f"rep {rep}: out has {len(errors)} mismatches" From 4d3eeea8559c2b6ad2109e5a5b60d4a252af87c3 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 09:48:24 -0600 Subject: [PATCH 13/22] Tests: lean on mlir-aie's verify and benchmark utilities verify_buffer now judges through aie.utils.verify.compare under a relative Tolerance rather than IRON's own nearly_equal. The signature and returned mismatch indices are unchanged. The old check let a NaN output pass; compare requires NaN to meet NaN and an infinity to meet the same infinity, and that holds regardless of max_error_rate. conftest adds a Provenance column to the metrics CSV (commit, Peano, mlir-aie, kernel sources and digest, plus device and power mode when the run used the NPU). Existing columns are unchanged, so the CI pretty scripts still read them. When device-gated tests are collected and there is no NPU runtime, it stops with mlir-aie's probe reason (exit 4, e.g. "xrt-smi not on PATH ...") instead of a wall of failures. The hand-rolled perf_counter timing in gemm and the swiglu tests now uses aie.utils.benchmark.run_iters. The infrastructure tests that only re-tested mlir-aie utilities (comparison, benchmark, sequence_subviews, sequence_output_sync) are removed; mlir-aie's own test suite covers them. Tests of IRON code stay. Non-extensive sweep of iron/operators + iron/tests on Strix: 283 passed, 4 skipped, 1 failed. The one failure is lazy_imports::test_lazy_catalog_does_not_import_mha, which asserts on sys.modules and so fails whenever mha runs earlier in the same session, also on the parent commit. No operator test flipped under the stricter comparison. Co-Authored-By: Claude --- AGENTS.md | 2 +- conftest.py | 17 ++++ iron/common/test_utils.py | 88 ++++++------------- iron/operators/gemm/test.py | 14 ++- iron/operators/swiglu_decode/test.py | 9 +- iron/operators/swiglu_prefill/test.py | 9 +- iron/operators/swiglu_prefill_stream/test.py | 13 +-- iron/tests/infrastructure/benchmark.py | 69 --------------- iron/tests/infrastructure/comparison.py | 47 ---------- .../infrastructure/sequence_output_sync.py | 83 ----------------- .../tests/infrastructure/sequence_subviews.py | 73 --------------- 11 files changed, 61 insertions(+), 363 deletions(-) delete mode 100644 iron/tests/infrastructure/benchmark.py delete mode 100644 iron/tests/infrastructure/comparison.py delete mode 100644 iron/tests/infrastructure/sequence_output_sync.py delete mode 100644 iron/tests/infrastructure/sequence_subviews.py diff --git a/AGENTS.md b/AGENTS.md index b74b256949..6794290f29 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -147,7 +147,7 @@ reuse lint - `device_manager.py`: XRT device initialization and management (singleton pattern) - `context.py`: `AIEContext` for operator compilation/execution - `utils.py`: Helper functions (`torch_to_numpy`, `numpy_to_torch`) - - `test_utils.py`: Test utilities (`verify_buffer`, `nearly_equal`) + - `test_utils.py`: Test utilities (`verify_buffer`, a wrapper over mlir-aie's `aie.utils.verify.compare`; `run_test`, timed with `aie.utils.benchmark.run_iters`) ### Key Concepts diff --git a/conftest.py b/conftest.py index 3cd72a36e2..e4b58bbff3 100644 --- a/conftest.py +++ b/conftest.py @@ -12,6 +12,8 @@ from iron.common import AIEContext import aie.utils as aie_utils +from aie.utils.benchmark import preflight, provenance +from aie.utils.probe import npu_unavailable_reason @pytest.fixture @@ -82,10 +84,21 @@ def add_result( def finalize_results(self): """Compute statistics for all collected metrics""" + # The commit alone does not say which toolchain and kernel sources + # produced a number; mlir-aie's provenance line does. Only a run that + # measured something has used the NPU, so only then is it described: + # opening it otherwise would contend for the single-tenant device. + measured = any(len(data) > 1 for data in self.test_metrics.values()) + if measured and aie_utils.DefaultNPURuntime is not None: + npu = preflight() + source = provenance(device=npu.device, pmode=npu.pmode) + else: + source = provenance() for (test_path, test_name), data in self.test_metrics.items(): row = { "Commit": self.commit, "Date": self.date, + "Provenance": source, "Test Path": test_path, "Test": test_name, "Checks": f"{sum(data['passed'])}/{len(data['passed'])}", @@ -185,6 +198,10 @@ def pytest_collection_modifyitems(config, items): # else holds it and erroring out when none is attached. return + if aie_utils.DefaultNPURuntime is None: + # Most often an unsourced XRT, which otherwise surfaces as a pile of + # failures that look like a toolchain regression. + raise pytest.UsageError(f"No NPU runtime: {npu_unavailable_reason()}") device = aie_utils.DefaultNPURuntime.device().resolve().name for item, marker in marked_items: if device not in marker.args: diff --git a/iron/common/test_utils.py b/iron/common/test_utils.py index afda7607f2..66db3a168e 100644 --- a/iron/common/test_utils.py +++ b/iron/common/test_utils.py @@ -7,6 +7,7 @@ import torch import aie.utils as aie_utils from aie.utils.benchmark import run_iters +from aie.utils.verify import Tolerance, compare, nearly_equal from ml_dtypes import bfloat16 from .base import AIEOperatorBase @@ -19,34 +20,6 @@ "i32": torch.int32, } -# TODO: Consider upstreaming generic buffer utilities to mlir-aie once operator abstractions stabilize. - - -def nearly_equal( - a: float, - b: float, - rel_tol: float = 128 * np.finfo(np.float32).eps, - abs_tol: float = np.finfo(np.float32).tiny, -) -> bool: - """ - Compare two floating point numbers for approximate equality. - - Adapted from Stack Overflow, License CC BY-SA 4.0 - Original author: P-Gn - Source: https://stackoverflow.com/a/32334103 - """ - if np.finfo(np.float32).eps > rel_tol: - raise ValueError(f"rel_tol {rel_tol!r} must be >= machine epsilon") - if rel_tol >= 1.0: - raise ValueError(f"rel_tol {rel_tol!r} must be < 1.0") - - if a == b: - return True - - diff = abs(float(a) - float(b)) - norm = min(abs(float(a)) + abs(float(b)), np.finfo(np.float32).max) - return diff < max(abs_tol, rel_tol * norm) - def verify_buffer( output: np.ndarray | torch.Tensor, @@ -59,6 +32,11 @@ def verify_buffer( """ Verify buffer contents match reference within tolerances. + The comparison is mlir-aie's ``aie.utils.verify.compare`` under a relative + ``Tolerance``: an element passes at ``|a - b| < max(abs_tol, rel_tol * (|a| + |b|))``, + so ``rel_tol=abs_tol=0`` demands exact equality, and a NaN or infinity must + meet the same value in the reference whatever ``max_error_rate`` allows. + Args: output: Output buffer to verify buf_name: Name of buffer for error messages @@ -71,7 +49,6 @@ def verify_buffer( Returns: List of error indices. Empty if verification passes. """ - errors = [] def _to_numpy(x): if isinstance(x, torch.Tensor): @@ -89,41 +66,34 @@ def _to_numpy(x): print( f"Buffer size mismatch for {buf_name}: expected {len(expected_np)}, got {len(output)}" ) - errors.extend(i for i in range(abs(len(output) - len(expected_np)))) - compare_len = min(len(output), len(expected_np)) - diff = np.abs( - output[:compare_len].astype(float) - expected_np[:compare_len].astype(float) - ) - norm = np.minimum( - np.abs(output[:compare_len].astype(float)) - + np.abs(expected_np[:compare_len].astype(float)), - np.finfo(np.float32).max, + return list(range(len(output), len(expected_np))) + output = output[: len(expected_np)] + + tolerance = Tolerance.relative(rel_tol, abs_tol, max_mismatch_frac=max_error_rate) + verdict = compare(output, expected_np, tolerance) + if verdict.n_mismatch and max_error_rate > 0.0: + within = "within" if verdict else "exceeds" + print( + f"{buf_name}: {verdict.n_mismatch} errors " + f"({verdict.n_mismatch / verdict.n_checked * 100:.2f}%) {within} allowed " + f"rate of {max_error_rate * 100:.2f}%" + ) + if verdict: + return [] + + print(f"{buf_name}: {verdict.detail}") + # compare() judges; it does not list the elements. nearly_equal is the same + # per-element test, except that it also rejects a NaN that meets a NaN. + bad = ~nearly_equal(output, expected_np, rtol=rel_tol, atol=abs_tol) + bad &= ~( + np.isnan(output.astype(np.float32)) & np.isnan(expected_np.astype(np.float32)) ) - # Use `>`, not `>=`, here, so that a user can pass rel_tol=abs_tol=0 - # check exact equality. - mask = diff > np.maximum(abs_tol, rel_tol * norm) - error_indices = np.where(mask)[0].tolist() + error_indices = np.flatnonzero(bad).tolist() for i in error_indices[:10]: print( f"Mismatch in {buf_name}[{i}]: expected {float(expected_np[i]):.6f}, got {float(output[i]):.6f}" ) - errors.extend(error_indices) - - # Check if error rate is acceptable - if max_error_rate > 0.0 and len(errors) > 0: - error_rate = len(errors) / compare_len - max_allowed_errors = int(compare_len * max_error_rate) - if len(errors) <= max_allowed_errors: - print( - f"{buf_name}: {len(errors)} errors ({error_rate*100:.2f}%) within allowed rate of {max_error_rate*100:.2f}% ({max_allowed_errors} errors)" - ) - return [] # Pass - within allowed error rate - else: - print( - f"{buf_name}: {len(errors)} errors ({error_rate*100:.2f}%) exceeds allowed rate of {max_error_rate*100:.2f}% ({max_allowed_errors} errors)" - ) - - return errors + return error_indices def _nbytes(buf) -> int: diff --git a/iron/operators/gemm/test.py b/iron/operators/gemm/test.py index 100b9c2ca9..7bc92691be 100755 --- a/iron/operators/gemm/test.py +++ b/iron/operators/gemm/test.py @@ -2,13 +2,12 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import time - import numpy as np import pytest import aie.utils as aie_utils import torch import ml_dtypes +from aie.utils.benchmark import run_iters from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor from iron.operators.gemm.op import GEMM @@ -188,12 +187,11 @@ def test_gemm( B_bufs.append(XRTTensor.from_torch(b_torch)) C_bufs.append(XRTTensor(c_shape, dtype=c_dtype)) - # Run each partition - start_time = time.perf_counter() - for i in range(partition_N): - op_func(A_buf, B_bufs[i], C_bufs[i]) - end_time = time.perf_counter() - latency_us = (end_time - start_time) * 1e6 + def run_partitions(): + for i in range(partition_N): + op_func(A_buf, B_bufs[i], C_bufs[i]) + + latency_us = run_iters(run_partitions).e2e.avg_us # Read back and concatenate C partitions along the column dimension C_parts_torch = [buf.to_torch().reshape(c_shape) for buf in C_bufs] diff --git a/iron/operators/swiglu_decode/test.py b/iron/operators/swiglu_decode/test.py index 45eb541140..6956749511 100755 --- a/iron/operators/swiglu_decode/test.py +++ b/iron/operators/swiglu_decode/test.py @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import time import pytest +from aie.utils.benchmark import run_iters from iron.operators.swiglu_decode.op import SwiGLUDecode from iron.operators.swiglu_decode.reference import generate_golden_reference @@ -52,12 +52,7 @@ def test_swiglu_decode(embedding_dim, hidden_dim, aie_context): # Set the per-invocation input. fc.get_buffer("in").torch_view()[:] = golden_ref["input"].reshape(-1) - # Warmup - fc() - - start = time.perf_counter() - fc() - elapsed_us = (time.perf_counter() - start) * 1e6 + elapsed_us = run_iters(fc, warmup=1, iters=1).e2e.avg_us total_bytes = (golden_ref["input"].numel() + embedding_dim) * 2 # bf16 bandwidth_gbps = total_bytes / (elapsed_us * 1e-6) / 1e9 diff --git a/iron/operators/swiglu_prefill/test.py b/iron/operators/swiglu_prefill/test.py index 538650010e..a89a2c221c 100755 --- a/iron/operators/swiglu_prefill/test.py +++ b/iron/operators/swiglu_prefill/test.py @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import time import pytest +from aie.utils.benchmark import run_iters from iron.operators.gemm.op import GEMM from iron.operators.swiglu_prefill.op import SwiGLUPrefill @@ -64,12 +64,7 @@ def _as_stored(w): # Set the per-invocation input. fc.get_buffer("in").torch_view()[:] = golden_ref["input"].reshape(-1) - # Warmup - fc() - - start = time.perf_counter() - fc() - elapsed_us = (time.perf_counter() - start) * 1e6 + elapsed_us = run_iters(fc, warmup=1, iters=1).e2e.avg_us total_bytes = (golden_ref["input"].numel() + seq_len * embedding_dim) * 2 # bf16 bandwidth_gbps = total_bytes / (elapsed_us * 1e-6) / 1e9 diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index 89ea3c2397..e0152a69a8 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -2,10 +2,9 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import time - import pytest import torch +from aie.utils.benchmark import run_iters from iron.common.tracing_utils import dump_traces @@ -95,16 +94,12 @@ def test_swiglu_prefill_stream(k, aie_context): # The first dispatch on a callable pays for its hardware context, so time the # ones after it. - latencies = [] - for _ in range(TIMED_RUNS): - start = time.perf_counter() - run() - latencies.append((time.perf_counter() - start) * 1e6) - elapsed_us = min(latencies) + latency = run_iters(run, iters=TIMED_RUNS).e2e + elapsed_us = latency.min_us total_bytes = 4 * SEQ_LEN * EMBEDDING_DIM # bf16 in + out print(f"Latency (us): {elapsed_us:.2f}") print( f"Latency min/mean/max (us): {elapsed_us:.2f} / " - f"{sum(latencies) / len(latencies):.2f} / {max(latencies):.2f}" + f"{latency.avg_us:.2f} / {latency.max_us:.2f}" ) print(f"Effective Bandwidth: {total_bytes / (elapsed_us * 1e-6) / 1e9:.4f} GB/s") diff --git a/iron/tests/infrastructure/benchmark.py b/iron/tests/infrastructure/benchmark.py deleted file mode 100644 index e1fdb48979..0000000000 --- a/iron/tests/infrastructure/benchmark.py +++ /dev/null @@ -1,69 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""The operator test adapter keeps reporting NPU, not host, latency.""" - -from types import SimpleNamespace - -import pytest - -torch = pytest.importorskip("torch") - -from aie.utils.hostruntime.tensor_class import CPUOnlyTensor -from iron.common.base import AIEOperatorBase, AIERuntimeArgSpec -from iron.common import test_utils - - -class _Operator(AIEOperatorBase): - def __init__(self, results): - self.results = iter(results) - self.calls = 0 - - def set_up_artifacts(self): - pass - - def compile(self): - return self - - def get_arg_spec(self): - return [ - AIERuntimeArgSpec("in", (32,)), - AIERuntimeArgSpec("out", (32,)), - ] - - def get_callable(self): - def run(source, target): - self.calls += 1 - target[:] = source.numpy() - return next(self.results) - - return run - - -@pytest.mark.parametrize("tuple_result", [False, True]) -def test_run_test_uses_upstream_npu_timing(monkeypatch, tuple_result): - monkeypatch.setattr(test_utils.aie_utils, "DEFAULT_TENSOR_CLASS", CPUOnlyTensor) - results = [SimpleNamespace(npu_time=ns) for ns in (1000000, 2000, 4000)] - if tuple_result: - results = [(None, result) for result in results] - op = _Operator(results) - data = torch.ones(32, dtype=torch.bfloat16) - - errors, latency_us, bandwidth = test_utils.run_test( - op, {"in": data}, {"out": data}, warmup_iters=1, timed_iters=2 - ) - - assert op.calls == 3 - assert errors == {} - assert latency_us == 3.0 - assert bandwidth == pytest.approx(128 / (3e-6) / 1e9) - - -def test_missing_npu_timing_is_rejected(monkeypatch): - monkeypatch.setattr(test_utils.aie_utils, "DEFAULT_TENSOR_CLASS", CPUOnlyTensor) - op = _Operator([None]) - data = torch.ones(32, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="NPU execution time"): - test_utils.run_test( - op, {"in": data}, {"out": data}, warmup_iters=0, timed_iters=1 - ) diff --git a/iron/tests/infrastructure/comparison.py b/iron/tests/infrastructure/comparison.py deleted file mode 100644 index c5a8f550d5..0000000000 --- a/iron/tests/infrastructure/comparison.py +++ /dev/null @@ -1,47 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""What the tolerances in verify_buffer are required to mean. - -An operator that does no arithmetic (transpose, mem_copy) should be gated on exact -equality, not on a tolerance that would also accept a wrong answer. That is -rel_tol=abs_tol=0, so the zero case has to behave -- and it is the case a -threshold comparison is easiest to get backwards, since the threshold is then the -same value as the difference between two identical buffers. -""" - -import numpy as np -import pytest -import torch - -from iron.common.test_utils import verify_buffer - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_zero_tolerance_accepts_an_identical_buffer(dtype): - buf = (torch.arange(64, dtype=torch.float32) / 8).to(dtype) - - assert verify_buffer(buf, "out", buf.clone(), rel_tol=0.0, abs_tol=0.0) == [] - - -@pytest.mark.parametrize("rel_tol,abs_tol", [(0.0, 0.0), (0.04, 1e-6)]) -def test_a_single_wrong_element_is_reported_alone(rel_tol, abs_tol): - reference = torch.arange(64, dtype=torch.float32) - output = reference.clone() - output[ - 17 - ] += 10.0 # past the 4% relative tolerance at this magnitude, not just past 0 - - assert verify_buffer(output, "out", reference, rel_tol, abs_tol) == [17] - - -def test_zero_tolerance_still_rejects_a_one_ulp_error(): - """The point of the zero case is that it is exact, not that it is lenient.""" - reference = torch.full((32,), 1.0, dtype=torch.float32) - output = reference.clone() - output[5] = float(np.nextafter(np.float32(1.0), np.float32(2.0))) - - assert verify_buffer(output, "out", reference, rel_tol=0.0, abs_tol=0.0) == [5] - # The default tolerance is meant to absorb exactly this. - assert verify_buffer(output, "out", reference) == [] diff --git a/iron/tests/infrastructure/sequence_output_sync.py b/iron/tests/infrastructure/sequence_output_sync.py deleted file mode 100644 index 2b2e58344d..0000000000 --- a/iron/tests/infrastructure/sequence_output_sync.py +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Device-free tests for the per-buffer output sync. - -A dispatch writes its output buffers on the device, which the host-side coherence -map does not observe. ``to("cpu")`` transfers only the ranges the map holds as -device-resident, so a range left marked ``cpu`` by an earlier read is skipped and -the next dispatch hands back the previous one's output. -""" - -import pytest - -from aie.utils.hostruntime.coherence import _CoherenceMap - -from iron.common.sequence import SequenceXclbinCallable - - -def test_a_pull_is_skipped_while_the_range_reads_as_host_resident(): - """The hazard the output sync has to defeat, at the layer that decides it.""" - coherence = _CoherenceMap(64, _CoherenceMap.DEVICE) - assert coherence.ranges(0, 64, _CoherenceMap.DEVICE) == [(0, 64)] - - coherence.set(0, 64, _CoherenceMap.HOST) - assert coherence.ranges(0, 64, _CoherenceMap.DEVICE) == [] - - coherence.set(0, 64, _CoherenceMap.DEVICE) - assert coherence.ranges(0, 64, _CoherenceMap.DEVICE) == [(0, 64)] - - -class _RecordingBuffer: - def __init__(self): - self.calls = [] - self._device = "cpu" - - @property - def device(self): - return self._device - - @device.setter - def device(self, value): - self._device = value - self.calls.append(("device", value)) - - def to(self, target): - self._device = target - self.calls.append(("to", target)) - - -class _Op: - def __init__(self, names, inputs): - self.subbuffer_layout = {n: (None, None, 8) for n in names} - self.input_args = set(inputs) - - -def _callable(names, inputs): - """A SequenceXclbinCallable with recording buffers and no XRT behind it.""" - call = object.__new__(SequenceXclbinCallable) - call.op = _Op(names, inputs) - call._buffers = {n: _RecordingBuffer() for n in names} - return call - - -def test_output_sync_claims_the_device_before_pulling(): - call = _callable(["a", "out"], inputs=["a"]) - call._sync_outputs() - assert call._buffers["out"].calls == [("device", "npu"), ("to", "cpu")] - - -@pytest.mark.parametrize("reps", [2, 3]) -def test_every_dispatch_pulls_again(reps): - call = _callable(["out"], inputs=[]) - for _ in range(reps): - call._sync_outputs() - assert call._buffers["out"].calls.count(("to", "cpu")) == reps - assert call._buffers["out"].calls.count(("device", "npu")) == reps - - -def test_inputs_are_left_alone(): - call = _callable(["a", "out"], inputs=["a"]) - call._sync_outputs() - assert call._buffers["a"].calls == [] diff --git a/iron/tests/infrastructure/sequence_subviews.py b/iron/tests/infrastructure/sequence_subviews.py deleted file mode 100644 index 64ef417b62..0000000000 --- a/iron/tests/infrastructure/sequence_subviews.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Host-only coverage of the shared upstream tensor subview path.""" - -from types import SimpleNamespace - -import numpy as np -import pytest -from ml_dtypes import bfloat16 - -from iron.common.sequence import SequenceReferenceCallable - - -@pytest.fixture -def run(): - op = SimpleNamespace( - subbuffer_layout={"packed": ("output", 0, 1024)}, - slice_info={ - "first": ("packed", 0, 512), - "second": ("packed", 512, 1024), - }, - ) - return SequenceReferenceCallable(op) - - -@pytest.mark.parametrize("name, start", [("first", 0), ("second", 256)]) -def test_slices_alias_the_parent_and_are_cached(run, name, start): - parent = run.get_buffer("packed") - parent.fill_(0) - view = run.get_buffer(name) - - assert view is run.get_buffer(name) - assert view is run._resolve_buffer(name) - assert view.dtype == np.dtype(bfloat16) - assert view.shape == (256,) - assert np.shares_memory(view.data, parent.data) - - view.fill_(3) - expected = np.zeros(512, dtype=bfloat16) - expected[start : start + 256] = 3 - np.testing.assert_array_equal(parent.numpy(), expected) - - parent.fill_(7) - np.testing.assert_array_equal(view.numpy(), np.full(256, 7, dtype=bfloat16)) - - -def test_unknown_buffer_is_rejected(run): - with pytest.raises(ValueError, match="Unknown buffer"): - run.get_buffer("missing") - - -def test_out_of_bounds_slice_is_rejected_by_upstream(run): - run.op.slice_info["invalid"] = ("packed", 512, 1536) - with pytest.raises(ValueError): - run.get_buffer("invalid") - - -def test_input_slices_resolve_during_reference_dispatch(run, monkeypatch): - run.op.input_args = ["packed"] - parent = run.get_buffer("packed") - - def evaluate(): - assert parent.device == "cpu" - for name in ("first", "second"): - view = run._resolve_buffer(name) - assert view.device == "cpu" - np.testing.assert_array_equal(view.numpy(), parent.numpy()[:256]) - - monkeypatch.setattr(run, "_run", evaluate) - for value in (3, 7): - parent.fill_(value) - run() From afbc276b3a65c1890092ad28435a7b0398c5a20d Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 10:17:00 -0600 Subject: [PATCH 14/22] Tests: judge 1:1 operators by their kernel's tolerance contract verify_buffer and run_test take a Tolerance, replacing rel_tol, abs_tol and max_error_rate when given. Only per-element kinds are accepted, since the failing elements are listed one by one. Tolerances only tighten. Each was measured on npu2 first. - relu, leaky_relu, elementwise_add, elementwise_mul, axpy and dequant now use their kernel's contract. That is exact for relu and 1 bf16 ulp for the rest, in place of rel 0.04 (0.01 for dequant). - axpy's golden rounds s * A + B once, as the kernel does. The old bf16 expression rounded the product too. - rope's golden uses the bf16 cos/sin tables the operator is actually given. Measured against those it is within 1 ulp. Its tolerance goes from rel 0.05 / abs 0.5 to rel 0.05 with no absolute floor. - layer_norm goes from rel 0.1 / abs 0.1 to rel 0.1 / abs 0.05, the tighter of the test's and the contract's bounds. Co-Authored-By: Claude --- iron/common/test_utils.py | 64 ++++++++++++++++++++------ iron/operators/axpy/reference.py | 5 +- iron/operators/axpy/test.py | 5 +- iron/operators/dequant/test.py | 5 +- iron/operators/elementwise_add/test.py | 5 +- iron/operators/elementwise_mul/test.py | 5 +- iron/operators/layer_norm/test.py | 8 +++- iron/operators/leaky_relu/test.py | 5 +- iron/operators/relu/test.py | 5 +- iron/operators/rope/reference.py | 4 ++ iron/operators/rope/test.py | 9 +++- 11 files changed, 97 insertions(+), 23 deletions(-) diff --git a/iron/common/test_utils.py b/iron/common/test_utils.py index 66db3a168e..e1a4709462 100644 --- a/iron/common/test_utils.py +++ b/iron/common/test_utils.py @@ -3,6 +3,8 @@ from __future__ import annotations +from dataclasses import replace + import numpy as np import torch import aie.utils as aie_utils @@ -28,12 +30,13 @@ def verify_buffer( rel_tol: float = 0.04, abs_tol: float = 1e-6, max_error_rate: float = 0.0, + tolerance: Tolerance | None = None, ) -> list[int]: """ Verify buffer contents match reference within tolerances. - The comparison is mlir-aie's ``aie.utils.verify.compare`` under a relative - ``Tolerance``: an element passes at ``|a - b| < max(abs_tol, rel_tol * (|a| + |b|))``, + The comparison is mlir-aie's ``aie.utils.verify.compare``, by default under + a relative ``Tolerance``: an element passes at ``|a - b| < max(abs_tol, rel_tol * (|a| + |b|))``, so ``rel_tol=abs_tol=0`` demands exact equality, and a NaN or infinity must meet the same value in the reference whatever ``max_error_rate`` allows. @@ -45,10 +48,24 @@ def verify_buffer( abs_tol: Absolute tolerance for comparison max_error_rate: Maximum fraction of elements allowed to exceed tolerances (0.0 to 1.0) For example, 0.01 allows up to 1% of elements to fail + tolerance: A ``Tolerance`` to judge by instead of ``rel_tol``, ``abs_tol`` + and ``max_error_rate``; typically the contract of the kernel + the operator runs (``ExternalFunction.contract.tolerance``). + It must be judgeable element by element: no ``range_frac`` + and not a bound. Returns: List of error indices. Empty if verification passes. """ + if tolerance is None: + tolerance = Tolerance.relative( + rel_tol, abs_tol, max_mismatch_frac=max_error_rate + ) + elif tolerance.kind == "bound" or tolerance.range_frac is not None: + raise ValueError( + f"{buf_name}: a {tolerance.kind} tolerance with range_frac=" + f"{tolerance.range_frac} depends on more than the element it judges" + ) def _to_numpy(x): if isinstance(x, torch.Tensor): @@ -69,26 +86,38 @@ def _to_numpy(x): return list(range(len(output), len(expected_np))) output = output[: len(expected_np)] - tolerance = Tolerance.relative(rel_tol, abs_tol, max_mismatch_frac=max_error_rate) verdict = compare(output, expected_np, tolerance) - if verdict.n_mismatch and max_error_rate > 0.0: + allowed = tolerance.max_mismatch_frac + if verdict.n_mismatch and allowed > 0.0: within = "within" if verdict else "exceeds" print( f"{buf_name}: {verdict.n_mismatch} errors " f"({verdict.n_mismatch / verdict.n_checked * 100:.2f}%) {within} allowed " - f"rate of {max_error_rate * 100:.2f}%" + f"rate of {allowed * 100:.2f}%" ) if verdict: return [] print(f"{buf_name}: {verdict.detail}") - # compare() judges; it does not list the elements. nearly_equal is the same - # per-element test, except that it also rejects a NaN that meets a NaN. - bad = ~nearly_equal(output, expected_np, rtol=rel_tol, atol=abs_tol) - bad &= ~( - np.isnan(output.astype(np.float32)) & np.isnan(expected_np.astype(np.float32)) - ) - error_indices = np.flatnonzero(bad).tolist() + # compare() judges; it does not list the elements. + if tolerance.kind == "relative": + # nearly_equal is the same per-element test, except that it also + # rejects a NaN that meets a NaN. + bad = ~nearly_equal( + output, expected_np, rtol=tolerance.rtol or 0.0, atol=tolerance.atol + ) + bad &= ~( + np.isnan(output.astype(np.float32)) + & np.isnan(expected_np.astype(np.float32)) + ) + error_indices = np.flatnonzero(bad).tolist() + else: + each = replace(tolerance, max_mismatch_frac=0.0) + error_indices = [ + i + for i in range(len(output)) + if not compare(output[i : i + 1], expected_np[i : i + 1], each) + ] for i in error_indices[:10]: print( f"Mismatch in {buf_name}[{i}]: expected {float(expected_np[i]):.6f}, got {float(output[i]):.6f}" @@ -116,6 +145,7 @@ def run_test( max_error_rate: float = 0.0, warmup_iters: int = 1, timed_iters: int = 1, + tolerance: Tolerance | None = None, ) -> tuple[dict[str, list[int]], float, float]: """ Run operator test with specified input/output buffers. @@ -129,6 +159,8 @@ def run_test( max_error_rate: Maximum fraction of elements allowed to exceed tolerances (0.0 to 1.0) warmup_iters: Number of warmup iterations before timing timed_iters: Number of timed iterations for latency/bandwidth measurement + tolerance: Judge the outputs by this ``Tolerance`` instead; see + ``verify_buffer`` Returns: (errors: dict, latency_us: float, bandwidth_gbps: float) @@ -201,7 +233,13 @@ def run_test( buf = output_map[buf_name] output_torch = buf.to_torch() buf_errors = verify_buffer( - output_torch, buf_name, expected, rel_tol, abs_tol, max_error_rate + output_torch, + buf_name, + expected, + rel_tol, + abs_tol, + max_error_rate, + tolerance=tolerance, ) if buf_errors: errors[buf_name] = buf_errors diff --git a/iron/operators/axpy/reference.py b/iron/operators/axpy/reference.py index 43bbc92a85..aa12a747be 100644 --- a/iron/operators/axpy/reference.py +++ b/iron/operators/axpy/reference.py @@ -13,8 +13,9 @@ def generate_golden_reference(input_length: int, scalar=3.0, dtype="bf16", seed= B = torch.rand(input_length, dtype=dtype_torch) * val_range s = torch.tensor(scalar, dtype=dtype_torch) - # Generate golden outputs - C = s * A + B + # Generate golden outputs: the kernel computes s * A + B in fp32 and rounds + # once, where bf16 arithmetic would round the product as well. + C = (s.float() * A.float() + B.float()).to(dtype_torch) return { "A": A, diff --git a/iron/operators/axpy/test.py b/iron/operators/axpy/test.py index 9aba1c94bf..afcb7e6515 100755 --- a/iron/operators/axpy/test.py +++ b/iron/operators/axpy/test.py @@ -63,7 +63,10 @@ def test_axpy(input_length, num_aie_columns, tile_size, scalar_factor, aie_conte output_buffers = {"output": golden_ref["C"]} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 + operator, + input_buffers, + output_buffers, + tolerance=operator._kernel().contract.tolerance, ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/dequant/test.py b/iron/operators/dequant/test.py index a0831d65c5..76134c39cb 100644 --- a/iron/operators/dequant/test.py +++ b/iron/operators/dequant/test.py @@ -77,7 +77,10 @@ def test_dequant( output_buffers = {"output": golden_ref["output"].flatten()} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.01, abs_tol=1e-6 + operator, + input_buffers, + output_buffers, + tolerance=operator._kernel().contract.tolerance, ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/elementwise_add/test.py b/iron/operators/elementwise_add/test.py index 4414b53036..10e54f0cd6 100755 --- a/iron/operators/elementwise_add/test.py +++ b/iron/operators/elementwise_add/test.py @@ -38,7 +38,10 @@ def test_elementwise_add(input_length, num_aie_columns, tile_size, aie_context): output_buffers = {"output": golden_ref["C"]} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 + operator, + input_buffers, + output_buffers, + tolerance=operator._kernel().contract.tolerance, ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/elementwise_mul/test.py b/iron/operators/elementwise_mul/test.py index 8d2c638b4f..0c4663d6e4 100755 --- a/iron/operators/elementwise_mul/test.py +++ b/iron/operators/elementwise_mul/test.py @@ -40,7 +40,10 @@ def test_elementwise_mul(input_length, num_aie_columns, tile_size, aie_context): output_buffers = {"output": golden_ref["C"]} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 + operator, + input_buffers, + output_buffers, + tolerance=operator._kernel().contract.tolerance, ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/layer_norm/test.py b/iron/operators/layer_norm/test.py index 9d85ee8919..29c512d6f6 100755 --- a/iron/operators/layer_norm/test.py +++ b/iron/operators/layer_norm/test.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance from iron.operators.layer_norm.op import LayerNorm from iron.operators.layer_norm.reference import generate_golden_reference @@ -46,7 +47,12 @@ def test_layer_norm( output_buffers = {"output": golden_ref["output"]} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.1, abs_tol=0.1 + operator, + input_buffers, + output_buffers, + # The tighter of this test's former rel_tol (0.1) and the kernel + # contract's atol (0.05). + tolerance=Tolerance.relative(0.1, 0.05), ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/leaky_relu/test.py b/iron/operators/leaky_relu/test.py index cc80547622..a9055905ff 100755 --- a/iron/operators/leaky_relu/test.py +++ b/iron/operators/leaky_relu/test.py @@ -51,7 +51,10 @@ def test_leaky_relu( output_buffers = {"output": golden_ref["output"]} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 + operator, + input_buffers, + output_buffers, + tolerance=operator._kernel().contract.tolerance, ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/relu/test.py b/iron/operators/relu/test.py index 6c9628334c..e23c9d4c3e 100755 --- a/iron/operators/relu/test.py +++ b/iron/operators/relu/test.py @@ -41,7 +41,10 @@ def test_relu(input_length, num_aie_columns, num_channels, tile_size, aie_contex output_buffers = {"output": golden_ref["output"]} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 + operator, + input_buffers, + output_buffers, + tolerance=operator._kernel().contract.tolerance, ) print(f"\nLatency (us): {latency_us:.1f}") diff --git a/iron/operators/rope/reference.py b/iron/operators/rope/reference.py index 147ea9c31d..43f46755cf 100644 --- a/iron/operators/rope/reference.py +++ b/iron/operators/rope/reference.py @@ -182,6 +182,10 @@ def generate_golden_reference( method_type=method_type, freq_config=freq_config, ) + # The operator is handed the tables in bf16, so the golden output is computed + # from those rounded values rather than from the fp32 ones. + cos = cos.to(torch.bfloat16).to(torch.float32) + sin = sin.to(torch.bfloat16).to(torch.float32) val_range = 4 # Head count is inferred from rows and context_len. This logic assumes rows is either # smaller than context_len (1 head, seq_len == rows) or an exact multiple of context_len diff --git a/iron/operators/rope/test.py b/iron/operators/rope/test.py index 9e4820ca10..2dc8c5dcf9 100755 --- a/iron/operators/rope/test.py +++ b/iron/operators/rope/test.py @@ -4,6 +4,7 @@ import pytest import aie.utils as aie_utils +from aie.utils.verify import Tolerance from iron.operators.rope.op import RoPE from iron.operators.rope.reference import generate_golden_reference from iron.common.test_utils import run_test @@ -83,7 +84,13 @@ def test_rope(rows, cols, angle_rows, aie_columns, method_type, aie_context): output_buffers = {"output": golden_ref["C"].transpose(0, 1).contiguous()} errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.05, abs_tol=0.5 + operator, + input_buffers, + output_buffers, + # The tighter of this test's former rel_tol and the kernel contract's + # atol (none): an output that cancels to near zero is judged + # relatively like any other. + tolerance=Tolerance.relative(0.05), ) print(f"\nLatency (us): {latency_us:.1f}") From 4f4c061fa33b627ce52087bc7e88a7d31fe80f55 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 10:49:00 -0600 Subject: [PATCH 15/22] Tests: one reference per operator, judged by its kernel's contract Each simple 1:1 operator now has exactly one reference, the op's reference(). Its reference.py keeps the math as reference(...) and draws inputs with generate_inputs(...), same seeds and draws as before; the golden output is no longer a separate code path that can drift from what dispatch="compare" checks against. test_utils.assert_matches_reference(op, *inputs) shapes the inputs to the arg spec, computes op.reference(...), dispatches once and judges the output by op.reference_tolerance() -- the declared contract of the kernel the operator runs -- unless the test passes a tighter tolerance. The tests of relu, tanh, sigmoid, gelu, silu, leaky_relu, layer_norm, rms_norm, softmax, elementwise_add/mul, axpy, dequant, mem_copy, transpose, repeat and rope become thin wrappers over it; names, parameters and metrics are unchanged. Tests that were stricter than the contract keep their tolerance explicitly. The new references are bit-identical to the old goldens on every test parameter set, except axpy (now rounds once, as the kernel does: the scalar to bf16, fp32 multiply-add, one rounding; 0 differences on the test sets) and dequant (the fp32 golden rounded once to bf16, the same verdict). Also drops torch_dtype_map and the dtype= parameters that only ever took "bf16". Co-Authored-By: Claude --- iron/common/base.py | 14 ++ iron/common/test_utils.py | 69 +++++++-- iron/operators/axpy/op.py | 5 + iron/operators/axpy/reference.py | 30 ++-- iron/operators/axpy/test.py | 23 +-- iron/operators/dequant/op.py | 6 + iron/operators/dequant/reference.py | 31 ++-- iron/operators/dequant/test.py | 23 +-- iron/operators/elementwise_add/reference.py | 10 +- iron/operators/elementwise_add/test.py | 24 +-- iron/operators/elementwise_mul/reference.py | 10 +- iron/operators/elementwise_mul/test.py | 24 +-- iron/operators/flm/gemm/reference.py | 7 +- iron/operators/gelu/op.py | 5 + iron/operators/gelu/reference.py | 12 +- iron/operators/gelu/test.py | 24 ++- iron/operators/gemm/reference.py | 13 +- iron/operators/layer_norm/op.py | 6 + iron/operators/layer_norm/reference.py | 22 +-- iron/operators/layer_norm/test.py | 25 +--- iron/operators/leaky_relu/op.py | 5 + iron/operators/leaky_relu/reference.py | 15 +- iron/operators/leaky_relu/test.py | 24 +-- iron/operators/mem_copy/op.py | 5 + iron/operators/mem_copy/reference.py | 16 +- iron/operators/mem_copy/test.py | 27 +--- iron/operators/relu/reference.py | 10 +- iron/operators/relu/test.py | 24 +-- iron/operators/repeat/reference.py | 7 +- iron/operators/repeat/test.py | 20 +-- iron/operators/rms_norm/reference.py | 18 +-- iron/operators/rms_norm/test.py | 27 ++-- iron/operators/rope/op.py | 11 +- iron/operators/rope/reference.py | 139 +++++------------- iron/operators/rope/test.py | 27 +--- iron/operators/sigmoid/op.py | 5 + iron/operators/sigmoid/reference.py | 12 +- iron/operators/sigmoid/test.py | 24 ++- iron/operators/silu/reference.py | 7 +- iron/operators/silu/test.py | 24 ++- iron/operators/softmax/reference.py | 15 +- iron/operators/softmax/test.py | 23 +-- iron/operators/strided_copy/reference.py | 7 +- iron/operators/tanh/op.py | 5 + iron/operators/tanh/reference.py | 12 +- iron/operators/tanh/test.py | 24 ++- iron/operators/transpose/op.py | 4 +- iron/operators/transpose/reference.py | 26 +--- iron/operators/transpose/test.py | 27 +--- .../operators/rope_reference_convention.py | 4 +- 50 files changed, 401 insertions(+), 576 deletions(-) diff --git a/iron/common/base.py b/iron/common/base.py index 1e3dd68d77..635a588f8b 100644 --- a/iron/common/base.py +++ b/iron/common/base.py @@ -14,6 +14,7 @@ from ml_dtypes import bfloat16 import aie.utils as aie_utils from aie.utils.npukernel import NPUKernel +from aie.utils.verify import Tolerance from . import compilation as comp from .context import AIEContext @@ -157,6 +158,19 @@ def get_mlir_artifact(self) -> CompilationArtifact: def get_kernel_artifacts(self) -> list[CompilationArtifact]: pass + def reference_tolerance(self) -> Tolerance | None: + """How close the NPU output must come to ``reference()``. + + This is the declared contract of the one kernel the operator runs, its + ``_kernel()``. ``None`` when it runs several kernels or none (a + ``_kernel()`` that returns ``None``), or when that kernel declares no + tolerance. + """ + kernel = getattr(self, "_kernel", lambda: None)() + if kernel is None or kernel.contract is None: + return None + return kernel.contract.tolerance + def get_artifacts( self, prefix: str = "" ) -> tuple[XclbinArtifact, InstsBinArtifact]: diff --git a/iron/common/test_utils.py b/iron/common/test_utils.py index e1a4709462..da2b1c8b96 100644 --- a/iron/common/test_utils.py +++ b/iron/common/test_utils.py @@ -13,15 +13,6 @@ from ml_dtypes import bfloat16 from .base import AIEOperatorBase -torch_dtype_map = { - "bf16": torch.bfloat16, - "f32": torch.float32, - "i8": torch.int8, - "ui8": torch.uint8, - "i16": torch.int16, - "i32": torch.int32, -} - def verify_buffer( output: np.ndarray | torch.Tensor, @@ -100,17 +91,19 @@ def _to_numpy(x): print(f"{buf_name}: {verdict.detail}") # compare() judges; it does not list the elements. + both_nan = np.isnan(output.astype(np.float32)) & np.isnan( + expected_np.astype(np.float32) + ) if tolerance.kind == "relative": # nearly_equal is the same per-element test, except that it also # rejects a NaN that meets a NaN. bad = ~nearly_equal( output, expected_np, rtol=tolerance.rtol or 0.0, atol=tolerance.atol ) - bad &= ~( - np.isnan(output.astype(np.float32)) - & np.isnan(expected_np.astype(np.float32)) - ) - error_indices = np.flatnonzero(bad).tolist() + error_indices = np.flatnonzero(bad & ~both_nan).tolist() + elif tolerance.kind == "exact": + bad = output != expected_np.astype(output.dtype) + error_indices = np.flatnonzero(bad & ~both_nan).tolist() else: each = replace(tolerance, max_mismatch_frac=0.0) error_indices = [ @@ -254,6 +247,54 @@ def run_test( return errors, latency_us, bandwidth_gbps +def assert_matches_reference( + operator: AIEOperatorBase, + *inputs: torch.Tensor, + tolerance: Tolerance | None = None, +) -> None: + """Dispatch ``operator`` once and assert its output matches ``reference()``. + + The expected output is ``operator.reference(*inputs)``, each input shaped + as its argument spec, so the test and a ``dispatch="compare"`` sequence + hold the operator to the same reference. Latency and bandwidth are printed + in the form the CI metrics parse. + + Args: + operator: An operator with a single output and a ``reference()`` + inputs: Its ``"in"`` arguments, in argument-spec order + tolerance: How close the output must come; defaults to + ``operator.reference_tolerance()``, the contract of the + kernel it runs + """ + in_specs = [s for s in operator.get_arg_spec() if s.direction == "in"] + if len(inputs) != len(in_specs): + raise ValueError( + f"{type(operator).__name__} takes {len(in_specs)} inputs, " + f"got {len(inputs)}" + ) + expected = operator.reference( + *(x.reshape(spec.shape) for x, spec in zip(inputs, in_specs)) + ) + if tolerance is None: + tolerance = operator.reference_tolerance() + if tolerance is None: + raise ValueError( + f"{type(operator).__name__} declares no tolerance; pass tolerance=" + ) + + errors, latency_us, bandwidth_gbps = run_test( + operator, + {f"input{i}": x for i, x in enumerate(inputs)}, + {"output": expected}, + tolerance=tolerance, + ) + + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") + + assert not errors, f"Test failed with errors: {errors}" + + def make_channeled_unary_params(input_lengths, tile_size_cap, num_channels_choices): """Generate parameter tuples for channeled unary operator tests. diff --git a/iron/operators/axpy/op.py b/iron/operators/axpy/op.py index 87ec4d2823..2d9f50d236 100644 --- a/iron/operators/axpy/op.py +++ b/iron/operators/axpy/op.py @@ -24,6 +24,11 @@ class AXPY(BinaryElementwiseOperator): def _kernel(self): return datamovement.axpy(self._tile_elements) + def reference(self, x, y): + from iron.operators.axpy.reference import reference + + return reference(x, y, self.scalar_factor) + def _mlir_callback_args(self): return super()._mlir_callback_args() + [self.scalar_factor, self._kernel()] diff --git a/iron/operators/axpy/reference.py b/iron/operators/axpy/reference.py index aa12a747be..ee37697abf 100644 --- a/iron/operators/axpy/reference.py +++ b/iron/operators/axpy/reference.py @@ -2,23 +2,21 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map -def generate_golden_reference(input_length: int, scalar=3.0, dtype="bf16", seed=42): - torch.manual_seed(seed) - val_range = 4 - dtype_torch = torch_dtype_map[dtype] - A = torch.rand(input_length, dtype=dtype_torch) * val_range - B = torch.rand(input_length, dtype=dtype_torch) * val_range - s = torch.tensor(scalar, dtype=dtype_torch) +def reference(x, y, scalar): + """CPU reference: ``scalar * x + y`` in fp32, rounded once (ground truth). + + The kernel takes ``scalar`` as bf16 and rounds only the result, where bf16 + arithmetic would round the product as well. + """ + a = torch.tensor(scalar, dtype=torch.bfloat16).float() + return (a * x.float() + y.float()).to(x.dtype) - # Generate golden outputs: the kernel computes s * A + B in fp32 and rounds - # once, where bf16 arithmetic would round the product as well. - C = (s.float() * A.float() + B.float()).to(dtype_torch) - return { - "A": A, - "B": B, - "C": C, - } +def generate_inputs(input_length: int, seed=42): + torch.manual_seed(seed) + val_range = 4 + x = torch.rand(input_length, dtype=torch.bfloat16) * val_range + y = torch.rand(input_length, dtype=torch.bfloat16) * val_range + return x, y diff --git a/iron/operators/axpy/test.py b/iron/operators/axpy/test.py index afcb7e6515..aa48e9cfc0 100755 --- a/iron/operators/axpy/test.py +++ b/iron/operators/axpy/test.py @@ -6,8 +6,8 @@ import aie.utils as aie_utils from iron.operators.axpy.op import AXPY -from iron.operators.axpy.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.axpy.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_params(): @@ -47,9 +47,7 @@ def get_params(): get_params(), ) def test_axpy(input_length, num_aie_columns, tile_size, scalar_factor, aie_context): - golden_ref = generate_golden_reference( - input_length=input_length, scalar=scalar_factor - ) + x, y = generate_inputs(input_length=input_length) operator = AXPY( size=input_length, @@ -59,17 +57,4 @@ def test_axpy(input_length, num_aie_columns, tile_size, scalar_factor, aie_conte context=aie_context, ) - input_buffers = {"x": golden_ref["A"], "y": golden_ref["B"]} - output_buffers = {"output": golden_ref["C"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - tolerance=operator._kernel().contract.tolerance, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, x, y) diff --git a/iron/operators/dequant/op.py b/iron/operators/dequant/op.py index 94c0bef221..d92befb83a 100644 --- a/iron/operators/dequant/op.py +++ b/iron/operators/dequant/op.py @@ -65,6 +65,12 @@ def get_mlir_artifact(self): def _kernel(self): return datamovement.expand(self.tile_size, self.group_size) + def reference(self, payload): + """CPU reference: each uint4 value times its group's bf16 scale.""" + from iron.operators.dequant.reference import reference + + return reference(payload, self.tile_size, self.group_size) + def get_kernel_artifacts(self): return [KernelObjectArtifact.from_extern(self._kernel())] diff --git a/iron/operators/dequant/reference.py b/iron/operators/dequant/reference.py index 4ab72f4951..e13829c619 100644 --- a/iron/operators/dequant/reference.py +++ b/iron/operators/dequant/reference.py @@ -2,12 +2,26 @@ # SPDX-License-Identifier: Apache-2.0 import torch -import numpy as np -from ml_dtypes import bfloat16 +from aie.iron.kernels.datamovement import expand_ref -def generate_golden_reference(input_length, tile_size, group_size): - torch.manual_seed(42) +def reference(payload, tile_size, group_size): + """CPU reference: each uint4 value times its group's bf16 scale. + + ``payload`` holds, per tile, ``tile_size`` packed uint4 values + (``tile_size // 2`` bytes, low nibble first) followed by one bf16 scale per + ``group_size`` elements. The product is exact in fp32 and rounded once. + """ + tile_bytes = tile_size // 2 + 2 * (tile_size // group_size) + tiles = payload.reshape(-1, tile_bytes).numpy() + out = expand_ref(tiles, tile_size=tile_size, group_size=group_size) + return torch.from_numpy(out).to(torch.bfloat16).reshape(-1) + + +def generate_inputs(input_length, tile_size, group_size, seed=42): + """Random bf16 values quantized to uint4 per group and packed as the + operator takes them (see ``reference``).""" + torch.manual_seed(seed) if input_length % tile_size != 0: raise ValueError("Input length must be a multiple of tile size.") @@ -23,8 +37,7 @@ def generate_golden_reference(input_length, tile_size, group_size): ) # Total bytes (uint8 elements) after processing each tile val_range = 3.75 # Values in [0, 3.75) - # Generate golden output with uniform distribution between 0 and val_range - # This output will be quantized to be used as the input + # Uniform values in [0, val_range), quantized below to make the input A = ( torch.rand(num_tiles * num_scale_factors, group_size, dtype=torch.bfloat16) * val_range @@ -46,7 +59,6 @@ def generate_golden_reference(input_length, tile_size, group_size): axis=0, dtype=torch.quint8, ) - B = torch.dequantize(A) # Convert A from a quantized tensor type to regular tensor type for data packing # We do the data packing here instead of the host to show how the data would need to be @@ -79,7 +91,4 @@ def generate_golden_reference(input_length, tile_size, group_size): 0xFF, ) - return { - "input": A_concat, - "output": B, - } + return A_concat.flatten() diff --git a/iron/operators/dequant/test.py b/iron/operators/dequant/test.py index 76134c39cb..222a1c77f8 100644 --- a/iron/operators/dequant/test.py +++ b/iron/operators/dequant/test.py @@ -6,8 +6,8 @@ import aie.utils as aie_utils from iron.operators.dequant.op import Dequant -from iron.operators.dequant.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.dequant.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_params(): @@ -56,7 +56,7 @@ def get_params(): def test_dequant( input_length, num_aie_columns, num_channels, tile_size, group_size, aie_context ): - golden_ref = generate_golden_reference( + payload = generate_inputs( input_length=input_length, tile_size=tile_size, group_size=group_size, @@ -71,19 +71,4 @@ def test_dequant( context=aie_context, ) - input_buffers = { - "input": golden_ref["input"].flatten(), - } - output_buffers = {"output": golden_ref["output"].flatten()} - - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - tolerance=operator._kernel().contract.tolerance, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, payload) diff --git a/iron/operators/elementwise_add/reference.py b/iron/operators/elementwise_add/reference.py index c34089853b..9ad76e874c 100644 --- a/iron/operators/elementwise_add/reference.py +++ b/iron/operators/elementwise_add/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(a, b): @@ -10,10 +9,9 @@ def reference(a, b): return a + b -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - dtype_torch = torch_dtype_map[dtype] - input_a = torch.rand(input_length, dtype=dtype_torch) * val_range - input_b = torch.rand(input_length, dtype=dtype_torch) * val_range - return {"A": input_a, "B": input_b, "C": reference(input_a, input_b)} + a = torch.rand(input_length, dtype=torch.bfloat16) * val_range + b = torch.rand(input_length, dtype=torch.bfloat16) * val_range + return a, b diff --git a/iron/operators/elementwise_add/test.py b/iron/operators/elementwise_add/test.py index 10e54f0cd6..745172d187 100755 --- a/iron/operators/elementwise_add/test.py +++ b/iron/operators/elementwise_add/test.py @@ -5,8 +5,11 @@ import pytest from iron.operators.elementwise_add.op import ElementwiseAdd -from iron.operators.elementwise_add.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_binary_elementwise_params +from iron.operators.elementwise_add.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_binary_elementwise_params, +) def get_params(): @@ -25,7 +28,7 @@ def get_params(): get_params(), ) def test_elementwise_add(input_length, num_aie_columns, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + a, b = generate_inputs(input_length=input_length) operator = ElementwiseAdd( size=input_length, @@ -34,17 +37,4 @@ def test_elementwise_add(input_length, num_aie_columns, tile_size, aie_context): context=aie_context, ) - input_buffers = {"input1": golden_ref["A"], "input2": golden_ref["B"]} - output_buffers = {"output": golden_ref["C"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - tolerance=operator._kernel().contract.tolerance, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, a, b) diff --git a/iron/operators/elementwise_mul/reference.py b/iron/operators/elementwise_mul/reference.py index f27e717f9c..de31421110 100644 --- a/iron/operators/elementwise_mul/reference.py +++ b/iron/operators/elementwise_mul/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(a, b): @@ -10,10 +9,9 @@ def reference(a, b): return a * b -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - dtype_torch = torch_dtype_map[dtype] - input_a = torch.rand(input_length, dtype=dtype_torch) * val_range - input_b = torch.rand(input_length, dtype=dtype_torch) * val_range - return {"A": input_a, "B": input_b, "C": reference(input_a, input_b)} + a = torch.rand(input_length, dtype=torch.bfloat16) * val_range + b = torch.rand(input_length, dtype=torch.bfloat16) * val_range + return a, b diff --git a/iron/operators/elementwise_mul/test.py b/iron/operators/elementwise_mul/test.py index 0c4663d6e4..27a1b4e0e9 100755 --- a/iron/operators/elementwise_mul/test.py +++ b/iron/operators/elementwise_mul/test.py @@ -5,8 +5,11 @@ import pytest from iron.operators.elementwise_mul.op import ElementwiseMul -from iron.operators.elementwise_mul.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_binary_elementwise_params +from iron.operators.elementwise_mul.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_binary_elementwise_params, +) def get_params(): @@ -27,7 +30,7 @@ def get_params(): get_params(), ) def test_elementwise_mul(input_length, num_aie_columns, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + a, b = generate_inputs(input_length=input_length) operator = ElementwiseMul( size=input_length, @@ -36,17 +39,4 @@ def test_elementwise_mul(input_length, num_aie_columns, tile_size, aie_context): context=aie_context, ) - input_buffers = {"input1": golden_ref["A"], "input2": golden_ref["B"]} - output_buffers = {"output": golden_ref["C"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - tolerance=operator._kernel().contract.tolerance, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, a, b) diff --git a/iron/operators/flm/gemm/reference.py b/iron/operators/flm/gemm/reference.py index 24eda2335f..6639a06034 100644 --- a/iron/operators/flm/gemm/reference.py +++ b/iron/operators/flm/gemm/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map from iron.operators.flm.gemm.design import Epilogue @@ -62,7 +61,6 @@ def generate_golden_reference( M: int, K: int, N: int, - dtype="bf16", seed=42, epilogue=Epilogue.NONE, clamp=None, @@ -77,8 +75,7 @@ def generate_golden_reference( range where the curve is actually interesting. """ torch.manual_seed(seed) - dtype_torch = torch_dtype_map[dtype] - input_a = torch.randn(M, K, dtype=dtype_torch) * scale - input_b = torch.rand(K, N, dtype=dtype_torch) * scale + input_a = torch.randn(M, K, dtype=torch.bfloat16) * scale + input_b = torch.rand(K, N, dtype=torch.bfloat16) * scale output = reference(input_a, input_b, epilogue, clamp) return {"input": input_a, "input_b": input_b, "output": output} diff --git a/iron/operators/gelu/op.py b/iron/operators/gelu/op.py index 1644f56dfb..426ce98e12 100644 --- a/iron/operators/gelu/op.py +++ b/iron/operators/gelu/op.py @@ -18,3 +18,8 @@ class GELU(ChanneledUnaryOperator): def _kernel(self): return activation.gelu_sized(self._line_size) + + def reference(self, x): + from iron.operators.gelu.reference import reference + + return reference(x) diff --git a/iron/operators/gelu/reference.py b/iron/operators/gelu/reference.py index 991d67bab0..6640ed757a 100644 --- a/iron/operators/gelu/reference.py +++ b/iron/operators/gelu/reference.py @@ -2,12 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def reference(x): + """CPU reference: GELU, tanh approximation (ground truth).""" + return torch.nn.functional.gelu(x, approximate="tanh") + + +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(input_length, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = torch.nn.functional.gelu(input_tensor, approximate="tanh") - return {"input": input_tensor, "output": output_tensor} + return torch.rand(input_length, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/gelu/test.py b/iron/operators/gelu/test.py index d2c7cb4bbc..bd3ea18ae1 100755 --- a/iron/operators/gelu/test.py +++ b/iron/operators/gelu/test.py @@ -3,10 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance from iron.operators.gelu.op import GELU -from iron.operators.gelu.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.gelu.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -30,7 +34,7 @@ def _marks(ext): get_params(), ) def test_gelu(input_length, num_aie_columns, num_channels, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + x = generate_inputs(input_length=input_length) operator = GELU( size=input_length, @@ -40,14 +44,6 @@ def test_gelu(input_length, num_aie_columns, num_channels, tile_size, aie_contex context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # The torch reference at this test's tolerance, tighter than the kernel + # contract's. + assert_matches_reference(operator, x, tolerance=Tolerance.relative(0.04, 1e-6)) diff --git a/iron/operators/gemm/reference.py b/iron/operators/gemm/reference.py index 4cc9fb4c33..84bc39b023 100644 --- a/iron/operators/gemm/reference.py +++ b/iron/operators/gemm/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(input_a, input_b, b_col_maj=False, c_col_maj=False): @@ -23,7 +22,6 @@ def generate_golden_reference( M: int, K: int, N: int, - dtype="bf16", seed=42, b_col_maj=False, c_col_maj=False, @@ -31,9 +29,8 @@ def generate_golden_reference( ): torch.manual_seed(seed) val_range = 4 - dtype_torch = torch_dtype_map[dtype] - input_a = torch.randn(M, K, dtype=dtype_torch) * val_range - input_b_full = torch.rand(K, N, dtype=dtype_torch) * val_range + input_a = torch.randn(M, K, dtype=torch.bfloat16) * val_range + input_b_full = torch.rand(K, N, dtype=torch.bfloat16) * val_range if False: # The following inputs are useful for debugging; # the A matrix becomes a matrix where each element encodes its row and column index, @@ -42,10 +39,10 @@ def generate_golden_reference( factor = 10 ** (col_digits + 1) row_indices = torch.arange(M, dtype=torch.int64).unsqueeze(1) col_indices = torch.arange(K, dtype=torch.int64).unsqueeze(0) - input_a = (row_indices * factor + col_indices).to(dtype=dtype_torch) - input_b_full = torch.zeros(K, N, dtype=dtype_torch) + input_a = (row_indices * factor + col_indices).to(dtype=torch.bfloat16) + input_b_full = torch.zeros(K, N, dtype=torch.bfloat16) diag_dim = min(K, N) - input_b_full[:diag_dim, :diag_dim] = torch.eye(diag_dim, dtype=dtype_torch) + input_b_full[:diag_dim, :diag_dim] = torch.eye(diag_dim, dtype=torch.bfloat16) # Store B in the operator's expected layout, then compute the output via the # shared reference so the test golden and the operator reference agree. if b_col_maj: diff --git a/iron/operators/layer_norm/op.py b/iron/operators/layer_norm/op.py index 0d398d395d..affb0e95a9 100644 --- a/iron/operators/layer_norm/op.py +++ b/iron/operators/layer_norm/op.py @@ -25,6 +25,12 @@ def __post_init__(self, trace_size): def _kernel(self): return norm.layer_norm(self._line_size) + def reference(self, x): + """CPU reference: layer normalization of each line the kernel sees.""" + from iron.operators.layer_norm.reference import reference + + return reference(x.reshape(-1, self._line_size)) + def _mlir_callback_args(self): return [ aie_utils.get_current_device(), diff --git a/iron/operators/layer_norm/reference.py b/iron/operators/layer_norm/reference.py index 86fb5855de..f23b9e8b1b 100644 --- a/iron/operators/layer_norm/reference.py +++ b/iron/operators/layer_norm/reference.py @@ -2,17 +2,19 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map -def generate_golden_reference(rows: int, cols: int, dtype="bf16", seed=42): +def reference(x): + """CPU reference: layer normalization of each row over its last dim, with no + learnable affine parameters (ground truth). + + The AIE kernel normalizes one line at a time, computing mean and variance + over that line alone. + """ + return torch.nn.functional.layer_norm(x, normalized_shape=(x.shape[-1],)) + + +def generate_inputs(rows: int, cols: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(rows, cols, dtype=torch_dtype_map[dtype]) * val_range - # normalized_shape=(cols,) normalizes each row independently over its `cols` elements. - # This matches the AIE kernel behavior, which processes one tile (one row) at a time - # and computes mean and variance per row (no learnable affine parameters). - output_tensor = torch.nn.functional.layer_norm( - input_tensor, normalized_shape=(cols,) - ) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(rows, cols, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/layer_norm/test.py b/iron/operators/layer_norm/test.py index 29c512d6f6..5214bfeed9 100755 --- a/iron/operators/layer_norm/test.py +++ b/iron/operators/layer_norm/test.py @@ -6,8 +6,11 @@ from aie.utils.verify import Tolerance from iron.operators.layer_norm.op import LayerNorm -from iron.operators.layer_norm.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.layer_norm.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -30,10 +33,7 @@ def get_params(): def test_layer_norm( input_length, num_aie_columns, num_channels, tile_size, aie_context ): - - rows = input_length // tile_size - cols = tile_size - golden_ref = generate_golden_reference(rows=rows, cols=cols) + x = generate_inputs(rows=input_length // tile_size, cols=tile_size) operator = LayerNorm( size=input_length, @@ -43,19 +43,10 @@ def test_layer_norm( context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( + assert_matches_reference( operator, - input_buffers, - output_buffers, + x, # The tighter of this test's former rel_tol (0.1) and the kernel # contract's atol (0.05). tolerance=Tolerance.relative(0.1, 0.05), ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" diff --git a/iron/operators/leaky_relu/op.py b/iron/operators/leaky_relu/op.py index 00d0b18458..6faf9162da 100644 --- a/iron/operators/leaky_relu/op.py +++ b/iron/operators/leaky_relu/op.py @@ -48,6 +48,11 @@ def __post_init__(self) -> None: def _kernel(self): return activation.leaky_relu(self._line_size) + def reference(self, x): + from iron.operators.leaky_relu.reference import reference + + return reference(x, self.alpha) + def _mlir_callback_args(self): return super()._mlir_callback_args() + [self.alpha, self._kernel()] diff --git a/iron/operators/leaky_relu/reference.py b/iron/operators/leaky_relu/reference.py index 8c23041cc1..30efaa3593 100644 --- a/iron/operators/leaky_relu/reference.py +++ b/iron/operators/leaky_relu/reference.py @@ -2,15 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map -def generate_golden_reference(input_length: int, alpha=0.01, dtype="bf16", seed=42): +def reference(x, alpha=0.01): + """CPU reference: leaky ReLU with negative slope ``alpha`` (ground truth).""" + return torch.nn.functional.leaky_relu(x, negative_slope=alpha) + + +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = ( - torch.rand(input_length, dtype=torch_dtype_map[dtype]) * val_range - - val_range / 2 - ) - output_tensor = torch.nn.functional.leaky_relu(input_tensor, negative_slope=alpha) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(input_length, dtype=torch.bfloat16) * val_range - val_range / 2 diff --git a/iron/operators/leaky_relu/test.py b/iron/operators/leaky_relu/test.py index a9055905ff..0fb944e919 100755 --- a/iron/operators/leaky_relu/test.py +++ b/iron/operators/leaky_relu/test.py @@ -5,8 +5,11 @@ import pytest from iron.operators.leaky_relu.op import LeakyReLU -from iron.operators.leaky_relu.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.leaky_relu.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -36,7 +39,7 @@ def get_params(): def test_leaky_relu( input_length, num_aie_columns, num_channels, tile_size, alpha, aie_context ): - golden_ref = generate_golden_reference(input_length=input_length, alpha=alpha) + x = generate_inputs(input_length=input_length) operator = LeakyReLU( size=input_length, @@ -47,17 +50,4 @@ def test_leaky_relu( context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - tolerance=operator._kernel().contract.tolerance, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, x) diff --git a/iron/operators/mem_copy/op.py b/iron/operators/mem_copy/op.py index bd4234fc00..e02f9d90e3 100644 --- a/iron/operators/mem_copy/op.py +++ b/iron/operators/mem_copy/op.py @@ -63,6 +63,11 @@ def _kernel(self): return None return eltwise.passthrough(mem_copy_line_size(self.tile_size), np.int16) + def reference(self, x): + from iron.operators.mem_copy.reference import reference + + return reference(x) + def get_kernel_artifacts(self): if self.bypass: return [] diff --git a/iron/operators/mem_copy/reference.py b/iron/operators/mem_copy/reference.py index 948a09ab15..40d15f036d 100644 --- a/iron/operators/mem_copy/reference.py +++ b/iron/operators/mem_copy/reference.py @@ -4,14 +4,12 @@ import torch -def generate_golden_reference(input_length): - torch.manual_seed(42) +def reference(x): + """CPU reference: a copy (ground truth).""" + return x.clone() - # Generate random input data - val_range = 4 - A = torch.rand(input_length, dtype=torch.bfloat16) * val_range - return { - "input": A, - "output": A.clone(), - } +def generate_inputs(input_length): + torch.manual_seed(42) + val_range = 4 + return torch.rand(input_length, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/mem_copy/test.py b/iron/operators/mem_copy/test.py index 07541141a6..6fb14ff303 100644 --- a/iron/operators/mem_copy/test.py +++ b/iron/operators/mem_copy/test.py @@ -3,11 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance import aie.utils as aie_utils from iron.operators.mem_copy.op import MemCopy -from iron.operators.mem_copy.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.mem_copy.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_params(): @@ -62,8 +63,9 @@ def get_params(): def test_mem_copy( input_length, num_cores, num_channels, bypass, tile_size, aie_context ): - golden_ref = generate_golden_reference(input_length=input_length) + x = generate_inputs(input_length=input_length) + # num_cores >= num_channels is required: each channel must have at least one core assigned operator = MemCopy( size=input_length, num_cores=num_cores, @@ -73,20 +75,5 @@ def test_mem_copy( context=aie_context, ) - # num_cores >= num_channels is required: each channel must have at least one core assigned - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - # A copy that alters a value is a broken copy, so gate it exactly. - operator, - input_buffers, - output_buffers, - rel_tol=0.0, - abs_tol=0.0, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # A copy that alters a value is a broken copy, so gate it exactly. + assert_matches_reference(operator, x, tolerance=Tolerance.exact()) diff --git a/iron/operators/relu/reference.py b/iron/operators/relu/reference.py index 2494ab6042..05319f45db 100644 --- a/iron/operators/relu/reference.py +++ b/iron/operators/relu/reference.py @@ -2,19 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(x): return torch.nn.functional.relu(x) -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = ( - torch.rand(input_length, dtype=torch_dtype_map[dtype]) * val_range - - val_range / 2 - ) - output_tensor = torch.nn.functional.relu(input_tensor) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(input_length, dtype=torch.bfloat16) * val_range - val_range / 2 diff --git a/iron/operators/relu/test.py b/iron/operators/relu/test.py index e23c9d4c3e..f3a034ecc2 100755 --- a/iron/operators/relu/test.py +++ b/iron/operators/relu/test.py @@ -5,8 +5,11 @@ import pytest from iron.operators.relu.op import ReLU -from iron.operators.relu.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.relu.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -27,7 +30,7 @@ def get_params(): get_params(), ) def test_relu(input_length, num_aie_columns, num_channels, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + x = generate_inputs(input_length=input_length) operator = ReLU( size=input_length, @@ -37,17 +40,4 @@ def test_relu(input_length, num_aie_columns, num_channels, tile_size, aie_contex context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - tolerance=operator._kernel().contract.tolerance, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, x) diff --git a/iron/operators/repeat/reference.py b/iron/operators/repeat/reference.py index 9952752c75..7f4e068970 100644 --- a/iron/operators/repeat/reference.py +++ b/iron/operators/repeat/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(x, repeat): @@ -10,9 +9,7 @@ def reference(x, repeat): return x.repeat_interleave(repeat, dim=0) -def generate_golden_reference(rows: int, cols: int, repeat: int, dtype="bf16", seed=42): +def generate_inputs(rows: int, cols: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(rows, cols, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = reference(input_tensor, repeat) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(rows, cols, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/repeat/test.py b/iron/operators/repeat/test.py index 499ec42424..1f80f95758 100644 --- a/iron/operators/repeat/test.py +++ b/iron/operators/repeat/test.py @@ -3,10 +3,11 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance from iron.operators.repeat.op import Repeat -from iron.operators.repeat.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.repeat.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_params(): @@ -38,7 +39,7 @@ def test_repeat(rows, cols, repeat, transfer_size, aie_context): is the whole failure mode here, since the only caller uses this to expand KV groups to attention heads and a misrouted group is numerically plausible. """ - golden_ref = generate_golden_reference(rows=rows, cols=cols, repeat=repeat) + x = generate_inputs(rows=rows, cols=cols) operator = Repeat( rows=rows, @@ -48,18 +49,7 @@ def test_repeat(rows, cols, repeat, transfer_size, aie_context): context=aie_context, ) - errors, latency_us, bandwidth_gbps = run_test( - operator, - {"input": golden_ref["input"]}, - {"output": golden_ref["output"]}, - rel_tol=0.0, - abs_tol=0.0, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + assert_matches_reference(operator, x, tolerance=Tolerance.exact()) @pytest.mark.parametrize( diff --git a/iron/operators/rms_norm/reference.py b/iron/operators/rms_norm/reference.py index 184ed7da9d..cc823f536b 100644 --- a/iron/operators/rms_norm/reference.py +++ b/iron/operators/rms_norm/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(x, w=None, weighted=False, eps=1e-5): @@ -17,16 +16,11 @@ def reference(x, w=None, weighted=False, eps=1e-5): return out -def generate_golden_reference( - rows: int, cols: int, dtype="bf16", seed=42, weighted=False, eps=1e-5 -): +def generate_inputs(rows: int, cols: int, seed=42, weighted=False): + """The input, and with ``weighted`` the weights too.""" torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(rows, cols, dtype=torch_dtype_map[dtype]) * val_range - if weighted: - weights = torch.rand(cols, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = reference(input_tensor, weights, weighted=True, eps=eps) - return {"input": input_tensor, "weight": weights, "output": output_tensor} - else: - output_tensor = reference(input_tensor, eps=eps) - return {"input": input_tensor, "output": output_tensor} + x = torch.rand(rows, cols, dtype=torch.bfloat16) * val_range + if not weighted: + return (x,) + return x, torch.rand(cols, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/rms_norm/test.py b/iron/operators/rms_norm/test.py index 26b5c7090d..ce8e6c1849 100755 --- a/iron/operators/rms_norm/test.py +++ b/iron/operators/rms_norm/test.py @@ -3,11 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance import aie.utils as aie_utils from iron.operators.rms_norm.op import RMSNorm -from iron.operators.rms_norm.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.rms_norm.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference from iron.common.utils import get_shim_dma_limit @@ -73,9 +74,9 @@ def get_params(): def test_rms_norm( input_length, num_aie_columns, num_channels, tile_size, weighted, aie_context ): - rows = input_length // tile_size - cols = tile_size - golden_ref = generate_golden_reference(rows=rows, cols=cols, weighted=weighted) + inputs = generate_inputs( + rows=input_length // tile_size, cols=tile_size, weighted=weighted + ) operator = RMSNorm( size=input_length, @@ -86,16 +87,8 @@ def test_rms_norm( context=aie_context, ) - input_buffers = {"input1": golden_ref["input"]} - if weighted: - input_buffers["weight"] = golden_ref["weight"] - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 + # The torch reference at this test's tolerance, tighter than the kernel + # contract's. + assert_matches_reference( + operator, *inputs, tolerance=Tolerance.relative(0.04, 1e-6) ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" diff --git a/iron/operators/rope/op.py b/iron/operators/rope/op.py index 27bc702b38..36a1b91c91 100644 --- a/iron/operators/rope/op.py +++ b/iron/operators/rope/op.py @@ -86,14 +86,7 @@ def get_arg_spec(self): ] def reference(self, x, angles): - """CPU reference for RoPE. - - Assumes ``angles`` holds interleaved [cos, sin, cos, sin, ...] pairs - along the last dim (length ``cols``). Only ``method_type == 0`` - (TWO_HALVES) is currently supported. - - ``angles`` may have fewer rows than ``x``; in that case the angles - are tiled along the row dimension to match ``x``.""" + """CPU reference for RoPE; see ``iron.operators.rope.reference``.""" from iron.operators.rope.reference import reference - return reference(x, angles, self.method_type, self.rows, self.cols) + return reference(x, angles, self.method_type) diff --git a/iron/operators/rope/reference.py b/iron/operators/rope/reference.py index 43f46755cf..f81c18709b 100644 --- a/iron/operators/rope/reference.py +++ b/iron/operators/rope/reference.py @@ -2,15 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 import torch -import numpy as np -from ml_dtypes import bfloat16 def compute_rope_params( head_dim, theta_base=10_000, context_length=4096, - method_type=0, freq_config=None, dtype=torch.float32, ): @@ -69,96 +66,46 @@ def compute_rope_params( return cos, sin -def apply_rope(x, cos, sin, method_type=0): - """Apply rotary position embedding to input tensor.""" - if method_type == 0: # For the two-halves method used in HF transformers - # x: (n_heads, seq_len, head_dim) - n_heads, seq_len, head_dim = x.shape - assert head_dim % 2 == 0, "Head dimension must be even" - - # Split x into first half and second half - x1 = x[..., : head_dim // 2] # First half - x2 = x[..., head_dim // 2 :] # Second half - - # Adjust sin and cos shapes - cos = cos[:seq_len, :] # Shape: (seq_len, head_dim / 2) - sin = sin[:seq_len, :] - - # Apply the rotary transformation - x_rotated = torch.empty_like(x) - x_rotated[..., : head_dim // 2] = (x1 * cos) + (-x2 * sin) - x_rotated[..., head_dim // 2 :] = (x2 * cos) + (x1 * sin) - - # It's ok to use lower-precision after applying cos and sin rotation - return x_rotated.to(dtype=x.dtype) - elif method_type == 1: # For the interleaved method used in the Llama paper - # x: (n_heads, seq_len, head_dim) - n_heads, seq_len, head_dim = x.shape - assert head_dim % 2 == 0, "Head dimension must be even" - - # Split x into even and odd columns - x_even = x[..., ::2] # Even columns - x_odd = x[..., 1::2] # Odd columns - - # Adjust sin and cos shapes - cos = cos[:seq_len, :] # Shape: (seq_len, head_dim / 2) - sin = sin[:seq_len, :] - - # Apply the rotary transformation and interleave the even and odd outputs - x_rotated = torch.empty_like(x) - x_rotated[..., ::2] = (x_even * cos) - (x_odd * sin) - x_rotated[..., 1::2] = (x_even * sin) + (x_odd * cos) - - # It's ok to use lower-precision after applying cos and sin rotation - return x_rotated.to(dtype=x.dtype) - else: - raise ValueError("Invalid method_type. Must be 0 or 1.") - - -def reference(x, angles, method_type=0, rows=None, cols=None): +def reference(x, angles, method_type=0): """CPU reference for RoPE from the operator's packed ``angles`` buffer. ``angles`` holds interleaved [cos, sin, cos, sin, ...] pairs along the last - dim (length ``cols``). Only ``method_type == 0`` (TWO_HALVES) is supported - here; the golden-data generator uses :func:`apply_rope`, which additionally - supports the interleaved method and works from the full-precision cos/sin - tables. ``angles`` may have fewer rows than ``x``; in that case each angle - row is repeated for ``rows / angles.shape[0]`` *consecutive* rows of ``x``, - matching the device kernel (design.py's ``core_body`` acquires one angle - row and applies it to that many consecutive input rows before moving on). + dim. ``method_type`` 0 rotates the two halves of each row (HF + transformers); 1 rotates its interleaved even/odd pairs (the Llama paper). + The rotation is computed in fp32 and rounded once. + + ``angles`` may have fewer rows than ``x``; each angle row then applies to + ``rows / angles.shape[0]`` *consecutive* rows of ``x``, matching the device + kernel (design.py's ``core_body`` acquires one angle row and applies it to + that many consecutive input rows before moving on). """ - if method_type != 0: - raise NotImplementedError( - f"RoPE reference only supports method_type=0 (TWO_HALVES), " - f"got {method_type}" + rows = x.shape[0] + if rows % angles.shape[0] != 0: + raise ValueError( + f"{rows} rows cannot share {angles.shape[0]} angle rows evenly" ) - if cols is None: - cols = x.shape[-1] - if rows is None: - rows = x.shape[0] - half = cols // 2 - cos = angles[..., 0::2].to(torch.float32) - sin = angles[..., 1::2].to(torch.float32) - if cos.shape[0] != rows: - if rows % cos.shape[0] == 0: - rep = rows // cos.shape[0] - cos = cos.repeat_interleave(rep, dim=0) - sin = sin.repeat_interleave(rep, dim=0) - else: - cos = cos[:rows] - sin = sin[:rows] + rep = rows // angles.shape[0] + cos = angles[..., 0::2].to(torch.float32).repeat_interleave(rep, dim=0) + sin = angles[..., 1::2].to(torch.float32).repeat_interleave(rep, dim=0) x32 = x.to(torch.float32) - x1, x2 = x32[..., :half], x32[..., half:] - y1 = x1 * cos - x2 * sin - y2 = x2 * cos + x1 * sin - return torch.cat([y1, y2], dim=-1).to(torch.bfloat16) + if method_type == 0: + half = x.shape[-1] // 2 + x1, x2 = x32[..., :half], x32[..., half:] + y = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + elif method_type == 1: + xe, xo = x32[..., 0::2], x32[..., 1::2] + y = torch.empty_like(x32) + y[..., 0::2] = xe * cos - xo * sin + y[..., 1::2] = xe * sin + xo * cos + else: + raise ValueError(f"method_type must be 0 or 1, got {method_type}") + return y.to(torch.bfloat16) -def generate_golden_reference( +def generate_inputs( rows=4096, cols=64, context_len=131072, - method_type=0, rope_theta_base=500000.0, rope_freq_factor=32.0, rope_freq_low_factor=1.0, @@ -166,9 +113,11 @@ def generate_golden_reference( rope_freq_orig_ctx_len=8192, seed=42, ): + """Random input rows and their cos/sin table, laid out as the operator + takes them: ``x`` is ``(rows, cols)`` with a sequence position's heads on + consecutive rows, and the table is one row per position.""" torch.manual_seed(seed) - # Generate golden inputs freq_config = { "factor": rope_freq_factor, "low_freq_factor": rope_freq_low_factor, @@ -179,13 +128,8 @@ def generate_golden_reference( head_dim=cols, theta_base=rope_theta_base, context_length=context_len, - method_type=method_type, freq_config=freq_config, ) - # The operator is handed the tables in bf16, so the golden output is computed - # from those rounded values rather than from the fp32 ones. - cos = cos.to(torch.bfloat16).to(torch.float32) - sin = sin.to(torch.bfloat16).to(torch.float32) val_range = 4 # Head count is inferred from rows and context_len. This logic assumes rows is either # smaller than context_len (1 head, seq_len == rows) or an exact multiple of context_len @@ -196,18 +140,11 @@ def generate_golden_reference( ) n_heads = rows // context_len if context_len < rows else 1 seq_len = rows // n_heads - A = torch.rand(n_heads, seq_len, cols, dtype=torch.bfloat16) * val_range - - # Create the lut by interleaving cos and sin - B = torch.zeros((seq_len, cols), dtype=torch.bfloat16) - B[:, ::2] = cos[:seq_len, : cols // 2] - B[:, 1::2] = sin[:seq_len, : cols // 2] + x = torch.rand(n_heads, seq_len, cols, dtype=torch.bfloat16) * val_range - # Generate golden outputs - C = apply_rope(A, cos, sin, method_type) + # The lut interleaves cos and sin + angles = torch.zeros((seq_len, cols), dtype=torch.bfloat16) + angles[:, ::2] = cos[:seq_len, : cols // 2] + angles[:, 1::2] = sin[:seq_len, : cols // 2] - return { - "A": A, - "B": B, - "C": C, - } + return x.transpose(0, 1).reshape(rows, cols).contiguous(), angles diff --git a/iron/operators/rope/test.py b/iron/operators/rope/test.py index 2dc8c5dcf9..9ccf13e642 100755 --- a/iron/operators/rope/test.py +++ b/iron/operators/rope/test.py @@ -6,8 +6,8 @@ import aie.utils as aie_utils from aie.utils.verify import Tolerance from iron.operators.rope.op import RoPE -from iron.operators.rope.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.rope.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_params(): @@ -62,9 +62,7 @@ def get_params(): get_params(), ) def test_rope(rows, cols, angle_rows, aie_columns, method_type, aie_context): - golden_ref = generate_golden_reference( - rows=rows, cols=cols, context_len=angle_rows, method_type=method_type - ) + x, angles = generate_inputs(rows=rows, cols=cols, context_len=angle_rows) operator = RoPE( rows=rows, @@ -75,25 +73,12 @@ def test_rope(rows, cols, angle_rows, aie_columns, method_type, aie_context): context=aie_context, ) - # golden reference produces tensors of shape (n_heads, seq_len, cols); - # NPU design expects (seq_len, n_heads, cols), so we transpose inputs/outputs - input_buffers = { - "in": golden_ref["A"].transpose(0, 1).contiguous(), - "angles": golden_ref["B"], - } - output_buffers = {"output": golden_ref["C"].transpose(0, 1).contiguous()} - - errors, latency_us, bandwidth_gbps = run_test( + assert_matches_reference( operator, - input_buffers, - output_buffers, + x, + angles, # The tighter of this test's former rel_tol and the kernel contract's # atol (none): an output that cancels to near zero is judged # relatively like any other. tolerance=Tolerance.relative(0.05), ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" diff --git a/iron/operators/sigmoid/op.py b/iron/operators/sigmoid/op.py index ea4005afd6..cc0a302822 100644 --- a/iron/operators/sigmoid/op.py +++ b/iron/operators/sigmoid/op.py @@ -17,3 +17,8 @@ class Sigmoid(ChanneledUnaryOperator): def _kernel(self): return activation.sigmoid(self._line_size) + + def reference(self, x): + from iron.operators.sigmoid.reference import reference + + return reference(x) diff --git a/iron/operators/sigmoid/reference.py b/iron/operators/sigmoid/reference.py index 753ee806c6..0c597a8d76 100644 --- a/iron/operators/sigmoid/reference.py +++ b/iron/operators/sigmoid/reference.py @@ -2,12 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def reference(x): + """CPU reference: sigmoid (ground truth).""" + return torch.sigmoid(x) + + +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(input_length, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = torch.sigmoid(input_tensor) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(input_length, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/sigmoid/test.py b/iron/operators/sigmoid/test.py index d3723a4a57..00deec81bd 100755 --- a/iron/operators/sigmoid/test.py +++ b/iron/operators/sigmoid/test.py @@ -3,10 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance from iron.operators.sigmoid.op import Sigmoid -from iron.operators.sigmoid.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.sigmoid.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -27,7 +31,7 @@ def get_params(): get_params(), ) def test_sigmoid(input_length, num_aie_columns, num_channels, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + x = generate_inputs(input_length=input_length) operator = Sigmoid( size=input_length, @@ -37,14 +41,6 @@ def test_sigmoid(input_length, num_aie_columns, num_channels, tile_size, aie_con context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # The torch reference at this test's tolerance, tighter than the kernel + # contract's. + assert_matches_reference(operator, x, tolerance=Tolerance.relative(0.04, 1e-6)) diff --git a/iron/operators/silu/reference.py b/iron/operators/silu/reference.py index 87b78140bb..0ab1caf82b 100644 --- a/iron/operators/silu/reference.py +++ b/iron/operators/silu/reference.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(x): @@ -10,9 +9,7 @@ def reference(x): return torch.nn.functional.silu(x) -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(input_length, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = reference(input_tensor) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(input_length, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/silu/test.py b/iron/operators/silu/test.py index bb989315bc..2d526f2248 100755 --- a/iron/operators/silu/test.py +++ b/iron/operators/silu/test.py @@ -3,10 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance from iron.operators.silu.op import SiLU -from iron.operators.silu.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.silu.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -27,7 +31,7 @@ def get_params(): get_params(), ) def test_silu(input_length, num_aie_columns, num_channels, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + x = generate_inputs(input_length=input_length) operator = SiLU( size=input_length, @@ -36,14 +40,6 @@ def test_silu(input_length, num_aie_columns, num_channels, tile_size, aie_contex context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # The torch reference at this test's tolerance, tighter than the kernel + # contract's. + assert_matches_reference(operator, x, tolerance=Tolerance.relative(0.04, 1e-6)) diff --git a/iron/operators/softmax/reference.py b/iron/operators/softmax/reference.py index 6e5660e7e2..3a8470b0d6 100644 --- a/iron/operators/softmax/reference.py +++ b/iron/operators/softmax/reference.py @@ -1,10 +1,9 @@ # SPDX-FileCopyrightText: Copyright (C) 2025 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Golden reference generator for softmax operator.""" +"""Reference for the softmax operator.""" import torch -from iron.common.test_utils import torch_dtype_map def reference(x): @@ -12,15 +11,7 @@ def reference(x): return torch.softmax(x, dim=-1) -def generate_golden_reference(rows: int, cols: int, dtype="bf16", seed=42): - """ - Generate golden reference data for softmax. - - Returns: - dict: Dictionary with tensors for inputs and outputs - """ +def generate_inputs(rows: int, cols: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(rows, cols, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = reference(input_tensor) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(rows, cols, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/softmax/test.py b/iron/operators/softmax/test.py index 066d230932..7b9644152a 100755 --- a/iron/operators/softmax/test.py +++ b/iron/operators/softmax/test.py @@ -3,11 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance import aie.utils as aie_utils from iron.operators.softmax.op import Softmax -from iron.operators.softmax.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.softmax.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_optimal_columns_channels(input_length, tile_size, max_columns): @@ -60,11 +61,9 @@ def get_params(): get_params(), ) def test_softmax(input_length, num_aie_columns, num_channels, tile_size, aie_context): - rows = input_length // tile_size cols = tile_size - - golden_ref = generate_golden_reference(rows=rows, cols=cols) + x = generate_inputs(rows=rows, cols=cols) operator = Softmax( rows=rows, @@ -74,14 +73,6 @@ def test_softmax(input_length, num_aie_columns, num_channels, tile_size, aie_con context=aie_context, ) - input_buffers = {"in": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # The torch reference at this test's tolerance, tighter than the kernel + # contract's. + assert_matches_reference(operator, x, tolerance=Tolerance.relative(0.04, 1e-6)) diff --git a/iron/operators/strided_copy/reference.py b/iron/operators/strided_copy/reference.py index 2f02878456..623071f3bb 100644 --- a/iron/operators/strided_copy/reference.py +++ b/iron/operators/strided_copy/reference.py @@ -4,8 +4,6 @@ import numpy as np import torch -from iron.common.test_utils import torch_dtype_map - def _pad_to_4d(sizes, strides): """design.py pads access patterns to 4D before building the taps; the reference @@ -89,14 +87,11 @@ def generate_golden_reference( num_aie_channels=1, input_offset_addend=0, output_offset_addend=0, - dtype="bf16", seed=42, ): torch.manual_seed(seed) val_range = 4 - input_tensor = ( - torch.rand(int(input_buffer_size), dtype=torch_dtype_map[dtype]) * val_range - ) + input_tensor = torch.rand(int(input_buffer_size), dtype=torch.bfloat16) * val_range output_tensor = reference( input_tensor, input_sizes, diff --git a/iron/operators/tanh/op.py b/iron/operators/tanh/op.py index 541303472f..931560c783 100644 --- a/iron/operators/tanh/op.py +++ b/iron/operators/tanh/op.py @@ -17,3 +17,8 @@ class Tanh(ChanneledUnaryOperator): def _kernel(self): return activation.tanh(self._line_size) + + def reference(self, x): + from iron.operators.tanh.reference import reference + + return reference(x) diff --git a/iron/operators/tanh/reference.py b/iron/operators/tanh/reference.py index 17e591ca61..5d68afd289 100644 --- a/iron/operators/tanh/reference.py +++ b/iron/operators/tanh/reference.py @@ -2,12 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map -def generate_golden_reference(input_length: int, dtype="bf16", seed=42): +def reference(x): + """CPU reference: tanh (ground truth).""" + return torch.tanh(x) + + +def generate_inputs(input_length: int, seed=42): torch.manual_seed(seed) val_range = 4 - input_tensor = torch.rand(input_length, dtype=torch_dtype_map[dtype]) * val_range - output_tensor = torch.tanh(input_tensor) - return {"input": input_tensor, "output": output_tensor} + return torch.rand(input_length, dtype=torch.bfloat16) * val_range diff --git a/iron/operators/tanh/test.py b/iron/operators/tanh/test.py index 6337b5484a..18ba30a123 100755 --- a/iron/operators/tanh/test.py +++ b/iron/operators/tanh/test.py @@ -3,10 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance from iron.operators.tanh.op import Tanh -from iron.operators.tanh.reference import generate_golden_reference -from iron.common.test_utils import run_test, make_channeled_unary_params +from iron.operators.tanh.reference import generate_inputs +from iron.common.test_utils import ( + assert_matches_reference, + make_channeled_unary_params, +) def get_params(): @@ -27,7 +31,7 @@ def get_params(): get_params(), ) def test_tanh(input_length, num_aie_columns, num_channels, tile_size, aie_context): - golden_ref = generate_golden_reference(input_length=input_length) + x = generate_inputs(input_length=input_length) operator = Tanh( size=input_length, @@ -37,14 +41,6 @@ def test_tanh(input_length, num_aie_columns, num_channels, tile_size, aie_contex context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - operator, input_buffers, output_buffers, rel_tol=0.04, abs_tol=1e-6 - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # The torch reference at this test's tolerance, tighter than the kernel + # contract's. + assert_matches_reference(operator, x, tolerance=Tolerance.relative(0.04, 1e-6)) diff --git a/iron/operators/transpose/op.py b/iron/operators/transpose/op.py index 56d4e2fa38..86fadf02bb 100644 --- a/iron/operators/transpose/op.py +++ b/iron/operators/transpose/op.py @@ -114,7 +114,7 @@ def get_arg_spec(self): ] def reference(self, x): - """CPU reference: 2D transpose of an (M, N) matrix stored row-major.""" + """CPU reference: transpose of each (M, N) matrix, stored row-major.""" from iron.operators.transpose.reference import reference - return reference(x.reshape(self.M, self.N)) + return reference(x.reshape(-1, self.M, self.N)) diff --git a/iron/operators/transpose/reference.py b/iron/operators/transpose/reference.py index 86e9c24ffe..3fba73eab6 100644 --- a/iron/operators/transpose/reference.py +++ b/iron/operators/transpose/reference.py @@ -2,28 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 import torch -from iron.common.test_utils import torch_dtype_map def reference(x): - """CPU reference: 2D transpose of an ``(rows, cols)`` matrix (ground truth).""" - return torch.transpose(x, 0, 1) + """CPU reference: transpose of the last two dims (ground truth), so each of + a batch of ``(rows, cols)`` matrices is transposed on its own.""" + return x.transpose(-2, -1) -def generate_golden_reference( - rows: int, cols: int, dtype="bf16", seed=42, num_batches=1 -): +def generate_inputs(rows: int, cols: int, seed=42, num_batches=1): + """``num_batches`` independent ``(rows, cols)`` matrices laid back-to-back; + the batch dim is dropped when there is one.""" torch.manual_seed(seed) val_range = 4 - # num_batches>1: B independent (rows,cols) matrices laid back-to-back; each is - # transposed independently and the results concatenated in the same order. - input_tensor = ( - torch.rand(num_batches, rows, cols, dtype=torch_dtype_map[dtype]) * val_range - ) - output_tensor = torch.stack( - [reference(input_tensor[b]) for b in range(num_batches)] - ) - # drop batch dimension if num_batches == 1 - input_tensor = torch.squeeze(input_tensor, 0) - output_tensor = torch.squeeze(output_tensor, 0) - return {"input": input_tensor, "output": output_tensor} + x = torch.rand(num_batches, rows, cols, dtype=torch.bfloat16) * val_range + return torch.squeeze(x, 0) diff --git a/iron/operators/transpose/test.py b/iron/operators/transpose/test.py index aebe4814b1..1d81debc7f 100755 --- a/iron/operators/transpose/test.py +++ b/iron/operators/transpose/test.py @@ -3,11 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from aie.utils.verify import Tolerance import aie.utils as aie_utils from iron.operators.transpose.op import Transpose -from iron.operators.transpose.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.operators.transpose.reference import generate_inputs +from iron.common.test_utils import assert_matches_reference def get_params(): @@ -79,7 +80,7 @@ def get_params(): ) @pytest.mark.parametrize("M,N,aie_columns,channels,m,n,s,num_batches", get_params()) def test_transpose(M, N, aie_columns, channels, m, n, s, num_batches, aie_context): - golden_ref = generate_golden_reference(rows=M, cols=N, num_batches=num_batches) + x = generate_inputs(rows=M, cols=N, num_batches=num_batches) operator = Transpose( M=M, @@ -93,23 +94,9 @@ def test_transpose(M, N, aie_columns, channels, m, n, s, num_batches, aie_contex context=aie_context, ) - input_buffers = {"input": golden_ref["input"]} - output_buffers = {"output": golden_ref["output"]} - - errors, latency_us, bandwidth_gbps = run_test( - # A transpose is a permutation. Any tolerance here also accepts some class of - # wrong permutation, so gate it exactly. - operator, - input_buffers, - output_buffers, - rel_tol=0.0, - abs_tol=0.0, - ) - - print(f"\nLatency (us): {latency_us:.1f}") - print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") - - assert not errors, f"Test failed with errors: {errors}" + # A transpose is a permutation. Any tolerance here also accepts some class of + # wrong permutation, so gate it exactly. + assert_matches_reference(operator, x, tolerance=Tolerance.exact()) # Shapes whose M*N is divisible by every factor while one per-dimension quotient floors diff --git a/iron/tests/operators/rope_reference_convention.py b/iron/tests/operators/rope_reference_convention.py index e199f915a3..6731d42eb7 100644 --- a/iron/tests/operators/rope_reference_convention.py +++ b/iron/tests/operators/rope_reference_convention.py @@ -50,7 +50,7 @@ def test_reference_matches_device_convention_for_batched_angle_rows(): rows, angle_rows = 6, 3 x, angles = _make_inputs(rows, angle_rows) expected = _block_major_expected(x, angles, rows, angle_rows) - got = reference(x, angles, rows=rows, cols=x.shape[-1]) + got = reference(x, angles) assert torch.equal(expected, got) @@ -58,7 +58,7 @@ def test_reference_matches_device_convention_across_shapes(): for rows, angle_rows in [(8, 2), (1024, 1), (4, 4), (13, 13), (12, 4)]: x, angles = _make_inputs(rows, angle_rows) expected = _block_major_expected(x, angles, rows, angle_rows) - got = reference(x, angles, rows=rows, cols=x.shape[-1]) + got = reference(x, angles) assert torch.equal( expected, got ), f"mismatch at rows={rows} angle_rows={angle_rows}" From 44a5667358b1ca820d518cdddd9a7ddfeb57ea7e Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 10:49:17 -0600 Subject: [PATCH 16/22] Sequences: judge compare mode by each step's kernel contract dispatch="compare" used one global rule for every step: flag the step if both max_abs and max_rel exceeded fixed thresholds. It now judges each step with aie.utils.verify.compare under that step's op.reference_tolerance(), the same contract its operator test holds it to. CompareDispatch(tolerance=...) overrides it for every step. A step whose operator declares no element-wise tolerance (none, a bound, or one with range_frac) falls back to relative(0.025, 1e-2) per element, which is stricter than the old max_abs-AND-max_rel rule. The logged stats are unchanged; a mismatch now reports the verdict and the tolerance. The infrastructure test checks both directions on real hardware: tanh passes compare mode under its contract and is flagged under exact(). Co-Authored-By: Claude --- iron/common/sequence.py | 47 +++++++++++++------ iron/tests/infrastructure/sequence.py | 67 +++++++++++++-------------- 2 files changed, 66 insertions(+), 48 deletions(-) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 7c449e79ad..69352f53e1 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -13,6 +13,7 @@ from aie.iron.device import NPU2 from aie.utils.hostruntime.tensor_class import CPUOnlyTensor from aie.utils.npukernel import NPUKernel +from aie.utils.verify import Tolerance, compare try: import pyxrt @@ -245,20 +246,33 @@ class CompareDispatch(SeparateDispatch): per-step deviation. Args: - rel_tol / abs_tol: Per-step tolerances; a step counts as a mismatch - only when it exceeds both. + tolerance: How close every step's output must come to its reference. + By default each step is held to its operator's + ``reference_tolerance()``, the contract of the kernel it runs, as + its own test holds it; a step without one that can be judged + element by element falls back to ``FALLBACK_TOLERANCE``. raise_on_mismatch: When True (default), raise ``RuntimeError`` on the first mismatching step instead of only logging it. """ name = "compare" - def __init__(self, rel_tol=0.05, abs_tol=1e-2, raise_on_mismatch=True): + FALLBACK_TOLERANCE = Tolerance.relative(0.025, 1e-2) + + def __init__(self, tolerance=None, raise_on_mismatch=True): super().__init__() - self.rel_tol = rel_tol - self.abs_tol = abs_tol + self.tolerance = tolerance self.raise_on_mismatch = raise_on_mismatch + def step_tolerance(self, op): + """The tolerance ``op``'s step is judged by.""" + if self.tolerance is not None: + return self.tolerance + tol = op.reference_tolerance() if isinstance(op, MLIROperator) else None + if tol is None or tol.kind == "bound" or tol.range_frac is not None: + return self.FALLBACK_TOLERANCE + return tol + def make_callable(self, seq): return SequenceCompareCallable(seq, self) @@ -305,7 +319,7 @@ class OperatorSequence(AIEOperatorBase): runs the ``"separate"`` xclbin path and, after each NPU step, also runs the operator's CPU reference on the NPU-produced inputs and logs the deviation for testing/debugging. Pass a - :class:`CompareDispatch` instance to tune the compare tolerances. + :class:`CompareDispatch` instance to set the compare tolerance. """ def __init__( @@ -855,8 +869,7 @@ class SequenceCompareCallable(SequenceXclbinCallable): def __init__(self, op, dispatch): super().__init__(op, dispatch) - self.rel_tol = dispatch.rel_tol - self.abs_tol = dispatch.abs_tol + self.dispatch = dispatch self.raise_on_mismatch = dispatch.raise_on_mismatch self.last_step_stats = [] @@ -882,7 +895,8 @@ def _run_step(self, step_idx, kernel, args, step): kernel(*args) torch = _torch() - npu_out = self._read_to_cpu(out_name, out_spec).to(torch.float32) + npu_raw = self._read_to_cpu(out_name, out_spec) + npu_out = npu_raw.to(torch.float32) ref_out = step_op.reference(*cpu_inputs) stats = { @@ -907,7 +921,13 @@ def _run_step(self, step_idx, kernel, args, step): max_rel=rel, ref_max=ref_max, ) - fail = (max_abs > self.abs_tol) and (rel > self.rel_tol) + tol = self.dispatch.step_tolerance(step_op) + if npu_raw.dtype == torch.bfloat16: + npu_np = npu_raw.view(torch.uint16).numpy().view(ml_dtypes.bfloat16) + else: + npu_np = npu_raw.numpy() + verdict = compare(npu_np, ref_flat.numpy(), tol) + fail = not verdict stats["mismatch"] = fail level = logging.ERROR if fail else logging.INFO logger.log( @@ -920,14 +940,13 @@ def _run_step(self, step_idx, kernel, args, step): mean_abs, rel, ref_max, - " MISMATCH" if fail else "", + f" MISMATCH: {verdict.detail}" if fail else "", ) if fail and self.raise_on_mismatch: raise RuntimeError( f"[compare step {step_idx}] {stats['op']} (name={stats['op_name']}) " f"-> {out_name}: NPU output deviates from reference " - f"(max_abs={max_abs:.4g}, max_rel={rel:.4g}, " - f"ref_max={ref_max:.4g}; inputs={list(in_names)}; " - f"tolerances abs_tol={self.abs_tol}, rel_tol={self.rel_tol})" + f"({verdict.detail}; max_abs={max_abs:.4g}, max_rel={rel:.4g}, " + f"ref_max={ref_max:.4g}; inputs={list(in_names)}; tolerance {tol})" ) self.last_step_stats.append(stats) diff --git a/iron/tests/infrastructure/sequence.py b/iron/tests/infrastructure/sequence.py index 3bdeac3cb8..22b3418424 100644 --- a/iron/tests/infrastructure/sequence.py +++ b/iron/tests/infrastructure/sequence.py @@ -27,12 +27,14 @@ import aie.utils as aie_utils from aie.iron.device import NPU2 +from aie.utils.verify import Tolerance -from iron.common.sequence import OperatorSequence +from iron.common.sequence import CompareDispatch, OperatorSequence from iron.common.compilation.sequence import fuse_mlir from iron.common.test_utils import verify_buffer from iron.operators.elementwise_add.op import ElementwiseAdd from iron.operators.relu.op import ReLU +from iron.operators.tanh.op import Tanh def _set_input(run, name, data): @@ -208,7 +210,7 @@ def test_dispatch_modes_bit_identical(dispatch, aie_context): # rather than a hand-rolled numpy view. Not covered by # test_dispatch_modes_bit_identical above, since reference() is a CPU # re-implementation and only expected to match the NPU output within -# tolerance, not bit-for-bit (see CompareDispatch's rel_tol/abs_tol). +# tolerance, not bit-for-bit (see CompareDispatch's tolerance). # --------------------------------------------------------------------------- _SLICE_SIZE = 1024 @@ -270,57 +272,54 @@ def test_reference_dispatch_resolves_sliced_buffer(aie_context): # --------------------------------------------------------------------------- -# 4. Compare mode flags (and by default raises on) a per-step reference/NPU -# mismatch on its own. +# 4. Compare mode holds each step to its kernel's contract, and flags (and by +# default raises on) a step that falls outside the tolerance it is judged by. # -# Normally the reference is trusted and the NPU kernel is the suspect; here -# we invert that (keep the NPU correct, vary the reference) because it is -# easier to inject a known-wrong reference than a known-wrong kernel. +# Tanh's kernel approximates torch.tanh: within its contract, but not +# bit-exact. So the same NPU output must pass under the default tolerance +# and fail under an exact one. # --------------------------------------------------------------------------- -@pytest.mark.parametrize("reference_is_correct", [True, False]) -def test_compare_mode_detects_wrong_reference(reference_is_correct, aie_context): +@pytest.mark.parametrize("exact", [False, True]) +def test_compare_mode_judges_each_step_by_its_tolerance(exact, aie_context): """dispatch="compare" runs the NPU pipeline and, per step, re-runs the - operator's ``reference()`` on the same NPU inputs. A correct reference must - run cleanly (no flagged step); a wrong one must make compare mode raise on - its own (``compare_raise_on_mismatch`` defaults to True).""" - size = 256 + operator's ``reference()`` on the same NPU inputs. Under its kernel's + contract the step runs cleanly (no flagged step); held to exact equality it + makes compare mode raise on its own (``raise_on_mismatch`` defaults to + True).""" + size = 1024 torch.manual_seed(0) - a = torch.rand(size, dtype=torch.bfloat16) - b = torch.rand(size, dtype=torch.bfloat16) + x = torch.rand(size, dtype=torch.bfloat16) * 4 - op = ElementwiseAdd( - size=size, tile_size=256, num_aie_columns=1, context=aie_context + op = Tanh( + size=size, + num_aie_columns=1, + num_channels=1, + tile_size=size, + context=aie_context, ) - if not reference_is_correct: - # Override the reference on this instance to disagree with the NPU - # kernel (which computes a + b). Keeping the real ElementwiseAdd class - # leaves its name/compilation intact for the xclbin compare path. - op.reference = lambda a, b: a + b + 1.0 - seq = OperatorSequence( - name="infra_compare_add", - runlist=[(op, "a", "b", "out")], - input_args=["a", "b"], + name="infra_compare_tanh", + runlist=[(op, "x", "out")], + input_args=["x"], output_args=["out"], - dispatch="compare", + dispatch=CompareDispatch(tolerance=Tolerance.exact() if exact else None), context=aie_context, ) seq.compile() assert seq._dispatch.name == "compare" run = seq.get_callable() - _set_input(run, "a", a) - _set_input(run, "b", b) + _set_input(run, "x", x) - if reference_is_correct: + if exact: + with pytest.raises(RuntimeError, match="deviates from reference"): + run() + else: run() # must not raise flagged = any(step.get("mismatch") for step in run.last_step_stats) - assert not flagged, "compare mode should not flag a matching reference" - else: - with pytest.raises(RuntimeError): - run() # compare mode reports the wrong reference by itself + assert not flagged, "compare mode flagged a step within its kernel contract" # --------------------------------------------------------------------------- From f2f0bc171f58f5743df333f42a1e32f26518acb9 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 10:49:24 -0600 Subject: [PATCH 17/22] Tests and docs: small cleanups - gemm test: drop trace_size, which every case set to 0 and nothing read. This changes the gemm test IDs. - lazy_imports: check in a fresh interpreter. In the session's own, any operator collected earlier is already imported, so the check depended on test order. iron is installed into the environment, so no path setup. - AGENTS.md: replace the torch_to_numpy/numpy_to_torch section (neither exists) with the bf16 view and DEFAULT_TENSOR_CLASS.from_torch/to_torch, and describe assert_matches_reference and the one-reference layout. Co-Authored-By: Claude --- AGENTS.md | 64 +++++++++++++---------- iron/operators/gemm/test.py | 58 ++++++++++---------- iron/tests/infrastructure/lazy_imports.py | 21 ++++++-- 3 files changed, 82 insertions(+), 61 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 6794290f29..6b7ada346b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -126,8 +126,10 @@ reuse lint - Each operator directory contains: - `op.py`: Python interface (inherits from `MLIROperator`) - defines operator parameters, compilation artifacts, and runtime argument specs - `design.py`: NPU implementation using MLIR-AIE Python API - defines ObjectFIFOs, Workers, and Runtime sequences - - `reference.py`: CPU reference implementation for validation - - `test.py`: End-to-end test (build, run, verify against reference) + - `reference.py`: `reference()`, the CPU ground truth the NPU output is + judged against (exposed as the operator's `reference()` method), and + `generate_inputs()`, the test's random inputs + - `test.py`: End-to-end test (build, run once, check against `reference()`) 2. **AIE Kernels** ([mlir-aie `aie_kernels/`](https://github.com/Xilinx/mlir-aie/tree/main/aie_kernels)) - Architecture-specific C++ compute kernels, sourced from the installed @@ -146,8 +148,8 @@ reuse lint - `fusion.py`: Operator sequencing framework (`OperatorSequence`) - `device_manager.py`: XRT device initialization and management (singleton pattern) - `context.py`: `AIEContext` for operator compilation/execution - - `utils.py`: Helper functions (`torch_to_numpy`, `numpy_to_torch`) - - `test_utils.py`: Test utilities (`verify_buffer`, a wrapper over mlir-aie's `aie.utils.verify.compare`; `run_test`, timed with `aie.utils.benchmark.run_iters`) + - `utils.py`: Helper functions (`float_to_name`, `get_shim_dma_limit`, `split_run`) + - `test_utils.py`: Test utilities (`assert_matches_reference`, the one-call operator check; `verify_buffer`, a wrapper over mlir-aie's `aie.utils.verify.compare`; `run_test`, timed with `aie.utils.benchmark.run_iters`) ### Key Concepts @@ -266,10 +268,12 @@ Data movement pattern: L3 → Shim DMA → L2 → L1 (tile local) → Compute - Choose appropriate directory: `generic/`, `aie2/`, or `aie2p/` - Use AIE API for portable vectorization when possible - Add `event0()` and `event1()` for performance profiling -5. Implement `reference.py` with CPU reference +5. Implement `reference.py` with the CPU reference and `generate_inputs()`, + and a `reference()` method on the operator that calls it 6. Implement `test.py` with pytest tests - Use `@pytest.mark.extensive` for slower/larger tests - - Use `verify_buffer()` from `iron.common.test_utils` + - Check the output with `assert_matches_reference()` from + `iron.common.test_utils` 7. Register operator in `iron/operators/__init__.py` ## Operator Sequences @@ -367,33 +371,38 @@ void my_kernel(bfloat16* in, bfloat16* out, int32_t size) { ### Test Verification Pattern ```python -from iron.common.test_utils import verify_buffer - -# Compare NPU output against CPU reference -errors = verify_buffer( - output=npu_output, - buf_name="output", - reference=cpu_reference, - rel_tol=0.04, # 4% relative tolerance - abs_tol=1e-6, # Absolute tolerance for small values - max_error_rate=0.0 # 0% of elements can fail (strict) -) -assert len(errors) == 0, f"Found {len(errors)} mismatches" +from aie.utils.verify import Tolerance +from iron.common.test_utils import assert_matches_reference + +x = generate_inputs(input_length=2048) +op = Tanh(size=2048, num_aie_columns=1, num_channels=1, tile_size=2048) + +# Dispatch once and compare with op.reference(x), under the declared +# tolerance contract of the kernel the operator runs +# (op.reference_tolerance()) ... +assert_matches_reference(op, x) + +# ... or under an explicit one, e.g. exact for pure data movement. +assert_matches_reference(op, x, tolerance=Tolerance.relative(0.04, 1e-6)) ``` -### Datatype Conversion Helpers +`verify_buffer()` compares a single buffer the same way, for tests that +dispatch by hand. -```python -from iron.common.utils import torch_to_numpy, numpy_to_torch +### bfloat16 between torch and numpy + +numpy has no bfloat16 of its own; use `ml_dtypes.bfloat16` and move the bits, +never going through float32: -# Convert torch tensor to numpy (preserves bfloat16) -np_array = torch_to_numpy(torch_tensor) +```python +import ml_dtypes, torch -# Convert numpy array to torch (preserves bfloat16) -torch_tensor = numpy_to_torch(np_array) +np_array = torch_tensor.view(torch.uint16).numpy().view(ml_dtypes.bfloat16) +torch_tensor = torch.from_numpy(np_array.view("uint16")).view(torch.bfloat16) ``` -These utilities handle bfloat16 conversion correctly (avoiding float32 intermediate). +Runtime tensors take and return torch tensors directly +(`aie.utils.DEFAULT_TENSOR_CLASS.from_torch()`, `.to_torch()`). ## Debugging and Performance @@ -479,7 +488,8 @@ logging.basicConfig(level=logging.DEBUG) - Check datatype consistency (bfloat16 has limited precision) - Verify reference implementation matches NPU kernel exactly - Look for memory alignment issues in C++ kernel -- Adjust tolerances in `verify_buffer()` if needed (`rel_tol`, `abs_tol`) +- Check which tolerance the test judges by: the kernel's contract + (`op.reference_tolerance()`) unless the test passes `tolerance=` **Dimension mismatch errors** diff --git a/iron/operators/gemm/test.py b/iron/operators/gemm/test.py index 7bc92691be..6e787d5944 100755 --- a/iron/operators/gemm/test.py +++ b/iron/operators/gemm/test.py @@ -20,39 +20,39 @@ def get_params(): max_aie_columns = dev.cols device_type = dev.resolve().name # fmt: off - # M, K, N, num_aie_columns, b_col_maj, c_col_maj, m, k, n, trace_size, partition_N + # M, K, N, num_aie_columns, b_col_maj, c_col_maj, m, k, n, partition_N regular_params = [ - (2048, 2048, 2048, 1, False, False, 64, 64, 64, 0, 1), - (2048, 2048, 2048, 2, True, False, 64, 64, 64, 0, 1), - (2048, 2048, 2048, 8, True, True, 64, 64, 64, 0, 1), - ( 384, 1536, 1792, 4, True, False, 32, 48, 64, 0, 1), - (1792, 896, 1152, 8, False, True, 64, 32, 48, 0, 1), - ( 896, 1792, 640, 8, False, True, 32, 64, 80, 0, 1), - ( 192, 384, 64, 4, False, False, 48, 96, 16, 0, 1), - ( 192, 384, 64, 4, True, True, 48, 96, 16, 0, 1), - ( 64, 512, 256, 4, True, False, 16, 64, 64, 0, 4), + (2048, 2048, 2048, 1, False, False, 64, 64, 64, 1), + (2048, 2048, 2048, 2, True, False, 64, 64, 64, 1), + (2048, 2048, 2048, 8, True, True, 64, 64, 64, 1), + ( 384, 1536, 1792, 4, True, False, 32, 48, 64, 1), + (1792, 896, 1152, 8, False, True, 64, 32, 48, 1), + ( 896, 1792, 640, 8, False, True, 32, 64, 80, 1), + ( 192, 384, 64, 4, False, False, 48, 96, 16, 1), + ( 192, 384, 64, 4, True, True, 48, 96, 16, 1), + ( 64, 512, 256, 4, True, False, 16, 64, 64, 4), ] extensive_params = [ - (2048, 2048, 2048, 8, False, False, 32, 32, 128, 0, 1), - (2048, 2048, 8192, 2, False, False, 64, 64, 64, 0, 1), - (2048, 8192, 2048, 2, False, False, 64, 64, 64, 0, 1), - (2048, 64, 2048, 2, False, False, 64, 64, 64, 0, 1), - (2048, 64, 8192, 2, False, False, 64, 64, 64, 0, 1), - (2048, 2048, 2048, 8, True, False, 128, 32, 32, 0, 1), - (2048, 2048, 8192, 2, True, False, 64, 64, 64, 0, 1), - (2048, 8192, 2048, 2, True, False, 64, 64, 64, 0, 1), - (2048, 64, 2048, 2, True, False, 64, 64, 64, 0, 1), - (2048, 64, 8192, 2, True, False, 64, 64, 64, 0, 1), - (2048, 2048, 2048, 2, False, True, 8, 16, 32, 0, 1), - (2048, 2048, 8192, 2, False, True, 64, 64, 64, 0, 1), - (2048, 8192, 2048, 2, False, True, 64, 64, 64, 0, 1), - (2048, 64, 2048, 2, False, True, 64, 64, 64, 0, 1), - (2048, 64, 8192, 2, False, True, 64, 64, 64, 0, 1), + (2048, 2048, 2048, 8, False, False, 32, 32, 128, 1), + (2048, 2048, 8192, 2, False, False, 64, 64, 64, 1), + (2048, 8192, 2048, 2, False, False, 64, 64, 64, 1), + (2048, 64, 2048, 2, False, False, 64, 64, 64, 1), + (2048, 64, 8192, 2, False, False, 64, 64, 64, 1), + (2048, 2048, 2048, 8, True, False, 128, 32, 32, 1), + (2048, 2048, 8192, 2, True, False, 64, 64, 64, 1), + (2048, 8192, 2048, 2, True, False, 64, 64, 64, 1), + (2048, 64, 2048, 2, True, False, 64, 64, 64, 1), + (2048, 64, 8192, 2, True, False, 64, 64, 64, 1), + (2048, 2048, 2048, 2, False, True, 8, 16, 32, 1), + (2048, 2048, 8192, 2, False, True, 64, 64, 64, 1), + (2048, 8192, 2048, 2, False, True, 64, 64, 64, 1), + (2048, 64, 2048, 2, False, True, 64, 64, 64, 1), + (2048, 64, 8192, 2, False, True, 64, 64, 64, 1), # N wide enough that C's row stride (mem_tile_m_C * N) overflows the # shim BD's 20-bit iteration step, so the drain is issued as one # descriptor per row-block. Cover for that split. - (1024, 2560, 10240, 8, False, False, 64, 64, 64, 0, 1), - (2048, 2560, 10240, 8, False, False, 64, 64, 64, 0, 1), + (1024, 2560, 10240, 8, False, False, 64, 64, 64, 1), + (2048, 2560, 10240, 8, False, False, 64, 64, 64, 1), ] # fmt: on @@ -71,7 +71,6 @@ def add_params(param_list, is_extensive): m, k, n, - trace_size, partition_N, ) = p @@ -99,7 +98,7 @@ def add_params(param_list, is_extensive): Throughput=r"Throughput: (?P[\d\.e\+-]+) GFLOP/s", ) @pytest.mark.parametrize( - "M,K,N,num_aie_columns,b_col_maj,c_col_maj,m,k,n,trace_size,partition_N", + "M,K,N,num_aie_columns,b_col_maj,c_col_maj,m,k,n,partition_N", get_params(), ) def test_gemm( @@ -112,7 +111,6 @@ def test_gemm( m, k, n, - trace_size, partition_N, aie_context, ): diff --git a/iron/tests/infrastructure/lazy_imports.py b/iron/tests/infrastructure/lazy_imports.py index a40bb2a558..84a29c3c12 100644 --- a/iron/tests/infrastructure/lazy_imports.py +++ b/iron/tests/infrastructure/lazy_imports.py @@ -1,14 +1,27 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Importing one operator must not import the rest of the catalog.""" +"""Importing one operator must not import the rest of the catalog. +The check runs in a fresh interpreter: in this one, whatever the session +collected before it has already imported the operators it looks for. +""" + +import subprocess import sys +_CHECK = """\ +import sys from iron.operators import ElementwiseAdd +assert ElementwiseAdd.__name__ == "ElementwiseAdd" +loaded = [m for m in ("iron.operators.mha.op", "iron.operators.swiglu_decode.op") + if m in sys.modules] +assert not loaded, f"importing ElementwiseAdd also imported {loaded}" +""" def test_lazy_catalog_does_not_import_mha(): - assert ElementwiseAdd.__name__ == "ElementwiseAdd" - assert "iron.operators.mha.op" not in sys.modules - assert "iron.operators.swiglu_decode.op" not in sys.modules + result = subprocess.run( + [sys.executable, "-c", _CHECK], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr From ff3aaf55e55ca284d12af7c1eca7051d03e67ee2 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 10:51:08 -0600 Subject: [PATCH 18/22] Tracing: take the trace buffer from mlir-aie; real tests, no mocks The full-ELF sequence callable now asks mlir-aie's get_trace_buffer() for the fused trace buffer's argument index and size, instead of assuming argument 3 and summing the slices itself. IRON's trace_buffer_size() and its test (trace_layout.py) go; the layout is mlir-aie's and tested there. The tests that stubbed their way around the hardware now run it: - tracing: dispatch a traced one-step layer_norm sequence, check tracing leaves the result bit-identical, and that dump_traces writes the buffer's raw 32-bit words and the Perfetto JSON. - conftest_lazy_device: run the root conftest's device gating in a real pytest subprocess instead of calling the hook on fake items against a stubbed runtime. The no-runtime cases hide pyxrt from the child's PYTHONPATH, which is what an unsourced XRT amounts to. - relative_build_dir: pass a relative build_dir instead of monkeypatching the working directory. Co-Authored-By: Claude --- iron/common/compilation/__init__.py | 1 - iron/common/compilation/sequence.py | 12 -- iron/common/sequence.py | 17 +- iron/tests/compilation/relative_build_dir.py | 6 +- .../infrastructure/conftest_lazy_device.py | 149 ++++++++---------- iron/tests/infrastructure/trace_layout.py | 35 ---- iron/tests/infrastructure/tracing.py | 83 ++++++---- 7 files changed, 136 insertions(+), 167 deletions(-) delete mode 100644 iron/tests/infrastructure/trace_layout.py diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index c1fb11855d..0819d31c8c 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -33,5 +33,4 @@ from .sequence import ( SequenceMLIRArtifact, FusePythonGeneratedMLIRCompilationRule, - trace_buffer_size, ) diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index 6a1b6858f1..b580cf1eb5 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -14,7 +14,6 @@ from aie import ir from aie.dialects import aie, aiex, memref from aie.extras.context import mlir_mod_ctx -from aie.utils.trace import get_trace_slices import ml_dtypes from typing import Any @@ -35,17 +34,6 @@ # ########################################################################## -def trace_buffer_size(mlir_text: str) -> int: - """Bytes of the fused trace buffer the dispatched sequence takes. - - `-aie-fuse-trace-buffers` gives the sequence one buffer covering every design - it configures, and records the split on the sequence. Returns 0 for an - untraced build. - """ - slices = get_trace_slices(mlir_text) - return max((s["offset"] + s["size"] for s in slices), default=0) - - class SequenceMLIRArtifact(MLIRArtifact): def __init__( self, diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 69352f53e1..6eec6e9319 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -13,6 +13,7 @@ from aie.iron.device import NPU2 from aie.utils.hostruntime.tensor_class import CPUOnlyTensor from aie.utils.npukernel import NPUKernel +from aie.utils.trace import get_trace_buffer from aie.utils.verify import Tolerance, compare try: @@ -633,7 +634,7 @@ def __init__(self, op, device_name="main", sequence_name="sequence"): self.run_handle.set_arg(1, self.output_buffer.buffer_object()) self.run_handle.set_arg(2, self.scratch_buffer.buffer_object()) if self.trace_buffer is not None: - self.run_handle.set_arg(3, self.trace_buffer.buffer_object()) + self.run_handle.set_arg(self._trace_arg, self.trace_buffer.buffer_object()) self._params = None @@ -672,13 +673,17 @@ def _allocate_buffers(self): (_n_elements(scratch_sz),), dtype=ml_dtypes.bfloat16 ) # Trace lowering appends one buffer covering every configured design, after - # the consolidated three. Its size depends on how many channels and - # sub-designs claim a share, so read it from the lowered module. + # the consolidated three. Its argument and size depend on how many channels + # and sub-designs claim a share, so read them from the lowered module. self.trace_buffer = None + self._trace_arg = None if self.op.trace_size: - total = comp.trace_buffer_size(self.lowered_mlir_text()) - if total: - self.trace_buffer = XRTTensor((total,), dtype=np.int8) + layout = get_trace_buffer( + self.lowered_mlir_text(), f"{self.device_name}:{self.sequence_name}" + ) + if layout: + self._trace_arg = layout["arg_index"] + self.trace_buffer = XRTTensor((layout["size"],), dtype=np.int8) def lowered_mlir_text(self) -> str: """aiecc's post-lowering module, which carries the trace buffer layout.""" diff --git a/iron/tests/compilation/relative_build_dir.py b/iron/tests/compilation/relative_build_dir.py index b5e4c4906f..764634ea80 100644 --- a/iron/tests/compilation/relative_build_dir.py +++ b/iron/tests/compilation/relative_build_dir.py @@ -11,6 +11,7 @@ tests never did because their build_dir is absolute. """ +import os from pathlib import Path import aie.utils as aie_utils @@ -21,10 +22,9 @@ from iron.operators.elementwise_mul.op import ElementwiseMul -def test_factory_kernel_compiles_with_a_relative_build_dir(tmp_path, monkeypatch): - monkeypatch.chdir(tmp_path) +def test_factory_kernel_compiles_with_a_relative_build_dir(tmp_path): aie_utils.set_current_device(NPU2()) - ctx = AIEContext(build_dir="build_rel") + ctx = AIEContext(build_dir=os.path.relpath(tmp_path / "build_rel")) op = ElementwiseMul(size=4096, tile_size=4096, num_aie_columns=1, context=ctx) op.compile() diff --git a/iron/tests/infrastructure/conftest_lazy_device.py b/iron/tests/infrastructure/conftest_lazy_device.py index 3ea80fcc42..87072f1a7f 100644 --- a/iron/tests/infrastructure/conftest_lazy_device.py +++ b/iron/tests/infrastructure/conftest_lazy_device.py @@ -2,103 +2,94 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The root conftest.py's pytest_collection_modifyitems must not resolve a -device unless some collected test restricts itself to specific devices via -@pytest.mark.supported_devices. Resolving one unconditionally opens the +"""The root conftest.py's device gating, run as a real pytest session. + +Its pytest_collection_modifyitems must not resolve a device unless some collected +test restricts itself via @pytest.mark.supported_devices: resolving one opens the single-tenant NPU on every plain `pytest` in this tree, whatever was selected. +When a test does restrict itself, it skips the tests this device is not listed +for, and stops with the reason when there is no NPU runtime at all. -pytest loads the root conftest.py for these tests too, so the hook under test -is imported by path instead and called directly, against fake items and a -stubbed aie_utils.DefaultNPURuntime that raises if .device() is reached. +Each case runs pytest in a subprocess, over a directory holding a copy of the +root conftest and one test module. The no-runtime cases hide pyxrt from it, +which is what an unsourced XRT amounts to and the setup the laziness exists for. """ import importlib.util +import os +import shutil +import subprocess import sys from pathlib import Path -from types import SimpleNamespace import pytest _ROOT_CONFTEST = Path(__file__).resolve().parents[3] / "conftest.py" +_INI = """\ +[pytest] +markers = + supported_devices(*devices): only supported on the given devices +""" -def _load_root_conftest(): - spec = importlib.util.spec_from_file_location( - "_root_conftest_under_test", _ROOT_CONFTEST - ) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -class _FakeMarker: - def __init__(self, *args): - self.args = args - - -class _FakeItem: - def __init__(self, marker=None): - self._marker = marker - self.markers_added = [] - - def get_closest_marker(self, name): - assert name == "supported_devices" - return self._marker - - def add_marker(self, marker): - self.markers_added.append(marker) - - -class _DeviceCalledError(Exception): - pass - - -def _stub_runtime_that_forbids_device_calls(root_conftest, monkeypatch): - def _raise(): - raise _DeviceCalledError( - "DefaultNPURuntime.device() was called with no marked test collected" - ) - monkeypatch.setattr( - root_conftest.aie_utils, - "DefaultNPURuntime", - SimpleNamespace(device=_raise), +def _pytest(tmp_path, test_source, without_xrt=False): + """Run pytest over one test module under the root conftest.""" + shutil.copy(_ROOT_CONFTEST, tmp_path / "conftest.py") + (tmp_path / "pytest.ini").write_text(_INI) + (tmp_path / "test_gated.py").write_text(test_source) + env = dict(os.environ) + pyxrt = importlib.util.find_spec("pyxrt") + if without_xrt and pyxrt is not None: + hidden = os.path.dirname(pyxrt.origin) + entries = env.get("PYTHONPATH", "").split(os.pathsep) + if hidden not in entries: + pytest.skip(f"pyxrt is installed, not on PYTHONPATH: {pyxrt.origin}") + env["PYTHONPATH"] = os.pathsep.join(p for p in entries if p != hidden) + return subprocess.run( + [sys.executable, "-m", "pytest", "-p", "no:cacheprovider"] + + ["--iterations", "1", "-v", "test_gated.py"], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, ) -def test_no_device_probe_when_nothing_is_device_restricted(monkeypatch): - root_conftest = _load_root_conftest() - _stub_runtime_that_forbids_device_calls(root_conftest, monkeypatch) - - items = [_FakeItem(), _FakeItem(), _FakeItem()] - root_conftest.pytest_collection_modifyitems(config=None, items=items) - assert all(item.markers_added == [] for item in items) - - -def test_device_probed_and_unsupported_items_skipped_when_a_test_is_restricted( - monkeypatch, -): - root_conftest = _load_root_conftest() - - class _FakeDevice: - def resolve(self): - return SimpleNamespace(name="npu2") - - monkeypatch.setattr( - root_conftest.aie_utils, - "DefaultNPURuntime", - SimpleNamespace(device=lambda: _FakeDevice()), +def test_unrestricted_tests_need_no_npu_runtime(tmp_path): + result = _pytest( + tmp_path, + "def test_plain():\n pass\n", + without_xrt=True, ) + assert result.returncode == 0, result.stdout + result.stderr + assert "1 passed" in result.stdout - unrestricted = _FakeItem() - matches_device = _FakeItem(_FakeMarker("npu1", "npu2")) - excludes_device = _FakeItem(_FakeMarker("npu1")) - root_conftest.pytest_collection_modifyitems( - config=None, items=[unrestricted, matches_device, excludes_device] +def test_restricted_test_without_npu_runtime_stops_with_the_reason(tmp_path): + result = _pytest( + tmp_path, + "import pytest\n" + "@pytest.mark.supported_devices('npu1', 'npu2')\n" + "def test_gated():\n pass\n", + without_xrt=True, ) - - assert unrestricted.markers_added == [] - assert matches_device.markers_added == [] - assert len(excludes_device.markers_added) == 1 - assert excludes_device.markers_added[0].name == "skip" + assert result.returncode == pytest.ExitCode.USAGE_ERROR, result.stdout + assert "No NPU runtime: " in result.stderr + assert "xrt" in result.stderr.lower(), result.stderr + + +@pytest.mark.supported_devices("npu1", "npu2") +def test_restricted_tests_skip_where_the_device_is_not_listed(tmp_path): + result = _pytest( + tmp_path, + "import pytest\n" + "def test_plain():\n pass\n" + "@pytest.mark.supported_devices('npu1', 'npu2')\n" + "def test_any_npu():\n pass\n" + "@pytest.mark.supported_devices('no_such_npu')\n" + "def test_elsewhere():\n pass\n", + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "2 passed, 1 skipped" in result.stdout, result.stdout + assert "test_elsewhere SKIPPED (Not supported on" in result.stdout, result.stdout diff --git a/iron/tests/infrastructure/trace_layout.py b/iron/tests/infrastructure/trace_layout.py deleted file mode 100644 index 8b556b5bd7..0000000000 --- a/iron/tests/infrastructure/trace_layout.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Reading back the trace buffer size the compiler recorded on the sequence.""" - -from iron.common.compilation import trace_buffer_size - -LOWERED = """ -module { - aie.device(npu1_1col) @main { - aie.runtime_sequence @sequence(%arg0: memref<4xi32>, %arg1: memref<12288xi8>) - attributes {trace_slices = [ - #aie.trace_slice, - #aie.trace_slice]} { - } - } -} -""" - -UNTRACED = """ -module { - aie.device(npu1_1col) @main { - aie.runtime_sequence @sequence(%arg0: memref<4xi32>) { - } - } -} -""" - - -def test_size_spans_every_slice(): - assert trace_buffer_size(LOWERED) == 12288 - - -def test_untraced_build_has_no_trace_buffer(): - assert trace_buffer_size(UNTRACED) == 0 diff --git a/iron/tests/infrastructure/tracing.py b/iron/tests/infrastructure/tracing.py index cd2de35662..64efdd1fd8 100644 --- a/iron/tests/infrastructure/tracing.py +++ b/iron/tests/infrastructure/tracing.py @@ -1,48 +1,69 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Trace dumps use the upstream tensor's host interface without torch.""" +"""IRON's side of tracing: the full-ELF sequence callable binds, fills and syncs +the trace buffer mlir-aie's lowering asks for, and dump_traces writes it out. + +Trace insertion, the buffer layout and event decoding are mlir-aie's, and tested +there. +""" import json -from types import SimpleNamespace import numpy as np import pytest -from aie.utils.hostruntime.tensor_class import CPUOnlyTensor +import torch + +from iron.common.sequence import OperatorSequence +from iron.common.tracing_utils import dump_traces +from iron.operators.layer_norm.op import LayerNorm -from iron.common import tracing_utils +SIZE = 2048 +TRACE_SIZE = 8192 -@pytest.mark.parametrize("dtype", [np.int8, np.uint8]) -def test_dump_preserves_raw_trace_bits(monkeypatch, tmp_path, dtype): - words = np.array([0xFFFFFFFF, 0x80000000, 0x12345678, 0], dtype=np.uint32) - buffer = CPUOnlyTensor(words.view(dtype), dtype=dtype) - run = SimpleNamespace(trace_buffer=buffer) - monkeypatch.setattr( - tracing_utils, "lowered_mlir", lambda run: (tmp_path / "test.mlir", "mlir") +def _layer_norm_run(context, trace_size): + """A dispatched one-step sequence, and its output.""" + layer_norm = LayerNorm( + size=SIZE, + num_aie_columns=1, + num_channels=1, + tile_size=SIZE, + trace_size=trace_size, + context=context, ) - events = [{"name": "event"}] + seq = OperatorSequence( + name="infra_trace_layer_norm", + runlist=[(layer_norm, "x", "y")], + input_args=["x"], + output_args=["y"], + dispatch="fused", + trace_size=trace_size, + context=context, + ) + seq.compile() + run = seq.get_callable() + torch.manual_seed(0) + run.get_buffer("x").torch_view()[:] = torch.randn(SIZE, dtype=torch.bfloat16) + run() + return run, run.get_buffer("y").torch_view()[:SIZE].clone() - def parse(actual, mlir_text, colshift): - np.testing.assert_array_equal(actual, words) - assert mlir_text == "mlir" - assert colshift == 2 - return [(None, events)] - monkeypatch.setattr(tracing_utils, "parse_trace_buffer", parse) - written = tracing_utils.dump_traces( - run, "test", out_dir=tmp_path, colshift=2, summary=False - ) +@pytest.mark.supported_devices("npu2") +def test_dump_writes_raw_words_and_perfetto_json(aie_context, tmp_path): + run, traced = _layer_norm_run(aie_context, TRACE_SIZE) + _, untraced = _layer_norm_run(aie_context, 0) + assert torch.equal(traced, untraced), "tracing changed the result" - assert written == [tmp_path / "test_trace.json"] - assert json.loads(written[0].read_text()) == events - assert (tmp_path / "test.txt").read_text().splitlines() == [ - "ffffffff", - "80000000", - "12345678", - "00000000", - ] + written = dump_traces(run, "layer_norm", out_dir=tmp_path, summary=False) + words = run.trace_buffer.numpy().view(np.uint32).reshape(-1) + assert words.any(), "the traced dispatch captured no trace data" + # The raw text is the buffer's 32-bit words, unchanged by the int8 buffer. + raw = (tmp_path / "layer_norm.txt").read_text().split() + assert [int(w, 16) for w in raw] == words.tolist() -def test_untraced_run_needs_no_buffer(tmp_path): - assert tracing_utils.dump_traces(SimpleNamespace(), "test", tmp_path) == [] + assert written, "a buffer with trace data produced no Perfetto file" + for path in written: + assert path.parent == tmp_path and path.name.startswith("layer_norm_") + assert json.loads(path.read_text()), f"{path.name} holds no events" From b8a909e596a8f48135c89925710cfcf71dd18edb Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 11:01:20 -0600 Subject: [PATCH 19/22] Tracing: write and decode dumps with mlir-aie's TraceConfig dump_traces wrote the raw words, split the buffer into slices, decoded each, named the files and warned of truncation itself, which duplicated TraceConfig.write_trace/trace_to_json. trace_to_json now does the per-slice part upstream (mlir-aie branch trace-to-json-slices), so dump_traces writes the text with write_trace and the JSON with trace_to_json, and keeps only IRON's knobs: the output directory, the column shift, the MLIR override and the cycles summary. parse_trace_buffer existed to turn the parser's SystemExit into an exception; the parser raises ValueError now, so it goes. lowered_mlir duplicated the callable's lookup of aiecc's lowered module, which the callable now exposes as lowered_mlir_path. The raw text drops trailing zero words, as TraceConfig's always has, so the test reads it back with read_trace. A buffer without slices now writes .json rather than _trace.json. Co-Authored-By: Claude --- iron/common/sequence.py | 12 +-- iron/common/tracing_utils.py | 123 +++++++-------------------- iron/tests/infrastructure/tracing.py | 12 ++- 3 files changed, 47 insertions(+), 100 deletions(-) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 6eec6e9319..7882bc3484 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -679,17 +679,19 @@ def _allocate_buffers(self): self._trace_arg = None if self.op.trace_size: layout = get_trace_buffer( - self.lowered_mlir_text(), f"{self.device_name}:{self.sequence_name}" + self.lowered_mlir_path.read_text(), + f"{self.device_name}:{self.sequence_name}", ) if layout: self._trace_arg = layout["arg_index"] self.trace_buffer = XRTTensor((layout["size"],), dtype=np.int8) - def lowered_mlir_text(self) -> str: - """aiecc's post-lowering module, which carries the trace buffer layout.""" + @property + def lowered_mlir_path(self) -> Path: + """aiecc's post-lowering module, which carries the trace configuration and + the trace buffer layout. A traced build asks aiecc to keep it.""" mlir_filename = self.op.artifacts[0].mlir_input.filename - path = comp._aiecc_work_dir(mlir_filename) / "input_with_addresses.mlir" - return path.read_text() + return comp._aiecc_work_dir(mlir_filename) / "input_with_addresses.mlir" def get_buffer(self, buffer_name): if buffer_name in self._buffer_cache: diff --git a/iron/common/tracing_utils.py b/iron/common/tracing_utils.py index 99dedb2e78..1667dfde1f 100644 --- a/iron/common/tracing_utils.py +++ b/iron/common/tracing_utils.py @@ -16,12 +16,10 @@ On an untraced build the call returns an empty list, so a test can call it unconditionally. -A dump writes the raw 32-bit words as hex text, plus one JSON file per traced -design for https://ui.perfetto.dev. Keep the text: :func:`parse_trace_buffer` -reparses it with a different column shift for the price of no further dispatch. - -:func:`dump_traces` also prints mlir-aie's per-tile cycles summary for each file it -writes. +The writing and decoding are mlir-aie's ``TraceConfig``: a dump is its raw trace +text, which ``TraceConfig.read_trace`` reads back to reparse without a further +dispatch, plus one JSON file per traced design for https://ui.perfetto.dev. +:func:`dump_traces` also prints mlir-aie's per-tile cycles summary for each. Environment: * ``IRON_TRACE_DIR`` where to write (default ``outputs/traces``) @@ -31,74 +29,18 @@ from __future__ import annotations -import json import os from pathlib import Path import numpy as np -from aie.utils.trace import parse_trace_slices, print_cycles_summary - -from . import compilation as comp +from aie.utils.trace import TraceConfig, print_cycles_summary -__all__ = [ - "dump_traces", - "parse_trace_buffer", - "lowered_mlir", -] +__all__ = ["dump_traces"] DEFAULT_TRACE_DIR = "outputs/traces" -def lowered_mlir(run) -> tuple[Path, str]: - """The post-lowering MLIR for a callable, as ``(path, text)``. - - mlir-aie's trace parser matches ``aiex.npu.write32`` ops against the trace unit's - config addresses. ``aie-insert-trace-flows`` emits those writes inside aiecc, so - the parser needs aiecc's lowered module. A traced build requests it with - ``--get-input-with-addresses``, which lands it in the work dir beside the source - (``.mlir.d/``). - """ - override = os.environ.get("IRON_TRACE_MLIR") - if override: - path = Path(override) - return path, path.read_text() - - source = Path(run.op.artifacts[0].mlir_input.filename) - path = comp._aiecc_work_dir(str(source)) / "input_with_addresses.mlir" - if not path.exists(): - raise FileNotFoundError( - f"{path} is missing; a traced build passes --get-input-with-addresses " - "to aiecc. Point IRON_TRACE_MLIR at a lowered module to override." - ) - return path, path.read_text() - - -def parse_trace_buffer(words, mlir_text: str, colshift: int | None = None): - """A trace buffer's words as ``(slice_info, events)`` per traced design. - - The parser splits the buffer by the layout the compiler recorded on the - dispatched sequence, and decodes each region against the device that wrote it. - - ``colshift`` of None lets the parser align the columns itself, which is what you - want by default: a design configured for one column may be loaded into another. - Override it when that alignment picks the wrong columns. - - The parser calls ``sys.exit`` on some malformed input, so SystemExit becomes a - RuntimeError here: a visualisation failure must not fail a test. - """ - try: - return parse_trace_slices( - np.asarray(words, dtype=np.uint32), mlir_text, colshift - ) - except SystemExit as exc: - raise RuntimeError( - "mlir-aie's trace parser exited; the usual cause is an MLIR without the " - "trace register writes, or a column shift that does not match the data. " - "Run with logging at DEBUG to see the tiles it found." - ) from exc - - def _slug(text: str) -> str: keep = "-_." return "".join(c if c.isalnum() or c in keep else "_" for c in text) @@ -111,15 +53,20 @@ def dump_traces( colshift: int | None = None, summary: bool = True, ) -> list[Path]: - """Write a completed run's trace buffer as hex text and Perfetto JSON. + """Write a completed run's trace buffer as trace text and Perfetto JSON. Call it after ``run()``: the callable syncs its trace buffer device->host as part of the dispatch, so this only reads host memory. Returns the JSON paths written, empty on an untraced build. ``tag`` distinguishes one dump from another - a test name or parameter id. The - layout the compiler recorded on the dispatched sequence splits the buffer, so a - fused sequence yields one JSON file per configured design. + text goes to ``.txt``. A fused sequence shares the buffer between the + designs it configures, and each gets its own + ``___.json``; otherwise the JSON is ``.json``. + + ``colshift`` of None lets the parser align the columns itself, which is what you + want by default: a design configured for one column may be loaded into another. + Override it when that alignment picks the wrong columns. """ buffer = getattr(run, "trace_buffer", None) if buffer is None: @@ -137,38 +84,32 @@ def dump_traces( env = os.environ.get("IRON_TRACE_COLSHIFT") colshift = int(env) if env else None - mlir_path, mlir_text = lowered_mlir(run) - print(f"[trace] parsing against {mlir_path}") - words = buffer.numpy().view(np.uint32).reshape(-1) tag = _slug(tag) - raw = (out_dir / tag).with_suffix(".txt") - raw.write_text("\n".join(f"{w:08x}" for w in words) + "\n") + config = TraceConfig( + trace_size=words.nbytes, trace_file=str(out_dir / f"{tag}.txt") + ) + config.write_trace(words) if not words.any(): print("[trace] buffer is all zeros, no trace data captured") return [] + mlir = os.environ.get("IRON_TRACE_MLIR") or run.lowered_mlir_path + print(f"[trace] parsing against {mlir}") try: - parsed = parse_trace_buffer(words, mlir_text, colshift) + written = config.trace_to_json( + str(mlir), + str(out_dir / f"{tag}.json"), + colshift=colshift, + kernel=f"{run.device_name}:{run.sequence_name}", + ) except Exception as exc: # a visualisation failure must not fail a run - print(f"[trace] parse failed ({exc}); raw words kept at {raw}") + print(f"[trace] parse failed ({exc}); raw words kept at {config.trace_file}") return [] - written = [] - for index, (entry, events) in enumerate(parsed): - # A device may hold several runtime sequences, so both names identify a slice. - name = f"{index}_{entry['device']}_{entry['sequence']}" if entry else "trace" - if entry and words[(entry["offset"] + entry["size"]) // 4 - 1]: - print( - f"[trace] {name}: slice full ({entry['size']} B), trace is likely " - "truncated - raise IRON_TRACE_SIZE" - ) - - target = (out_dir / f"{tag}_{_slug(name)}").with_suffix(".json") - target.write_text(json.dumps(events)) - print(f"[trace] {target} ({len(events)} events)") - written.append(target) - + paths = [Path(p) for p in written] + for path in paths: + print(f"[trace] {path}") if summary: - print_cycles_summary(target) - return written + print_cycles_summary(path) + return paths diff --git a/iron/tests/infrastructure/tracing.py b/iron/tests/infrastructure/tracing.py index 64efdd1fd8..f9d0b37940 100644 --- a/iron/tests/infrastructure/tracing.py +++ b/iron/tests/infrastructure/tracing.py @@ -13,6 +13,7 @@ import numpy as np import pytest import torch +from aie.utils.trace import TraceConfig from iron.common.sequence import OperatorSequence from iron.common.tracing_utils import dump_traces @@ -59,11 +60,14 @@ def test_dump_writes_raw_words_and_perfetto_json(aie_context, tmp_path): words = run.trace_buffer.numpy().view(np.uint32).reshape(-1) assert words.any(), "the traced dispatch captured no trace data" - # The raw text is the buffer's 32-bit words, unchanged by the int8 buffer. - raw = (tmp_path / "layer_norm.txt").read_text().split() - assert [int(w, 16) for w in raw] == words.tolist() + # The text reads back as the buffer's 32-bit words, unchanged by the int8 + # buffer, so it can be reparsed without another dispatch. + raw = TraceConfig( + trace_size=words.nbytes, trace_file=str(tmp_path / "layer_norm.txt") + ) + assert np.array_equal(raw.read_trace(), words) assert written, "a buffer with trace data produced no Perfetto file" for path in written: - assert path.parent == tmp_path and path.name.startswith("layer_norm_") + assert path.parent == tmp_path and path.name.startswith("layer_norm") assert json.loads(path.read_text()), f"{path.name} holds no events" From a60d1701c6b3841a9ae1fc577d5a56324f9cfcf3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:37:21 +0000 Subject: [PATCH 20/22] Make missing Llama weights fail in CI Co-authored-by: hunhoffe <54562339+hunhoffe@users.noreply.github.com> --- iron/applications/llama_3.2_1b/test.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/iron/applications/llama_3.2_1b/test.py b/iron/applications/llama_3.2_1b/test.py index 9545019420..85ecbad224 100644 --- a/iron/applications/llama_3.2_1b/test.py +++ b/iron/applications/llama_3.2_1b/test.py @@ -29,11 +29,12 @@ def generate_test_params(): params, names = generate_test_params() requires_weights = pytest.mark.skipif( - not ( + not os.environ.get("CI") + and not ( (weights_dir / "llama3.2-1b" / "model.safetensors").exists() and (weights_dir / "llama3.2-1b" / "tokenizer.model").exists() ), - reason="llama3.2-1b weights not found", + reason="llama3.2-1b weights not found outside CI", ) From a31d9a6922ad371812f85bfeaa9d0bbe40c4b284 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 25 Sep 2026 22:39:50 +0000 Subject: [PATCH 21/22] Clarify AXPY coefficient rounding and add non-integer regressions Co-authored-by: hunhoffe <54562339+hunhoffe@users.noreply.github.com> --- iron/operators/axpy/reference.py | 5 +++-- iron/operators/axpy/test.py | 4 ++-- iron/tests/operators/axpy_reference.py | 29 ++++++++++++++++++++++++++ 3 files changed, 34 insertions(+), 4 deletions(-) create mode 100644 iron/tests/operators/axpy_reference.py diff --git a/iron/operators/axpy/reference.py b/iron/operators/axpy/reference.py index ee37697abf..7fb48d5bda 100644 --- a/iron/operators/axpy/reference.py +++ b/iron/operators/axpy/reference.py @@ -7,8 +7,9 @@ def reference(x, y, scalar): """CPU reference: ``scalar * x + y`` in fp32, rounded once (ground truth). - The kernel takes ``scalar`` as bf16 and rounds only the result, where bf16 - arithmetic would round the product as well. + The vectorized kernel accepts ``scalar`` as fp32 but broadcasts + ``bfloat16(a)`` internally. The product stays in an fp32 accumulator until + after adding ``y``; only the coefficient and final result round to bf16. """ a = torch.tensor(scalar, dtype=torch.bfloat16).float() return (a * x.float() + y.float()).to(x.dtype) diff --git a/iron/operators/axpy/test.py b/iron/operators/axpy/test.py index aa48e9cfc0..c05541b367 100755 --- a/iron/operators/axpy/test.py +++ b/iron/operators/axpy/test.py @@ -13,7 +13,7 @@ def get_params(): max_aie_columns = aie_utils.get_current_device().cols input_lengths = [1024, 2048, 4096, 8192] - scalar_factors = [3.0, 10.0] + scalar_factors = [3.0, 10.0, 1.003] params = [] for input_length in input_lengths: @@ -23,7 +23,7 @@ def get_params(): continue for scalar in scalar_factors: # Determine if this is a regular test case - is_regular = input_length == 2048 and scalar == 3.0 + is_regular = input_length == 2048 and scalar in (3.0, 1.003) marks = [] if is_regular else [pytest.mark.extensive] params.append( diff --git a/iron/tests/operators/axpy_reference.py b/iron/tests/operators/axpy_reference.py new file mode 100644 index 0000000000..b55a0fed16 --- /dev/null +++ b/iron/tests/operators/axpy_reference.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from iron.operators.axpy.reference import reference + + +@pytest.mark.parametrize("scalar", [1.003, -1.003]) +def test_reference_rounds_scalar_to_bf16(scalar): + x = torch.tensor([1.0], dtype=torch.bfloat16) + y = torch.tensor([-1.0 if scalar > 0 else 1.0], dtype=torch.bfloat16) + + unrounded = (torch.tensor(scalar, dtype=torch.float32) + y.float()).to(x.dtype) + actual = reference(x, y, scalar) + + assert unrounded.item() != 0.0 + assert actual.dtype == x.dtype + assert torch.equal(actual, torch.zeros_like(x)) + + +def test_reference_does_not_round_intermediate_product(): + x = torch.tensor([1.0078125], dtype=torch.bfloat16) + y = torch.tensor([-1.015625], dtype=torch.bfloat16) + + actual = reference(x, y, 1.0078125) + + assert torch.equal(actual, torch.tensor([2**-14], dtype=torch.bfloat16)) From c4bdfd29892097e30b0e8380717eeece33531729 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 25 Sep 2026 17:54:57 -0600 Subject: [PATCH 22/22] Follow mlir-aie #3801's aie_kernels layout mlir-aie #3801 (f119c6947) groups aie_kernels by family instead of by architecture: generic/, aie2/ and aie2p/ are gone, and architecture-specific code moved into *_aie2.h / *_aie2p.h headers the family's .cc selects. - flm.GEMM compiles fused/fused_mm_tile.cc: mm_fused.cc became the header fused/mm_fused.h, and activations.h / zero.cc moved to common/. - Stream ops: mm.cc is linalg/mm.cc, zero.cc is zero/zero.cc, silu and mul live in activation/ and eltwise/. StreamKernel.subdir now names the family directory instead of defaulting to the device one, which no longer exists. - Comments, docstrings, README kernel links and AGENTS.md follow the new paths. Co-Authored-By: Claude --- AGENTS.md | 21 ++++---- README.md | 49 ++++++++++--------- iron/common/stream/ops.py | 20 ++++---- iron/operators/flm/gemm/design.py | 2 +- iron/operators/flm/gemm/op.py | 19 ++++--- iron/operators/gemm/op.py | 8 +-- iron/operators/gemv/op.py | 4 +- iron/operators/gemv/reference.py | 2 +- iron/operators/rope/README.md | 2 +- .../operators/swiglu_prefill_stream/README.md | 2 +- .../kernel_object_arch_isolation.py | 4 +- .../tests/operators/gemm_tile_divisibility.py | 2 +- 12 files changed, 72 insertions(+), 63 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 6b7ada346b..a9f3f3e9d6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -132,13 +132,13 @@ reuse lint - `test.py`: End-to-end test (build, run once, check against `reference()`) 2. **AIE Kernels** ([mlir-aie `aie_kernels/`](https://github.com/Xilinx/mlir-aie/tree/main/aie_kernels)) - - Architecture-specific C++ compute kernels, sourced from the installed - mlir-aie package, not from this repo. Operators get them from mlir-aie's - kernel factories (`aie.iron.kernels`), each of which returns an - `ExternalFunction` carrying its source, flags, symbol and argument types: - - `generic/`: Works on both AIE2 and AIE2P - - `aie2/`: AIE2-specific (NPU1) - - `aie2p/`: AIE2P-specific (NPU2) + - C++ compute kernels, sourced from the installed mlir-aie package, not + from this repo. Operators get them from mlir-aie's kernel factories + (`aie.iron.kernels`), each of which returns an `ExternalFunction` + carrying its source, flags, symbol and argument types + - Grouped by family (`activation/`, `eltwise/`, `linalg/`, `norm/`, + `fused/`, `common/`, ...), not by architecture: a kernel's `.cc` includes + its `*_aie2.h` or `*_aie2p.h` header, chosen by `aie_arch.h` - Use AIE API for vectorization (e.g., `aie::mmul`, `aie::add`, `aie::mul`) - Compiled to `.o` files and linked into operator `.xclbin` @@ -265,7 +265,8 @@ Data movement pattern: L3 → Shim DMA → L2 → L1 (tile local) → Compute 4. If a new C++ compute kernel is needed, add it to the [mlir-aie kernel library](https://github.com/Xilinx/mlir-aie/tree/main/aie_kernels) with a factory in `aie.iron.kernels`; IRON no longer hosts kernels - - Choose appropriate directory: `generic/`, `aie2/`, or `aie2p/` + - Choose the family directory (`activation/`, `eltwise/`, `linalg/`, ...); + put architecture-specific code in `*_aie2.h` / `*_aie2p.h` headers - Use AIE API for portable vectorization when possible - Add `event0()` and `event1()` for performance profiling 5. Implement `reference.py` with the CPU reference and `generate_inputs()`, @@ -472,7 +473,7 @@ logging.basicConfig(level=logging.DEBUG) **"Kernel not found" or "Symbol not defined"** - Verify the kernel `.cc` exists under the installed mlir-aie package's - `include/aie_kernels//` (`AIEContext.kernels_dir`, overridden by + `include/aie_kernels//` (`AIEContext.kernels_dir`, overridden by `MLIR_AIE_KERNEL_SOURCES`) - Check `get_kernel_artifacts()` in `op.py` returns every factory the design uses - Ensure the C signature matches the factory's (or `bind()`'s) argument types @@ -505,7 +506,7 @@ logging.basicConfig(level=logging.DEBUG) **Kernel compilation failures** -- Check kernel is in correct architecture directory (`generic/`, `aie2/`, `aie2p/`) +- Check the kernel's `.cc` includes the right `*_aie2.h` / `*_aie2p.h` header for the target - Verify `#include ` for AIE API kernels - Ensure template parameters match function signature - Check for syntax errors in vectorization code diff --git a/README.md b/README.md index adb53d1764..e99f59423f 100755 --- a/README.md +++ b/README.md @@ -42,33 +42,33 @@ The IRON Python API for Ryzen™ AI NPUs is described in the following paper: | Section | Description | Datatype | AIE2 | AIE2P | Status | Design Example | |:--------|:------------|:---------|:-----|:------|:-------|:-------------| -| [Element-wise Add](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/add.cc) | Element-wise addition kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/elementwise_add/](./iron/operators/elementwise_add/) | -| [Element-wise Mul](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/mul.cc) | Element-wise multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/elementwise_mul/](./iron/operators/elementwise_mul/) | -| [GEMM](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/mm.cc) | General Matrix Multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gemm/](./iron/operators/gemm/) | -| [Alternative GEMM](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/mm_fused.cc) | General Matrix Multiplication with a fused activation epilogue, specialised for transformer projection shapes. M, K and N are runtime parameters, so one xclbin serves every shape. B is stored as bfp16 on AIE2P | bfloat16, bfp16 | ✓ | ✓ | 🟢 | [iron/operators/flm/gemm/](./iron/operators/flm/gemm/) | -| [GEMV](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/mv.cc) | General Matrix-Vector Multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gemv/](./iron/operators/gemv/) | -| [GQA](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/mha.cc) | Grouped Query Attention kernel (Single pipeline) | bfloat16 | | ✓ | 🟢 | [iron/operators/mha/](./iron/operators/mha/) | -| [MHA](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/mha.cc) | Multi-Head Attention kernel & Grouped Query Attention | bfloat16 | | ✓ | 🟢 | [iron/operators/mha/](./iron/operators/mha/) | -| [RMSNorm](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/rms_norm.cc) | RMSNorm kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rms_norm/](./iron/operators/rms_norm/) | -| [RoPE](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/rope.cc) | Rotary Positional Embedding kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rope/](./iron/operators/rope/) | -| [SiLU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/silu.cc) | Sigmoid Linear Unit activation kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/silu/](./iron/operators/silu/) | -| [Softmax](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/softmax.cc) | Softmax kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/softmax/](./iron/operators/softmax/) | -| [Weighted RMSNorm](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2p/rms_norm.cc) | Weighted RMSNorm kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rms_norm/](./iron/operators/rms_norm/) | -| [Copy](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/passThrough.cc) | Copy | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/mem_copy/](./iron/operators/mem_copy/) | -| [Transpose](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/transpose.cc) | Transpose | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/transpose/](./iron/operators/transpose/) | -| [AXPY](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/axpy.cc) | AXPY | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/axpy/](./iron/operators/axpy/) | +| [Element-wise Add](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/eltwise/add.cc) | Element-wise addition kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/elementwise_add/](./iron/operators/elementwise_add/) | +| [Element-wise Mul](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/eltwise/mul.cc) | Element-wise multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/elementwise_mul/](./iron/operators/elementwise_mul/) | +| [GEMM](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/linalg/mm.cc) | General Matrix Multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gemm/](./iron/operators/gemm/) | +| [Alternative GEMM](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/fused/mm_fused.h) | General Matrix Multiplication with a fused activation epilogue, specialised for transformer projection shapes. M, K and N are runtime parameters, so one xclbin serves every shape. B is stored as bfp16 on AIE2P | bfloat16, bfp16 | ✓ | ✓ | 🟢 | [iron/operators/flm/gemm/](./iron/operators/flm/gemm/) | +| [GEMV](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/linalg/mv_bf16.cc) | General Matrix-Vector Multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gemv/](./iron/operators/gemv/) | +| [GQA](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/linalg/mha.cc) | Grouped Query Attention kernel (Single pipeline) | bfloat16 | | ✓ | 🟢 | [iron/operators/mha/](./iron/operators/mha/) | +| [MHA](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/linalg/mha.cc) | Multi-Head Attention kernel & Grouped Query Attention | bfloat16 | | ✓ | 🟢 | [iron/operators/mha/](./iron/operators/mha/) | +| [RMSNorm](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/norm/rms_norm.cc) | RMSNorm kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rms_norm/](./iron/operators/rms_norm/) | +| [RoPE](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/datamovement/rope.cc) | Rotary Positional Embedding kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rope/](./iron/operators/rope/) | +| [SiLU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/activation/silu.cc) | Sigmoid Linear Unit activation kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/silu/](./iron/operators/silu/) | +| [Softmax](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/activation/softmax.cc) | Softmax kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/softmax/](./iron/operators/softmax/) | +| [Weighted RMSNorm](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/norm/rms_norm.cc) | Weighted RMSNorm kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rms_norm/](./iron/operators/rms_norm/) | +| [Copy](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/eltwise/passThrough.cc) | Copy | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/mem_copy/](./iron/operators/mem_copy/) | +| [Transpose](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/datamovement/transpose.cc) | Transpose | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/transpose/](./iron/operators/transpose/) | +| [AXPY](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/datamovement/axpy.cc) | AXPY | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/axpy/](./iron/operators/axpy/) | | [Reduction]() | Reduction | bfloat16 | | | 🟡 | | -| [Dequant](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/expand.cc) | Dequant Q4NX from [AWQ](https://github.com/mit-han-lab/llm-awq) to bfloat16 | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/dequant/](./iron/operators/dequant/) | -| [Dequant to bfp16](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/q4nx_dequant.cc) | Dequant Q4NX to bfp16, laid out for the B operand of the [alternative GEMM](./iron/operators/flm/gemm/) | q4nx → bfp16 | | ✓ | 🟢 | [iron/operators/flm/dequant/](./iron/operators/flm/dequant/) | -| [RELU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/relu.cc) | RELU | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/relu/](./iron/operators/relu/) | -| [Leaky RELU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/leaky_relu.cc) | Leaky RELU | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/leaky_relu/](./iron/operators/leaky_relu/) | -| [GELU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/gelu.cc) | GELU | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gelu/](./iron/operators/gelu/) | -| [LayerNorm](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/layer_norm.cc) | LayerNorm | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/layer_norm/](./iron/operators/layer_norm/) | +| [Dequant](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/datamovement/expand.cc) | Dequant Q4NX from [AWQ](https://github.com/mit-han-lab/llm-awq) to bfloat16 | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/dequant/](./iron/operators/dequant/) | +| [Dequant to bfp16](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/quant/q4nx_dequant.cc) | Dequant Q4NX to bfp16, laid out for the B operand of the [alternative GEMM](./iron/operators/flm/gemm/) | q4nx → bfp16 | | ✓ | 🟢 | [iron/operators/flm/dequant/](./iron/operators/flm/dequant/) | +| [RELU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/eltwise/relu.cc) | RELU | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/relu/](./iron/operators/relu/) | +| [Leaky RELU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/activation/leaky_relu.cc) | Leaky RELU | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/leaky_relu/](./iron/operators/leaky_relu/) | +| [GELU](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/activation/gelu.cc) | GELU | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gelu/](./iron/operators/gelu/) | +| [LayerNorm](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/norm/layer_norm.cc) | LayerNorm | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/layer_norm/](./iron/operators/layer_norm/) | | [Convolution]() | Convolution | bfloat16 | | | 🟡 | | | [MaxPool]() | MaxPool | bfloat16 | | | ⚪ | | | [AveragePool]() | AveragePool | bfloat16 | | | ⚪ | | -| [Tanh](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/tanh.cc) | Tanh kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/tanh/](./iron/operators/tanh/) | -| [Sigmoid](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/aie2/sigmoid.cc) | Sigmoid kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/sigmoid/](./iron/operators/sigmoid/) | +| [Tanh](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/activation/tanh.cc) | Tanh kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/tanh/](./iron/operators/tanh/) | +| [Sigmoid](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/activation/sigmoid.cc) | Sigmoid kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/sigmoid/](./iron/operators/sigmoid/) | > Use this dashboard to quickly check the status of each kernel and locate relevant setup, build, and usage information. @@ -199,7 +199,8 @@ IRON uses a three-layer architecture: - Each operator has: `op.py` (interface), `design.py` (MLIR-AIE implementation), `reference.py` (CPU reference), `test.py` (validation) 2. **AIE Kernels** ([mlir-aie `aie_kernels/`](https://github.com/Xilinx/mlir-aie/tree/main/aie_kernels)): Low-level C++ compute kernels - - Organized by architecture: `generic/`, `aie2/`, `aie2p/` + - Organized by family (`activation/`, `eltwise/`, `linalg/`, `norm/`, ...); architecture-specific + code lives in `*_aie2.h` / `*_aie2p.h` headers the family's `.cc` selects between - Vectorized using AIE API for optimal performance 3. **Common Infrastructure** (`iron/common/`): Compilation, device management, and utilities diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index da847c95bd..a6572bce67 100644 --- a/iron/common/stream/ops.py +++ b/iron/common/stream/ops.py @@ -13,7 +13,7 @@ the exporter emits them as a single node. Supporting a new op is one :class:`StreamKernel` plus one :data:`TORCH_OPS` entry -- -the kernel source is mlir-aie's ``aie_kernels//.cc``, exactly as the +the kernel source is mlir-aie's ``aie_kernels//.cc``, exactly as the hand-written operators use it. """ @@ -94,12 +94,12 @@ def _gemm_artifacts(kernels_dir, kernel_dir, m: int, k: int, n: int): from iron.common.compilation import KernelObjectArtifact, SourceArtifact suffix = f"{m}_{k}_{n}" - zero_source = kernels_dir / "generic" / "zero.cc" + zero_source = kernels_dir / "zero" / "zero.cc" return [ KernelObjectArtifact( f"mm_{suffix}.o", dependencies=[ - SourceArtifact(kernels_dir / kernel_dir / "mm.cc"), + SourceArtifact(kernels_dir / "linalg" / "mm.cc"), SourceArtifact(zero_source), ], extra_flags=[ @@ -129,8 +129,8 @@ class StreamKernel: """An AIE kernel: its stream-dse identity, its source, and its operand layouts. ``source``/``subdir`` name the file in mlir-aie's ``aie_kernels`` library the same - way the hand-written operators do (``subdir=None`` means the device directory, - e.g. ``aie2p``). The object name must equal the kernel's ``linkwith_name`` in + way the hand-written operators do (``subdir`` is the family directory, e.g. + ``activation``). The object name must equal the kernel's ``linkwith_name`` in stream-dse, since the generated MLIR links against it. """ @@ -146,12 +146,11 @@ def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs): return self.artifacts(kernels_dir, kernel_dir, **kwargs) from iron.common.compilation import KernelObjectArtifact, SourceArtifact - subdir = self.subdir or kernel_dir return [ KernelObjectArtifact( f"{self.source}.o", dependencies=[ - SourceArtifact(kernels_dir / subdir / f"{self.source}.cc") + SourceArtifact(kernels_dir / self.subdir / f"{self.source}.cc") ], ) ] @@ -159,13 +158,16 @@ def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs): GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts) SILU = StreamKernel( - key="silu", layouts=lambda: elementwise_layouts(2), source="silu", subdir="generic" + key="silu", + layouts=lambda: elementwise_layouts(2), + source="silu", + subdir="activation", ) ELTWISE_MUL = StreamKernel( key="eltwise_mul", layouts=lambda: elementwise_layouts(3), source="mul", - subdir="generic", + subdir="eltwise", ) Silu = custom_op("Silu") diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 162577b1b7..1e0b4a2cd3 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -438,7 +438,7 @@ def unit_rows(u): kernel_object, [ct_a_obj_ty, ct_b_ty, ct_acc_ty, np.int32], ) - # Same object as the mmul: the epilogue is compiled into mm_fused.cc, so + # Same object as the mmul: the epilogue is compiled into mm_fused.h, so # one -D flag set and one artifact cover both. epilogue_chunk = Kernel( EPILOGUE_SYMBOL, diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index 05f4aa1d22..10221a469e 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -374,11 +374,15 @@ def set_up_artifacts(self) -> None: def get_kernel_artifacts(self): # Built by hand rather than from aie.iron.kernels.fused_mm: that # factory compiles in one epilogue mode (this operator selects among - # several at runtime, from one xclbin), always rounds to nearest-even, - # and wraps mm_fused.cc in fused_mm_tile.cc. + # several at runtime, from one xclbin) and always rounds to + # nearest-even. Its translation unit, fused_mm_tile.cc, is still the + # one to compile, since mm_fused.h is a header. The whole-tile entry + # point it adds is never called, so the link drops it, but it also + # compiles out the per-step event0/event1 markers. kernel_dir = get_kernel_dir() kernels_dir = self.context.kernels_dir - generic = kernels_dir / "generic" + fused = kernels_dir / "fused" + common = kernels_dir / "common" # AIE2P lowers the 8x8x8 mmul onto two bfp16-emulated macs, which this # selects; AIE2 lowers it onto four native bf16 macs and ignores it. @@ -415,11 +419,12 @@ def get_kernel_artifacts(self): kernel_obj = KernelObjectArtifact( self._kernel_object, dependencies=[ - SourceArtifact(generic / "mm_fused.cc"), - SourceArtifact(generic / "mm_fused_mmul.h"), - SourceArtifact(generic / "activations.h"), + SourceArtifact(fused / "fused_mm_tile.cc"), + SourceArtifact(fused / "mm_fused.h"), + SourceArtifact(fused / "mm_fused_mmul.h"), + SourceArtifact(common / "activations.h"), SourceArtifact(kernels_dir / "aie_kernel_utils.h"), - SourceArtifact(generic / "zero.cc"), + SourceArtifact(common / "zero.h"), ], extra_flags=flags, ) diff --git a/iron/operators/gemm/op.py b/iron/operators/gemm/op.py index 3dbeb757b9..62b6142335 100644 --- a/iron/operators/gemm/op.py +++ b/iron/operators/gemm/op.py @@ -62,7 +62,7 @@ def __post_init__(self): raise ValueError(f"N ({self.N}) must be a multiple of {min_N}") # r, s, t are the aie::mmul tile dims the bf16 kernel is built from - # (aie_kernels/aie2p/mm.cc, matmul_vectorized_2x2_mmul) + # (aie_kernels/linalg/mm_aie2p.h, matmul_vectorized_2x2_mmul) if self.emulate_bf16_mmul_with_bfp16: r, s, t = 8, 8, 8 else: @@ -71,17 +71,17 @@ def __post_init__(self): if self.tile_m % min_tile_m != 0: raise ValueError( f"tile_m ({self.tile_m}) must be a multiple of {min_tile_m} " - f"(aie_kernels/aie2p/mm.cc requires m % (2*r) == 0, r={r})" + f"(aie_kernels/linalg/mm_aie2p.h requires m % (2*r) == 0, r={r})" ) if self.tile_k % min_tile_k != 0: raise ValueError( f"tile_k ({self.tile_k}) must be a multiple of {min_tile_k} " - f"(aie_kernels/aie2p/mm.cc requires k % s == 0, s={s})" + f"(aie_kernels/linalg/mm_aie2p.h requires k % s == 0, s={s})" ) if self.tile_n % min_tile_n != 0: raise ValueError( f"tile_n ({self.tile_n}) must be a multiple of {min_tile_n} " - f"(aie_kernels/aie2p/mm.cc requires n % (2*t) == 0, t={t})" + f"(aie_kernels/linalg/mm_aie2p.h requires n % (2*t) == 0, t={t})" ) MLIROperator.__init__(self, context=self.context) diff --git a/iron/operators/gemv/op.py b/iron/operators/gemv/op.py index c9aea02459..2ead44793a 100644 --- a/iron/operators/gemv/op.py +++ b/iron/operators/gemv/op.py @@ -98,8 +98,8 @@ def _matvec(self): ) def _gelu(self): - # The epilogue is gelu.cc's in-place gelu_tile_bf16, which only aie2p's - # gelu.cc exports; it rides in the object the gelu factory builds. + # The epilogue is gelu.cc's in-place gelu_tile_bf16, which only + # gelu_aie2p.h provides; it rides in the object the gelu factory builds. if get_kernel_dir() != "aie2p": raise NotImplementedError( "gemv gelu epilogue is only available on NPU2 (aie2p); " diff --git a/iron/operators/gemv/reference.py b/iron/operators/gemv/reference.py index 140d8a0de3..4d9af8d810 100644 --- a/iron/operators/gemv/reference.py +++ b/iron/operators/gemv/reference.py @@ -67,7 +67,7 @@ def generate_golden_reference_batched(M=128, K=128, num_batches=2, seed=42): def gelu_tanh_approx(x): - """Tanh-approximation GELU, matching aie_kernels/aie2p/gelu.cc. + """Tanh-approximation GELU, matching aie_kernels/activation/gelu_aie2p.h. 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). Computed in float32. """ diff --git a/iron/operators/rope/README.md b/iron/operators/rope/README.md index bdcbbe363b..a49158d5e4 100644 --- a/iron/operators/rope/README.md +++ b/iron/operators/rope/README.md @@ -5,7 +5,7 @@ SPDX-License-Identifier: Apache-2.0 # RoPE Example for Ryzen NPU -This repository contains an example implementation of the **RoPE (Rotary Position Embedding)** algorithm for the **Ryzen Neural Processing Unit (NPU)**. The implementation utilizes the [RoPE](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/generic/rope.cc) kernel to demonstrate the capabilities of the NPU in handling advanced embedding techniques. +This repository contains an example implementation of the **RoPE (Rotary Position Embedding)** algorithm for the **Ryzen Neural Processing Unit (NPU)**. The implementation utilizes the [RoPE](https://github.com/Xilinx/mlir-aie/blob/main/aie_kernels/datamovement/rope.cc) kernel to demonstrate the capabilities of the NPU in handling advanced embedding techniques. ## Notes The RoPE kernel offers two methods for applying embedding: a two-halves method (default) and an interleaved method. The two-halves method is what's used in Hugging Face's `transformer` library, while the interleaved method is used in Meta's official repo. Using the two-halves method is necessary when using Llama weights from Hugging Face, as the parameters of some layers are re-permuted while converting the Llama weights to Hugging Face. See [Issue #25199](https://github.com/huggingface/transformers/issues/25199) for reference. \ No newline at end of file diff --git a/iron/operators/swiglu_prefill_stream/README.md b/iron/operators/swiglu_prefill_stream/README.md index a75ac5afb3..bb10b2d2c1 100644 --- a/iron/operators/swiglu_prefill_stream/README.md +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -116,7 +116,7 @@ pytest iron/operators/swiglu_prefill_stream/test.py ## Adding another operator One `StreamKernel` plus one `TORCH_OPS` entry in `iron/common/stream/ops.py`, pointing -at mlir-aie's `aie_kernels//.cc`, plus that operator's own placement. The +at mlir-aie's `aie_kernels//.cc`, plus that operator's own placement. The kernel entry carries both the compile flags and the operand layouts, so the layout the generated DMAs produce and the layout the compiled object expects come from one place. diff --git a/iron/tests/compilation/kernel_object_arch_isolation.py b/iron/tests/compilation/kernel_object_arch_isolation.py index 031e6829a4..c3051c4bcc 100644 --- a/iron/tests/compilation/kernel_object_arch_isolation.py +++ b/iron/tests/compilation/kernel_object_arch_isolation.py @@ -6,7 +6,7 @@ KernelCompilationRule.compile() passes a different --target and aie_runtime_lib -I per arch for the same output filename (e.g. "mul.o"), and -for aie_kernels/generic/ sources the very same input file compiles to +for aie_kernels/ sources shared by both arches the very same input file compiles to different machine code per arch. CompilationArtifact.is_available_in_filesystem() only ever compares mtimes and never records which arch an object was built for, so if two arches' objects resolve to the same build_dir path, whichever @@ -44,7 +44,7 @@ def _mul_kernel_object(build_dir, device): def test_two_arches_do_not_resolve_the_same_kernel_object_path(tmp_path): - """aie_kernels/generic/mul.cc is one source shared by aie2 and aie2p + """aie_kernels/eltwise/mul.cc is one source shared by aie2 and aie2p (eltwise.mul_sized); its object must not collide in build_dir.""" aie2 = _mul_kernel_object(tmp_path, NPU1()) aie2p = _mul_kernel_object(tmp_path, NPU2()) diff --git a/iron/tests/operators/gemm_tile_divisibility.py b/iron/tests/operators/gemm_tile_divisibility.py index 6eacfbede5..b8f72a3fbf 100644 --- a/iron/tests/operators/gemm_tile_divisibility.py +++ b/iron/tests/operators/gemm_tile_divisibility.py @@ -4,7 +4,7 @@ """GEMM.__post_init__ must reject tile sizes the kernel cannot build. -aie_kernels/aie2p/mm.cc's matmul_vectorized_*x*x*_bf16_* wrappers static_assert +aie_kernels/linalg/mm_aie2p.h's matmul_vectorized_*x*x*_bf16_* wrappers static_assert m % (2*r) == 0, k % s == 0 and n % (2*t) == 0 for the (r, s, t) triple selected by emulate_bf16_mmul_with_bfp16 (r=t=8 when enabled, the default; r=4, t=8 otherwise). A tile size that meets a lower bound without dividing evenly