diff --git a/.github/actions/prereqs/action.yaml b/.github/actions/prereqs/action.yaml index 7363b0fbf..bf1616a01 100644 --- a/.github/actions/prereqs/action.yaml +++ b/.github/actions/prereqs/action.yaml @@ -21,4 +21,11 @@ runs: source ${{ inputs.env_name }}/bin/activate pip install --upgrade pip pip install -r requirements.txt + pip install -r requirements_stream.txt + # stream-dse's pure-Python AIE codegen dialects (xdsl-aie, snax-mlir) are + # git installs, not PyPI dependencies; this console script installs them so + # the stream-backed operators can generate their MLIR at build time. It does + # NOT install mlir_aie/llvm-aie -- those stay pinned by requirements.txt above + # (stream-dse>=1.13.7; older releases reinstalled/downgraded them here). + stream-setup-aie echo "Prerequisites installed into ${{ inputs.env_name }}" diff --git a/.gitignore b/.gitignore index d19c375d6..ec6f4f799 100755 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,4 @@ id_ed25519.pub .cline_storage *.egg-info **/*.prj/** +/outputs/ diff --git a/ci/docker-based/Dockerfile b/ci/docker-based/Dockerfile index ccffadcf2..38081e1c8 100644 --- a/ci/docker-based/Dockerfile +++ b/ci/docker-based/Dockerfile @@ -18,6 +18,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libssl-dev libboost-dev libboost-filesystem-dev libboost-program-options-dev uuid-dev \ # Required by MHA kernel llvm-18 \ + # graphviz provides the `dot` binary that stream-dse's codegen shells out to + # (via pydot) to render its workload graph; without it the stream-backed + # operators fail at MLIR-generation time with: "dot" not found in path. + graphviz \ # GitHub Actions runner requirements git jq curl tar \ && rm -rf /var/lib/apt/lists/* diff --git a/ci/docker-based/test_docker_ci.sh b/ci/docker-based/test_docker_ci.sh index 5d44c63a2..6eef50173 100755 --- a/ci/docker-based/test_docker_ci.sh +++ b/ci/docker-based/test_docker_ci.sh @@ -8,7 +8,9 @@ GITHUB_OWNER="amd" GITHUB_REPO="IRON" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -GITHUB_PAT=$(cat "${SCRIPT_DIR}/secret_github_token") +# This script drops into an interactive shell (CMD override below), which does not +# register a runner, so the PAT is optional here -- don't fail if it is absent. +GITHUB_PAT=$(cat "${SCRIPT_DIR}/secret_github_token" 2>/dev/null || true) DATE=$(printf '%(%Y_%m_%d_%H_%M_%S)T') NAME="ci-run-${DATE}" diff --git a/iron/common/__init__.py b/iron/common/__init__.py index 9f5a8f7ae..cb2ff31be 100644 --- a/iron/common/__init__.py +++ b/iron/common/__init__.py @@ -18,3 +18,4 @@ PythonGeneratedMLIRArtifact, DesignGenerator, ) +from .layout import Stride, TiledStride, TiledStridedLayout, tiled_2d diff --git a/iron/common/base.py b/iron/common/base.py index e8dea2535..701e90dfe 100644 --- a/iron/common/base.py +++ b/iron/common/base.py @@ -115,6 +115,15 @@ class MLIROperator(AIEOperatorBase): def operator_dir(self) -> Path: return Path(inspect.getfile(type(self))).parent + def design_key(self) -> str | None: + """Identifies the design this operator compiles to, for sharing it. + + Two operators returning the same key must produce byte-identical MLIR before + the fused build prefixes their kernel symbols, and must take the same runtime + argument shapes. ``None`` means the design is never shared. + """ + return None + @property def name(self) -> str: """Unique name for this operator instance, derived from its parameters. diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index da15df038..8b06de537 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -318,10 +318,12 @@ def __init__( filename: str, mlir_input: CompilationArtifact, dependencies: list[CompilationArtifact], + extra_flags: list[str] | None = None, ) -> None: if mlir_input not in dependencies: dependencies = dependencies + [mlir_input] super().__init__(filename, dependencies) + self.extra_flags = extra_flags if extra_flags is not None else [] class XclbinArtifact(_MLIRInputMixin, CompilationArtifact): @@ -535,6 +537,7 @@ def compile(self, graph): "--generate-full-elf", "--full-elf-name", os.path.abspath(artifact.filename), + *artifact.extra_flags, os.path.abspath(artifact.mlir_input.filename), ] commands.append( diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index c941a8572..52382b805 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -27,6 +27,9 @@ MLIRArtifact, ) +RESET_DEVICE = "reset_device" + + # Compilation Artifacts # ########################################################################## @@ -89,6 +92,26 @@ def get_child_mlir_module(mlir_artifact: PythonGeneratedMLIRArtifact) -> Any: return callback_function(*gen.args, **gen.kwargs) +def needs_additional_reset(runlist: list[Any]) -> bool: + """Whether the sequence must configure one more device than the runlist asks for. + + ``aiecc --expand-load-pdis`` marks each configure point by loading one of two + otherwise empty PDIs, alternating between them from a fixed start. A load of the + PDI already loaded has no effect, so a sequence with an odd number of configure + points ends on the one the next dispatch starts with, and that dispatch + reconfigures over the state the last design left. Configuring one more device + makes the count even. Consecutive entries running the same operator share a + configure point. + """ + points = 0 + previous = None + for op_name, *_ in runlist: + if op_name != previous: + points += 1 + previous = op_name + return points % 2 == 1 + + def fuse_mlir(artifact: SequenceMLIRArtifact) -> None: """Fuse multiple MLIR modules by inlining their device operations and adding a new main device and runtime sequence that call into sequence of operations based on a runlist.""" @@ -170,6 +193,17 @@ def fuse_mlir(artifact: SequenceMLIRArtifact) -> None: dev_op.sym_name = ir.StringAttr.get(op_name) ctx.module.body.append(dev_op) + needs_reset = needs_additional_reset(artifact.runlist) + if needs_reset: + + @aie.device(device_ty) + def reset(): + @aiex.runtime_sequence() + def sequence(): + pass + + reset.operation.attributes["sym_name"] = ir.StringAttr.get(RESET_DEVICE) + # Create the main device -- this contains the runtime sequence calling into the other devices @aie.device(device_ty) def main(): @@ -274,6 +308,10 @@ def sequence(input_buf, output_buf, scratch_buf): sequence_sym_ref_attr = ir.FlatSymbolRefAttr.get("sequence") run_op = aiex.RunOp(sequence_sym_ref_attr, buffer_ssa_values) + if needs_reset: + reset_op = aiex.ConfigureOp(ir.FlatSymbolRefAttr.get(RESET_DEVICE)) + reset_op.body.blocks.append() + # Write the fused MLIR to file with open(artifact.filename, "w") as f: f.write(str(ctx.module)) diff --git a/iron/common/layout.py b/iron/common/layout.py new file mode 100644 index 000000000..09770524c --- /dev/null +++ b/iron/common/layout.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tiled-strided memory layouts for IRON operators. + +A tiled-strided layout describes how a logical multi-dimensional tensor is laid +out in memory as a hierarchy of tiles, each level carrying its own ``(step, +bound)`` stride. It is the layout model AIE kernels are written against: a GEMM +microkernel, for example, reads its ``MxK`` operand as ``mt x kt`` tiles of +``r x s`` elements, which is exactly a two-level tiled-strided layout. + +The types here mirror ``snaxc.ir.tsl`` (``Stride`` -> ``TiledStride`` -> +``TiledStridedLayout``) so an IRON-authored layout can be handed to stream-dse's +code generation verbatim via :meth:`TiledStridedLayout.to_snaxc`. They carry no +stream-dse / snaxc / xdsl dependency themselves -- the snaxc import is lazy and +confined to ``to_snaxc`` -- so they are usable (and testable) in a plain IRON +install with no AIE codegen toolchain present. + +This is a common primitive: it is meant to be shared across operators as the one +place a kernel's operand layouts are defined, rather than re-derived per operator +or hand-copied into stream-dse. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Stride: + """One stride level: ``bound`` elements spaced ``step`` apart. + + ``step``/``bound`` may be ``None`` to denote a dynamic (run-time) value, + matching snaxc's convention. + """ + + step: int | None + bound: int | None + + +@dataclass +class TiledStride: + """The strides of a single tensor dimension, outermost tile first. + + A simple (untiled) dimension has one stride; one level of tiling has two + (the outer tile stride followed by the inner element stride), and so on. + """ + + strides: tuple[Stride, ...] + + def __post_init__(self) -> None: + self.strides = tuple(self.strides) + + +@dataclass +class TiledStridedLayout: + """A tiled-strided layout: one :class:`TiledStride` per tensor dimension.""" + + tstrides: tuple[TiledStride, ...] + offset: int = 0 + + def __post_init__(self) -> None: + self.tstrides = tuple(self.tstrides) + + def to_snaxc(self): + """Return the equivalent ``snaxc.ir.tsl.TiledStridedLayout``. + + The snaxc import is deferred to here so this module stays usable without + the AIE codegen toolchain installed. Used to feed IRON-authored layouts + into stream-dse code generation. + """ + from snaxc.ir.tsl import ( + Stride as SnaxStride, + TiledStride as SnaxTiledStride, + TiledStridedLayout as SnaxTiledStridedLayout, + ) + + return SnaxTiledStridedLayout( + [ + SnaxTiledStride([SnaxStride(s.step, s.bound) for s in ts.strides]) + for ts in self.tstrides + ], + offset=self.offset, + ) + + +def tiled_2d(rows: int, cols: int, row_unit: int, col_unit: int) -> TiledStridedLayout: + """Two-level tiled-strided layout for a ``rows x cols`` tensor. + + The tensor is tiled into ``(rows // row_unit) x (cols // col_unit)`` tiles of + ``row_unit x col_unit`` elements, the tiles laid out row-major and each tile + stored row-major internally. This reproduces stream-dse's GEMM/elementwise + operand layouts (the intrinsic ``row_unit``/``col_unit`` are the kernel's MAC + tile dimensions). + """ + rows_t, cols_t = rows // row_unit, cols // col_unit + return TiledStridedLayout( + ( + TiledStride( + ( + Stride(row_unit * col_unit * cols_t, rows_t), + Stride(col_unit, row_unit), + ) + ), + TiledStride((Stride(row_unit * col_unit, cols_t), Stride(1, col_unit))), + ) + ) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 747f3244d..c4034cad4 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -85,6 +85,7 @@ def set_up_artifacts(self, seq): f"{seq.name}.elf", mlir_input=mlir_artifact, dependencies=[mlir_artifact] + kernel_objects, + extra_flags=seq.extra_flags, ) seq.add_artifacts([full_elf_artifact]) @@ -97,18 +98,19 @@ def build_fused_mlir(self, seq): """ operator_mlir_map = {} comp_runlist = [] - op_names = {} # id(op) -> op_name + designs, design_of = seq.unique_designs() + design_names = [] - for idx, op in enumerate(seq.unique_operators()): + for idx, op in enumerate(designs): mlir_artifact = op.get_mlir_artifact() if len(op.get_kernel_artifacts()) > 0: mlir_artifact.generator.kwargs["func_prefix"] = f"op{idx}_" op_name = f"op{idx}_{op.__class__.__name__}" - op_names[id(op)] = op_name + design_names.append(op_name) operator_mlir_map[op_name] = mlir_artifact for op, *bufs in seq.runlist: - comp_runlist.append((op_names[id(op)], *bufs)) + comp_runlist.append((design_names[design_of[id(op)]], *bufs)) return comp.SequenceMLIRArtifact( seq.name + "_fused.mlir", @@ -122,7 +124,7 @@ def build_fused_mlir(self, seq): def _collect_kernel_artifacts(self, seq): """Kernel artifacts from all child operators, prefixed per operator index.""" kernel_artifacts = [] - for idx, op in enumerate(seq.unique_operators()): + for idx, op in enumerate(seq.unique_designs()[0]): objs = op.get_kernel_artifacts() for obj in objs: obj.filename = f"op{idx}_{obj.filename}" @@ -262,6 +264,8 @@ def __init__( output_args, buffer_sizes=None, dispatch="auto", + extra_flags=None, + share_designs=False, *args, **kwargs, ): @@ -276,12 +280,19 @@ def __init__( ) super().__init__(*args, **kwargs) self.runlist = runlist - self.name = name + # Sharing changes which designs are built, so it belongs in the name that + # keys the build artifacts. + self.name = name + "_shared" if share_designs else name self.input_args = input_args self.output_args = output_args self.explicit_buffer_sizes = ( buffer_sizes or {} ) # Optional dict: buffer_name -> size_in_bytes + # Extra aiecc flags forwarded to the full-ELF build (e.g. --dynamic-objFifos + # for placed/routed whole-array designs that would otherwise overflow AIE2p + # program memory). Empty by default, so other sequences are unaffected. + self.extra_flags = extra_flags or [] + self.share_designs = share_designs self._dispatch = dispatch @staticmethod @@ -300,6 +311,32 @@ def unique_operators(self): seen.setdefault(id(op), op) return list(seen.values()) + def unique_designs(self): + """The designs to build, and which design each operator uses. + + With ``share_designs`` set, operators reporting the same ``design_key`` + collapse onto one design, so it is built, prefixed and configured once. + """ + designs = [] + design_of = {} + first_with_key = {} + for op in self.unique_operators(): + key = op.design_key() if self.share_designs else None + if key is not None and key in first_with_key: + shared = designs[first_with_key[key]] + if op.get_arg_spec() != shared.get_arg_spec(): + raise ValueError( + f"{op.name} and {shared.name} report the same design_key but " + "different runtime arguments, so the design cannot be shared" + ) + design_of[id(op)] = first_with_key[key] + continue + if key is not None: + first_with_key[key] = len(designs) + design_of[id(op)] = len(designs) + designs.append(op) + return designs, design_of + def calculate_buffer_layout(self): args = {} # base_buffer_name -> args_spec sliced_buffers = ( diff --git a/iron/common/stream/__init__.py b/iron/common/stream/__init__.py new file mode 100644 index 000000000..1e35c3fae --- /dev/null +++ b/iron/common/stream/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Building blocks for stream-dse-backed operators. + +An operator supplies a reference ``nn.Module`` and a placement; these modules turn +that into everything stream-dse needs: + +* :mod:`~iron.common.stream.ops` -- the registry binding a torch ATen op to its ONNX + form, its stream-dse kernel and IRON's ``aie_kernels`` source. +* :mod:`~iron.common.stream.workload` -- ``torch.export`` of the module into the ONNX + workload stream-dse optimizes. +* :mod:`~iron.common.stream.mapping` -- the mapping YAML, named from that same graph. + +The submodules are not re-exported here: they need ``onnx``/``pyyaml`` (installed +with stream-dse, see ``requirements_stream.txt``), so importing an operator must not +pull them in. Import them directly from the module that builds the design. +""" diff --git a/iron/common/stream/hardware.py b/iron/common/stream/hardware.py new file mode 100644 index 000000000..532e0fefd --- /dev/null +++ b/iron/common/stream/hardware.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Address the device's compute tiles the way the array is laid out. + +stream identifies a tile by an integer core id, which says nothing on its own. +IRON already describes the device: mlir-aie's :class:`~aie.iron.device.Device` +knows the grid and the type of every tile in it. This turns that into columns of +core ids, so a placement is written in columns and rows and no other IRON module +has to know what a stream core id means. + +An operator that gives every layer the whole array (the layer-by-layer designs, +and a single layer sent to stream for a performance estimate) asks for +:attr:`ComputeArray.all_columns`. An operator that pipelines layers across the +array asks :meth:`ComputeArray.allocate` for a column budget per layer. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable, Sequence + +from aie.iron.device.device import AIETileType + + +@dataclass(frozen=True) +class ComputeArray: + """The device's compute tiles, as stream core ids grouped by column.""" + + columns: tuple[tuple[int, ...], ...] + + @classmethod + def from_device(cls, device) -> ComputeArray: + """Read the compute grid from an mlir-aie ``Device``. + + stream numbers tiles ``column * rows + row`` across the device's whole + grid, so the stride is the device's row count -- shim and memory rows + included -- and not the number of compute rows. + """ + return cls( + tuple( + tuple( + column * device.rows + row + for row in range(device.rows) + if device.get_tile_type(column, row) == AIETileType.CoreTile + ) + for column in range(device.cols) + ) + ) + + @property + def num_columns(self) -> int: + return len(self.columns) + + @property + def num_rows(self) -> int: + return len(self.columns[0]) if self.columns else 0 + + @property + def all_columns(self) -> tuple[int, ...]: + return tuple(range(self.num_columns)) + + def cores( + self, columns: Iterable[int], rows: Sequence[int] | None = None + ) -> tuple[int, ...]: + """The core ids of ``columns``, column by column and row by row. + + ``rows`` takes only some rows of each column, for a layer that wants the + array's width but not its full depth -- an elementwise layer with one + worker per column, as IRON's own channeled operators place it. + """ + selected = range(self.num_rows) if rows is None else rows + return tuple( + self.columns[column][row] for column in columns for row in selected + ) + + def allocate(self, budgets: Sequence[int]) -> tuple[tuple[int, ...], ...]: + """Consecutive, disjoint column ranges, one per budget. + + Layers placed on disjoint columns pipeline across steady-state + iterations instead of taking turns on the same tiles. + """ + if sum(budgets) > self.num_columns: + raise ValueError( + f"column budgets {list(budgets)} need {sum(budgets)} columns, " + f"the array has {self.num_columns}" + ) + ranges, first = [], 0 + for budget in budgets: + ranges.append(tuple(range(first, first + budget))) + first += budget + return tuple(ranges) diff --git a/iron/common/stream/mapping.py b/iron/common/stream/mapping.py new file mode 100644 index 000000000..68b3f7251 --- /dev/null +++ b/iron/common/stream/mapping.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Emit the stream-dse mapping for an exported workload. + +An operator declares *where* each node runs (:class:`Placement`) and how nodes +are fused (:class:`FusedGroup`); this module turns that into the mapping YAML +stream-dse consumes. Node names are taken from the +:class:`~iron.common.stream.workload.StreamWorkload` the ONNX was generated from, +and every placement is checked against it, so a mapping can never refer to a node +the workload does not contain. + +The placement is deliberately explicit -- which compute tiles a layer occupies is +a hardware decision worth reading in one place -- while the repetitive YAML +plumbing is generated. Placements carry no absolute tensor dimensions, only tile +sizes and core sets, so they hold across problem sizes. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Sequence + +from iron.common.stream.hardware import ComputeArray +from iron.common.stream.workload import StreamWorkload + + +@dataclass(frozen=True) +class Placement: + """Where one workload node runs. + + ``columns`` are the array columns it occupies, resolved to core ids against + the :class:`~iron.common.stream.hardware.ComputeArray`; ``rows`` narrows that + to some rows of each column (all of them by default); ``splits`` is the + inter-core tiling as ``(dim, split)`` pairs; ``kernel_kwargs`` are the + arguments of the node's stream-dse kernel (e.g. a GEMM's tile shape). + """ + + columns: Sequence[int] + splits: Sequence[tuple[str, int]] = () + kernel_kwargs: dict = field(default_factory=dict) + rows: Sequence[int] | None = None + + +@dataclass(frozen=True) +class FusedGroup: + """A set of nodes fused into one design, with its layer-fusion tiling. + + ``intra_core_tiling`` entries are ``(node_name, dim, tile)``. Splitting a + workload into several groups makes stream-dse emit one design per group. + """ + + name: str + layers: Sequence[str] + intra_core_tiling: Sequence[tuple[str, str, int]] = () + + +def group_boundaries( + workload: StreamWorkload, group_layers: Sequence[Sequence[str]] +) -> list[tuple[tuple[str, ...], tuple[str, ...]]]: + """Per group, the tensors it consumes from outside and produces for outside. + + These are the group's runtime arguments: what an operator built from several + fused groups has to hand from one to the next. Both sequences follow the order + the group's nodes use them, which is the order stream-dse gives the generated + design its arguments in. + """ + graph = workload.model.graph + produced_by = {out: node.name for node in graph.node for out in node.output} + consumers: dict[str, list[str]] = {} + for node in graph.node: + for tensor in node.input: + consumers.setdefault(tensor, []).append(node.name) + graph_outputs = {out.name for out in graph.output} + + boundaries = [] + for layers in group_layers: + members = set(layers) + inputs: list[str] = [] + outputs: list[str] = [] + for node in graph.node: + if node.name not in members: + continue + for tensor in node.input: + if produced_by.get(tensor) not in members and tensor not in inputs: + inputs.append(tensor) + for tensor in node.output: + escapes = tensor in graph_outputs or any( + consumer not in members for consumer in consumers.get(tensor, []) + ) + if escapes and tensor not in outputs: + outputs.append(tensor) + boundaries.append((tuple(inputs), tuple(outputs))) + return boundaries + + +def _layer_entry( + name: str, placement: Placement, kernel_key: str, array: ComputeArray +) -> dict: + return { + "name": name, + "core_allocation": [list(array.cores(placement.columns, placement.rows))], + "inter_core_tiling": [ + [{"dim": dim, "split": split} for dim, split in placement.splits] + ], + "kernel": {"name": kernel_key, "kwargs": dict(placement.kernel_kwargs)}, + } + + +def build_mapping( + workload: StreamWorkload, + placements: dict[str, Placement], + groups: Sequence[FusedGroup], + array: ComputeArray, +) -> dict: + """The mapping for ``workload`` as a plain dict (validated against it).""" + kernel_of = dict(workload.nodes) + unknown = set(placements) - set(kernel_of) + if unknown: + raise ValueError( + f"placements refer to nodes absent from the workload: {sorted(unknown)}" + ) + + grouped = [name for group in groups for name in group.layers] + missing = [name for name in grouped if name not in placements] + if missing: + raise ValueError(f"fused groups refer to nodes without a placement: {missing}") + + layers = [ + _layer_entry(name, placements[name], kernel, array) + for name, kernel in workload.nodes + if name in placements + ] + return { + "layers": layers, + "fused_groups": [ + { + "name": group.name, + "layers": list(group.layers), + "intra_core_tiling": [ + {"dim": f"{node}.{dim}", "tile": tile} + for node, dim, tile in group.intra_core_tiling + ], + } + for group in groups + ], + "runtime_args": {buffer: {} for buffer in workload.buffers}, + } + + +def emit_mapping( + workload: StreamWorkload, + placements: dict[str, Placement], + groups: Sequence[FusedGroup], + array: ComputeArray, + path, +) -> str: + """Write the mapping YAML to ``path`` and return it.""" + import yaml + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + yaml.dump( + build_mapping(workload, placements, groups, array), + f, + default_flow_style=False, + sort_keys=False, + ) + return str(path) diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py new file mode 100644 index 000000000..825527ef4 --- /dev/null +++ b/iron/common/stream/ops.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Registry binding torch operators to their ONNX form and their AIE kernel. + +One :class:`StreamOp` entry per supported torch operator is all a stream-dse-backed +operator needs: how the op is emitted by the ONNX exporter, which stream-dse kernel +implements it, which ``aie_kernels`` source that kernel is compiled from, and what +operand layouts the generated DMAs must use. + +Ops stream-dse implements with a fused kernel but ONNX has no operator for are +declared with :func:`custom_op`, which gives them a schema in a private domain so +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 IRON's existing ``aie_kernels//.cc``, exactly as the +hand-written operators use it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch +from onnx import defs +from onnxscript import opset18 +from onnxscript.values import Op, Opset + +from iron.common.layout import TiledStridedLayout, tiled_2d + +# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets. The +# operand layouts are the contract the generated DMAs and the compiled kernel +# objects agree on. +# mm.cc takes an 8-row MAC tile when bf16 matmuls run on the bfp16 MACs and a +# 4-row one when they do not. +R, S, T = 4, 8, 8 +MAC_ROWS_BFP16 = 8 + +# Element tile the stream-dse elementwise kernels are written against. +ELEMENTWISE_TILE = (32, 64) + +# Private domain for ops that exist as an AIE kernel but not as an ONNX operator. +CUSTOM_DOMAIN = Opset("com.example", 1) + +_ELEMENT_TYPES = ["tensor(bfloat16)", "tensor(float)"] + + +def custom_op(name: str, arity: int = 1) -> Op: + """An operator in :data:`CUSTOM_DOMAIN`, emitted by the exporter as one node.""" + schema = defs.OpSchema( + name, + CUSTOM_DOMAIN.domain, + CUSTOM_DOMAIN.version, + inputs=[defs.OpSchema.FormalParameter(f"X{i}", "T") for i in range(arity)], + outputs=[defs.OpSchema.FormalParameter("Y", "T")], + type_constraints=[("T", _ELEMENT_TYPES, "")], + ) + return Op(CUSTOM_DOMAIN, name, schema) + + +def mac_rows(bfp16_mmul: bool) -> int: + """Rows of the MAC tile a kernel object compiled this way takes.""" + return MAC_ROWS_BFP16 if bfp16_mmul else R + + +def gemm_layouts( + m: int, k: int, n: int, bfp16_mmul: bool = False +) -> tuple[TiledStridedLayout, ...]: + """Layouts of a GEMM's ``A[m,k]``, ``B[k,n]`` and ``C[m,n]`` operands.""" + rows = mac_rows(bfp16_mmul) + return (tiled_2d(m, k, rows, S), tiled_2d(k, n, S, T), tiled_2d(m, n, rows, T)) + + +def elementwise_layouts( + nb_operands: int, bfp16_mmul: bool = False +) -> tuple[TiledStridedLayout, ...]: + """Identical tiled layout for each operand of an elementwise kernel.""" + return (tiled_2d(*ELEMENTWISE_TILE, mac_rows(bfp16_mmul), T),) * nb_operands + + +def _gemm_artifacts(base_dir, kernel_dir, m: int, k: int, n: int): + """The ``mm.cc`` object specialized for one tile shape. + + stream-dse emits dimension-suffixed symbols so GEMMs of different tile shapes + coexist in one design (``GemmKernel.function_name``/``zero_name``); rename + ``mm.cc``'s unsuffixed symbols to match. + """ + from iron.common.compilation import KernelObjectArtifact, SourceArtifact + + suffix = f"{m}_{k}_{n}" + return [ + KernelObjectArtifact( + f"mm_{suffix}.o", + dependencies=[ + SourceArtifact(base_dir / "aie_kernels" / kernel_dir / "mm.cc") + ], + extra_flags=[ + f"-DDIM_M={m}", + f"-DDIM_K={k}", + f"-DDIM_N={n}", + "-Dbf16_bf16_ONLY", + # Emulating the matmul on the bfp16 MACs is what makes the 8-row + # MAC tile available, so it and the layouts move together. + "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", + "-DROUND_CONV_EVEN", + ], + rename_symbols={ + "matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}", + "zero_bf16": f"zero_bf16_{suffix}", + }, + ) + ] + + +@dataclass(frozen=True) +class StreamKernel: + """An AIE kernel: its stream-dse identity, its source, and its operand layouts. + + ``source``/``subdir`` name the file in IRON'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 + stream-dse, since the generated MLIR links against it. + """ + + key: str # stream-dse AIEKernels key + layouts: Callable[..., tuple[TiledStridedLayout, ...]] + source: str | None = None + subdir: str | None = None + artifacts: Callable | None = None # overrides source/subdir when tile-specialized + + def kernel_artifacts(self, base_dir, kernel_dir, **kwargs): + """Compilation artifacts building this kernel's object file.""" + if self.artifacts is not None: + return self.artifacts(base_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( + base_dir / "aie_kernels" / subdir / f"{self.source}.cc" + ) + ], + ) + ] + + +GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts) +SILU = StreamKernel(key="silu", layouts=lambda: elementwise_layouts(2), source="silu") +ELTWISE_MUL = StreamKernel( + key="eltwise_mul", + layouts=lambda: elementwise_layouts(3), + source="mul", + subdir="generic", +) + +Silu = custom_op("Silu") + + +def _to_gemm(a, b): + return opset18.Gemm(a, b) + + +def _to_silu(x): + return Silu(x) + + +def _to_mul(a, b): + return opset18.Mul(a, b) + + +@dataclass(frozen=True) +class StreamOp: + """How one torch operator is exported, and which kernel runs it. + + ``translation`` overrides how the exporter lowers the operator, and is needed + only when its default lowering is not what stream-dse parses. Leaving it unset + keeps the exporter's own lowering and just binds the resulting ONNX operator to + a kernel. + """ + + onnx_type: str + kernel: StreamKernel + translation: Callable | None = None + + +# torch operator -> its ONNX form and AIE kernel. Gemm rather than the exporter's +# default MatMul because stream-dse's Gemm parser iterates (m, k, n), which is the +# order the mappings address as D0/D1/D2. +TORCH_OPS: dict[Callable, StreamOp] = { + torch.ops.aten.matmul.default: StreamOp("Gemm", GEMM, _to_gemm), + torch.ops.aten.silu.default: StreamOp("Silu", SILU, _to_silu), + torch.ops.aten.mul.Tensor: StreamOp("Mul", ELTWISE_MUL, _to_mul), +} + +_BY_ONNX_TYPE = {op.onnx_type: op for op in TORCH_OPS.values()} + + +def translation_table() -> dict[Callable, Callable]: + """The ``custom_translation_table`` for :func:`torch.onnx.export`.""" + return { + target: op.translation + for target, op in TORCH_OPS.items() + if op.translation is not None + } + + +def op_for_onnx_type(onnx_type: str) -> StreamOp: + """The :class:`StreamOp` an exported node's operator type belongs to.""" + try: + return _BY_ONNX_TYPE[onnx_type] + except KeyError: + raise NotImplementedError( + f"ONNX operator '{onnx_type}' has no stream-dse mapping; " + f"add it to iron.common.stream.ops.TORCH_OPS" + ) from None diff --git a/iron/common/stream/workload.py b/iron/common/stream/workload.py new file mode 100644 index 000000000..d573dc849 --- /dev/null +++ b/iron/common/stream/workload.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Export a reference ``nn.Module`` into the ONNX workload stream-dse optimizes. + +:func:`torch.onnx.export` captures the module and lowers it through the +translation table in :mod:`~iron.common.stream.ops`, so every operator is emitted +in the form stream-dse's parsers expect. Because the operator's reference module is +the only description of the computation, the generated design cannot drift from the +reference the operator is tested against. + +The exported graph is then adjusted for what stream-dse consumes: weight payloads +are dropped (it needs shapes, not values) and nodes and their results are given the +operator's names, which the mapping refers to. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from iron.common.stream.ops import op_for_onnx_type, translation_table + +_WEIGHT_DATA_FIELDS = ( + "float_data", + "double_data", + "int32_data", + "int64_data", + "uint64_data", + "raw_data", +) + + +@dataclass(frozen=True) +class StreamWorkload: + """An exported workload: the ONNX model and the names the mapping refers to.""" + + model: object # onnx.ModelProto + nodes: tuple[tuple[str, str], ...] # (node name, kernel key), topological order + buffers: tuple[str, ...] # runtime buffer names, in argument order + + @property + def shapes(self) -> dict[str, tuple[int, ...]]: + """Every named tensor's shape, as the exported graph declares it.""" + graph = self.model.graph + declared = list(graph.input) + list(graph.value_info) + list(graph.output) + shapes = { + value.name: tuple(d.dim_value for d in value.type.tensor_type.shape.dim) + for value in declared + } + shapes.update( + { + initializer.name: tuple(initializer.dims) + for initializer in graph.initializer + } + ) + return shapes + + def write(self, path) -> str: + """Write the ONNX model to ``path`` and return it. + + stream-dse's parser loads the workload from a file, so the model is + materialized at build time (under the build directory, never in the tree). + """ + import onnx + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + onnx.save(self.model, str(path)) + return str(path) + + +def _drop_weight_data(model) -> None: + """Keep each initializer's shape and type, drop its values.""" + for initializer in model.graph.initializer: + for field in _WEIGHT_DATA_FIELDS: + initializer.ClearField(field) + + +def _rename(model, node_names, result_names, output_name: str) -> None: + """Give nodes and the tensors they produce the operator's names.""" + nodes = list(model.graph.node) + if node_names is not None and len(node_names) != len(nodes): + raise ValueError( + f"node_names has {len(node_names)} entries but the exported graph has " + f"{len(nodes)} nodes" + ) + + graph_output = model.graph.output[0].name + renamed: dict[str, str] = {} + for position, node in enumerate(nodes): + if node_names is not None: + node.name = node_names[position] + produced = node.output[0] + target = ( + output_name + if produced == graph_output + else result_names.get(node.name, f"out_{node.name}") + ) + renamed[produced] = target + + for node in nodes: + node.input[:] = [renamed.get(name, name) for name in node.input] + node.output[:] = [renamed.get(name, name) for name in node.output] + for value in model.graph.value_info: + value.name = renamed.get(value.name, value.name) + for output in model.graph.output: + output.name = renamed.get(output.name, output.name) + + +def export_workload( + module, + example_inputs, + node_names=None, + result_names=None, + output_name: str = "output", +) -> StreamWorkload: + """Export ``module`` into a :class:`StreamWorkload`. + + ``example_inputs`` fixes the shapes: re-exporting with different ones is all + that is needed for a different problem size, since the mapping carries tile + sizes and placement but no absolute dimensions. + + Tensor names come from the module -- its ``forward`` argument names and its + parameter names. The exporter names the computation nodes and their results + after the operators it captured, which the mapping would then have to refer to, + so both can be renamed: + + * ``node_names`` -- computation nodes, in topological order. + * ``result_names`` -- node name -> the name of the tensor it produces + (default ``out_{node_name}``). The final result is always ``output_name``. + """ + import torch + + program = torch.onnx.export( + module, + tuple(example_inputs), + dynamo=True, + custom_translation_table=translation_table(), + optimize=False, + verbose=False, + ) + model = program.model_proto + _drop_weight_data(model) + _rename(model, node_names, result_names or {}, output_name) + + nodes = tuple( + (node.name, op_for_onnx_type(node.op_type).kernel.key) + for node in model.graph.node + ) + # Runtime buffer order: activations, then weights in module order, then the + # output -- the order the operator's argument spec and the mapping follow. + buffers = ( + tuple(i.name for i in model.graph.input) + + tuple(i.name for i in model.graph.initializer) + + (output_name,) + ) + return StreamWorkload(model, nodes, buffers) diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 4a6c56044..6d62e215b 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -12,6 +12,7 @@ from .softmax.op import Softmax from .swiglu_decode.op import SwiGLUDecode from .swiglu_prefill.op import SwiGLUPrefill +from .swiglu_prefill_stream.op import SwiGLUPrefillStream from .transpose.op import Transpose from .strided_copy.op import StridedCopy from .repeat.op import Repeat diff --git a/iron/operators/swiglu_prefill_stream/README.md b/iron/operators/swiglu_prefill_stream/README.md new file mode 100644 index 000000000..d39b6de1f --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -0,0 +1,130 @@ + + +# SwiGLU prefill (stream-dse codegen) + +The whole SwiGLU-prefill block, `(SiLU(x @ gate) * (x @ up)) @ down`, deployed as one +full ELF. Unlike the hand-written operators the MLIR is generated at build time by +[`stream-dse`](https://github.com/KULeuven-MICAS/stream), +which solves where every layer runs and emits the AIE design. + +## What IRON provides and what stream-dse returns + +stream-dse needs two inputs, and IRON writes both from one source. + +| | Source | Built by | +| --- | --- | --- | +| Workload (ONNX) | [`reference.py`](./reference.py), the `SwiGLU` `nn.Module` | `torch.export` via [`iron/common/stream/workload.py`](../../common/stream/workload.py) | +| Mapping (YAML) | the placement in [`stream_design.py`](./stream_design.py) | [`iron/common/stream/mapping.py`](../../common/stream/mapping.py) | +| Kernels (`.cc`) | IRON's `aie_kernels` library | the registry in [`iron/common/stream/ops.py`](../../common/stream/ops.py) | + +`reference.py` is the single source of truth. Running it produces the golden output the +test compares against; exporting it produces the workload the design is generated from. +The mapping reads its layer names back from the exported graph rather than restating +them, so workload and mapping cannot disagree. Both files are written into the +experiment's output directory at build time; nothing is committed. + +stream-dse returns one MLIR design per fusion group. IRON takes it from there: +`iron/common/sequence.py` fuses the designs into a single module and compiles it with +`aiecc` into one full ELF. + +Nothing crosses the boundary except those files, which is why stream-dse can be an +optional dependency: a checkout without it imports and tests everything else. + +## Fusion granularity: the `k` modes + +`k` is how many fused groups the block is split into, and so how many designs the ELF +holds. The groups are `stream_design.GROUP_LAYERS`. + +| `k` | Groups | Shape on the array | +| --- | --- | --- | +| 1 | `gate, up, silu, mul, down` | the whole block at once, layers on disjoint columns, pipelined | +| 2 | `gate, up, silu, mul` + `down` | as above, split after the multiply | +| 5 | one per layer | layer by layer, each taking the whole array in turn | + +k=1 and k=2 fuse several layers onto each core, so intermediates stay on chip. k=5 is +the shape [`swiglu_prefill`](../swiglu_prefill) uses, every layer its own design. + +A core holds the operands of every layer in its group, so the kernel tile a group can +afford shrinks as more layers fuse onto it. That is why the tile is chosen per `k` +(`stream_design.FUSED_TILES` and `LAYER_TILES`). A tile that does not fit is rejected at +build time by stream-dse, naming the core and the shortfall. + +The tile also bounds the problem sizes: `seq_len` must be a multiple of the array rows, +and `embedding_dim` and `hidden_dim` multiples of their tile times the column split. +`stream_design._check_shapes` enforces this and names the offending dimension. + +Designs that come out byte-identical are built and configured once: at k=5 the gate and +up projections are the same design, so the ELF holds four rather than five. Set +`share_designs=False` on the operator to switch that off. + +## Expected performance + +Warm dispatch on one callable, 20 dispatches, seq 256 / embedding 512 / hidden 2048, on +an idle NPU2. `swiglu_prefill` is the hand-written operator at the same shape. + +| Design | Median (us) | Relative | +| --- | --- | --- | +| `swiglu_prefill` | 1348 | 1.00x | +| `k=1` | 1165 | 0.86x | +| `k=2` | 1452 | 1.08x | +| `k=5` | 1399 | 1.04x | + +k=1 is fastest: fusing the whole block keeps the intermediates on chip instead of +returning them to memory between layers. + +Medians from one machine, so treat them as orders of magnitude rather than exact. +Accuracy tracks the hand-written operator: 0.139 of output elements more than 8 percent +from the golden reference, against 0.139 to 0.144 here. + +## Runtime buffers + +Named by the reference module: `input`, `w_gate`, `w_up`, `w_down`, `output`. + +```python +run = operator.get_callable() +run.get_buffer("w_gate").torch_view()[:] = weights.flatten() +``` + +## Installing + +stream-dse is an optional, separately installed dependency and is not in IRON's +`requirements.txt`. + +```bash +pip install -r requirements_stream.txt +stream-setup-aie # required: installs stream-dse's AIE codegen dialects +``` + +`stream-setup-aie` installs the pure-Python codegen dialects (`xdsl-aie`, `snax-mlir`) +that cannot be plain PyPI dependencies because they are direct git installs. It does not +install `mlir_aie` / `llvm-aie`: IRON's `requirements.txt` pins those, and the codegen +only emits text MLIR through xdsl. Solving uses the licence-free OR-Tools GSCIP backend, +so no Gurobi licence is needed. + +Importing the operator does not require stream-dse; only building it does. + +## Build and run + +```bash +source /opt/xilinx/xrt/setup.sh +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 IRON'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. + +## Notes + +The hardware description (`whole_array_strix.yaml`) is resolved from the installed +`stream` package, where it ships as package data; nothing is vendored here. + +Node names are set explicitly, naming the role each layer plays rather than the ATen +op the exporter captured (`matmul`, `matmul_1`, ...). The mapping and the generated +design are both read by those names. diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py new file mode 100644 index 000000000..68de6187a --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from typing import Any + +import aie.utils as aie_utils + +from iron.common import ( + MLIROperator, + AIERuntimeArgSpec, + PythonGeneratedMLIRArtifact, + DesignGenerator, +) +from iron.common.device_utils import get_kernel_dir +from iron.common.sequence import OperatorSequence +from iron.common.stream.ops import ELTWISE_MUL, GEMM, SILU + + +@dataclass +class _SwiGLUStreamGroup(MLIROperator): + """One stream-dse design, used as an ``OperatorSequence`` child. + + ``k`` is how many fused groups the block is split into and ``group_index`` + which of them this is, in the order + :data:`~iron.operators.swiglu_prefill_stream.stream_design.GROUP_LAYERS` + lists them. + """ + + seq_len: int + embedding_dim: int + hidden_dim: int + k: int + group_index: int + context: Any = field(default=None, repr=False, compare=False) + + def __post_init__(self): + MLIROperator.__init__(self, context=self.context) + + @property + def _design(self): + from iron.operators.swiglu_prefill_stream import stream_design + + return stream_design + + def get_mlir_artifact(self): + return PythonGeneratedMLIRArtifact( + f"{self.name}.mlir", + DesignGenerator( + self.operator_dir / "stream_design.py", + "load_group", + (self.group_index,), + { + "k": self.k, + "seq_len": self.seq_len, + "embedding_dim": self.embedding_dim, + "hidden_dim": self.hidden_dim, + "npu": aie_utils.get_current_device().resolve().name, + }, + ), + ) + + def get_kernel_artifacts(self): + # The registry is the single place a kernel's source, compile flags and + # symbol names are declared, so the object and the design agree. + design = self._design + gemm_tiles = design.gemm_tiles(self.k) + per_layer = { + design.GATE: (GEMM, gemm_tiles[design.GATE]), + design.UP: (GEMM, gemm_tiles[design.UP]), + design.DOWN: (GEMM, gemm_tiles[design.DOWN]), + design.SILU: (SILU, None), + design.MUL: (ELTWISE_MUL, None), + } + layers = design.GROUP_LAYERS[self.k][self.group_index] + base_dir, kernel_dir = self.context.base_dir, get_kernel_dir() + return [ + artifact + for kernel, tiles in dict.fromkeys(per_layer[layer] for layer in layers) + for artifact in kernel.kernel_artifacts( + base_dir, kernel_dir, **(dict(zip("mkn", tiles)) if tiles else {}) + ) + ] + + def design_key(self): + """Groups whose generated design is byte-identical share it.""" + return self._design.group_digest( + self.group_index, + k=self.k, + seq_len=self.seq_len, + embedding_dim=self.embedding_dim, + hidden_dim=self.hidden_dim, + npu=aie_utils.get_current_device().resolve().name, + ) + + def get_arg_spec(self): + """The group's runtime arguments, shaped by the exported graph. + + Both the names and their order come from the workload, which is also the + order the generated design takes its arguments in. + """ + dims = (self.seq_len, self.embedding_dim, self.hidden_dim) + shapes = self._design.workload_for(*dims).shapes + inputs, outputs = self._design.group_ports(*dims, k=self.k)[self.group_index] + return [AIERuntimeArgSpec("in", shapes[name]) for name in inputs] + [ + AIERuntimeArgSpec("out", shapes[name]) for name in outputs + ] + + +def _wiring(seq_len, embedding_dim, hidden_dim, k): + """Each group's arguments, and the operator's own inputs and outputs. + + A group's arguments are the tensors it consumes and produces, in the order the + exported graph uses them. The operator's own arguments are what no group + produces and what none consumes. Imported lazily, so only building the operator + needs stream-dse, not importing it. + """ + from iron.operators.swiglu_prefill_stream.stream_design import group_ports + + boundaries = group_ports(seq_len, embedding_dim, hidden_dim, k) + produced = {name for _, outputs in boundaries for name in outputs} + consumed = {name for inputs, _ in boundaries for name in inputs} + ports = [inputs + outputs for inputs, outputs in boundaries] + external_inputs = dict.fromkeys( + name for inputs, _ in boundaries for name in inputs if name not in produced + ) + external_outputs = dict.fromkeys( + name for _, outputs in boundaries for name in outputs if name not in consumed + ) + return ports, list(external_inputs), list(external_outputs) + + +class SwiGLUPrefillStream(OperatorSequence): + """SwiGLU-prefill block generated by stream-dse and deployed as one full ELF. + + ``k`` is how many fused groups the block is split into: 1 keeps the whole block + on the array at once, 2 splits after the elementwise multiply, and 5 runs layer + by layer, each layer taking the whole array in turn as + :mod:`iron.operators.swiglu_prefill` does. The split is decided by the mapping's + fused groups, and the external buffers are the same either way. + + Runtime buffers (``get_callable().get_buffer(name)``) are named by the reference + module: ``input``, ``w_gate``, ``w_up``, ``w_down``, ``output``. Building + requires ``stream-dse`` (``pip install stream-dse`` + ``stream-setup-aie``); + importing this module does not. + """ + + def __init__( + self, seq_len, embedding_dim, hidden_dim, k=1, context=None, share_designs=True + ): + ports, inputs, outputs = _wiring(seq_len, embedding_dim, hidden_dim, k) + groups = [ + _SwiGLUStreamGroup( + seq_len=seq_len, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + k=k, + group_index=index, + context=context, + ) + for index in range(len(ports)) + ] + super().__init__( + name=f"swiglu_prefill_stream_k{k}_m{seq_len}_e{embedding_dim}_h{hidden_dim}", + runlist=[ + (group, *group_ports) for group, group_ports in zip(groups, ports) + ], + input_args=inputs, + output_args=outputs, + extra_flags=["--dynamic-objFifos"], + share_designs=share_designs, + context=context, + ) diff --git a/iron/operators/swiglu_prefill_stream/reference.py b/iron/operators/swiglu_prefill_stream/reference.py new file mode 100644 index 000000000..92caf8eb1 --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/reference.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reference SwiGLU-prefill block: ``(SiLU(x @ gate) * (x @ up)) @ down``. + +Running this module produces the golden output; exporting it produces the +workload stream-dse generates the design from. + +The names below are the block's vocabulary, and they are the ones +:mod:`iron.operators.swiglu_decode.reference` -- the golden reference this +operator shares -- gives the same tensors. Everything downstream is named from +here: the ONNX tensors, the mapping's layers and runtime arguments, the runtime +buffers, and the tensor handed between fusion groups. +""" + +import torch +from torch import nn + +INPUT = "input" +OUTPUT = "output" +GATE_PROJECTION = "left" +UP_PROJECTION = "right" +ACTIVATION = "left_swished" +HIDDEN = "intermediate" + +WEIGHTS = ("w_gate", "w_up", "w_down") + +# Every name this module exports, for the correspondence check in iron/tests/stream. +TENSOR_NAMES = ( + INPUT, + OUTPUT, + GATE_PROJECTION, + UP_PROJECTION, + ACTIVATION, + HIDDEN, + *WEIGHTS, +) + + +class SwiGLU(nn.Module): + """SwiGLU prefill block over a ``[seq_len, embedding_dim]`` activation.""" + + def __init__(self, embedding_dim: int, hidden_dim: int, dtype=torch.bfloat16): + super().__init__() + gate_up = (embedding_dim, hidden_dim) + self.w_gate = nn.Parameter(torch.zeros(gate_up, dtype=dtype)) + self.w_up = nn.Parameter(torch.zeros(gate_up, dtype=dtype)) + self.w_down = nn.Parameter( + torch.zeros((hidden_dim, embedding_dim), dtype=dtype) + ) + + def forward(self, input): + gate = input @ self.w_gate + up = input @ self.w_up + return (torch.nn.functional.silu(gate) * up) @ self.w_down + + +def swiglu_module(embedding_dim, hidden_dim, golden_reference=None) -> SwiGLU: + """A :class:`SwiGLU`, optionally holding ``golden_reference``'s weights. + + Weight *values* are irrelevant to the exported graph (only shapes and the + topology are), so the operator builds its design from a zero-filled module. + """ + module = SwiGLU(embedding_dim, hidden_dim).eval() + if golden_reference is not None: + with torch.no_grad(): + for name in WEIGHTS: + getattr(module, name).copy_(golden_reference[name]) + return module diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py new file mode 100644 index 000000000..b4edae178 --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -0,0 +1,400 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate the SwiGLU-prefill design with stream-dse. + +Both inputs stream-dse needs are produced here, from IRON: + +* the **workload**, exported from :mod:`~iron.operators.swiglu_prefill_stream.reference`, + the same module the test checks the result against; +* the **mapping**, from the placement below. + +Both are written into the experiment's output directory at build time, never into +the source tree, and the mapping's node names come from the exported workload, so +the two cannot disagree. stream-dse then solves the allocation and emits the MLIR. + +This module is imported lazily (by ``DesignGenerator`` at compile time), so +importing the operator does not require ``stream-dse`` to be installed, only +building it does. +""" + +import hashlib +import os +import re +from functools import lru_cache +from pathlib import Path + +import stream +import torch +from stream.api import optimize_allocation_co + +from iron.common.stream.hardware import ComputeArray +from iron.common.stream.mapping import ( + FusedGroup, + Placement, + emit_mapping, + group_boundaries, +) +from iron.common.stream.workload import export_workload +from iron.operators.swiglu_prefill_stream import reference +from iron.operators.swiglu_prefill_stream.reference import swiglu_module + +# Hardware description for the whole-array Strix (npu2) target, shipped as package +# data inside the installed stream package. +ACCELERATOR = os.path.join( + os.path.dirname(stream.__file__), + "inputs", + "aie", + "hardware", + "whole_array_strix.yaml", +) + +BACKEND = "ortools_gscip" # license-free OR-Tools GSCIP, no Gurobi needed +OUTPUT_ROOT = "outputs" + +# Names for the exported graph's computation nodes, in topological order, and for +# the tensors they produce. They name the roles rather than the ATen ops the +# exporter captured, and they are what the mapping and the generated design are +# read by. +GATE, UP, SILU, MUL, DOWN = "Gemm_Left", "Gemm_Right", "Silu", "Elt_Mul", "Gemm_Down" +NODE_NAMES = [GATE, UP, SILU, MUL, DOWN] +RESULT_NAMES = { + GATE: reference.GATE_PROJECTION, + UP: reference.UP_PROJECTION, + SILU: reference.ACTIVATION, + MUL: reference.HIDDEN, +} + +# The kernel tile each layer is compiled and mapped for, as (sequence, embedding, +# hidden). A core holds the operands of every layer in its group, so the tile a group +# can afford shrinks as more layers fuse onto it. Carrying the tile and no absolute +# dimension is what lets one mapping hold across problem sizes. +FUSED_TILES = (32, 32, 64) # k=1, k=2: several layers share a core +LAYER_TILES = (64, 64, 64) # k=5: one layer per core + +# Sequence positions an elementwise layer works at a time when it reads from and +# writes to memory. Its tile is then this many whole rows, which is contiguous in +# a row-major tensor, so each transfer runs the length of the rows rather than one +# MAC tile at a time. +ELEMENTWISE_ROWS = 1 + +# Which layers each fused group contains, per number of groups ``k``. Splitting +# makes stream-dse emit one design per group; the tensor handed from one group to +# the next comes from the exported graph. +LAYER_BY_LAYER = 5 +GROUP_LAYERS = { + 1: [[GATE, UP, SILU, MUL, DOWN]], + 2: [[GATE, UP, SILU, MUL], [DOWN]], + LAYER_BY_LAYER: [[GATE], [UP], [SILU], [MUL], [DOWN]], +} + + +def tiles_for(k): + """The kernel tile, as (sequence, embedding, hidden), for ``k`` fused groups.""" + return LAYER_TILES if k == LAYER_BY_LAYER else FUSED_TILES + + +def gemm_tiles(k): + """Each GEMM layer's kernel tile, in the (m, k, n) order the kernel takes.""" + sequence, embedding, hidden = tiles_for(k) + return { + GATE: (sequence, embedding, hidden), + UP: (sequence, embedding, hidden), + DOWN: (sequence, hidden, embedding), + } + + +@lru_cache(maxsize=None) +def array() -> ComputeArray: + """The compute grid of the device being built for.""" + import aie.utils as aie_utils + + return ComputeArray.from_device(aie_utils.get_current_device()) + + +def _placements(k, hidden_dim): + """Where each layer runs. + + Fused (k=1, k=2): the layers sit on disjoint columns, two per GEMM and one per + elementwise layer, so they pipeline across steady-state iterations. Each splits + over the array's rows (D0, the sequence dimension) and a GEMM over its two + columns as well (D2, the output dimension). + + Layer by layer (k=5): the layers run in turn, so each takes the whole array. + The GEMMs use every row; the elementwise layers take one core per column and + split the sequence across them, the shape IRON's channeled operators use. They + also read whole rows, so their transfers to and from memory are contiguous. + """ + grid = array() + sequence_tile, _, hidden_tile = tiles_for(k) + tiles = gemm_tiles(k) + + def gemm(tiles): + return dict( + zip("mkn", tiles), utilization=61.8, layout="default", bfp16_mmul=True + ) + + def elementwise(rows, columns, layout, bfp16_mmul=False): + return { + "utilization": 50.0, + "layout": layout, + "m": rows, + "n": columns, + "bfp16_mmul": bfp16_mmul, + } + + if k == LAYER_BY_LAYER: + wide = grid.all_columns + gemm_split = (("D0", grid.num_rows), ("D2", grid.num_columns)) + elementwise_split = (("D0", grid.num_columns),) + rows_wide = elementwise(ELEMENTWISE_ROWS, hidden_dim, "contiguous") + return { + GATE: Placement(wide, gemm_split, gemm(tiles[GATE])), + UP: Placement(wide, gemm_split, gemm(tiles[UP])), + SILU: Placement(wide, elementwise_split, rows_wide, rows=[0]), + MUL: Placement(wide, elementwise_split, rows_wide, rows=[0]), + DOWN: Placement(wide, gemm_split, gemm(tiles[DOWN])), + } + + columns = dict(zip(NODE_NAMES, grid.allocate([2, 2, 1, 1, 2]))) + gemm_split = (("D0", grid.num_rows), ("D2", 2)) + elementwise_split = (("D0", grid.num_rows),) + # Fused behind a GEMM, so the operands take the layout that GEMM writes. + fused = elementwise(sequence_tile, hidden_tile, "default", bfp16_mmul=True) + return { + GATE: Placement(columns[GATE], gemm_split, gemm(tiles[GATE])), + UP: Placement(columns[UP], gemm_split, gemm(tiles[UP])), + SILU: Placement(columns[SILU], elementwise_split, fused), + MUL: Placement(columns[MUL], elementwise_split, fused), + DOWN: Placement(columns[DOWN], gemm_split, gemm(tiles[DOWN])), + } + + +def _layer_tiling(layer, hidden_dim, k): + """Intra-core tiling of one layer, over the dimensions it iterates.""" + if layer not in (GATE, UP, DOWN): + return [(layer, "D1", hidden_dim), (layer, "D0", ELEMENTWISE_ROWS)] + sequence, contraction, output = gemm_tiles(k)[layer] + return [ + (layer, "D1", contraction), + (layer, "D2", output), + (layer, "D0", sequence), + ] + + +def _groups(k, hidden_dim): + """The fused groups, each with the intra-core tiling of its leading GEMM. + + Splitting makes stream-dse emit one design per group, which IRON then deploys + as a single full ELF. D0 is the sequence dimension, D1 the contraction and D2 + the output dimension. + """ + sequence_tile, embedding_tile, hidden_tile = tiles_for(k) + if k == LAYER_BY_LAYER: + tiling = [_layer_tiling(layers[0], hidden_dim, k) for layers in GROUP_LAYERS[k]] + elif k == 2: + tiling = [ + _layer_tiling(GATE, hidden_dim, k), + _layer_tiling(DOWN, hidden_dim, k), + ] + else: + tiling = [ + [ + (GATE, "D1", embedding_tile), + (DOWN, "D2", embedding_tile), + (GATE, "D2", hidden_tile), + (GATE, "D0", sequence_tile), + ] + ] + return [ + FusedGroup(f"Fused_Group_{index + 1}", layers, group_tiling) + for index, (layers, group_tiling) in enumerate(zip(GROUP_LAYERS[k], tiling)) + ] + + +def _check_shapes(seq_len, embedding_dim, hidden_dim, k): + """Reject a problem size the placement and the kernel tiles cannot divide.""" + grid = array() + sequence_tile, embedding_tile, hidden_tile = tiles_for(k) + gemm_split = grid.num_columns if k == LAYER_BY_LAYER else 2 + if seq_len % grid.num_rows or seq_len < sequence_tile * grid.num_rows: + raise ValueError( + f"seq_len ({seq_len}) must be a multiple of {grid.num_rows} and at " + f"least {sequence_tile * grid.num_rows}" + ) + if embedding_dim % (embedding_tile * gemm_split): + raise ValueError( + f"embedding_dim ({embedding_dim}) must be a multiple of " + f"{embedding_tile * gemm_split}" + ) + if hidden_dim % (hidden_tile * gemm_split): + raise ValueError( + f"hidden_dim ({hidden_dim}) must be a multiple of " + f"{hidden_tile * gemm_split}" + ) + + +@lru_cache(maxsize=None) +def workload_for(seq_len, embedding_dim, hidden_dim): + """The exported workload for one problem size.""" + return export_workload( + swiglu_module(embedding_dim, hidden_dim), + (torch.zeros(seq_len, embedding_dim, dtype=torch.bfloat16),), + node_names=NODE_NAMES, + result_names=RESULT_NAMES, + ) + + +def group_ports(seq_len, embedding_dim, hidden_dim, k=1): + """Per fused group, the tensor names it takes in and hands on. + + These are the operator's runtime arguments, including the tensors a split + design passes from one group to the next. + """ + return group_boundaries( + workload_for(seq_len, embedding_dim, hidden_dim), GROUP_LAYERS[k] + ) + + +def build_inputs(seq_len, embedding_dim, hidden_dim, output_dir, k=1): + """Write the workload and mapping for one configuration; return their paths.""" + _check_shapes(seq_len, embedding_dim, hidden_dim, k) + workload = workload_for(seq_len, embedding_dim, hidden_dim) + output_dir = Path(output_dir) + return ( + workload.write(output_dir / "workload.onnx"), + emit_mapping( + workload, + _placements(k, hidden_dim), + _groups(k, hidden_dim), + array(), + output_dir / "mapping.yaml", + ), + ) + + +def _experiment_id(seq_len, embedding_dim, hidden_dim, k): + grid = array() + hardware = os.path.splitext(os.path.basename(ACCELERATOR))[0] + suffix = f"_k{k}" if k > 1 else "" + return ( + f"{hardware}-swiglu{suffix}_{seq_len}_{embedding_dim}_{hidden_dim}" + f"-{grid.num_rows}_row_{grid.num_columns}_col" + ) + + +def _design_paths(seq_len, embedding_dim, hidden_dim, k): + """Where stream-dse writes each group's MLIR. + + A single fused group goes through stream-dse's single-design pipeline and lands + in ``codegen/``; several groups each land in their own ``group_i/codegen/``. + """ + output_dir = os.path.join( + OUTPUT_ROOT, _experiment_id(seq_len, embedding_dim, hidden_dim, k) + ) + if k == 1: + return [os.path.join(output_dir, "codegen", "final.mlir")] + return [ + os.path.join(output_dir, f"group_{index}", "codegen", "final.mlir") + for index in range(len(GROUP_LAYERS[k])) + ] + + +def _run_codegen(seq_len, embedding_dim, hidden_dim, npu, k): + """Run stream-dse's constraint optimization and code generation once.""" + grid = array() + experiment_id = _experiment_id(seq_len, embedding_dim, hidden_dim, k) + workload_path, mapping_path = build_inputs( + seq_len, + embedding_dim, + hidden_dim, + os.path.join(OUTPUT_ROOT, experiment_id), + k=k, + ) + optimize_allocation_co( + hardware=ACCELERATOR, + workload=workload_path, + mapping=mapping_path, + experiment_id=experiment_id, + output_path=OUTPUT_ROOT, + skip_if_exists=False, + enable_codegen=True, + trace_size=0, + nb_cols_to_use=grid.num_columns, + npu=npu, + backend=BACKEND, + ) + + +def _prefixed(mlir_text: str, func_prefix: str) -> str: + """Apply a fused-operator ``func_prefix`` (``op_``) to a group's MLIR. + + ``OperatorSequence`` renames each child's kernel object files and symbols to + ``op_...`` so the groups stay distinct inside one ELF; the group's MLIR + must reference the same prefixed names. Prefix the ``link_with`` object files + and every privately declared kernel symbol, and its call sites. + """ + if not func_prefix: + return mlir_text + mlir_text = re.sub( + r'link_with\s*=\s*"([^"]+)"', + lambda m: f'link_with = "{func_prefix}{m.group(1)}"', + mlir_text, + ) + symbols = sorted( + set(re.findall(r"func\.func\s+private\s+@([A-Za-z0-9_]+)", mlir_text)), + key=len, + reverse=True, + ) + for symbol in symbols: + mlir_text = re.sub( + rf"@{re.escape(symbol)}\b", f"@{func_prefix}{symbol}", mlir_text + ) + return mlir_text + + +def region_module(mlir_text: str, func_prefix: str = ""): + """Parse a group's MLIR text into an ``aie`` module for fusion. + + ``OperatorSequence`` consumes ``aie.DeviceOp`` objects, so the xDSL-emitted + group text is re-parsed with the mlir-aie bindings, after ``func_prefix`` + rewriting. + """ + from aie import ir + from aie.extras.context import mlir_mod_ctx + + with mlir_mod_ctx(): + return ir.Module.parse(_prefixed(mlir_text, func_prefix)) + + +def _group_text(group_index, *, k, seq_len, embedding_dim, hidden_dim, npu) -> str: + """One group's generated MLIR, before any ``func_prefix`` rewriting.""" + finals = _design_paths(seq_len, embedding_dim, hidden_dim, k) + if not all(os.path.exists(final) for final in finals): + _run_codegen(seq_len, embedding_dim, hidden_dim, npu, k) + return Path(finals[group_index]).read_text() + + +def group_digest(group_index, **dims) -> str: + """Digest of a group's design, for recognising groups that share one.""" + return hashlib.sha256(_group_text(group_index, **dims).encode()).hexdigest() + + +def load_group( + group_index, func_prefix="", *, k, seq_len, embedding_dim, hidden_dim, npu +): + """Generate the ``k``-group design once and return one group's aie module. + + ``group_index`` selects the group, in the order :data:`GROUP_LAYERS` lists them. + ``func_prefix`` is injected by ``OperatorSequence``. Every group loader calls + this; the first generates the design and the rest reuse the files on disk. + """ + text = _group_text( + group_index, + k=k, + seq_len=seq_len, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + npu=npu, + ) + return region_module(text, func_prefix) diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py new file mode 100644 index 000000000..c0a9deaa8 --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time + +import pytest +import torch + +# The design is generated by stream-dse at compile() time. stream-dse is an +# optional dependency (see requirements_stream.txt) absent from the default CI +# image, so skip this whole module when it is unavailable. +pytest.importorskip( + "stream", reason="stream-dse not installed (see requirements_stream.txt)" +) + +from iron.operators.swiglu_prefill_stream.op import SwiGLUPrefillStream + +# The operator's design is generated from this module; the values it is checked +# against come from swiglu_decode's reference, which it shares. +from iron.operators.swiglu_decode.reference import generate_golden_reference +from iron.operators.swiglu_prefill_stream.reference import INPUT, OUTPUT, WEIGHTS +from iron.common.test_utils import verify_buffer + +# The MILP-feasible shape on the whole-array Strix (npu2) target. +SEQ_LEN, EMBEDDING_DIM, HIDDEN_DIM = 256, 512, 2048 + +# Fused groups to deploy the block as: one design, a front end plus the down +# projection, or one design per layer. +FUSION_GROUPS = [1, 2, 5] + +# Timed dispatches per test; the reported latency is the fastest of them. +TIMED_RUNS = 3 + + +def _staged(operator, golden_ref): + """A callable with its inputs staged. + + Inputs are named by the golden reference, and the design consumes the weights in + their natural (K, N) layout. + """ + run = operator.get_callable() + for name in (INPUT, *WEIGHTS): + buffer = run.get_buffer(name) + buffer.torch_view()[:] = golden_ref[name].to(torch.bfloat16).flatten() + buffer.to("npu") + return run + + +@pytest.mark.supported_devices("npu2") +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", +) +@pytest.mark.parametrize("k", FUSION_GROUPS) +def test_swiglu_prefill_stream(k, aie_context): + golden_ref = generate_golden_reference(M=SEQ_LEN, K=EMBEDDING_DIM, N=HIDDEN_DIM) + operator = SwiGLUPrefillStream( + seq_len=SEQ_LEN, + embedding_dim=EMBEDDING_DIM, + hidden_dim=HIDDEN_DIM, + k=k, + context=aie_context, + ) + operator.compile() + + # The whole block is computed in bf16, so rounding accumulates and reorders + # across the chain and the K=hidden_dim down-projection sum (near-cancellation + # amplifies relative error): ~20% of elements drift past the 8% bound, so allow + # up to 25%. Tolerances are local to this test. + run = _staged(operator, golden_ref) + run() + output = run.get_buffer(OUTPUT).to_torch().reshape((SEQ_LEN, EMBEDDING_DIM)) + errors = verify_buffer( + output, + OUTPUT, + golden_ref[OUTPUT], + rel_tol=0.08, + abs_tol=0.7, + max_error_rate=0.25, + ) + assert not errors, f"Test failed with errors: {errors}" + + # 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) + 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}" + ) + print(f"Effective Bandwidth: {total_bytes / (elapsed_us * 1e-6) / 1e9:.4f} GB/s") diff --git a/iron/tests/stream/kernel_layouts.py b/iron/tests/stream/kernel_layouts.py new file mode 100644 index 000000000..d6efdad5b --- /dev/null +++ b/iron/tests/stream/kernel_layouts.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The operand layouts IRON's kernels are compiled for must match stream-dse's. + +stream-dse generates the DMAs that feed the kernel objects IRON compiles from +``aie_kernels``; both sides must agree on how an operand is tiled in memory. The +layouts declared in :mod:`iron.common.stream.ops` are that contract. They happen +to coincide with stream-dse's built-in kernel layouts today, so no override is +needed -- this test fails if a future stream-dse release changes them, which would +otherwise corrupt results silently. +""" + +import pytest + +pytest.importorskip( + "stream", reason="stream-dse not installed (see requirements_stream.txt)" +) + +from stream.compiler.kernels import AIEKernels # noqa: E402 + +from iron.common.stream.ops import ( # noqa: E402 + ELTWISE_MUL, + GEMM, + SILU, + elementwise_layouts, + gemm_layouts, +) + + +def _assert_same(iron_layouts, stream_kernel): + expected = [layout.to_snaxc() for layout in iron_layouts] + actual = list(stream_kernel.operand_layouts()) + assert [str(layout) for layout in expected] == [str(layout) for layout in actual] + + +@pytest.mark.parametrize("bfp16_mmul", [False, True]) +@pytest.mark.parametrize("m,k,n", [(32, 32, 64), (32, 64, 32), (64, 64, 64)]) +def test_gemm_layouts_match_stream(m, k, n, bfp16_mmul): + _assert_same( + gemm_layouts(m, k, n, bfp16_mmul), + AIEKernels[GEMM.key](61.8, m, k, n, "default", bfp16_mmul), + ) + + +@pytest.mark.parametrize("bfp16_mmul", [False, True]) +def test_silu_layouts_match_stream(bfp16_mmul): + _assert_same( + elementwise_layouts(2, bfp16_mmul), + AIEKernels[SILU.key](50.0, "default", bfp16_mmul=bfp16_mmul), + ) + + +@pytest.mark.parametrize("bfp16_mmul", [False, True]) +def test_eltwise_mul_layouts_match_stream(bfp16_mmul): + _assert_same( + elementwise_layouts(3, bfp16_mmul), + AIEKernels[ELTWISE_MUL.key](50.0, "default", bfp16_mmul=bfp16_mmul), + ) + + +@pytest.mark.parametrize("kernel", [GEMM, SILU, ELTWISE_MUL]) +def test_kernel_keys_exist_in_stream(kernel): + assert kernel.key in AIEKernels diff --git a/iron/tests/stream/names.py b/iron/tests/stream/names.py new file mode 100644 index 000000000..cdb58d57c --- /dev/null +++ b/iron/tests/stream/names.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The names the design is built from must be the golden reference's names. + +Every tensor the operator exports, maps and wires is named once, in +:mod:`iron.operators.swiglu_prefill_stream.reference`, using the vocabulary the +shared golden reference uses for the same tensors. These tests pin that +correspondence, and pin that the mapping and the fused-group wiring take their +names from the exported graph rather than restating them. +""" + +import pytest + +from iron.operators.swiglu_decode.reference import generate_golden_reference +from iron.operators.swiglu_prefill_stream import reference + +SHAPE = dict(M=8, K=8, N=16) + + +@pytest.fixture(scope="module") +def golden_keys(): + return set(generate_golden_reference(**SHAPE)) + + +@pytest.mark.parametrize("name", reference.TENSOR_NAMES) +def test_name_is_a_golden_reference_key(name, golden_keys): + assert name in golden_keys + + +def test_module_parameters_cover_the_golden_weights(): + module = reference.swiglu_module(SHAPE["K"], SHAPE["N"]) + assert set(dict(module.named_parameters())) == set(reference.WEIGHTS) + + +def test_golden_weights_load_into_the_module(): + golden = generate_golden_reference(**SHAPE) + module = reference.swiglu_module(SHAPE["K"], SHAPE["N"], golden) + for name in reference.WEIGHTS: + assert getattr(module, name).equal(golden[name]) + + +def test_module_computes_the_golden_reference(): + """The design is generated from this module and the result is checked against + the golden reference, so the two have to be the same computation.""" + golden = generate_golden_reference(**SHAPE) + module = reference.swiglu_module(SHAPE["K"], SHAPE["N"], golden) + assert module(golden[reference.INPUT]).equal(golden[reference.OUTPUT]) + + +stream = pytest.importorskip( + "stream", reason="stream-dse not installed (see requirements_stream.txt)" +) + +from iron.operators.swiglu_prefill_stream import stream_design # noqa: E402 + +DIMS = (256, 512, 2048) + + +@pytest.mark.parametrize("k", [1, 2, 5]) +def test_group_ports_are_named_by_the_exported_graph(k): + workload = stream_design.workload_for(*DIMS) + known = set(workload.buffers) | set(reference.TENSOR_NAMES) + for inputs, outputs in stream_design.group_ports(*DIMS, k): + assert set(inputs) | set(outputs) <= known + + +def test_split_design_hands_on_the_hidden_state(): + (_, front_outputs), (down_inputs, _) = stream_design.group_ports(*DIMS, k=2) + assert front_outputs == (reference.HIDDEN,) + assert down_inputs[0] == reference.HIDDEN + + +def test_external_arguments_match_the_runtime_buffers(): + workload = stream_design.workload_for(*DIMS) + boundaries = stream_design.group_ports(*DIMS, k=2) + produced = {name for _, outputs in boundaries for name in outputs} + consumed = {name for inputs, _ in boundaries for name in inputs} + external = [ + name for inputs, _ in boundaries for name in inputs if name not in produced + ] + [name for _, outputs in boundaries for name in outputs if name not in consumed] + assert set(external) == set(workload.buffers) diff --git a/iron/tests/stream/placement.py b/iron/tests/stream/placement.py new file mode 100644 index 000000000..0615715ae --- /dev/null +++ b/iron/tests/stream/placement.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Placements are written in array columns and resolved to stream's core ids. + +:class:`~iron.common.stream.hardware.ComputeArray` reads the grid from the +mlir-aie device IRON is building for, and is the only place that knows what a +stream core id means, so operators never spell one out. +""" + +import pytest + +pytest.importorskip( + "stream", reason="stream-dse not installed (see requirements_stream.txt)" +) + +import aie.utils as aie_utils # noqa: E402 +from aie.iron.device import NPU2 # noqa: E402 + +aie_utils.set_current_device(NPU2()) + +from iron.common.stream.hardware import ComputeArray # noqa: E402 +from iron.operators.swiglu_prefill_stream import stream_design # noqa: E402 + +ARRAY = stream_design.array() + +DIMS = (256, 512, 2048) + +# The core ids the fused placement resolves to on the whole-array Strix target. +MEASURED_ALLOCATION = { + "Gemm_Left": [2, 3, 4, 5, 8, 9, 10, 11], + "Gemm_Right": [14, 15, 16, 17, 20, 21, 22, 23], + "Silu": [26, 27, 28, 29], + "Elt_Mul": [32, 33, 34, 35], + "Gemm_Down": [38, 39, 40, 41, 44, 45, 46, 47], +} + + +def test_array_matches_the_device(): + assert (ARRAY.num_columns, ARRAY.num_rows) == (8, 4) + assert ARRAY.all_columns == tuple(range(8)) + + +def test_columns_hold_only_compute_tiles(): + ids = [core for column in ARRAY.columns for core in column] + assert len(ids) == ARRAY.num_columns * ARRAY.num_rows + assert len(set(ids)) == len(ids) + + +def test_cores_follow_column_then_row_order(): + assert ARRAY.cores([0]) == ARRAY.columns[0] + assert ARRAY.cores([0, 1]) == ARRAY.columns[0] + ARRAY.columns[1] + + +def test_rows_narrow_a_placement_to_one_worker_per_column(): + """The shape IRON's channeled operators give an elementwise layer.""" + assert ARRAY.cores(ARRAY.all_columns, rows=[0]) == tuple( + column[0] for column in ARRAY.columns + ) + + +def test_allocate_gives_disjoint_consecutive_columns(): + ranges = ARRAY.allocate([2, 2, 1, 1, 2]) + assert ranges == ((0, 1), (2, 3), (4,), (5,), (6, 7)) + flat = [column for group in ranges for column in group] + assert len(set(flat)) == len(flat) + + +def test_allocate_rejects_an_oversubscribed_array(): + with pytest.raises(ValueError): + ARRAY.allocate([ARRAY.num_columns, 1]) + + +def test_ids_agree_with_the_accelerator_stream_solves_against(): + """IRON derives core ids from the device; stream-dse reads them from its own + accelerator description. A design is only correct while the two agree.""" + import os + + import stream + import yaml + + path = os.path.join( + os.path.dirname(stream.__file__), + "inputs", + "aie", + "hardware", + "whole_array_strix.yaml", + ) + description = yaml.safe_load(open(path)) + by_column: dict[int, list[tuple[int, int]]] = {} + for core_id, coordinate in description["core_coordinates"].items(): + if description["cores"][core_id].endswith("aie_tile.yaml"): + column, row = coordinate + by_column.setdefault(column, []).append((row, core_id)) + expected = tuple( + tuple(core_id for _, core_id in sorted(rows)) + for _, rows in sorted(by_column.items()) + ) + assert ARRAY.columns == expected + + +def test_emitted_allocation_resolves_to_the_expected_cores(tmp_path): + import yaml + + _, mapping_path = stream_design.build_inputs(*DIMS, tmp_path / "design") + emitted = { + layer["name"]: layer["core_allocation"][0] + for layer in yaml.safe_load(open(mapping_path))["layers"] + } + assert emitted == MEASURED_ALLOCATION + + +def test_layer_by_layer_gives_every_layer_the_whole_array(tmp_path): + import yaml + + _, mapping_path = stream_design.build_inputs( + *DIMS, tmp_path / "design_k5", k=stream_design.LAYER_BY_LAYER + ) + mapping = yaml.safe_load(open(mapping_path)) + columns = { + layer["name"]: { + core // (ARRAY.num_rows + 2) for core in layer["core_allocation"][0] + } + for layer in mapping["layers"] + } + assert all(used == set(ARRAY.all_columns) for used in columns.values()) + assert len(mapping["fused_groups"]) == stream_design.LAYER_BY_LAYER + + +def test_devices_other_than_the_default_resolve(): + from aie.iron.device import NPU1 + + array = ComputeArray.from_device(NPU1()) + assert array.num_columns and array.num_rows + assert array.cores(array.all_columns) diff --git a/requirements_stream.txt b/requirements_stream.txt new file mode 100644 index 000000000..092dc823c --- /dev/null +++ b/requirements_stream.txt @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Optional dependencies for the stream-dse-backed fused SwiGLU-prefill operator +# (iron/operators/swiglu_prefill_stream). +# +# Not installed by the default CI (requirements.txt); the operator's test skips +# itself (pytest.importorskip) when stream-dse is absent. Install this file to +# build and run the operator and its test: +# +# pip install -r requirements_stream.txt +# stream-setup-aie # REQUIRED: installs stream-dse's pure-Python AIE codegen +# # dialects (xdsl-aie, snax-mlir), which cannot be PyPI +# # dependencies. It leaves mlir_aie/llvm-aie alone; IRON +# # pins those in requirements.txt. +# +# stream-dse generates the fused MLIR design at build time with the license-free +# OR-Tools GSCIP solver, and writes its workload and mapping into its own installed +# package directory, so that environment must be writable. + +onnxscript>=0.7 +stream-dse>=1.13.11