From 71f92d07adfd2a0ac64aa9f6b675c1f839cf8c4e Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 01/21] iron/common: forward full-ELF aiecc flags FullElfArtifact/OperatorSequence gain an extra_flags list forwarded into the aiecc --generate-full-elf command (mirrors XclbinArtifact), so placed/routed whole-array designs can pass --dynamic-objFifos and not overflow AIE2p program memory; empty by default, so other sequences are unaffected. --- iron/common/compilation/base.py | 3 +++ iron/common/sequence.py | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index da15df03..8b06de53 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/sequence.py b/iron/common/sequence.py index 747f3244..a7a33cef 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]) @@ -262,6 +263,7 @@ def __init__( output_args, buffer_sizes=None, dispatch="auto", + extra_flags=None, *args, **kwargs, ): @@ -282,6 +284,10 @@ def __init__( self.explicit_buffer_sizes = ( buffer_sizes or {} ) # Optional dict: buffer_name -> size_in_bytes + # Extra aiecc flags forwarded to the full-ELF build (e.g. --dynamic-objFifos + # for placed/routed whole-array designs that would otherwise overflow AIE2p + # program memory). Empty by default, so other sequences are unaffected. + self.extra_flags = extra_flags or [] self._dispatch = dispatch @staticmethod From 51e637a98cd48697a10bead693c477daa5fb5609 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 02/21] iron/common: add TiledStridedLayout for tiled-strided operand layouts Small IRON-side tiled-strided layout (Stride / TiledStride / TiledStridedLayout + tiled_2d) with a to_snaxc() bridge, used to author kernel operand layouts for stream-dse-backed operators. --- iron/common/__init__.py | 1 + iron/common/layout.py | 107 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 iron/common/layout.py diff --git a/iron/common/__init__.py b/iron/common/__init__.py index 9f5a8f7a..cb2ff31b 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/layout.py b/iron/common/layout.py new file mode 100644 index 00000000..09770524 --- /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))), + ) + ) From c6503a0662d0010eb6bb7095d01baf742c2ffdec Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 03/21] ci: package stream-dse for stream-backed operators Optional stream-dse dependency set (requirements_stream.txt) plus the CI bits to build stream-backed operators: run stream-setup-aie in prereqs, install graphviz in the docker image (stream-dse codegen shells out to dot), and tolerate a missing GitHub token in test_docker_ci.sh. --- .github/actions/prereqs/action.yaml | 5 +++++ ci/docker-based/Dockerfile | 4 ++++ ci/docker-based/test_docker_ci.sh | 4 +++- requirements_stream.txt | 30 +++++++++++++++++++++++++++++ 4 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 requirements_stream.txt diff --git a/.github/actions/prereqs/action.yaml b/.github/actions/prereqs/action.yaml index 7363b0fb..2b09153b 100644 --- a/.github/actions/prereqs/action.yaml +++ b/.github/actions/prereqs/action.yaml @@ -21,4 +21,9 @@ 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 AIE codegen deps (snax-mlir/snaxc, xdsl-aie, aie-python-extras) + # are not PyPI dependencies; this console script installs them so the + # stream-backed operators can generate their MLIR at build time. + stream-setup-aie echo "Prerequisites installed into ${{ inputs.env_name }}" diff --git a/ci/docker-based/Dockerfile b/ci/docker-based/Dockerfile index ccffadcf..38081e1c 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 5d44c63a..6eef5017 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/requirements_stream.txt b/requirements_stream.txt new file mode 100644 index 00000000..26a8029a --- /dev/null +++ b/requirements_stream.txt @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Optional dependency for the stream-dse-backed fused SwiGLU-prefill operator +# (iron/operators/swiglu_prefill_stream). +# +# It is 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/run the operator and its test: +# +# pip install -r requirements_stream.txt +# stream-setup-aie # REQUIRED: installs stream-dse's AIE codegen deps that +# # cannot be PyPI dependencies (snax-mlir/snaxc, xdsl-aie, +# # aie-python-extras); also installs the mlir_aie/llvm-aie +# # wheels, skipping any already provided by requirements.txt. +# +# Notes: +# - stream-dse generates the fused MLIR design at build time (license-free +# OR-Tools GSCIP solver; no Gurobi needed) and writes its generated workload/ +# mapping files into its own installed package directory, so that environment +# must be writable. +# - >=1.13.6 is required: stream_design.py feeds IRON-authored operand layouts +# into code generation via optimize_allocation_co(kernels=...), the override +# hook added in stream-dse 1.13.4; the k=2 variant (op_k2.py) additionally needs +# the two-fusion-group support (make_swiglu_mapping(split_groups=...) + the +# multi-group AIE codegen pipeline) added in stream-dse 1.13.5; and 1.13.6 makes +# the debug workload-graph visualization non-fatal, so code generation no longer +# requires graphviz (`dot`) to be installed. + +stream-dse>=1.13.6 From c242e905505a12d55a81adfb3d3fd9f62ffc1d89 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 04/21] swiglu_prefill_stream: stream-dse-backed SwiGLU-prefill operator (k=1 and k=2) Fused SwiGLU-prefill whose MLIR is generated by stream-dse and deployed as a single full-ELF via OperatorSequence. Two variants share one _SwiGLUStreamGroup child: SwiGLUPrefillStream fuses the whole block into one design; SwiGLUPrefillStreamK2 splits it into a gate/up/SiLU/mul group and a down-projection group, with the hidden state h kept on device between them. Passes --dynamic-objFifos so the whole-array design does not overflow program memory. Verified on npu2 (Strix). --- .gitignore | 1 + iron/operators/__init__.py | 1 + .../operators/swiglu_prefill_stream/README.md | 53 ++++ iron/operators/swiglu_prefill_stream/op.py | 245 ++++++++++++++++ .../swiglu_prefill_stream/stream_design.py | 265 ++++++++++++++++++ .../swiglu_prefill_stream/stream_kernels.py | 94 +++++++ iron/operators/swiglu_prefill_stream/test.py | 154 ++++++++++ 7 files changed, 813 insertions(+) create mode 100644 iron/operators/swiglu_prefill_stream/README.md create mode 100644 iron/operators/swiglu_prefill_stream/op.py create mode 100644 iron/operators/swiglu_prefill_stream/stream_design.py create mode 100644 iron/operators/swiglu_prefill_stream/stream_kernels.py create mode 100644 iron/operators/swiglu_prefill_stream/test.py diff --git a/.gitignore b/.gitignore index d19c375d..ec6f4f79 100755 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,4 @@ id_ed25519.pub .cline_storage *.egg-info **/*.prj/** +/outputs/ diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 4a6c5604..62841c7f 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, SwiGLUPrefillStreamK2 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 00000000..2b7c779c --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -0,0 +1,53 @@ + + +# SwiGLU prefill (stream-dse codegen) + +This operator is **fused**: the whole SwiGLU-prefill block (both GEMMs + SiLU + +elementwise-mul) is emitted as a **single MLIR design generated by +[`stream-dse`](https://github.com/KULeuven-MICAS/stream)**, then compiled by IRON's normal +flow into one xclbin. Unlike the other operators, its MLIR is not written by hand — it is +produced at build time by [`stream_design.py`](./stream_design.py), which calls the installed +`stream` package (`stream.api.optimize_allocation_co(..., enable_codegen=True)`). + +## Enabling stream codegen + +`stream-dse` is an **optional, separately-installed** dependency (it is *not* in IRON's +`requirements.txt`). Install it into the **same environment** as IRON via the extra +requirements file: + +```bash +pip install -r requirements_stream.txt +stream-setup-aie # required: installs stream-dse's AIE codegen deps +``` + +Notes: +- MLIR generation uses the open-source **OR-Tools GSCIP** solver (`backend="ortools_gscip"`), + so **no Gurobi license** is required. +- `stream-setup-aie` is **required**: it installs the AIE codegen packages stream-dse needs + that cannot be plain PyPI dependencies (`snax-mlir`/`snaxc`, `xdsl-aie`, `aie-python-extras`), + since they are direct git/URL installs. It also installs the `mlir_aie` / `llvm-aie` wheels, + but skips those if IRON's `requirements.txt` already provided them. +- Importing the operator does **not** require `stream-dse` (the launcher is imported lazily); + only **building** (`operator.compile()` / running the test) does. + +## Build & run + +```bash +# build + run on an NPU2 (Strix) device +source /opt/xilinx/xrt/setup.sh # XRT on PATH (provides pyxrt + xclbinutil) +pytest iron/operators/swiglu_prefill_stream/test.py +``` + +The feasible/verified shape is **seq 256 / embedding 512 / hidden 2048**, tiles +**32 / 32 / 64**, target **npu2**. + +## Caveats (stream-dse packaging) + +- The hardware-description YAML (`whole_array_strix.yaml` + `hardware/cores/*.yaml`) is + resolved from the **installed `stream` package**, where it ships as package data + (stream-dse >= 1.13.3); nothing is vendored in this operator. +- `stream-dse` writes its generated ONNX workload / mapping YAML **into its installed package + directory**, so that environment must be writable. diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py new file mode 100644 index 00000000..3e6cdca3 --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from typing import Any, Optional + +import aie.utils as aie_utils + +from iron.common import ( + MLIROperator, + AIERuntimeArgSpec, + KernelObjectArtifact, + SourceArtifact, + PythonGeneratedMLIRArtifact, + DesignGenerator, +) +from iron.common.device_utils import get_kernel_dir +from iron.common.sequence import OperatorSequence + + +@dataclass +class _SwiGLUStreamGroup(MLIROperator): + """One stream-dse design, used as an ``OperatorSequence`` child. + + ``group_index`` selects the design: ``None`` is the whole SwiGLU block + (``x, w1, w2, w3 -> y``); ``0`` is the gate/up/SiLU/mul front end + (``x, w1, w2 -> h``); ``1`` is the down projection (``h, w3 -> y``). + """ + + seq_len: int + embedding_dim: int + hidden_dim: int + group_index: Optional[int] = None + seq_len_tile_size: int = 32 + embedding_tile_size: int = 32 + hidden_tile_size: int = 64 + in_dtype: str = field(default="bf16", repr=False) + out_dtype: str = field(default="bf16", repr=False) + rows: int = field(default=4, repr=False) + num_aie_columns: int = field(default=8, repr=False) + backend: str = field(default="ortools_gscip", repr=False) + context: Any = field(default=None, repr=False, compare=False) + + def __post_init__(self): + MLIROperator.__init__(self, context=self.context) + + def get_mlir_artifact(self): + npu = aie_utils.get_current_device().resolve().name + kwargs = { + "seq_len": self.seq_len, + "embedding_dim": self.embedding_dim, + "hidden_dim": self.hidden_dim, + "in_dtype": self.in_dtype, + "out_dtype": self.out_dtype, + "rows": self.rows, + "cols": self.num_aie_columns, + "npu": npu, + "seq_len_tile_size": self.seq_len_tile_size, + "embedding_tile_size": self.embedding_tile_size, + "hidden_tile_size": self.hidden_tile_size, + "backend": self.backend, + } + if self.group_index is None: + fn, args = "run_main_aie_codegen_swiglu", () + kwargs["last_gemm_down"] = True + else: + fn, args = "load_swiglu_k2_group", (self.group_index,) + return PythonGeneratedMLIRArtifact( + f"{self.name}.mlir", + DesignGenerator(self.operator_dir / "stream_design.py", fn, args, kwargs), + ) + + def _mm_kernel(self, tile_m, tile_k, tile_n): + # stream-dse emits dimension-suffixed kernel symbols so the gate/up and + # down GEMMs (different tile shapes) coexist; rename mm.cc's unsuffixed + # symbols to match. + base_dir = self.context.base_dir + suffix = f"{tile_m}_{tile_k}_{tile_n}" + return KernelObjectArtifact( + f"mm_{suffix}.o", + dependencies=[ + SourceArtifact(base_dir / "aie_kernels" / get_kernel_dir() / "mm.cc") + ], + extra_flags=[ + f"-DDIM_M={tile_m}", + f"-DDIM_K={tile_k}", + f"-DDIM_N={tile_n}", + "-Dbf16_bf16_ONLY", + ], + rename_symbols={ + "matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}", + "zero_bf16": f"zero_bf16_{suffix}", + }, + ) + + def get_kernel_artifacts(self): + base_dir = self.context.base_dir + kernel_dir = get_kernel_dir() + gate_up = self._mm_kernel( + self.seq_len_tile_size, self.embedding_tile_size, self.hidden_tile_size + ) + down = self._mm_kernel( + self.seq_len_tile_size, self.hidden_tile_size, self.embedding_tile_size + ) + if self.group_index == 1: + return [down] + silu = KernelObjectArtifact( + "silu.o", + dependencies=[ + SourceArtifact(base_dir / "aie_kernels" / kernel_dir / "silu.cc") + ], + ) + mul = KernelObjectArtifact( + "mul.o", + dependencies=[ + SourceArtifact(base_dir / "aie_kernels" / "generic" / "mul.cc") + ], + ) + if self.group_index == 0: + return [gate_up, silu, mul] + return [gate_up, down, silu, mul] + + def get_arg_spec(self): + m, e, h = self.seq_len, self.embedding_dim, self.hidden_dim + if self.group_index == 0: + return [ + AIERuntimeArgSpec("in", (m, e)), # x + AIERuntimeArgSpec("in", (e, h)), # w_gate + AIERuntimeArgSpec("in", (e, h)), # w_up + AIERuntimeArgSpec("out", (m, h)), # h + ] + if self.group_index == 1: + return [ + AIERuntimeArgSpec("in", (m, h)), # h + AIERuntimeArgSpec("in", (h, e)), # w_down + AIERuntimeArgSpec("out", (m, e)), # y + ] + return [ + AIERuntimeArgSpec("in", (m, e)), # x + AIERuntimeArgSpec("in", (e, h)), # w_gate + AIERuntimeArgSpec("in", (e, h)), # w_up + AIERuntimeArgSpec("in", (h, e)), # w_down + AIERuntimeArgSpec("out", (m, e)), # y + ] + + +def _name(kind, m, e, h, st, et, ht): + return f"{kind}_m{m}_e{e}_h{h}_st{st}_et{et}_ht{ht}" + + +class SwiGLUPrefillStream(OperatorSequence): + """Fused SwiGLU-prefill block generated by stream-dse and deployed as a + single full-ELF (``OperatorSequence`` default dispatch). + + Runtime buffers (``get_callable().get_buffer(name)``): ``input``, ``weights_1`` + (gate), ``weights_2`` (up), ``weights_3`` (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, + seq_len_tile_size=32, + embedding_tile_size=32, + hidden_tile_size=64, + context=None, + ): + block = _SwiGLUStreamGroup( + seq_len=seq_len, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + seq_len_tile_size=seq_len_tile_size, + embedding_tile_size=embedding_tile_size, + hidden_tile_size=hidden_tile_size, + context=context, + ) + super().__init__( + name=_name( + "swiglu_prefill_stream", + seq_len, + embedding_dim, + hidden_dim, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + ), + runlist=[(block, "input", "weights_1", "weights_2", "weights_3", "output")], + input_args=["input", "weights_1", "weights_2", "weights_3"], + output_args=["output"], + extra_flags=["--dynamic-objFifos"], + context=context, + ) + + +class SwiGLUPrefillStreamK2(OperatorSequence): + """Two-fusion-group SwiGLU-prefill: a gate/up/SiLU/mul group and a separate + down-projection group fused into one full-ELF, with the hidden state ``h`` + kept on device between them. Same external buffers as + :class:`SwiGLUPrefillStream`; the split is decided by the stream mapping + (``make_swiglu_mapping(split_groups=True)``). + """ + + def __init__( + self, + seq_len, + embedding_dim, + hidden_dim, + seq_len_tile_size=32, + embedding_tile_size=32, + hidden_tile_size=64, + context=None, + ): + common = dict( + seq_len=seq_len, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + seq_len_tile_size=seq_len_tile_size, + embedding_tile_size=embedding_tile_size, + hidden_tile_size=hidden_tile_size, + context=context, + ) + front = _SwiGLUStreamGroup(group_index=0, **common) + down = _SwiGLUStreamGroup(group_index=1, **common) + super().__init__( + name=_name( + "swiglu_prefill_stream_k2", + seq_len, + embedding_dim, + hidden_dim, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + ), + runlist=[ + (front, "input", "weights_1", "weights_2", "h"), + (down, "h", "weights_3", "output"), + ], + input_args=["input", "weights_1", "weights_2", "weights_3"], + output_args=["output"], + extra_flags=["--dynamic-objFifos"], + context=context, + ) 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 00000000..967903ac --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Stream-dse MLIR generation launcher for the fused SwiGLU-prefill operator. + +This is the in-IRON replacement for the previously hardcoded +``/home/micas/stream_aie/main_swiglu.py`` entry point. It calls the *installed* +``stream-dse`` package (``pip install stream-dse`` followed by ``stream-setup-aie``) +to produce a single fused MLIR module for the whole SwiGLU-prefill block, which +IRON then fuses into a single full-ELF via `OperatorSequence`. + +The function signature mirrors ``run_main_aie_codegen_swiglu`` from stream-dse's +``scripts/main_swiglu.py`` reference entry point. Because ``scripts/`` is not +shipped in the stream-dse wheel, that logic is vendored here; the hardware- +description YAML is resolved from the installed ``stream`` package, where it ships +as package data (stream-dse >= 1.13.3). + +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 os +import re +from pathlib import Path + +import stream +from stream.api import optimize_allocation_co +from stream.inputs.aie.mapping.make_swiglu_mapping import make_swiglu_mapping +from stream.inputs.aie.workload.make_onnx_swiglu import make_swiglu_workload + +from iron.operators.swiglu_prefill_stream.stream_kernels import iron_kernels + +# Hardware description for the whole-array Strix (npu2) target, shipped as package +# data inside the installed stream package (stream-dse >= 1.13.3). +_ACCELERATOR = os.path.join( + os.path.dirname(stream.__file__), + "inputs", + "aie", + "hardware", + "whole_array_strix.yaml", +) + + +def run_main_aie_codegen_swiglu( + seq_len, + embedding_dim, + hidden_dim, + in_dtype="bf16", + out_dtype="bf16", + trace_size=0, + rows=4, + cols=8, + npu="npu2", + seq_len_tile_size=32, + embedding_tile_size=32, + hidden_tile_size=64, + last_gemm_down=True, + backend="ortools_gscip", + func_prefix="", +): + """Generate the fused SwiGLU-prefill MLIR module via stream-dse. + + Returns an ``aie`` MLIR module. ``func_prefix`` (injected by + ``OperatorSequence``) prefixes the kernel symbols / ``link_with`` objects so + the design can be deployed as one fusion group; see ``region_module``. + + The default ``ortools_gscip`` backend is the license-free OR-Tools GSCIP + solver, so no Gurobi license is required. + """ + workload_path = make_swiglu_workload( + seq_len, + embedding_dim, + hidden_dim, + in_dtype, + out_dtype, + last_gemm_down=last_gemm_down, + ) + mapping_path = make_swiglu_mapping( + seq_len, + embedding_dim, + hidden_dim, + last_gemm_down, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + ) + + hw_name = os.path.splitext(os.path.basename(_ACCELERATOR))[0] + wl_name = re.split(r"/|\.", workload_path)[-1] + if wl_name == "onnx": + wl_name = re.split(r"/|\.", workload_path)[-2] + experiment_id = f"{hw_name}-{wl_name}-{rows}_row_{cols}_col" + + ctx = optimize_allocation_co( + hardware=_ACCELERATOR, + workload=workload_path, + mapping=mapping_path, + experiment_id=experiment_id, + output_path="outputs", + skip_if_exists=False, + enable_codegen=True, + trace_size=trace_size, + nb_cols_to_use=cols, + npu=npu, + backend=backend, + kernels=iron_kernels(), # IRON-authored operand layouts drive the DMA tiling + ) + return region_module(str(ctx.get("module")), func_prefix) + + +# --------------------------------------------------------------------------- +# k=2 variant: two fusion groups (gate/up/SiLU/mul -> h, then down-projection) +# --------------------------------------------------------------------------- +# +# stream-dse emits a separate ``aie.device`` design per fusion group (under +# ``/group_i/codegen/final.mlir``). IRON fuses the two groups into one +# full-ELF via ``OperatorSequence``; each group is loaded below as a child +# design. The split itself is expressed entirely in the stream mapping +# (``make_swiglu_mapping(split_groups=True)``); see that function. + + +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 sym in symbols: + mlir_text = re.sub(rf"@{re.escape(sym)}\b", f"@{func_prefix}{sym}", mlir_text) + return mlir_text + + +def region_module(mlir_text: str, func_prefix: str = ""): + """Parse a stream 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 _swiglu_k2_experiment_id(seq_len, embedding_dim, hidden_dim, rows, cols): + hw_name = os.path.splitext(os.path.basename(_ACCELERATOR))[0] + return f"{hw_name}-swiglu_k2_{seq_len}_{embedding_dim}_{hidden_dim}-{rows}_row_{cols}_col" + + +def _run_swiglu_k2_codegen( + seq_len, + embedding_dim, + hidden_dim, + in_dtype, + out_dtype, + rows, + cols, + npu, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + backend, + output_root="outputs", +): + """Run the two-group SwiGLU codegen once (cached by output existence). + + Returns the experiment output directory containing ``group_0`` / ``group_1``. + Both group loaders call this; the first generates, the rest reuse the files. + """ + experiment_id = _swiglu_k2_experiment_id( + seq_len, embedding_dim, hidden_dim, rows, cols + ) + out_dir = os.path.join(output_root, experiment_id) + finals = [ + os.path.join(out_dir, f"group_{i}", "codegen", "final.mlir") for i in (0, 1) + ] + if all(os.path.exists(f) for f in finals): + return out_dir + + workload_path = make_swiglu_workload( + seq_len, embedding_dim, hidden_dim, in_dtype, out_dtype, last_gemm_down=True + ) + mapping_path = make_swiglu_mapping( + seq_len, + embedding_dim, + hidden_dim, + True, # last_gemm_down + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + split_groups=True, + ) + 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=cols, + npu=npu, + backend=backend, + kernels=iron_kernels(), + ) + return out_dir + + +def load_swiglu_k2_group( + group_index, + func_prefix="", + *, + seq_len, + embedding_dim, + hidden_dim, + in_dtype="bf16", + out_dtype="bf16", + rows=4, + cols=8, + npu="npu2", + seq_len_tile_size=32, + embedding_tile_size=32, + hidden_tile_size=64, + backend="ortools_gscip", +): + """Generate the two-group design (cached) and return one group's aie module. + + ``group_index`` 0 is the gate/up/SiLU/mul front end (``x, w1, w2 -> h``); 1 is + the down-projection (``h, w3 -> y``). ``func_prefix`` is injected by + ``OperatorSequence``. + """ + out_dir = _run_swiglu_k2_codegen( + seq_len, + embedding_dim, + hidden_dim, + in_dtype, + out_dtype, + rows, + cols, + npu, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + backend, + ) + text = Path(out_dir, f"group_{group_index}", "codegen", "final.mlir").read_text() + return region_module(text, func_prefix) diff --git a/iron/operators/swiglu_prefill_stream/stream_kernels.py b/iron/operators/swiglu_prefill_stream/stream_kernels.py new file mode 100644 index 00000000..17bb57aa --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/stream_kernels.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""IRON-authored operand layouts for the stream-dse SwiGLU-prefill kernels. + +stream-dse selects an AIE kernel per computation node and uses each kernel's +``operand_layouts()`` to drive the DMA tiling emitted into the design MLIR. +:func:`iron_kernels` returns the ``optimize_allocation_co(kernels=...)`` override +that keeps every kernel stream would build but replaces its operand layouts with +the ones defined here -- the single source of truth -- converted to stream's +tiled-strided layout via :meth:`iron.common.TiledStridedLayout.to_snaxc`. IRON +owns the layouts; stream owns construction, symbol names and the MLIR rewrite. + +Each override is stream's own kernel, re-typed to a subclass that overrides only +``operand_layouts()``: the kernel is still built by stream's ``AIEKernels`` +factory (so its constructor signature is inherited, not re-declared here), then +its ``__class__`` is swapped. The subclasses are module-level so the kernels stay +picklable (stream stores them on the mapping). + +stream / snaxc / xdsl are imported at module load, so this module is only +importable where the AIE codegen toolchain is installed; it is imported only from +``stream_design.py``. +""" + +from __future__ import annotations + +from typing import Any, Callable + +from stream.compiler.kernels.eltwise_mul import EltwiseMulKernel +from stream.compiler.kernels.gemm import GemmKernel +from stream.compiler.kernels.silu import SiluKernel + +from iron.common import TiledStridedLayout, tiled_2d + +# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets; the +# operand layouts below are the contract the generated DMAs and the compiled +# kernel objects must agree on. +R, S, T = 4, 8, 8 + + +def _gemm_layouts(m: int, k: int, n: int) -> tuple[TiledStridedLayout, ...]: + return (tiled_2d(m, k, R, S), tiled_2d(k, n, S, T), tiled_2d(m, n, R, T)) + + +def _elementwise_layouts( + count: int, tile: tuple[int, int] = (32, 64) +) -> tuple[TiledStridedLayout, ...]: + return (tiled_2d(*tile, R, T),) * count + + +class _IronGemmKernel(GemmKernel): + def operand_layouts(self): + return [tsl.to_snaxc() for tsl in _gemm_layouts(self.m, self.k, self.n)] + + +class _IronSiluKernel(SiluKernel): + def operand_layouts(self): + return [tsl.to_snaxc() for tsl in _elementwise_layouts(2)] + + +class _IronEltwiseMulKernel(EltwiseMulKernel): + def operand_layouts(self): + return [tsl.to_snaxc() for tsl in _elementwise_layouts(3)] + + +# stream AIEKernels name -> IRON subclass overriding operand_layouts(). +_OVERRIDES: dict[str, type] = { + "gemm": _IronGemmKernel, + "silu": _IronSiluKernel, + "eltwise_mul": _IronEltwiseMulKernel, +} + + +def iron_kernels() -> dict[str, Callable[..., Any]]: + """Return the ``optimize_allocation_co(kernels=...)`` override registry. + + Only kernels for which IRON defines layouts are overridden; any other kernel + stream needs falls through to its built-in ``AIEKernels`` entry. + """ + from stream.compiler.kernels import AIEKernels + + def override(factory: Callable[..., Any], cls: type) -> Callable[..., Any]: + def make(*args: Any, **kwargs: Any) -> Any: + kernel = factory(*args, **kwargs) + kernel.__class__ = cls + return kernel + + return make + + return { + name: override(AIEKernels[name], cls) + for name, cls in _OVERRIDES.items() + if name in AIEKernels + } diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py new file mode 100644 index 00000000..b6b7b05d --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time +import inspect + +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 stream.inputs.aie.mapping.make_swiglu_mapping import make_swiglu_mapping + +from iron.operators.swiglu_prefill_stream.op import ( + SwiGLUPrefillStream, + SwiGLUPrefillStreamK2, +) + +# swiglu_prefill_stream shares swiglu_decode's reference: W3 @ (SiLU(W1@x)*(W2@x)). +from iron.operators.swiglu_decode.reference import generate_golden_reference +from iron.common.test_utils import verify_buffer + +# The k=2 variant needs stream-dse's two-fusion-group support; skip it on older +# releases that lack the ``split_groups`` mapping option. +_STREAM_HAS_SPLIT_GROUPS = ( + "split_groups" in inspect.signature(make_swiglu_mapping).parameters +) + + +def get_params(): + # (seq_len, embedding_dim, hidden_dim, seq_tile, embedding_tile, hidden_tile): + # the MILP-feasible shape on the whole-array Strix (npu2) target. + return [pytest.param(256, 512, 2048, 32, 32, 64)] + + +def _run_and_verify(operator, seq_len, embedding_dim, golden_ref): + operator.compile() + run = operator.get_callable() + + # Populate inputs by name; the design consumes weights in their natural + # (K, N) layout (no transpose). + run.get_buffer("input").torch_view()[:] = ( + golden_ref["input"].to(torch.bfloat16).flatten() + ) + run.get_buffer("weights_1").torch_view()[:] = ( + golden_ref["w_gate"].to(torch.bfloat16).flatten() + ) + run.get_buffer("weights_2").torch_view()[:] = ( + golden_ref["w_up"].to(torch.bfloat16).flatten() + ) + run.get_buffer("weights_3").torch_view()[:] = ( + golden_ref["w_down"].to(torch.bfloat16).flatten() + ) + + # Correctness is checked on the FIRST run: creating the hw_context applies the + # design's init CDO, so the first run on the freshly loaded full-ELF computes + # correctly. The full-ELF path does not re-initialize the array between runs, + # so later runs on the same callable read stale state and are used only for + # timing below. + # + # The whole block is fused 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() + 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}" + + # Performance: time a run that reuses the hw_context, so it reflects the actual + # kernel latency without the one-off ELF/hw_context setup. + start = time.perf_counter() + run() + elapsed_us = (time.perf_counter() - start) * 1e6 + total_bytes = 2 * (seq_len * embedding_dim + seq_len * embedding_dim) # bf16 in+out + print(f"Latency (us): {elapsed_us:.2f}") + print(f"Effective Bandwidth: {total_bytes / (elapsed_us * 1e-6) / 1e9:.4f} GB/s") + + +@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( + "seq_len,embedding_dim,hidden_dim,seq_tile,embedding_tile,hidden_tile", get_params() +) +def test_swiglu_prefill_stream( + seq_len, + embedding_dim, + hidden_dim, + seq_tile, + embedding_tile, + hidden_tile, + 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, + seq_len_tile_size=seq_tile, + embedding_tile_size=embedding_tile, + hidden_tile_size=hidden_tile, + context=aie_context, + ) + _run_and_verify(operator, seq_len, embedding_dim, golden_ref) + + +@pytest.mark.skipif( + not _STREAM_HAS_SPLIT_GROUPS, + reason="installed stream-dse lacks two-fusion-group support (split_groups)", +) +@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( + "seq_len,embedding_dim,hidden_dim,seq_tile,embedding_tile,hidden_tile", get_params() +) +def test_swiglu_prefill_stream_k2( + seq_len, + embedding_dim, + hidden_dim, + seq_tile, + embedding_tile, + hidden_tile, + aie_context, +): + golden_ref = generate_golden_reference(M=seq_len, K=embedding_dim, N=hidden_dim) + operator = SwiGLUPrefillStreamK2( + seq_len=seq_len, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + seq_len_tile_size=seq_tile, + embedding_tile_size=embedding_tile, + hidden_tile_size=hidden_tile, + context=aie_context, + ) + _run_and_verify(operator, seq_len, embedding_dim, golden_ref) From 3ae6ab8ec52492d176826121c157818c25073618 Mon Sep 17 00:00:00 2001 From: asyms Date: Fri, 24 Jul 2026 19:38:32 +0200 Subject: [PATCH 05/21] ci: require stream-dse>=1.13.7 so stream-setup-aie stops clobbering mlir_aie stream-setup-aie (from stream-dse) runs after requirements.txt in the prereqs action and, up to 1.13.6, unconditionally reinstalled an older pinned mlir_aie/llvm-aie, downgrading the wheels requirements.txt had just installed (mlir_aie 1.3.5 -> 0.0.1, numpy 2.x -> 1.26.4). Every operator test then failed against the stale bindings (AnyComputeTile.col, Worker(tile=...), ObjectFifoHandle.split(tile=...), aie.iron.ScratchpadParameter): 525 failures. stream-dse 1.13.7 makes stream-setup-aie install only the pure-Python codegen dialects (xdsl-aie, snax-mlir) and leaves mlir_aie/llvm-aie to requirements.txt. Bump the pin and correct the now-inaccurate comments/README. --- .github/actions/prereqs/action.yaml | 8 ++++--- .../operators/swiglu_prefill_stream/README.md | 10 +++++---- requirements_stream.txt | 21 ++++++++++++------- 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/.github/actions/prereqs/action.yaml b/.github/actions/prereqs/action.yaml index 2b09153b..bf1616a0 100644 --- a/.github/actions/prereqs/action.yaml +++ b/.github/actions/prereqs/action.yaml @@ -22,8 +22,10 @@ runs: pip install --upgrade pip pip install -r requirements.txt pip install -r requirements_stream.txt - # stream-dse's AIE codegen deps (snax-mlir/snaxc, xdsl-aie, aie-python-extras) - # are not PyPI dependencies; this console script installs them so the - # stream-backed operators can generate their MLIR at build time. + # 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/iron/operators/swiglu_prefill_stream/README.md b/iron/operators/swiglu_prefill_stream/README.md index 2b7c779c..dc9ae7a5 100644 --- a/iron/operators/swiglu_prefill_stream/README.md +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -26,10 +26,12 @@ stream-setup-aie # required: installs stream-dse's AIE codegen deps Notes: - MLIR generation uses the open-source **OR-Tools GSCIP** solver (`backend="ortools_gscip"`), so **no Gurobi license** is required. -- `stream-setup-aie` is **required**: it installs the AIE codegen packages stream-dse needs - that cannot be plain PyPI dependencies (`snax-mlir`/`snaxc`, `xdsl-aie`, `aie-python-extras`), - since they are direct git/URL installs. It also installs the `mlir_aie` / `llvm-aie` wheels, - but skips those if IRON's `requirements.txt` already provided them. +- `stream-setup-aie` is **required**: it installs the pure-Python AIE codegen dialects + stream-dse needs that cannot be plain PyPI dependencies (`xdsl-aie`, `snax-mlir`), since they + are direct git installs. As of **stream-dse 1.13.7** it does **not** install `mlir_aie` / + `llvm-aie`: IRON's `requirements.txt` already pins those, and stream-dse's codegen only emits + text MLIR via xdsl (it never imports the `aie` bindings). Earlier releases reinstalled an + older `mlir_aie` here, which downgraded IRON's pinned wheel and broke the test suite. - Importing the operator does **not** require `stream-dse` (the launcher is imported lazily); only **building** (`operator.compile()` / running the test) does. diff --git a/requirements_stream.txt b/requirements_stream.txt index 26a8029a..d21b2842 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -9,22 +9,27 @@ # to build/run the operator and its test: # # pip install -r requirements_stream.txt -# stream-setup-aie # REQUIRED: installs stream-dse's AIE codegen deps that -# # cannot be PyPI dependencies (snax-mlir/snaxc, xdsl-aie, -# # aie-python-extras); also installs the mlir_aie/llvm-aie -# # wheels, skipping any already provided by requirements.txt. +# stream-setup-aie # REQUIRED: installs stream-dse's pure-Python AIE codegen +# # dialects (xdsl-aie, snax-mlir) that cannot be PyPI +# # dependencies. As of stream-dse 1.13.7 it does NOT install +# # mlir_aie/llvm-aie: IRON already pins those in +# # requirements.txt, and codegen only emits text MLIR via +# # xdsl (it never imports the aie bindings). Earlier releases +# # reinstalled an older mlir_aie here, clobbering IRON's wheel. # # Notes: # - stream-dse generates the fused MLIR design at build time (license-free # OR-Tools GSCIP solver; no Gurobi needed) and writes its generated workload/ # mapping files into its own installed package directory, so that environment # must be writable. -# - >=1.13.6 is required: stream_design.py feeds IRON-authored operand layouts +# - >=1.13.7 is required: stream_design.py feeds IRON-authored operand layouts # into code generation via optimize_allocation_co(kernels=...), the override # hook added in stream-dse 1.13.4; the k=2 variant (op_k2.py) additionally needs # the two-fusion-group support (make_swiglu_mapping(split_groups=...) + the -# multi-group AIE codegen pipeline) added in stream-dse 1.13.5; and 1.13.6 makes +# multi-group AIE codegen pipeline) added in stream-dse 1.13.5; 1.13.6 makes # the debug workload-graph visualization non-fatal, so code generation no longer -# requires graphviz (`dot`) to be installed. +# requires graphviz (`dot`); and 1.13.7 makes `stream-setup-aie` stop reinstalling +# mlir_aie/llvm-aie, which previously downgraded IRON's pinned mlir_aie (1.3.5 -> +# 0.0.1) and numpy (2.x -> 1.26.4) and broke the whole operator test suite. -stream-dse>=1.13.6 +stream-dse>=1.13.7 From b02f47db08bc32d877ac22483a415cb4d42fe2c0 Mon Sep 17 00:00:00 2001 From: asyms Date: Sat, 25 Jul 2026 18:38:38 +0200 Subject: [PATCH 06/21] swiglu_prefill_stream: generate the workload and mapping in IRON The design's inputs were built by stream-dse helpers (make_swiglu_workload / make_swiglu_mapping), so the graph the NPU ran was described separately from the reference the result was checked against, and the mapping referred to node names restated by hand. Both are now produced in IRON from one source: - reference.py holds the SwiGLU nn.Module. torch.onnx.export lowers it through the translation table in iron/common/stream/ops.py, which emits each operator in the form stream-dse parses; weight payloads are dropped and nodes are named by the operator afterwards. - iron/common/stream/mapping.py emits the mapping from the operator's placement, taking node names from that same exported graph and validating every placement against it, so workload and mapping cannot disagree. - iron/common/stream/ops.py binds each torch operator to its ONNX form, its stream-dse kernel and the aie_kernels source that kernel is built from, so op.py no longer hardcodes kernel objects, flags and symbol renames. Ops with a fused AIE kernel but no ONNX operator get a schema in a private domain, so the exporter emits them as a single node. Changing the problem size means re-exporting with different example shapes; the placement carries tile sizes and core sets but no absolute dimensions. stream_kernels.py is dropped: its operand-layout overrides reproduced stream-dse's built-in layouts exactly, so the design is unchanged without them. iron/tests/stream now asserts that equivalence instead, to catch a future divergence. --- iron/common/stream/__init__.py | 18 + iron/common/stream/mapping.py | 126 ++++++ iron/common/stream/ops.py | 192 +++++++++ iron/common/stream/workload.py | 141 +++++++ .../operators/swiglu_prefill_stream/README.md | 34 +- iron/operators/swiglu_prefill_stream/op.py | 68 +-- .../swiglu_prefill_stream/reference.py | 50 +++ .../swiglu_prefill_stream/stream_design.py | 388 ++++++++++++------ .../swiglu_prefill_stream/stream_kernels.py | 94 ----- iron/operators/swiglu_prefill_stream/test.py | 38 +- iron/tests/stream/kernel_layouts.py | 53 +++ requirements_stream.txt | 3 +- 12 files changed, 900 insertions(+), 305 deletions(-) create mode 100644 iron/common/stream/__init__.py create mode 100644 iron/common/stream/mapping.py create mode 100644 iron/common/stream/ops.py create mode 100644 iron/common/stream/workload.py create mode 100644 iron/operators/swiglu_prefill_stream/reference.py delete mode 100644 iron/operators/swiglu_prefill_stream/stream_kernels.py create mode 100644 iron/tests/stream/kernel_layouts.py diff --git a/iron/common/stream/__init__.py b/iron/common/stream/__init__.py new file mode 100644 index 00000000..0a55ae7c --- /dev/null +++ b/iron/common/stream/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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/mapping.py b/iron/common/stream/mapping.py new file mode 100644 index 00000000..dd2c3e70 --- /dev/null +++ b/iron/common/stream/mapping.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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.workload import StreamWorkload + + +@dataclass(frozen=True) +class Placement: + """Where one workload node runs. + + ``cores`` lists the compute tiles per allocation, ``splits`` the matching + inter-core tiling as ``(dim, split)`` pairs, and ``kernel_kwargs`` the + arguments of the node's stream-dse kernel (e.g. a GEMM's tile shape). + """ + + cores: Sequence[Sequence[int]] + splits: Sequence[Sequence[tuple[str, int]]] = () + kernel_kwargs: dict = field(default_factory=dict) + + +@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 _layer_entry(name: str, placement: Placement, kernel_key: str) -> dict: + entry = { + "name": name, + "core_allocation": [list(cores) for cores in placement.cores], + "inter_core_tiling": [ + [{"dim": dim, "split": split} for dim, split in group] + for group in placement.splits + ], + "kernel": {"name": kernel_key, "kwargs": dict(placement.kernel_kwargs)}, + } + return entry + + +def build_mapping( + workload: StreamWorkload, + placements: dict[str, Placement], + groups: Sequence[FusedGroup], +) -> 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) + 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], + 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), + 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 00000000..475fadff --- /dev/null +++ b/iron/common/stream/ops.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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. +R, S, T = 4, 8, 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 gemm_layouts(m: int, k: int, n: int) -> tuple[TiledStridedLayout, ...]: + """Layouts of a GEMM's ``A[m,k]``, ``B[k,n]`` and ``C[m,n]`` operands.""" + return (tiled_2d(m, k, R, S), tiled_2d(k, n, S, T), tiled_2d(m, n, R, T)) + + +def elementwise_layouts(nb_operands: int) -> tuple[TiledStridedLayout, ...]: + """Identical tiled layout for each operand of an elementwise kernel.""" + return (tiled_2d(*ELEMENTWISE_TILE, R, 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", + ], + 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.""" + + onnx_type: str + kernel: StreamKernel + translation: Callable + + +# torch operator -> its ONNX form and AIE kernel. Gemm rather than 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()} + + +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 00000000..488ad69d --- /dev/null +++ b/iron/common/stream/workload.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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 + + 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/swiglu_prefill_stream/README.md b/iron/operators/swiglu_prefill_stream/README.md index dc9ae7a5..828e227a 100644 --- a/iron/operators/swiglu_prefill_stream/README.md +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -12,6 +12,30 @@ flow into one xclbin. Unlike the other operators, its MLIR is not written by han produced at build time by [`stream_design.py`](./stream_design.py), which calls the installed `stream` package (`stream.api.optimize_allocation_co(..., enable_codegen=True)`). +## Where the design comes from + +Both inputs stream-dse needs are generated **by IRON**, from one source: + +| | Source | Built by | +|---|---|---| +| Workload (ONNX) | [`reference.py`](./reference.py) — the `SwiGLU` `nn.Module` | `torch.export` → [`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 | [`iron/common/stream/ops.py`](../../common/stream/ops.py) registry | + +The module that generates the design is the same one the test checks the result +against, so the two cannot drift. Mapping node names are taken from the exported +graph rather than restated, so workload and mapping cannot disagree either. Both +files are written into the experiment's output directory at build time — nothing +is committed to the tree. + +Changing the problem size means re-exporting with different example shapes: the +placement carries tile sizes and core sets but no absolute dimensions, so it is +unaffected (within the divisibility limits `stream_design._check_shapes` enforces). + +Supporting another op (say a softmax block) is one `StreamKernel` plus one +`ATEN_OPS` entry in `iron/common/stream/ops.py`, pointing at the existing +`aie_kernels//softmax.cc`, plus that operator's own placement. + ## Enabling stream codegen `stream-dse` is an **optional, separately-installed** dependency (it is *not* in IRON's @@ -51,5 +75,11 @@ The feasible/verified shape is **seq 256 / embedding 512 / hidden 2048**, tiles - The hardware-description YAML (`whole_array_strix.yaml` + `hardware/cores/*.yaml`) is resolved from the **installed `stream` package**, where it ships as package data (stream-dse >= 1.13.3); nothing is vendored in this operator. -- `stream-dse` writes its generated ONNX workload / mapping YAML **into its installed package - directory**, so that environment must be writable. +- Nodes and their results are renamed from the ATen names `torch.export` produces + (`matmul`, `matmul_1`, …). Those are prefixes of one another, which stream-dse's + tensor bookkeeping does not distinguish reliably (it raises a `KeyError` while + building the memory-capacity constraints). +- Naming does not change the problem posed to the solver — same nodes, tensors, + shapes and connections — but it does change SCIP's presolve, and with it how + long the allocation takes to solve. The names here are the ones this operator + was measured with. diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 3e6cdca3..d31efd6f 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -9,13 +9,12 @@ from iron.common import ( MLIROperator, AIERuntimeArgSpec, - KernelObjectArtifact, - SourceArtifact, 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 @@ -70,55 +69,28 @@ def get_mlir_artifact(self): DesignGenerator(self.operator_dir / "stream_design.py", fn, args, kwargs), ) - def _mm_kernel(self, tile_m, tile_k, tile_n): - # stream-dse emits dimension-suffixed kernel symbols so the gate/up and - # down GEMMs (different tile shapes) coexist; rename mm.cc's unsuffixed - # symbols to match. - base_dir = self.context.base_dir - suffix = f"{tile_m}_{tile_k}_{tile_n}" - return KernelObjectArtifact( - f"mm_{suffix}.o", - dependencies=[ - SourceArtifact(base_dir / "aie_kernels" / get_kernel_dir() / "mm.cc") - ], - extra_flags=[ - f"-DDIM_M={tile_m}", - f"-DDIM_K={tile_k}", - f"-DDIM_N={tile_n}", - "-Dbf16_bf16_ONLY", - ], - rename_symbols={ - "matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}", - "zero_bf16": f"zero_bf16_{suffix}", - }, - ) - def get_kernel_artifacts(self): - base_dir = self.context.base_dir - kernel_dir = get_kernel_dir() - gate_up = self._mm_kernel( + # Each kernel's source, compile flags and symbol names come from the + # stream op registry, so they stay in step with the design stream-dse + # generates against IRON's aie_kernels library. + from iron.operators.swiglu_prefill_stream.stream_design import gemm_tiles + + base_dir, kernel_dir = self.context.base_dir, get_kernel_dir() + gate_up, down = gemm_tiles( self.seq_len_tile_size, self.embedding_tile_size, self.hidden_tile_size ) - down = self._mm_kernel( - self.seq_len_tile_size, self.hidden_tile_size, self.embedding_tile_size - ) - if self.group_index == 1: - return [down] - silu = KernelObjectArtifact( - "silu.o", - dependencies=[ - SourceArtifact(base_dir / "aie_kernels" / kernel_dir / "silu.cc") - ], - ) - mul = KernelObjectArtifact( - "mul.o", - dependencies=[ - SourceArtifact(base_dir / "aie_kernels" / "generic" / "mul.cc") - ], - ) - if self.group_index == 0: - return [gate_up, silu, mul] - return [gate_up, down, silu, mul] + kernels = { + 0: [(GEMM, gate_up), (SILU, None), (ELTWISE_MUL, None)], + 1: [(GEMM, down)], + None: [(GEMM, gate_up), (GEMM, down), (SILU, None), (ELTWISE_MUL, None)], + }[self.group_index] + return [ + artifact + for kernel, tiles in kernels + for artifact in kernel.kernel_artifacts( + base_dir, kernel_dir, **(dict(zip("mkn", tiles)) if tiles else {}) + ) + ] def get_arg_spec(self): m, e, h = self.seq_len, self.embedding_dim, self.hidden_dim diff --git a/iron/operators/swiglu_prefill_stream/reference.py b/iron/operators/swiglu_prefill_stream/reference.py new file mode 100644 index 00000000..78b7718f --- /dev/null +++ b/iron/operators/swiglu_prefill_stream/reference.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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. Its parameter names are the ONNX +tensor names, hence the operator's runtime buffer names. + +Numerical values come from :mod:`iron.operators.swiglu_decode.reference`, which +this operator shares; :func:`swiglu_module` loads them into the module. +""" + +import torch +from torch import nn + +# Reference weight -> this module's parameter (and so the runtime buffer name). +WEIGHTS = {"w_gate": "weights_1", "w_up": "weights_2", "w_down": "weights_3"} + + +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.weights_1 = nn.Parameter(torch.zeros(gate_up, dtype=dtype)) # gate + self.weights_2 = nn.Parameter(torch.zeros(gate_up, dtype=dtype)) # up + self.weights_3 = nn.Parameter( + torch.zeros((hidden_dim, embedding_dim), dtype=dtype) + ) # down + + def forward(self, input): + gate = input @ self.weights_1 + up = input @ self.weights_2 + return (torch.nn.functional.silu(gate) * up) @ self.weights_3 + + +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 reference, parameter in WEIGHTS.items(): + getattr(module, parameter).copy_(golden_reference[reference]) + return module diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index 967903ac..33b02e9c 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -1,19 +1,17 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Stream-dse MLIR generation launcher for the fused SwiGLU-prefill operator. +"""Generate the fused SwiGLU-prefill design with stream-dse. -This is the in-IRON replacement for the previously hardcoded -``/home/micas/stream_aie/main_swiglu.py`` entry point. It calls the *installed* -``stream-dse`` package (``pip install stream-dse`` followed by ``stream-setup-aie``) -to produce a single fused MLIR module for the whole SwiGLU-prefill block, which -IRON then fuses into a single full-ELF via `OperatorSequence`. +Both inputs stream-dse needs are produced here, from IRON: -The function signature mirrors ``run_main_aie_codegen_swiglu`` from stream-dse's -``scripts/main_swiglu.py`` reference entry point. Because ``scripts/`` is not -shipped in the stream-dse wheel, that logic is vendored here; the hardware- -description YAML is resolved from the installed ``stream`` package, where it ships -as package data (stream-dse >= 1.13.3). +* 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 @@ -25,11 +23,12 @@ from pathlib import Path import stream +import torch from stream.api import optimize_allocation_co -from stream.inputs.aie.mapping.make_swiglu_mapping import make_swiglu_mapping -from stream.inputs.aie.workload.make_onnx_swiglu import make_swiglu_workload -from iron.operators.swiglu_prefill_stream.stream_kernels import iron_kernels +from iron.common.stream.mapping import FusedGroup, Placement, emit_mapping +from iron.common.stream.workload import export_workload +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 (stream-dse >= 1.13.3). @@ -41,6 +40,217 @@ "whole_array_strix.yaml", ) +# 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 +# torch.export 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: "left", + UP: "right", + SILU: "left_swished", + MUL: "intermediate", +} + +# Compute-tile placement on the 4-row x 8-column array: each GEMM gets two +# columns' worth of cores split 4-ways over rows (D0, the sequence dimension) and +# 2-ways over D2, the two elementwise layers one column each. This is a property +# of the array, not of the problem size, so it holds across shapes. +_GEMM_SPLIT = [[("D0", 4), ("D2", 2)]] +_ELEMENTWISE_SPLIT = [[("D0", 4)]] +_CORES = { + GATE: [[2, 3, 4, 5, 8, 9, 10, 11]], + UP: [[14, 15, 16, 17, 20, 21, 22, 23]], + SILU: [[26, 27, 28, 29]], + MUL: [[32, 33, 34, 35]], + DOWN: [[38, 39, 40, 41, 44, 45, 46, 47]], +} + + +def gemm_tiles(seq_tile, embedding_tile, hidden_tile): + """Kernel tile shape ``(m, k, n)`` of the gate/up GEMMs and of the down GEMM. + + The gate and up projections contract over the embedding dimension and produce + the hidden one; the down projection contracts the other way round. + """ + return (seq_tile, embedding_tile, hidden_tile), ( + seq_tile, + hidden_tile, + embedding_tile, + ) + + +def _placements(seq_tile, embedding_tile, hidden_tile): + gate_up, down = gemm_tiles(seq_tile, embedding_tile, hidden_tile) + + def gemm(tiles): + m, k, n = tiles + return {"utilization": 61.8, "m": m, "k": k, "n": n, "layout": "default"} + + elementwise = {"utilization": 50.0, "layout": "default"} + return { + GATE: Placement(_CORES[GATE], _GEMM_SPLIT, gemm(gate_up)), + UP: Placement(_CORES[UP], _GEMM_SPLIT, gemm(gate_up)), + SILU: Placement(_CORES[SILU], _ELEMENTWISE_SPLIT, elementwise), + MUL: Placement(_CORES[MUL], _ELEMENTWISE_SPLIT, elementwise), + DOWN: Placement(_CORES[DOWN], _GEMM_SPLIT, gemm(down)), + } + + +def _groups(seq_tile, embedding_tile, hidden_tile, split_groups): + """Fusion groups: one fully fused design, or a front end plus down projection. + + 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 of the group's leading GEMM. + """ + if split_groups: + return [ + FusedGroup( + "Fused_Group_1", + [GATE, UP, SILU, MUL], + [ + (GATE, "D1", embedding_tile), + (GATE, "D2", hidden_tile), + (GATE, "D0", seq_tile), + ], + ), + FusedGroup( + "Fused_Group_2", + [DOWN], + [ + (DOWN, "D1", hidden_tile), + (DOWN, "D2", embedding_tile), + (DOWN, "D0", seq_tile), + ], + ), + ] + return [ + FusedGroup( + "Fused_Group_1", + [GATE, UP, SILU, MUL, DOWN], + [ + (GATE, "D1", embedding_tile), + (DOWN, "D2", embedding_tile), + (GATE, "D2", hidden_tile), + (GATE, "D0", seq_tile), + ], + ) + ] + + +def _check_shapes( + seq_len, embedding_dim, hidden_dim, seq_tile, embedding_tile, hidden_tile +): + """Reject shapes the fixed 4x8 placement cannot tile evenly.""" + rows, gemm_split = 4, 2 + if seq_len % rows or seq_len < seq_tile * rows: + raise ValueError( + f"seq_len ({seq_len}) must be a multiple of {rows} and at least {seq_tile * 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} ({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} ({hidden_tile * gemm_split})" + ) + + +def build_inputs( + seq_len, + embedding_dim, + hidden_dim, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + output_dir, + split_groups=False, +): + """Write the workload and mapping for one configuration; return their paths.""" + _check_shapes( + seq_len, + embedding_dim, + hidden_dim, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + ) + module = swiglu_module(embedding_dim, hidden_dim) + workload = export_workload( + module, + (torch.zeros(seq_len, embedding_dim, dtype=torch.bfloat16),), + node_names=NODE_NAMES, + result_names=RESULT_NAMES, + ) + output_dir = Path(output_dir) + return ( + workload.write(output_dir / "workload.onnx"), + emit_mapping( + workload, + _placements(seq_len_tile_size, embedding_tile_size, hidden_tile_size), + _groups( + seq_len_tile_size, embedding_tile_size, hidden_tile_size, split_groups + ), + output_dir / "mapping.yaml", + ), + ) + + +def _experiment_id(seq_len, embedding_dim, hidden_dim, rows, cols, suffix=""): + hw_name = os.path.splitext(os.path.basename(_ACCELERATOR))[0] + return f"{hw_name}-swiglu{suffix}_{seq_len}_{embedding_dim}_{hidden_dim}-{rows}_row_{cols}_col" + + +def _run_codegen( + seq_len, + embedding_dim, + hidden_dim, + rows, + cols, + npu, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + backend, + trace_size=0, + split_groups=False, + output_root="outputs", +): + """Run stream-dse's constraint-optimization + code generation once.""" + experiment_id = _experiment_id( + seq_len, embedding_dim, hidden_dim, rows, cols, "_k2" if split_groups else "" + ) + output_dir = os.path.join(output_root, experiment_id) + workload_path, mapping_path = build_inputs( + seq_len, + embedding_dim, + hidden_dim, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + output_dir, + split_groups=split_groups, + ) + ctx = 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=trace_size, + nb_cols_to_use=cols, + npu=npu, + backend=backend, + ) + return ctx, output_dir + def run_main_aie_codegen_swiglu( seq_len, @@ -59,52 +269,27 @@ def run_main_aie_codegen_swiglu( backend="ortools_gscip", func_prefix="", ): - """Generate the fused SwiGLU-prefill MLIR module via stream-dse. + """Generate the fully fused SwiGLU-prefill MLIR module. Returns an ``aie`` MLIR module. ``func_prefix`` (injected by ``OperatorSequence``) prefixes the kernel symbols / ``link_with`` objects so - the design can be deployed as one fusion group; see ``region_module``. + the design can be deployed as one fusion group; see :func:`region_module`. The default ``ortools_gscip`` backend is the license-free OR-Tools GSCIP solver, so no Gurobi license is required. """ - workload_path = make_swiglu_workload( + ctx, _ = _run_codegen( seq_len, embedding_dim, hidden_dim, - in_dtype, - out_dtype, - last_gemm_down=last_gemm_down, - ) - mapping_path = make_swiglu_mapping( - seq_len, - embedding_dim, - hidden_dim, - last_gemm_down, + rows, + cols, + npu, seq_len_tile_size, embedding_tile_size, hidden_tile_size, - ) - - hw_name = os.path.splitext(os.path.basename(_ACCELERATOR))[0] - wl_name = re.split(r"/|\.", workload_path)[-1] - if wl_name == "onnx": - wl_name = re.split(r"/|\.", workload_path)[-2] - experiment_id = f"{hw_name}-{wl_name}-{rows}_row_{cols}_col" - - ctx = optimize_allocation_co( - hardware=_ACCELERATOR, - workload=workload_path, - mapping=mapping_path, - experiment_id=experiment_id, - output_path="outputs", - skip_if_exists=False, - enable_codegen=True, + backend, trace_size=trace_size, - nb_cols_to_use=cols, - npu=npu, - backend=backend, - kernels=iron_kernels(), # IRON-authored operand layouts drive the DMA tiling ) return region_module(str(ctx.get("module")), func_prefix) @@ -115,9 +300,7 @@ def run_main_aie_codegen_swiglu( # # stream-dse emits a separate ``aie.device`` design per fusion group (under # ``/group_i/codegen/final.mlir``). IRON fuses the two groups into one -# full-ELF via ``OperatorSequence``; each group is loaded below as a child -# design. The split itself is expressed entirely in the stream mapping -# (``make_swiglu_mapping(split_groups=True)``); see that function. +# full-ELF via ``OperatorSequence``; each group is loaded below as a child design. def _prefixed(mlir_text: str, func_prefix: str) -> str: @@ -159,71 +342,6 @@ def region_module(mlir_text: str, func_prefix: str = ""): return ir.Module.parse(_prefixed(mlir_text, func_prefix)) -def _swiglu_k2_experiment_id(seq_len, embedding_dim, hidden_dim, rows, cols): - hw_name = os.path.splitext(os.path.basename(_ACCELERATOR))[0] - return f"{hw_name}-swiglu_k2_{seq_len}_{embedding_dim}_{hidden_dim}-{rows}_row_{cols}_col" - - -def _run_swiglu_k2_codegen( - seq_len, - embedding_dim, - hidden_dim, - in_dtype, - out_dtype, - rows, - cols, - npu, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - backend, - output_root="outputs", -): - """Run the two-group SwiGLU codegen once (cached by output existence). - - Returns the experiment output directory containing ``group_0`` / ``group_1``. - Both group loaders call this; the first generates, the rest reuse the files. - """ - experiment_id = _swiglu_k2_experiment_id( - seq_len, embedding_dim, hidden_dim, rows, cols - ) - out_dir = os.path.join(output_root, experiment_id) - finals = [ - os.path.join(out_dir, f"group_{i}", "codegen", "final.mlir") for i in (0, 1) - ] - if all(os.path.exists(f) for f in finals): - return out_dir - - workload_path = make_swiglu_workload( - seq_len, embedding_dim, hidden_dim, in_dtype, out_dtype, last_gemm_down=True - ) - mapping_path = make_swiglu_mapping( - seq_len, - embedding_dim, - hidden_dim, - True, # last_gemm_down - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - split_groups=True, - ) - 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=cols, - npu=npu, - backend=backend, - kernels=iron_kernels(), - ) - return out_dir - - def load_swiglu_k2_group( group_index, func_prefix="", @@ -243,23 +361,29 @@ def load_swiglu_k2_group( ): """Generate the two-group design (cached) and return one group's aie module. - ``group_index`` 0 is the gate/up/SiLU/mul front end (``x, w1, w2 -> h``); 1 is - the down-projection (``h, w3 -> y``). ``func_prefix`` is injected by - ``OperatorSequence``. + ``group_index`` 0 is the gate/up/SiLU/mul front end (``x, w_gate, w_up -> h``); + 1 is the down projection (``h, w_down -> y``). ``func_prefix`` is injected by + ``OperatorSequence``. Both group loaders call this; the first generates the + design, the second reuses the files on disk. """ - out_dir = _run_swiglu_k2_codegen( - seq_len, - embedding_dim, - hidden_dim, - in_dtype, - out_dtype, - rows, - cols, - npu, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - backend, + output_dir = os.path.join( + "outputs", _experiment_id(seq_len, embedding_dim, hidden_dim, rows, cols, "_k2") ) - text = Path(out_dir, f"group_{group_index}", "codegen", "final.mlir").read_text() - return region_module(text, func_prefix) + finals = [ + os.path.join(output_dir, f"group_{i}", "codegen", "final.mlir") for i in (0, 1) + ] + if not all(os.path.exists(f) for f in finals): + _run_codegen( + seq_len, + embedding_dim, + hidden_dim, + rows, + cols, + npu, + seq_len_tile_size, + embedding_tile_size, + hidden_tile_size, + backend, + split_groups=True, + ) + return region_module(Path(finals[group_index]).read_text(), func_prefix) diff --git a/iron/operators/swiglu_prefill_stream/stream_kernels.py b/iron/operators/swiglu_prefill_stream/stream_kernels.py deleted file mode 100644 index 17bb57aa..00000000 --- a/iron/operators/swiglu_prefill_stream/stream_kernels.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""IRON-authored operand layouts for the stream-dse SwiGLU-prefill kernels. - -stream-dse selects an AIE kernel per computation node and uses each kernel's -``operand_layouts()`` to drive the DMA tiling emitted into the design MLIR. -:func:`iron_kernels` returns the ``optimize_allocation_co(kernels=...)`` override -that keeps every kernel stream would build but replaces its operand layouts with -the ones defined here -- the single source of truth -- converted to stream's -tiled-strided layout via :meth:`iron.common.TiledStridedLayout.to_snaxc`. IRON -owns the layouts; stream owns construction, symbol names and the MLIR rewrite. - -Each override is stream's own kernel, re-typed to a subclass that overrides only -``operand_layouts()``: the kernel is still built by stream's ``AIEKernels`` -factory (so its constructor signature is inherited, not re-declared here), then -its ``__class__`` is swapped. The subclasses are module-level so the kernels stay -picklable (stream stores them on the mapping). - -stream / snaxc / xdsl are imported at module load, so this module is only -importable where the AIE codegen toolchain is installed; it is imported only from -``stream_design.py``. -""" - -from __future__ import annotations - -from typing import Any, Callable - -from stream.compiler.kernels.eltwise_mul import EltwiseMulKernel -from stream.compiler.kernels.gemm import GemmKernel -from stream.compiler.kernels.silu import SiluKernel - -from iron.common import TiledStridedLayout, tiled_2d - -# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets; the -# operand layouts below are the contract the generated DMAs and the compiled -# kernel objects must agree on. -R, S, T = 4, 8, 8 - - -def _gemm_layouts(m: int, k: int, n: int) -> tuple[TiledStridedLayout, ...]: - return (tiled_2d(m, k, R, S), tiled_2d(k, n, S, T), tiled_2d(m, n, R, T)) - - -def _elementwise_layouts( - count: int, tile: tuple[int, int] = (32, 64) -) -> tuple[TiledStridedLayout, ...]: - return (tiled_2d(*tile, R, T),) * count - - -class _IronGemmKernel(GemmKernel): - def operand_layouts(self): - return [tsl.to_snaxc() for tsl in _gemm_layouts(self.m, self.k, self.n)] - - -class _IronSiluKernel(SiluKernel): - def operand_layouts(self): - return [tsl.to_snaxc() for tsl in _elementwise_layouts(2)] - - -class _IronEltwiseMulKernel(EltwiseMulKernel): - def operand_layouts(self): - return [tsl.to_snaxc() for tsl in _elementwise_layouts(3)] - - -# stream AIEKernels name -> IRON subclass overriding operand_layouts(). -_OVERRIDES: dict[str, type] = { - "gemm": _IronGemmKernel, - "silu": _IronSiluKernel, - "eltwise_mul": _IronEltwiseMulKernel, -} - - -def iron_kernels() -> dict[str, Callable[..., Any]]: - """Return the ``optimize_allocation_co(kernels=...)`` override registry. - - Only kernels for which IRON defines layouts are overridden; any other kernel - stream needs falls through to its built-in ``AIEKernels`` entry. - """ - from stream.compiler.kernels import AIEKernels - - def override(factory: Callable[..., Any], cls: type) -> Callable[..., Any]: - def make(*args: Any, **kwargs: Any) -> Any: - kernel = factory(*args, **kwargs) - kernel.__class__ = cls - return kernel - - return make - - return { - name: override(AIEKernels[name], cls) - for name, cls in _OVERRIDES.items() - if name in AIEKernels - } diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index b6b7b05d..778d6ced 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -3,7 +3,6 @@ # SPDX-License-Identifier: Apache-2.0 import time -import inspect import pytest import torch @@ -15,23 +14,17 @@ "stream", reason="stream-dse not installed (see requirements_stream.txt)" ) -from stream.inputs.aie.mapping.make_swiglu_mapping import make_swiglu_mapping - from iron.operators.swiglu_prefill_stream.op import ( SwiGLUPrefillStream, SwiGLUPrefillStreamK2, ) -# swiglu_prefill_stream shares swiglu_decode's reference: W3 @ (SiLU(W1@x)*(W2@x)). +# 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 WEIGHTS from iron.common.test_utils import verify_buffer -# The k=2 variant needs stream-dse's two-fusion-group support; skip it on older -# releases that lack the ``split_groups`` mapping option. -_STREAM_HAS_SPLIT_GROUPS = ( - "split_groups" in inspect.signature(make_swiglu_mapping).parameters -) - def get_params(): # (seq_len, embedding_dim, hidden_dim, seq_tile, embedding_tile, hidden_tile): @@ -43,20 +36,13 @@ def _run_and_verify(operator, seq_len, embedding_dim, golden_ref): operator.compile() run = operator.get_callable() - # Populate inputs by name; the design consumes weights in their natural - # (K, N) layout (no transpose). - run.get_buffer("input").torch_view()[:] = ( - golden_ref["input"].to(torch.bfloat16).flatten() - ) - run.get_buffer("weights_1").torch_view()[:] = ( - golden_ref["w_gate"].to(torch.bfloat16).flatten() - ) - run.get_buffer("weights_2").torch_view()[:] = ( - golden_ref["w_up"].to(torch.bfloat16).flatten() - ) - run.get_buffer("weights_3").torch_view()[:] = ( - golden_ref["w_down"].to(torch.bfloat16).flatten() - ) + # Populate inputs by name; buffers are the reference module's parameters (see + # reference.WEIGHTS) and the design consumes weights in their natural (K, N) + # layout, no transpose. + for reference, buffer in (("input", "input"), *WEIGHTS.items()): + run.get_buffer(buffer).torch_view()[:] = ( + golden_ref[reference].to(torch.bfloat16).flatten() + ) # Correctness is checked on the FIRST run: creating the hw_context applies the # design's init CDO, so the first run on the freshly loaded full-ELF computes @@ -120,10 +106,6 @@ def test_swiglu_prefill_stream( _run_and_verify(operator, seq_len, embedding_dim, golden_ref) -@pytest.mark.skipif( - not _STREAM_HAS_SPLIT_GROUPS, - reason="installed stream-dse lacks two-fusion-group support (split_groups)", -) @pytest.mark.supported_devices("npu2") @pytest.mark.metrics( Latency=r"Latency \(us\): (?P[\d\.]+)", diff --git a/iron/tests/stream/kernel_layouts.py b/iron/tests/stream/kernel_layouts.py new file mode 100644 index 00000000..a8f9444a --- /dev/null +++ b/iron/tests/stream/kernel_layouts.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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("m,k,n", [(32, 32, 64), (32, 64, 32)]) +def test_gemm_layouts_match_stream(m, k, n): + _assert_same(gemm_layouts(m, k, n), AIEKernels[GEMM.key](61.8, m, k, n, "default")) + + +def test_silu_layouts_match_stream(): + _assert_same(elementwise_layouts(2), AIEKernels[SILU.key](50.0, "default")) + + +def test_eltwise_mul_layouts_match_stream(): + _assert_same(elementwise_layouts(3), AIEKernels[ELTWISE_MUL.key](50.0, "default")) + + +@pytest.mark.parametrize("kernel", [GEMM, SILU, ELTWISE_MUL]) +def test_kernel_keys_exist_in_stream(kernel): + assert kernel.key in AIEKernels diff --git a/requirements_stream.txt b/requirements_stream.txt index d21b2842..c2e8c816 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# Optional dependency for the stream-dse-backed fused SwiGLU-prefill operator +# Optional dependencies for the stream-dse-backed fused SwiGLU-prefill operator # (iron/operators/swiglu_prefill_stream). # # It is NOT installed by the default CI (requirements.txt); the operator's test @@ -32,4 +32,5 @@ # mlir_aie/llvm-aie, which previously downgraded IRON's pinned mlir_aie (1.3.5 -> # 0.0.1) and numpy (2.x -> 1.26.4) and broke the whole operator test suite. +onnxscript>=0.7 stream-dse>=1.13.7 From bf0482a4998172a35c05cd0a1ab5abc61610989f Mon Sep 17 00:00:00 2001 From: asyms Date: Sat, 25 Jul 2026 18:40:45 +0200 Subject: [PATCH 07/21] iron/common/stream: make the exporter translation optional per op An operator whose default lowering already produces what stream-dse parses only needs its ONNX operator bound to a kernel. --- iron/common/stream/ops.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index 475fadff..5695c531 100644 --- a/iron/common/stream/ops.py +++ b/iron/common/stream/ops.py @@ -157,16 +157,22 @@ def _to_mul(a, b): @dataclass(frozen=True) class StreamOp: - """How one torch operator is exported, and which kernel runs it.""" + """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 + translation: Callable | None = None -# torch operator -> its ONNX form and AIE kernel. Gemm rather than MatMul because -# stream-dse's Gemm parser iterates (m, k, n), which is the order the mappings -# address as D0/D1/D2. +# 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), @@ -178,7 +184,11 @@ class StreamOp: 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()} + 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: From 845500ff59fe9cc6e25436b457e140406d86c517 Mon Sep 17 00:00:00 2001 From: asyms Date: Sun, 26 Jul 2026 04:44:47 +0200 Subject: [PATCH 08/21] swiglu_prefill_stream: name the design from the golden reference The tensor handed between fusion groups was called "h" in the operator's runlist and "intermediate" in the exported graph, and the sequence restated the runtime buffer names that the workload already carries. reference.py now holds the block's vocabulary, using the names the shared golden reference gives the same tensors, and iron/tests/stream/names.py pins that correspondence. mapping.group_boundaries derives each fused group's arguments from the exported graph, so the runlist and the operator's own arguments follow from the workload instead of being written out; the split design hands on "intermediate" by that derivation. The weights keep weights_1/2/3 rather than the golden w_gate/w_up/w_down: an otherwise identical model solves in ~115 s under these names and does not finish within 600 s under those, so the mapping in reference.WEIGHTS stays. --- iron/common/stream/mapping.py | 39 ++++++++ iron/operators/swiglu_prefill_stream/op.py | 68 ++++++++++--- .../swiglu_prefill_stream/reference.py | 34 ++++++- .../swiglu_prefill_stream/stream_design.py | 99 ++++++++++++------- iron/operators/swiglu_prefill_stream/test.py | 17 ++-- iron/tests/stream/names.py | 77 +++++++++++++++ 6 files changed, 271 insertions(+), 63 deletions(-) create mode 100644 iron/tests/stream/names.py diff --git a/iron/common/stream/mapping.py b/iron/common/stream/mapping.py index dd2c3e70..5042fd3b 100644 --- a/iron/common/stream/mapping.py +++ b/iron/common/stream/mapping.py @@ -52,6 +52,45 @@ class FusedGroup: 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) -> dict: entry = { "name": name, diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index d31efd6f..70f1d74d 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -120,12 +120,52 @@ def _name(kind, m, e, h, st, et, ht): return f"{kind}_m{m}_e{e}_h{h}_st{st}_et{et}_ht{ht}" +def _boundaries(seq_len, embedding_dim, hidden_dim, split_groups=False): + """Each fused group's (inputs, outputs), named as the exported graph names them. + + Imported lazily: the names come from the exported workload, so only building + the operator needs stream-dse, not importing it. + """ + from iron.operators.swiglu_prefill_stream.stream_design import group_ports + + return group_ports(seq_len, embedding_dim, hidden_dim, split_groups) + + +def _ports(seq_len, embedding_dim, hidden_dim, split_groups=False): + """Each fused group's arguments in the order its design takes them.""" + return [ + inputs + outputs + for inputs, outputs in _boundaries( + seq_len, embedding_dim, hidden_dim, split_groups + ) + ] + + +def _external(seq_len, embedding_dim, hidden_dim): + """The operator's own arguments: what no group produces, and what none consumes.""" + boundaries = _boundaries(seq_len, embedding_dim, hidden_dim, split_groups=True) + produced = {name for _, outputs in boundaries for name in outputs} + consumed = {name for inputs, _ in boundaries for name in inputs} + inputs = tuple( + dict.fromkeys( + name for group, _ in boundaries for name in group if name not in produced + ) + ) + outputs = tuple( + dict.fromkeys( + name for _, group in boundaries for name in group if name not in consumed + ) + ) + return inputs, outputs + + class SwiGLUPrefillStream(OperatorSequence): """Fused SwiGLU-prefill block generated by stream-dse and deployed as a single full-ELF (``OperatorSequence`` default dispatch). - Runtime buffers (``get_callable().get_buffer(name)``): ``input``, ``weights_1`` - (gate), ``weights_2`` (up), ``weights_3`` (down), ``output``. Building requires + Runtime buffers (``get_callable().get_buffer(name)``) are named by the + reference module: ``input``, ``weights_1`` (gate), ``weights_2`` (up), + ``weights_3`` (down), ``output``. Building requires ``stream-dse`` (``pip install stream-dse`` + ``stream-setup-aie``); importing this module does not. """ @@ -159,9 +199,9 @@ def __init__( embedding_tile_size, hidden_tile_size, ), - runlist=[(block, "input", "weights_1", "weights_2", "weights_3", "output")], - input_args=["input", "weights_1", "weights_2", "weights_3"], - output_args=["output"], + runlist=[(block, *_ports(seq_len, embedding_dim, hidden_dim)[0])], + input_args=list(_external(seq_len, embedding_dim, hidden_dim)[0]), + output_args=list(_external(seq_len, embedding_dim, hidden_dim)[1]), extra_flags=["--dynamic-objFifos"], context=context, ) @@ -169,10 +209,9 @@ def __init__( class SwiGLUPrefillStreamK2(OperatorSequence): """Two-fusion-group SwiGLU-prefill: a gate/up/SiLU/mul group and a separate - down-projection group fused into one full-ELF, with the hidden state ``h`` - kept on device between them. Same external buffers as - :class:`SwiGLUPrefillStream`; the split is decided by the stream mapping - (``make_swiglu_mapping(split_groups=True)``). + down-projection group fused into one full-ELF, with the hidden state kept on + device between them. Same external buffers as :class:`SwiGLUPrefillStream`; + the split is decided by the mapping's fused groups. """ def __init__( @@ -207,11 +246,14 @@ def __init__( hidden_tile_size, ), runlist=[ - (front, "input", "weights_1", "weights_2", "h"), - (down, "h", "weights_3", "output"), + (group, *ports) + for group, ports in zip( + (front, down), + _ports(seq_len, embedding_dim, hidden_dim, split_groups=True), + ) ], - input_args=["input", "weights_1", "weights_2", "weights_3"], - output_args=["output"], + input_args=list(_external(seq_len, embedding_dim, hidden_dim)[0]), + output_args=list(_external(seq_len, embedding_dim, hidden_dim)[1]), extra_flags=["--dynamic-objFifos"], context=context, ) diff --git a/iron/operators/swiglu_prefill_stream/reference.py b/iron/operators/swiglu_prefill_stream/reference.py index 78b7718f..0b8f5260 100644 --- a/iron/operators/swiglu_prefill_stream/reference.py +++ b/iron/operators/swiglu_prefill_stream/reference.py @@ -4,19 +4,43 @@ """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. Its parameter names are the ONNX -tensor names, hence the operator's runtime buffer names. +workload stream-dse generates the design from. -Numerical values come from :mod:`iron.operators.swiglu_decode.reference`, which -this operator shares; :func:`swiglu_module` loads them into the module. +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 -# Reference weight -> this module's parameter (and so the runtime buffer name). +INPUT = "input" +OUTPUT = "output" +GATE_PROJECTION = "left" +UP_PROJECTION = "right" +ACTIVATION = "left_swished" +HIDDEN = "intermediate" + +# Golden-reference weight -> this module's parameter, and so the ONNX tensor and +# runtime buffer name. The weights are the one place the two vocabularies differ: +# stream-dse's allocation solve is markedly slower under the golden names (>600 s +# versus ~115 s for an otherwise identical model), so the parameters keep the +# names the operator was measured with. WEIGHTS = {"w_gate": "weights_1", "w_up": "weights_2", "w_down": "weights_3"} +# 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.""" diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index 33b02e9c..2c77a2a9 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -20,14 +20,21 @@ 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.mapping import FusedGroup, Placement, emit_mapping +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 @@ -47,10 +54,18 @@ GATE, UP, SILU, MUL, DOWN = "Gemm_Left", "Gemm_Right", "Silu", "Elt_Mul", "Gemm_Down" NODE_NAMES = [GATE, UP, SILU, MUL, DOWN] RESULT_NAMES = { - GATE: "left", - UP: "right", - SILU: "left_swished", - MUL: "intermediate", + GATE: reference.GATE_PROJECTION, + UP: reference.UP_PROJECTION, + SILU: reference.ACTIVATION, + MUL: reference.HIDDEN, +} + +# Which layers each fused group contains. Splitting makes stream-dse emit one +# design per group; the tensor handed from one to the next is derived from the +# exported graph, not named here. +GROUP_LAYERS = { + False: [[GATE, UP, SILU, MUL, DOWN]], + True: [[GATE, UP, SILU, MUL], [DOWN]], } # Compute-tile placement on the 4-row x 8-column array: each GEMM gets two @@ -106,36 +121,31 @@ def _groups(seq_tile, embedding_tile, hidden_tile, split_groups): the output dimension of the group's leading GEMM. """ if split_groups: - return [ - FusedGroup( - "Fused_Group_1", - [GATE, UP, SILU, MUL], - [ - (GATE, "D1", embedding_tile), - (GATE, "D2", hidden_tile), - (GATE, "D0", seq_tile), - ], - ), - FusedGroup( - "Fused_Group_2", - [DOWN], - [ - (DOWN, "D1", hidden_tile), - (DOWN, "D2", embedding_tile), - (DOWN, "D0", seq_tile), - ], - ), + tiling = [ + [ + (GATE, "D1", embedding_tile), + (GATE, "D2", hidden_tile), + (GATE, "D0", seq_tile), + ], + [ + (DOWN, "D1", hidden_tile), + (DOWN, "D2", embedding_tile), + (DOWN, "D0", seq_tile), + ], ] - return [ - FusedGroup( - "Fused_Group_1", - [GATE, UP, SILU, MUL, DOWN], + else: + tiling = [ [ (GATE, "D1", embedding_tile), (DOWN, "D2", embedding_tile), (GATE, "D2", hidden_tile), (GATE, "D0", seq_tile), - ], + ] + ] + return [ + FusedGroup(f"Fused_Group_{index + 1}", layers, group_tiling) + for index, (layers, group_tiling) in enumerate( + zip(GROUP_LAYERS[split_groups], tiling) ) ] @@ -161,6 +171,29 @@ def _check_shapes( ) +@lru_cache(maxsize=None) +def workload_for(seq_len, embedding_dim, hidden_dim): + """The exported workload for one problem size.""" + module = swiglu_module(embedding_dim, hidden_dim) + return export_workload( + module, + (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, split_groups=False): + """Per fused group, the tensor names it takes in and hands on. + + These are the operator's runtime arguments, including the tensor a split + design passes from its front end to its down projection. + """ + return group_boundaries( + workload_for(seq_len, embedding_dim, hidden_dim), GROUP_LAYERS[split_groups] + ) + + def build_inputs( seq_len, embedding_dim, @@ -180,13 +213,7 @@ def build_inputs( embedding_tile_size, hidden_tile_size, ) - module = swiglu_module(embedding_dim, hidden_dim) - workload = export_workload( - module, - (torch.zeros(seq_len, embedding_dim, dtype=torch.bfloat16),), - node_names=NODE_NAMES, - result_names=RESULT_NAMES, - ) + workload = workload_for(seq_len, embedding_dim, hidden_dim) output_dir = Path(output_dir) return ( workload.write(output_dir / "workload.onnx"), diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index 778d6ced..12104726 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -22,7 +22,7 @@ # 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 WEIGHTS +from iron.operators.swiglu_prefill_stream.reference import INPUT, OUTPUT, WEIGHTS from iron.common.test_utils import verify_buffer @@ -36,12 +36,11 @@ def _run_and_verify(operator, seq_len, embedding_dim, golden_ref): operator.compile() run = operator.get_callable() - # Populate inputs by name; buffers are the reference module's parameters (see - # reference.WEIGHTS) and the design consumes weights in their natural (K, N) - # layout, no transpose. - for reference, buffer in (("input", "input"), *WEIGHTS.items()): + # Populate inputs by golden-reference name; the design consumes weights in + # their natural (K, N) layout, no transpose. + for name, buffer in ((INPUT, INPUT), *WEIGHTS.items()): run.get_buffer(buffer).torch_view()[:] = ( - golden_ref[reference].to(torch.bfloat16).flatten() + golden_ref[name].to(torch.bfloat16).flatten() ) # Correctness is checked on the FIRST run: creating the hw_context applies the @@ -55,11 +54,11 @@ def _run_and_verify(operator, seq_len, embedding_dim, golden_ref): # amplifies relative error): ~20% of elements drift past the 8% bound, so allow # up to 25%. Tolerances are local to this test. run() - output = run.get_buffer("output").to_torch().reshape((seq_len, embedding_dim)) + output = run.get_buffer(OUTPUT).to_torch().reshape((seq_len, embedding_dim)) errors = verify_buffer( output, - "output", - golden_ref["output"], + OUTPUT, + golden_ref[OUTPUT], rel_tol=0.08, abs_tol=0.7, max_error_rate=0.25, diff --git a/iron/tests/stream/names.py b/iron/tests/stream/names.py new file mode 100644 index 00000000..ed5367b0 --- /dev/null +++ b/iron/tests/stream/names.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. 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.values()) + + +def test_golden_weights_load_into_the_module(): + golden = generate_golden_reference(**SHAPE) + module = reference.swiglu_module(SHAPE["K"], SHAPE["N"], golden) + for name, parameter in reference.WEIGHTS.items(): + assert getattr(module, parameter).equal(golden[name]) + + +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("split_groups", [False, True]) +def test_group_ports_are_named_by_the_exported_graph(split_groups): + workload = stream_design.workload_for(*DIMS) + known = set(workload.buffers) | set(reference.TENSOR_NAMES) + for inputs, outputs in stream_design.group_ports(*DIMS, split_groups): + assert set(inputs) | set(outputs) <= known + + +def test_split_design_hands_on_the_hidden_state(): + (_, front_outputs), (down_inputs, _) = stream_design.group_ports( + *DIMS, split_groups=True + ) + 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, split_groups=True) + 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) From 06c09ffd39530adf170e42285d3d2b735dabafe3 Mon Sep 17 00:00:00 2001 From: asyms Date: Mon, 27 Jul 2026 03:36:58 +0200 Subject: [PATCH 09/21] iron/common/stream: place layers by array column, not by core id Placements spelled out stream's integer core ids, which say nothing in IRON about where a layer runs and silently assumed the 8-column Strix array. ComputeArray reads the accelerator description -- which already gives every compute tile a [column, row] coordinate and a type -- and is the only place that knows what an id means. A placement now names columns: allocate() hands out disjoint consecutive ranges from a per-layer budget, and all_columns gives a layer the whole array, which is what a layer-by-layer design or a single layer sent to stream for an estimate asks for. The SwiGLU placement becomes a budget of [2, 2, 1, 1, 2] columns and the row split comes from the array. The emitted mapping is byte-identical. (cherry picked from commit b589d7c597dbf85ea418e42a6e67a842becf0422) --- iron/common/stream/hardware.py | 103 ++++++++++++++++++ iron/common/stream/mapping.py | 28 +++-- .../swiglu_prefill_stream/stream_design.py | 36 +++--- iron/tests/stream/placement.py | 85 +++++++++++++++ 4 files changed, 221 insertions(+), 31 deletions(-) create mode 100644 iron/common/stream/hardware.py create mode 100644 iron/tests/stream/placement.py diff --git a/iron/common/stream/hardware.py b/iron/common/stream/hardware.py new file mode 100644 index 00000000..d36b3f3f --- /dev/null +++ b/iron/common/stream/hardware.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Address an accelerator'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: +which ids are compute tiles, and which column they sit in, is in the accelerator +description. This turns that description into columns of core ids, so a placement +is written in columns and no other IRON module has to know what an 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 functools import lru_cache +from pathlib import Path +from typing import Iterable, Sequence + +_COMPUTE_TYPE_SUFFIX = "compute" + + +@lru_cache(maxsize=None) +def _core_type(path: str) -> str: + import yaml + + return yaml.safe_load(Path(path).read_text())["type"] + + +def _core_description(accelerator: Path, reference: str) -> Path: + """Resolve a core reference, which is relative to the accelerator or to ``cores/``.""" + path = accelerator.parent / reference + return ( + path if path.exists() else accelerator.parent / "cores" / Path(reference).name + ) + + +@dataclass(frozen=True) +class ComputeArray: + """The accelerator's compute tiles, as core ids grouped by column.""" + + columns: tuple[tuple[int, ...], ...] + + @classmethod + def from_accelerator(cls, path) -> ComputeArray: + import yaml + + path = Path(path) + description = yaml.safe_load(path.read_text()) + coordinates = description.get("core_coordinates", {}) + by_column: dict[int, list[tuple[int, int]]] = {} + for core_id, reference in description["cores"].items(): + if core_id not in coordinates: + continue + if not _core_type(str(_core_description(path, reference))).endswith( + _COMPUTE_TYPE_SUFFIX + ): + continue + column, row = coordinates[core_id] + by_column.setdefault(column, []).append((row, core_id)) + return cls( + tuple( + tuple(core_id for _, core_id in sorted(rows)) + for _, rows in sorted(by_column.items()) + ) + ) + + @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]) -> tuple[int, ...]: + """The core ids of ``columns``, column by column and row by row.""" + return tuple(core for column in columns for core in self.columns[column]) + + 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 index 5042fd3b..10e9cdec 100644 --- a/iron/common/stream/mapping.py +++ b/iron/common/stream/mapping.py @@ -22,6 +22,7 @@ from pathlib import Path from typing import Sequence +from iron.common.stream.hardware import ComputeArray from iron.common.stream.workload import StreamWorkload @@ -29,13 +30,14 @@ class Placement: """Where one workload node runs. - ``cores`` lists the compute tiles per allocation, ``splits`` the matching - inter-core tiling as ``(dim, split)`` pairs, and ``kernel_kwargs`` the + ``columns`` are the array columns it occupies, resolved to core ids against + the :class:`~iron.common.stream.hardware.ComputeArray`; ``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). """ - cores: Sequence[Sequence[int]] - splits: Sequence[Sequence[tuple[str, int]]] = () + columns: Sequence[int] + splits: Sequence[tuple[str, int]] = () kernel_kwargs: dict = field(default_factory=dict) @@ -91,23 +93,24 @@ def group_boundaries( return boundaries -def _layer_entry(name: str, placement: Placement, kernel_key: str) -> dict: - entry = { +def _layer_entry( + name: str, placement: Placement, kernel_key: str, array: ComputeArray +) -> dict: + return { "name": name, - "core_allocation": [list(cores) for cores in placement.cores], + "core_allocation": [list(array.cores(placement.columns))], "inter_core_tiling": [ - [{"dim": dim, "split": split} for dim, split in group] - for group in placement.splits + [{"dim": dim, "split": split} for dim, split in placement.splits] ], "kernel": {"name": kernel_key, "kwargs": dict(placement.kernel_kwargs)}, } - return entry 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) @@ -123,7 +126,7 @@ def build_mapping( raise ValueError(f"fused groups refer to nodes without a placement: {missing}") layers = [ - _layer_entry(name, placements[name], kernel) + _layer_entry(name, placements[name], kernel, array) for name, kernel in workload.nodes if name in placements ] @@ -148,6 +151,7 @@ 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.""" @@ -157,7 +161,7 @@ def emit_mapping( path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w") as f: yaml.dump( - build_mapping(workload, placements, groups), + build_mapping(workload, placements, groups, array), f, default_flow_style=False, sort_keys=False, diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index 2c77a2a9..e36f95ef 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -27,6 +27,7 @@ import torch from stream.api import optimize_allocation_co +from iron.common.stream.hardware import ComputeArray from iron.common.stream.mapping import ( FusedGroup, Placement, @@ -68,19 +69,15 @@ True: [[GATE, UP, SILU, MUL], [DOWN]], } -# Compute-tile placement on the 4-row x 8-column array: each GEMM gets two -# columns' worth of cores split 4-ways over rows (D0, the sequence dimension) and -# 2-ways over D2, the two elementwise layers one column each. This is a property -# of the array, not of the problem size, so it holds across shapes. -_GEMM_SPLIT = [[("D0", 4), ("D2", 2)]] -_ELEMENTWISE_SPLIT = [[("D0", 4)]] -_CORES = { - GATE: [[2, 3, 4, 5, 8, 9, 10, 11]], - UP: [[14, 15, 16, 17, 20, 21, 22, 23]], - SILU: [[26, 27, 28, 29]], - MUL: [[32, 33, 34, 35]], - DOWN: [[38, 39, 40, 41, 44, 45, 46, 47]], -} +ARRAY = ComputeArray.from_accelerator(_ACCELERATOR) + +# Columns per layer: two for each GEMM, one for each elementwise layer. The five +# layers sit on disjoint columns so they pipeline across steady-state iterations. +# Each layer splits over the array's rows (D0, the sequence dimension), and a GEMM +# splits over its two columns as well (D2, the output dimension). +_COLUMNS = dict(zip(NODE_NAMES, ARRAY.allocate([2, 2, 1, 1, 2]))) +_GEMM_SPLIT = (("D0", ARRAY.num_rows), ("D2", 2)) +_ELEMENTWISE_SPLIT = (("D0", ARRAY.num_rows),) def gemm_tiles(seq_tile, embedding_tile, hidden_tile): @@ -105,11 +102,11 @@ def gemm(tiles): elementwise = {"utilization": 50.0, "layout": "default"} return { - GATE: Placement(_CORES[GATE], _GEMM_SPLIT, gemm(gate_up)), - UP: Placement(_CORES[UP], _GEMM_SPLIT, gemm(gate_up)), - SILU: Placement(_CORES[SILU], _ELEMENTWISE_SPLIT, elementwise), - MUL: Placement(_CORES[MUL], _ELEMENTWISE_SPLIT, elementwise), - DOWN: Placement(_CORES[DOWN], _GEMM_SPLIT, gemm(down)), + GATE: Placement(_COLUMNS[GATE], _GEMM_SPLIT, gemm(gate_up)), + UP: Placement(_COLUMNS[UP], _GEMM_SPLIT, gemm(gate_up)), + SILU: Placement(_COLUMNS[SILU], _ELEMENTWISE_SPLIT, elementwise), + MUL: Placement(_COLUMNS[MUL], _ELEMENTWISE_SPLIT, elementwise), + DOWN: Placement(_COLUMNS[DOWN], _GEMM_SPLIT, gemm(down)), } @@ -154,7 +151,7 @@ def _check_shapes( seq_len, embedding_dim, hidden_dim, seq_tile, embedding_tile, hidden_tile ): """Reject shapes the fixed 4x8 placement cannot tile evenly.""" - rows, gemm_split = 4, 2 + rows, gemm_split = ARRAY.num_rows, 2 if seq_len % rows or seq_len < seq_tile * rows: raise ValueError( f"seq_len ({seq_len}) must be a multiple of {rows} and at least {seq_tile * rows}" @@ -223,6 +220,7 @@ def build_inputs( _groups( seq_len_tile_size, embedding_tile_size, hidden_tile_size, split_groups ), + ARRAY, output_dir / "mapping.yaml", ), ) diff --git a/iron/tests/stream/placement.py b/iron/tests/stream/placement.py new file mode 100644 index 00000000..f051617b --- /dev/null +++ b/iron/tests/stream/placement.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Placements are written in array columns and resolved to stream's core ids. + +An accelerator description gives every compute tile a ``[column, row]`` +coordinate; :class:`~iron.common.stream.hardware.ComputeArray` is the only place +that reads it, so operators never spell out a core id. +""" + +import pytest + +pytest.importorskip( + "stream", reason="stream-dse not installed (see requirements_stream.txt)" +) + +from iron.common.stream.hardware import ComputeArray # noqa: E402 +from iron.operators.swiglu_prefill_stream import stream_design # noqa: E402 + +ARRAY = stream_design.ARRAY + +# The core ids this operator was measured with, 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_accelerator(): + 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_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_emitted_allocation_is_the_measured_one(tmp_path): + import yaml + + _, mapping_path = stream_design.build_inputs( + 256, 512, 2048, 32, 32, 64, tmp_path / "design" + ) + emitted = { + layer["name"]: layer["core_allocation"][0] + for layer in yaml.safe_load(open(mapping_path))["layers"] + } + assert emitted == MEASURED_ALLOCATION + + +@pytest.mark.parametrize("accelerator", ["whole_array.yaml", "single_col.yaml"]) +def test_other_accelerators_resolve(accelerator): + import os + + import stream + + path = os.path.join( + os.path.dirname(stream.__file__), "inputs", "aie", "hardware", accelerator + ) + array = ComputeArray.from_accelerator(path) + assert array.num_rows == 4 + assert array.cores(array.all_columns) From ff9c3b12c68f97c1022c3f89e91051cde806769c Mon Sep 17 00:00:00 2001 From: asyms Date: Mon, 27 Jul 2026 19:27:45 +0200 Subject: [PATCH 10/21] swiglu_prefill_stream: name the weights from the golden reference stream-dse 1.13.8 derives the generated design's runtime argument order from the exported graph instead of from tensor names, so the weights no longer need the weights_1/2/3 workaround that kept the allocation solve fast. They now carry their golden-reference names, and the reference's name mapping collapses to the tuple of names it shares with swiglu_decode. --- iron/operators/swiglu_prefill_stream/op.py | 4 +-- .../swiglu_prefill_stream/reference.py | 25 ++++++++----------- iron/operators/swiglu_prefill_stream/test.py | 4 +-- iron/tests/stream/names.py | 6 ++--- requirements_stream.txt | 11 +++++--- 5 files changed, 24 insertions(+), 26 deletions(-) diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 70f1d74d..79d79dd1 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -164,8 +164,8 @@ class SwiGLUPrefillStream(OperatorSequence): single full-ELF (``OperatorSequence`` default dispatch). Runtime buffers (``get_callable().get_buffer(name)``) are named by the - reference module: ``input``, ``weights_1`` (gate), ``weights_2`` (up), - ``weights_3`` (down), ``output``. Building requires + 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. """ diff --git a/iron/operators/swiglu_prefill_stream/reference.py b/iron/operators/swiglu_prefill_stream/reference.py index 0b8f5260..9056ae59 100644 --- a/iron/operators/swiglu_prefill_stream/reference.py +++ b/iron/operators/swiglu_prefill_stream/reference.py @@ -23,12 +23,7 @@ ACTIVATION = "left_swished" HIDDEN = "intermediate" -# Golden-reference weight -> this module's parameter, and so the ONNX tensor and -# runtime buffer name. The weights are the one place the two vocabularies differ: -# stream-dse's allocation solve is markedly slower under the golden names (>600 s -# versus ~115 s for an otherwise identical model), so the parameters keep the -# names the operator was measured with. -WEIGHTS = {"w_gate": "weights_1", "w_up": "weights_2", "w_down": "weights_3"} +WEIGHTS = ("w_gate", "w_up", "w_down") # Every name this module exports, for the correspondence check in iron/tests/stream. TENSOR_NAMES = ( @@ -48,16 +43,16 @@ class SwiGLU(nn.Module): def __init__(self, embedding_dim: int, hidden_dim: int, dtype=torch.bfloat16): super().__init__() gate_up = (embedding_dim, hidden_dim) - self.weights_1 = nn.Parameter(torch.zeros(gate_up, dtype=dtype)) # gate - self.weights_2 = nn.Parameter(torch.zeros(gate_up, dtype=dtype)) # up - self.weights_3 = nn.Parameter( + 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) - ) # down + ) def forward(self, input): - gate = input @ self.weights_1 - up = input @ self.weights_2 - return (torch.nn.functional.silu(gate) * up) @ self.weights_3 + 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: @@ -69,6 +64,6 @@ def swiglu_module(embedding_dim, hidden_dim, golden_reference=None) -> SwiGLU: module = SwiGLU(embedding_dim, hidden_dim).eval() if golden_reference is not None: with torch.no_grad(): - for reference, parameter in WEIGHTS.items(): - getattr(module, parameter).copy_(golden_reference[reference]) + for name in WEIGHTS: + getattr(module, name).copy_(golden_reference[name]) return module diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index 12104726..e694ba82 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -38,8 +38,8 @@ def _run_and_verify(operator, seq_len, embedding_dim, golden_ref): # Populate inputs by golden-reference name; the design consumes weights in # their natural (K, N) layout, no transpose. - for name, buffer in ((INPUT, INPUT), *WEIGHTS.items()): - run.get_buffer(buffer).torch_view()[:] = ( + for name in (INPUT, *WEIGHTS): + run.get_buffer(name).torch_view()[:] = ( golden_ref[name].to(torch.bfloat16).flatten() ) diff --git a/iron/tests/stream/names.py b/iron/tests/stream/names.py index ed5367b0..f2358ee5 100644 --- a/iron/tests/stream/names.py +++ b/iron/tests/stream/names.py @@ -31,14 +31,14 @@ def test_name_is_a_golden_reference_key(name, 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.values()) + 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, parameter in reference.WEIGHTS.items(): - assert getattr(module, parameter).equal(golden[name]) + for name in reference.WEIGHTS: + assert getattr(module, name).equal(golden[name]) stream = pytest.importorskip( diff --git a/requirements_stream.txt b/requirements_stream.txt index c2e8c816..be597fbb 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -22,15 +22,18 @@ # OR-Tools GSCIP solver; no Gurobi needed) and writes its generated workload/ # mapping files into its own installed package directory, so that environment # must be writable. -# - >=1.13.7 is required: stream_design.py feeds IRON-authored operand layouts +# - >=1.13.8 is required: stream_design.py feeds IRON-authored operand layouts # into code generation via optimize_allocation_co(kernels=...), the override # hook added in stream-dse 1.13.4; the k=2 variant (op_k2.py) additionally needs # the two-fusion-group support (make_swiglu_mapping(split_groups=...) + the # multi-group AIE codegen pipeline) added in stream-dse 1.13.5; 1.13.6 makes # the debug workload-graph visualization non-fatal, so code generation no longer -# requires graphviz (`dot`); and 1.13.7 makes `stream-setup-aie` stop reinstalling +# requires graphviz (`dot`); 1.13.7 makes `stream-setup-aie` stop reinstalling # mlir_aie/llvm-aie, which previously downgraded IRON's pinned mlir_aie (1.3.5 -> -# 0.0.1) and numpy (2.x -> 1.26.4) and broke the whole operator test suite. +# 0.0.1) and numpy (2.x -> 1.26.4) and broke the whole operator test suite; and +# 1.13.8 orders the generated design's runtime arguments by the exported graph +# instead of by tensor name, which is what lets the weights below carry their +# reference names. onnxscript>=0.7 -stream-dse>=1.13.7 +stream-dse>=1.13.8 From 1aacdc86971b03c4b8ee1a69ce329b37529f0349 Mon Sep 17 00:00:00 2001 From: asyms Date: Tue, 28 Jul 2026 19:44:40 +0200 Subject: [PATCH 11/21] swiglu_prefill_stream: deploy the block at any fusion granularity One SwiGLUPrefillStream now takes k, the number of fused groups the block is split into: 1 keeps the whole block on the array, 2 splits after the elementwise multiply, and 5 runs layer by layer, each layer taking the whole array in turn as swiglu_prefill does. This replaces the three near identical operator classes and the two design entry points behind them. ComputeArray reads the grid from mlir-aie's Device instead of parsing stream-dse's accelerator description, and a placement can now take some rows of a column, which is what an elementwise layer with one worker per column needs. The kernel tiles are constants of the design rather than arguments threaded through the operator, the test and the generator. The mapping carries them and no absolute dimension, so it holds across problem sizes within the limits _check_shapes enforces. Every dispatch takes a fresh callable: the full-ELF path does not reset array state between dispatches, so reusing one reads stale state, and a design of several fused groups deadlocks on the second dispatch. Requires stream-dse 1.13.8 or newer. --- iron/common/stream/__init__.py | 2 +- iron/common/stream/hardware.py | 81 ++-- iron/common/stream/mapping.py | 8 +- iron/common/stream/ops.py | 2 +- iron/common/stream/workload.py | 19 +- iron/operators/__init__.py | 2 +- .../operators/swiglu_prefill_stream/README.md | 27 +- iron/operators/swiglu_prefill_stream/op.py | 282 ++++-------- .../swiglu_prefill_stream/reference.py | 2 +- .../swiglu_prefill_stream/stream_design.py | 412 ++++++++---------- iron/operators/swiglu_prefill_stream/test.py | 151 +++---- iron/tests/stream/kernel_layouts.py | 2 +- iron/tests/stream/names.py | 22 +- iron/tests/stream/placement.py | 89 +++- requirements_stream.txt | 37 +- 15 files changed, 492 insertions(+), 646 deletions(-) diff --git a/iron/common/stream/__init__.py b/iron/common/stream/__init__.py index 0a55ae7c..1e35c3fa 100644 --- a/iron/common/stream/__init__.py +++ b/iron/common/stream/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Building blocks for stream-dse-backed operators. diff --git a/iron/common/stream/hardware.py b/iron/common/stream/hardware.py index d36b3f3f..532e0fef 100644 --- a/iron/common/stream/hardware.py +++ b/iron/common/stream/hardware.py @@ -1,12 +1,13 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Address an accelerator's compute tiles the way the array is laid out. +"""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: -which ids are compute tiles, and which column they sit in, is in the accelerator -description. This turns that description into columns of core ids, so a placement -is written in columns and no other IRON module has to know what an id means. +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 @@ -17,55 +18,33 @@ from __future__ import annotations from dataclasses import dataclass -from functools import lru_cache -from pathlib import Path from typing import Iterable, Sequence -_COMPUTE_TYPE_SUFFIX = "compute" - - -@lru_cache(maxsize=None) -def _core_type(path: str) -> str: - import yaml - - return yaml.safe_load(Path(path).read_text())["type"] - - -def _core_description(accelerator: Path, reference: str) -> Path: - """Resolve a core reference, which is relative to the accelerator or to ``cores/``.""" - path = accelerator.parent / reference - return ( - path if path.exists() else accelerator.parent / "cores" / Path(reference).name - ) +from aie.iron.device.device import AIETileType @dataclass(frozen=True) class ComputeArray: - """The accelerator's compute tiles, as core ids grouped by column.""" + """The device's compute tiles, as stream core ids grouped by column.""" columns: tuple[tuple[int, ...], ...] @classmethod - def from_accelerator(cls, path) -> ComputeArray: - import yaml - - path = Path(path) - description = yaml.safe_load(path.read_text()) - coordinates = description.get("core_coordinates", {}) - by_column: dict[int, list[tuple[int, int]]] = {} - for core_id, reference in description["cores"].items(): - if core_id not in coordinates: - continue - if not _core_type(str(_core_description(path, reference))).endswith( - _COMPUTE_TYPE_SUFFIX - ): - continue - column, row = coordinates[core_id] - by_column.setdefault(column, []).append((row, core_id)) + 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(core_id for _, core_id in sorted(rows)) - for _, rows in sorted(by_column.items()) + 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) ) ) @@ -81,9 +60,19 @@ def num_rows(self) -> int: def all_columns(self) -> tuple[int, ...]: return tuple(range(self.num_columns)) - def cores(self, columns: Iterable[int]) -> tuple[int, ...]: - """The core ids of ``columns``, column by column and row by row.""" - return tuple(core for column in columns for core in self.columns[column]) + 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. diff --git a/iron/common/stream/mapping.py b/iron/common/stream/mapping.py index 10e9cdec..68b3f725 100644 --- a/iron/common/stream/mapping.py +++ b/iron/common/stream/mapping.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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. @@ -31,7 +31,8 @@ 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`; ``splits`` is the + 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). """ @@ -39,6 +40,7 @@ class Placement: columns: Sequence[int] splits: Sequence[tuple[str, int]] = () kernel_kwargs: dict = field(default_factory=dict) + rows: Sequence[int] | None = None @dataclass(frozen=True) @@ -98,7 +100,7 @@ def _layer_entry( ) -> dict: return { "name": name, - "core_allocation": [list(array.cores(placement.columns))], + "core_allocation": [list(array.cores(placement.columns, placement.rows))], "inter_core_tiling": [ [{"dim": dim, "split": split} for dim, split in placement.splits] ], diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index 5695c531..242fae92 100644 --- a/iron/common/stream/ops.py +++ b/iron/common/stream/ops.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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. diff --git a/iron/common/stream/workload.py b/iron/common/stream/workload.py index 488ad69d..d573dc84 100644 --- a/iron/common/stream/workload.py +++ b/iron/common/stream/workload.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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. @@ -39,6 +39,23 @@ class StreamWorkload: 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. diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 62841c7f..6d62e215 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -12,7 +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, SwiGLUPrefillStreamK2 +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 index 828e227a..cfa71445 100644 --- a/iron/operators/swiglu_prefill_stream/README.md +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -1,5 +1,5 @@ @@ -8,7 +8,7 @@ SPDX-License-Identifier: Apache-2.0 This operator is **fused**: the whole SwiGLU-prefill block (both GEMMs + SiLU + elementwise-mul) is emitted as a **single MLIR design generated by [`stream-dse`](https://github.com/KULeuven-MICAS/stream)**, then compiled by IRON's normal -flow into one xclbin. Unlike the other operators, its MLIR is not written by hand — it is +flow into one full ELF. Unlike the other operators, its MLIR is not written by hand — it is produced at build time by [`stream_design.py`](./stream_design.py), which calls the installed `stream` package (`stream.api.optimize_allocation_co(..., enable_codegen=True)`). @@ -32,8 +32,14 @@ Changing the problem size means re-exporting with different example shapes: the placement carries tile sizes and core sets but no absolute dimensions, so it is unaffected (within the divisibility limits `stream_design._check_shapes` enforces). +`k` selects how many fused groups the block is split into, and so how many designs +the ELF holds: `1` keeps the whole block on the array, `2` splits after the +elementwise multiply, and `5` runs layer by layer like `swiglu_prefill` does. The +groups themselves are `stream_design.GROUP_LAYERS`; the kernel tiles every design +is built for are `SEQ_TILE`, `EMBEDDING_TILE` and `HIDDEN_TILE` in the same module. + Supporting another op (say a softmax block) is one `StreamKernel` plus one -`ATEN_OPS` entry in `iron/common/stream/ops.py`, pointing at the existing +`TORCH_OPS` entry in `iron/common/stream/ops.py`, pointing at the existing `aie_kernels//softmax.cc`, plus that operator's own placement. ## Enabling stream codegen @@ -52,10 +58,9 @@ Notes: so **no Gurobi license** is required. - `stream-setup-aie` is **required**: it installs the pure-Python AIE codegen dialects stream-dse needs that cannot be plain PyPI dependencies (`xdsl-aie`, `snax-mlir`), since they - are direct git installs. As of **stream-dse 1.13.7** it does **not** install `mlir_aie` / - `llvm-aie`: IRON's `requirements.txt` already pins those, and stream-dse's codegen only emits - text MLIR via xdsl (it never imports the `aie` bindings). Earlier releases reinstalled an - older `mlir_aie` here, which downgraded IRON's pinned wheel and broke the test suite. + are direct git installs. It does not install `mlir_aie` / `llvm-aie`: IRON's + `requirements.txt` already pins those, and stream-dse's codegen only emits text MLIR via + xdsl (it never imports the `aie` bindings). - Importing the operator does **not** require `stream-dse` (the launcher is imported lazily); only **building** (`operator.compile()` / running the test) does. @@ -73,13 +78,9 @@ The feasible/verified shape is **seq 256 / embedding 512 / hidden 2048**, tiles ## Caveats (stream-dse packaging) - The hardware-description YAML (`whole_array_strix.yaml` + `hardware/cores/*.yaml`) is - resolved from the **installed `stream` package**, where it ships as package data - (stream-dse >= 1.13.3); nothing is vendored in this operator. + resolved from the **installed `stream` package**, where it ships as package data; + nothing is vendored in this operator. - Nodes and their results are renamed from the ATen names `torch.export` produces (`matmul`, `matmul_1`, …). Those are prefixes of one another, which stream-dse's tensor bookkeeping does not distinguish reliably (it raises a `KeyError` while building the memory-capacity constraints). -- Naming does not change the problem posed to the solver — same nodes, tensors, - shapes and connections — but it does change SCIP's presolve, and with it how - long the allocation takes to solve. The names here are the ones this operator - was measured with. diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 79d79dd1..9e23b194 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -1,8 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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, Optional +from typing import Any import aie.utils as aie_utils @@ -21,239 +21,139 @@ class _SwiGLUStreamGroup(MLIROperator): """One stream-dse design, used as an ``OperatorSequence`` child. - ``group_index`` selects the design: ``None`` is the whole SwiGLU block - (``x, w1, w2, w3 -> y``); ``0`` is the gate/up/SiLU/mul front end - (``x, w1, w2 -> h``); ``1`` is the down projection (``h, w3 -> y``). + ``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 - group_index: Optional[int] = None - seq_len_tile_size: int = 32 - embedding_tile_size: int = 32 - hidden_tile_size: int = 64 - in_dtype: str = field(default="bf16", repr=False) - out_dtype: str = field(default="bf16", repr=False) - rows: int = field(default=4, repr=False) - num_aie_columns: int = field(default=8, repr=False) - backend: str = field(default="ortools_gscip", repr=False) + 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): - npu = aie_utils.get_current_device().resolve().name - kwargs = { - "seq_len": self.seq_len, - "embedding_dim": self.embedding_dim, - "hidden_dim": self.hidden_dim, - "in_dtype": self.in_dtype, - "out_dtype": self.out_dtype, - "rows": self.rows, - "cols": self.num_aie_columns, - "npu": npu, - "seq_len_tile_size": self.seq_len_tile_size, - "embedding_tile_size": self.embedding_tile_size, - "hidden_tile_size": self.hidden_tile_size, - "backend": self.backend, - } - if self.group_index is None: - fn, args = "run_main_aie_codegen_swiglu", () - kwargs["last_gemm_down"] = True - else: - fn, args = "load_swiglu_k2_group", (self.group_index,) return PythonGeneratedMLIRArtifact( f"{self.name}.mlir", - DesignGenerator(self.operator_dir / "stream_design.py", fn, args, kwargs), + 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): - # Each kernel's source, compile flags and symbol names come from the - # stream op registry, so they stay in step with the design stream-dse - # generates against IRON's aie_kernels library. - from iron.operators.swiglu_prefill_stream.stream_design import gemm_tiles - + # Each kernel's source, compile flags and symbol names come from the stream + # op registry, so they stay in step with the design stream-dse generates + # against IRON's aie_kernels library. Only this group's layers are built. + design = self._design + per_layer = { + design.GATE: (GEMM, design.GATE_UP_TILES), + design.UP: (GEMM, design.GATE_UP_TILES), + design.DOWN: (GEMM, design.DOWN_TILES), + 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() - gate_up, down = gemm_tiles( - self.seq_len_tile_size, self.embedding_tile_size, self.hidden_tile_size - ) - kernels = { - 0: [(GEMM, gate_up), (SILU, None), (ELTWISE_MUL, None)], - 1: [(GEMM, down)], - None: [(GEMM, gate_up), (GEMM, down), (SILU, None), (ELTWISE_MUL, None)], - }[self.group_index] return [ artifact - for kernel, tiles in kernels + 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 get_arg_spec(self): - m, e, h = self.seq_len, self.embedding_dim, self.hidden_dim - if self.group_index == 0: - return [ - AIERuntimeArgSpec("in", (m, e)), # x - AIERuntimeArgSpec("in", (e, h)), # w_gate - AIERuntimeArgSpec("in", (e, h)), # w_up - AIERuntimeArgSpec("out", (m, h)), # h - ] - if self.group_index == 1: - return [ - AIERuntimeArgSpec("in", (m, h)), # h - AIERuntimeArgSpec("in", (h, e)), # w_down - AIERuntimeArgSpec("out", (m, e)), # y - ] - return [ - AIERuntimeArgSpec("in", (m, e)), # x - AIERuntimeArgSpec("in", (e, h)), # w_gate - AIERuntimeArgSpec("in", (e, h)), # w_up - AIERuntimeArgSpec("in", (h, e)), # w_down - AIERuntimeArgSpec("out", (m, e)), # y + """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 _name(kind, m, e, h, st, et, ht): - return f"{kind}_m{m}_e{e}_h{h}_st{st}_et{et}_ht{ht}" - +def _wiring(seq_len, embedding_dim, hidden_dim, k): + """Each group's arguments, and the operator's own inputs and outputs. -def _boundaries(seq_len, embedding_dim, hidden_dim, split_groups=False): - """Each fused group's (inputs, outputs), named as the exported graph names them. - - Imported lazily: the names come from the exported workload, so only building - the operator needs stream-dse, not importing it. + 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 - return group_ports(seq_len, embedding_dim, hidden_dim, split_groups) - - -def _ports(seq_len, embedding_dim, hidden_dim, split_groups=False): - """Each fused group's arguments in the order its design takes them.""" - return [ - inputs + outputs - for inputs, outputs in _boundaries( - seq_len, embedding_dim, hidden_dim, split_groups - ) - ] - - -def _external(seq_len, embedding_dim, hidden_dim): - """The operator's own arguments: what no group produces, and what none consumes.""" - boundaries = _boundaries(seq_len, embedding_dim, hidden_dim, split_groups=True) + 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} - inputs = tuple( - dict.fromkeys( - name for group, _ in boundaries for name in group if name not in produced - ) + 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 ) - outputs = tuple( - dict.fromkeys( - name for _, group in boundaries for name in group if name not in consumed - ) + external_outputs = dict.fromkeys( + name for _, outputs in boundaries for name in outputs if name not in consumed ) - return inputs, outputs + return ports, list(external_inputs), list(external_outputs) class SwiGLUPrefillStream(OperatorSequence): - """Fused SwiGLU-prefill block generated by stream-dse and deployed as a - single full-ELF (``OperatorSequence`` default dispatch). - - 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. + """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, - seq_len_tile_size=32, - embedding_tile_size=32, - hidden_tile_size=64, - context=None, - ): - block = _SwiGLUStreamGroup( - seq_len=seq_len, - embedding_dim=embedding_dim, - hidden_dim=hidden_dim, - seq_len_tile_size=seq_len_tile_size, - embedding_tile_size=embedding_tile_size, - hidden_tile_size=hidden_tile_size, - context=context, - ) - super().__init__( - name=_name( - "swiglu_prefill_stream", - seq_len, - embedding_dim, - hidden_dim, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - ), - runlist=[(block, *_ports(seq_len, embedding_dim, hidden_dim)[0])], - input_args=list(_external(seq_len, embedding_dim, hidden_dim)[0]), - output_args=list(_external(seq_len, embedding_dim, hidden_dim)[1]), - extra_flags=["--dynamic-objFifos"], - context=context, - ) - - -class SwiGLUPrefillStreamK2(OperatorSequence): - """Two-fusion-group SwiGLU-prefill: a gate/up/SiLU/mul group and a separate - down-projection group fused into one full-ELF, with the hidden state kept on - device between them. Same external buffers as :class:`SwiGLUPrefillStream`; - the split is decided by the mapping's fused groups. - """ - - def __init__( - self, - seq_len, - embedding_dim, - hidden_dim, - seq_len_tile_size=32, - embedding_tile_size=32, - hidden_tile_size=64, - context=None, - ): - common = dict( - seq_len=seq_len, - embedding_dim=embedding_dim, - hidden_dim=hidden_dim, - seq_len_tile_size=seq_len_tile_size, - embedding_tile_size=embedding_tile_size, - hidden_tile_size=hidden_tile_size, - context=context, - ) - front = _SwiGLUStreamGroup(group_index=0, **common) - down = _SwiGLUStreamGroup(group_index=1, **common) + def __init__(self, seq_len, embedding_dim, hidden_dim, k=1, context=None): + 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=_name( - "swiglu_prefill_stream_k2", - seq_len, - embedding_dim, - hidden_dim, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - ), + name=f"swiglu_prefill_stream_k{k}_m{seq_len}_e{embedding_dim}_h{hidden_dim}", runlist=[ - (group, *ports) - for group, ports in zip( - (front, down), - _ports(seq_len, embedding_dim, hidden_dim, split_groups=True), - ) + (group, *group_ports) for group, group_ports in zip(groups, ports) ], - input_args=list(_external(seq_len, embedding_dim, hidden_dim)[0]), - output_args=list(_external(seq_len, embedding_dim, hidden_dim)[1]), + input_args=inputs, + output_args=outputs, extra_flags=["--dynamic-objFifos"], context=context, ) diff --git a/iron/operators/swiglu_prefill_stream/reference.py b/iron/operators/swiglu_prefill_stream/reference.py index 9056ae59..92caf8eb 100644 --- a/iron/operators/swiglu_prefill_stream/reference.py +++ b/iron/operators/swiglu_prefill_stream/reference.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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``. diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index e36f95ef..0f4c7ac7 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -1,20 +1,20 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Generate the fused SwiGLU-prefill design with stream-dse. +"""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 **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 +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 +importing the operator does not require ``stream-dse`` to be installed, only building it does. """ @@ -39,8 +39,8 @@ 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 (stream-dse >= 1.13.3). -_ACCELERATOR = os.path.join( +# data inside the installed stream package. +ACCELERATOR = os.path.join( os.path.dirname(stream.__file__), "inputs", "aie", @@ -48,10 +48,13 @@ "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 -# torch.export captured, and they are what the mapping and the generated design -# are read by. +# 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 = { @@ -61,271 +64,228 @@ MUL: reference.HIDDEN, } -# Which layers each fused group contains. Splitting makes stream-dse emit one -# design per group; the tensor handed from one to the next is derived from the -# exported graph, not named here. +# The kernel tile every layer is compiled and mapped for. The mapping carries these +# and no absolute dimension, so it holds across problem sizes; _check_shapes rejects +# a workload they cannot tile evenly. +SEQ_TILE, EMBEDDING_TILE, HIDDEN_TILE = 32, 32, 64 +GATE_UP_TILES = (SEQ_TILE, EMBEDDING_TILE, HIDDEN_TILE) +DOWN_TILES = (SEQ_TILE, HIDDEN_TILE, EMBEDDING_TILE) + +# 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 to the +# next is derived from the exported graph, not named here. k=5 is layer by layer, +# every layer its own design, as in :mod:`iron.operators.swiglu_prefill`. +LAYER_BY_LAYER = 5 GROUP_LAYERS = { - False: [[GATE, UP, SILU, MUL, DOWN]], - True: [[GATE, UP, SILU, MUL], [DOWN]], + 1: [[GATE, UP, SILU, MUL, DOWN]], + 2: [[GATE, UP, SILU, MUL], [DOWN]], + LAYER_BY_LAYER: [[GATE], [UP], [SILU], [MUL], [DOWN]], } -ARRAY = ComputeArray.from_accelerator(_ACCELERATOR) -# Columns per layer: two for each GEMM, one for each elementwise layer. The five -# layers sit on disjoint columns so they pipeline across steady-state iterations. -# Each layer splits over the array's rows (D0, the sequence dimension), and a GEMM -# splits over its two columns as well (D2, the output dimension). -_COLUMNS = dict(zip(NODE_NAMES, ARRAY.allocate([2, 2, 1, 1, 2]))) -_GEMM_SPLIT = (("D0", ARRAY.num_rows), ("D2", 2)) -_ELEMENTWISE_SPLIT = (("D0", ARRAY.num_rows),) +@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 gemm_tiles(seq_tile, embedding_tile, hidden_tile): - """Kernel tile shape ``(m, k, n)`` of the gate/up GEMMs and of the down GEMM. - The gate and up projections contract over the embedding dimension and produce - the hidden one; the down projection contracts the other way round. - """ - return (seq_tile, embedding_tile, hidden_tile), ( - seq_tile, - hidden_tile, - embedding_tile, - ) +def _placements(k): + """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). -def _placements(seq_tile, embedding_tile, hidden_tile): - gate_up, down = gemm_tiles(seq_tile, embedding_tile, hidden_tile) + 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. + """ + grid = array() def gemm(tiles): - m, k, n = tiles - return {"utilization": 61.8, "m": m, "k": k, "n": n, "layout": "default"} + return dict(zip("mkn", tiles), utilization=61.8, layout="default") elementwise = {"utilization": 50.0, "layout": "default"} + + 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),) + return { + GATE: Placement(wide, gemm_split, gemm(GATE_UP_TILES)), + UP: Placement(wide, gemm_split, gemm(GATE_UP_TILES)), + SILU: Placement(wide, elementwise_split, elementwise, rows=[0]), + MUL: Placement(wide, elementwise_split, elementwise, rows=[0]), + DOWN: Placement(wide, gemm_split, gemm(DOWN_TILES)), + } + + 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),) return { - GATE: Placement(_COLUMNS[GATE], _GEMM_SPLIT, gemm(gate_up)), - UP: Placement(_COLUMNS[UP], _GEMM_SPLIT, gemm(gate_up)), - SILU: Placement(_COLUMNS[SILU], _ELEMENTWISE_SPLIT, elementwise), - MUL: Placement(_COLUMNS[MUL], _ELEMENTWISE_SPLIT, elementwise), - DOWN: Placement(_COLUMNS[DOWN], _GEMM_SPLIT, gemm(down)), + GATE: Placement(columns[GATE], gemm_split, gemm(GATE_UP_TILES)), + UP: Placement(columns[UP], gemm_split, gemm(GATE_UP_TILES)), + SILU: Placement(columns[SILU], elementwise_split, elementwise), + MUL: Placement(columns[MUL], elementwise_split, elementwise), + DOWN: Placement(columns[DOWN], gemm_split, gemm(DOWN_TILES)), } -def _groups(seq_tile, embedding_tile, hidden_tile, split_groups): - """Fusion groups: one fully fused design, or a front end plus down projection. +def _layer_tiling(layer): + """Intra-core tiling of one layer, over the dimensions it iterates.""" + if layer is DOWN: + contraction, output = HIDDEN_TILE, EMBEDDING_TILE + elif layer in (GATE, UP): + contraction, output = EMBEDDING_TILE, HIDDEN_TILE + else: + return [(layer, "D1", HIDDEN_TILE), (layer, "D0", SEQ_TILE)] + return [ + (layer, "D1", contraction), + (layer, "D2", output), + (layer, "D0", SEQ_TILE), + ] + + +def _groups(k): + """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 of the group's leading GEMM. + as a single full ELF. D0 is the sequence dimension, D1 the contraction and D2 + the output dimension. """ - if split_groups: - tiling = [ - [ - (GATE, "D1", embedding_tile), - (GATE, "D2", hidden_tile), - (GATE, "D0", seq_tile), - ], - [ - (DOWN, "D1", hidden_tile), - (DOWN, "D2", embedding_tile), - (DOWN, "D0", seq_tile), - ], - ] + if k == LAYER_BY_LAYER: + tiling = [_layer_tiling(layers[0]) for layers in GROUP_LAYERS[k]] + elif k == 2: + tiling = [_layer_tiling(GATE), _layer_tiling(DOWN)] else: tiling = [ [ - (GATE, "D1", embedding_tile), - (DOWN, "D2", embedding_tile), - (GATE, "D2", hidden_tile), - (GATE, "D0", seq_tile), + (GATE, "D1", EMBEDDING_TILE), + (DOWN, "D2", EMBEDDING_TILE), + (GATE, "D2", HIDDEN_TILE), + (GATE, "D0", SEQ_TILE), ] ] return [ FusedGroup(f"Fused_Group_{index + 1}", layers, group_tiling) - for index, (layers, group_tiling) in enumerate( - zip(GROUP_LAYERS[split_groups], tiling) - ) + for index, (layers, group_tiling) in enumerate(zip(GROUP_LAYERS[k], tiling)) ] -def _check_shapes( - seq_len, embedding_dim, hidden_dim, seq_tile, embedding_tile, hidden_tile -): - """Reject shapes the fixed 4x8 placement cannot tile evenly.""" - rows, gemm_split = ARRAY.num_rows, 2 - if seq_len % rows or seq_len < seq_tile * rows: +def _check_shapes(seq_len, embedding_dim, hidden_dim, k): + """Reject a problem size the placement and the kernel tiles cannot divide.""" + grid = array() + gemm_split = grid.num_columns if k == LAYER_BY_LAYER else 2 + if seq_len % grid.num_rows or seq_len < SEQ_TILE * grid.num_rows: raise ValueError( - f"seq_len ({seq_len}) must be a multiple of {rows} and at least {seq_tile * rows}" + f"seq_len ({seq_len}) must be a multiple of {grid.num_rows} and at " + f"least {SEQ_TILE * grid.num_rows}" ) - if embedding_dim % (embedding_tile * gemm_split): + if embedding_dim % (EMBEDDING_TILE * gemm_split): raise ValueError( f"embedding_dim ({embedding_dim}) must be a multiple of " - f"embedding_tile * {gemm_split} ({embedding_tile * gemm_split})" + f"{EMBEDDING_TILE * gemm_split}" ) - if hidden_dim % (hidden_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} ({hidden_tile * gemm_split})" + 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.""" - module = swiglu_module(embedding_dim, hidden_dim) return export_workload( - module, + 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, split_groups=False): +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 tensor a split - design passes from its front end to its down projection. + 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[split_groups] + workload_for(seq_len, embedding_dim, hidden_dim), GROUP_LAYERS[k] ) -def build_inputs( - seq_len, - embedding_dim, - hidden_dim, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - output_dir, - split_groups=False, -): +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, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - ) + _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(seq_len_tile_size, embedding_tile_size, hidden_tile_size), - _groups( - seq_len_tile_size, embedding_tile_size, hidden_tile_size, split_groups - ), - ARRAY, + _placements(k), + _groups(k), + array(), output_dir / "mapping.yaml", ), ) -def _experiment_id(seq_len, embedding_dim, hidden_dim, rows, cols, suffix=""): - hw_name = os.path.splitext(os.path.basename(_ACCELERATOR))[0] - return f"{hw_name}-swiglu{suffix}_{seq_len}_{embedding_dim}_{hidden_dim}-{rows}_row_{cols}_col" - - -def _run_codegen( - seq_len, - embedding_dim, - hidden_dim, - rows, - cols, - npu, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - backend, - trace_size=0, - split_groups=False, - output_root="outputs", -): - """Run stream-dse's constraint-optimization + code generation once.""" - experiment_id = _experiment_id( - seq_len, embedding_dim, hidden_dim, rows, cols, "_k2" if split_groups else "" +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" ) - output_dir = os.path.join(output_root, experiment_id) + + +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, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - output_dir, - split_groups=split_groups, + os.path.join(OUTPUT_ROOT, experiment_id), + k=k, ) - ctx = optimize_allocation_co( - hardware=_ACCELERATOR, + optimize_allocation_co( + hardware=ACCELERATOR, workload=workload_path, mapping=mapping_path, experiment_id=experiment_id, - output_path=output_root, + output_path=OUTPUT_ROOT, skip_if_exists=False, enable_codegen=True, - trace_size=trace_size, - nb_cols_to_use=cols, + trace_size=0, + nb_cols_to_use=grid.num_columns, npu=npu, - backend=backend, - ) - return ctx, output_dir - - -def run_main_aie_codegen_swiglu( - seq_len, - embedding_dim, - hidden_dim, - in_dtype="bf16", - out_dtype="bf16", - trace_size=0, - rows=4, - cols=8, - npu="npu2", - seq_len_tile_size=32, - embedding_tile_size=32, - hidden_tile_size=64, - last_gemm_down=True, - backend="ortools_gscip", - func_prefix="", -): - """Generate the fully fused SwiGLU-prefill MLIR module. - - Returns an ``aie`` MLIR module. ``func_prefix`` (injected by - ``OperatorSequence``) prefixes the kernel symbols / ``link_with`` objects so - the design can be deployed as one fusion group; see :func:`region_module`. - - The default ``ortools_gscip`` backend is the license-free OR-Tools GSCIP - solver, so no Gurobi license is required. - """ - ctx, _ = _run_codegen( - seq_len, - embedding_dim, - hidden_dim, - rows, - cols, - npu, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - backend, - trace_size=trace_size, + backend=BACKEND, ) - return region_module(str(ctx.get("module")), func_prefix) - - -# --------------------------------------------------------------------------- -# k=2 variant: two fusion groups (gate/up/SiLU/mul -> h, then down-projection) -# --------------------------------------------------------------------------- -# -# stream-dse emits a separate ``aie.device`` design per fusion group (under -# ``/group_i/codegen/final.mlir``). IRON fuses the two groups into one -# full-ELF via ``OperatorSequence``; each group is loaded below as a child design. def _prefixed(mlir_text: str, func_prefix: str) -> str: @@ -334,7 +294,7 @@ def _prefixed(mlir_text: str, func_prefix: str) -> str: ``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). + and every privately declared kernel symbol, and its call sites. """ if not func_prefix: return mlir_text @@ -348,15 +308,17 @@ def _prefixed(mlir_text: str, func_prefix: str) -> str: key=len, reverse=True, ) - for sym in symbols: - mlir_text = re.sub(rf"@{re.escape(sym)}\b", f"@{func_prefix}{sym}", mlir_text) + 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 stream group's MLIR text into an ``aie`` module for fusion. + """Parse a group's MLIR text into an ``aie`` module for fusion. - ``OperatorSequence`` consumes ``aie.DeviceOp`` objects, so the (xDSL-emitted) + ``OperatorSequence`` consumes ``aie.DeviceOp`` objects, so the xDSL-emitted group text is re-parsed with the mlir-aie bindings, after ``func_prefix`` rewriting. """ @@ -367,48 +329,16 @@ def region_module(mlir_text: str, func_prefix: str = ""): return ir.Module.parse(_prefixed(mlir_text, func_prefix)) -def load_swiglu_k2_group( - group_index, - func_prefix="", - *, - seq_len, - embedding_dim, - hidden_dim, - in_dtype="bf16", - out_dtype="bf16", - rows=4, - cols=8, - npu="npu2", - seq_len_tile_size=32, - embedding_tile_size=32, - hidden_tile_size=64, - backend="ortools_gscip", +def load_group( + group_index, func_prefix="", *, k, seq_len, embedding_dim, hidden_dim, npu ): - """Generate the two-group design (cached) and return one group's aie module. + """Generate the ``k``-group design once and return one group's aie module. - ``group_index`` 0 is the gate/up/SiLU/mul front end (``x, w_gate, w_up -> h``); - 1 is the down projection (``h, w_down -> y``). ``func_prefix`` is injected by - ``OperatorSequence``. Both group loaders call this; the first generates the - design, the second reuses the files on disk. + ``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. """ - output_dir = os.path.join( - "outputs", _experiment_id(seq_len, embedding_dim, hidden_dim, rows, cols, "_k2") - ) - finals = [ - os.path.join(output_dir, f"group_{i}", "codegen", "final.mlir") for i in (0, 1) - ] - if not all(os.path.exists(f) for f in finals): - _run_codegen( - seq_len, - embedding_dim, - hidden_dim, - rows, - cols, - npu, - seq_len_tile_size, - embedding_tile_size, - hidden_tile_size, - backend, - split_groups=True, - ) + 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 region_module(Path(finals[group_index]).read_text(), func_prefix) diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index e694ba82..c3a762a9 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. # SPDX-License-Identifier: Apache-2.0 import time @@ -14,10 +14,7 @@ "stream", reason="stream-dse not installed (see requirements_stream.txt)" ) -from iron.operators.swiglu_prefill_stream.op import ( - SwiGLUPrefillStream, - SwiGLUPrefillStreamK2, -) +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. @@ -25,36 +22,58 @@ 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 -def get_params(): - # (seq_len, embedding_dim, hidden_dim, seq_tile, embedding_tile, hidden_tile): - # the MILP-feasible shape on the whole-array Strix (npu2) target. - return [pytest.param(256, 512, 2048, 32, 32, 64)] +# 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 _run_and_verify(operator, seq_len, embedding_dim, golden_ref): - operator.compile() - run = operator.get_callable() - # Populate inputs by golden-reference name; the design consumes weights in - # their natural (K, N) layout, no transpose. +def _dispatch(operator, golden_ref): + """A callable on a freshly configured array, with its inputs staged. + + Creating the hw_context applies the design's init CDO. The full-ELF path does + not reset array state between dispatches, so every dispatch takes a new + callable: reusing one reads stale state, and a design of several fused groups + deadlocks on the second dispatch. 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): - run.get_buffer(name).torch_view()[:] = ( - golden_ref[name].to(torch.bfloat16).flatten() - ) + 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() - # Correctness is checked on the FIRST run: creating the hw_context applies the - # design's init CDO, so the first run on the freshly loaded full-ELF computes - # correctly. The full-ELF path does not re-initialize the array between runs, - # so later runs on the same callable read stale state and are used only for - # timing below. - # - # The whole block is fused in bf16, so rounding accumulates and reorders across - # the chain and the K=hidden_dim down-projection sum (near-cancellation + # 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 = _dispatch(operator, golden_ref) run() - output = run.get_buffer(OUTPUT).to_torch().reshape((seq_len, embedding_dim)) + output = run.get_buffer(OUTPUT).to_torch().reshape((SEQ_LEN, EMBEDDING_DIM)) errors = verify_buffer( output, OUTPUT, @@ -65,71 +84,19 @@ def _run_and_verify(operator, seq_len, embedding_dim, golden_ref): ) assert not errors, f"Test failed with errors: {errors}" - # Performance: time a run that reuses the hw_context, so it reflects the actual - # kernel latency without the one-off ELF/hw_context setup. - start = time.perf_counter() - run() - elapsed_us = (time.perf_counter() - start) * 1e6 - total_bytes = 2 * (seq_len * embedding_dim + seq_len * embedding_dim) # bf16 in+out + # Performance: time dispatches of their own, each on a freshly configured + # array, so every timed run is one that computed the verified result. + latencies = [] + for _ in range(TIMED_RUNS): + run = _dispatch(operator, golden_ref) + 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"Effective Bandwidth: {total_bytes / (elapsed_us * 1e-6) / 1e9:.4f} GB/s") - - -@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( - "seq_len,embedding_dim,hidden_dim,seq_tile,embedding_tile,hidden_tile", get_params() -) -def test_swiglu_prefill_stream( - seq_len, - embedding_dim, - hidden_dim, - seq_tile, - embedding_tile, - hidden_tile, - 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, - seq_len_tile_size=seq_tile, - embedding_tile_size=embedding_tile, - hidden_tile_size=hidden_tile, - context=aie_context, + print( + f"Latency min/mean/max (us): {elapsed_us:.2f} / " + f"{sum(latencies) / len(latencies):.2f} / {max(latencies):.2f}" ) - _run_and_verify(operator, seq_len, embedding_dim, golden_ref) - - -@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( - "seq_len,embedding_dim,hidden_dim,seq_tile,embedding_tile,hidden_tile", get_params() -) -def test_swiglu_prefill_stream_k2( - seq_len, - embedding_dim, - hidden_dim, - seq_tile, - embedding_tile, - hidden_tile, - aie_context, -): - golden_ref = generate_golden_reference(M=seq_len, K=embedding_dim, N=hidden_dim) - operator = SwiGLUPrefillStreamK2( - seq_len=seq_len, - embedding_dim=embedding_dim, - hidden_dim=hidden_dim, - seq_len_tile_size=seq_tile, - embedding_tile_size=embedding_tile, - hidden_tile_size=hidden_tile, - context=aie_context, - ) - _run_and_verify(operator, seq_len, embedding_dim, golden_ref) + 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 index a8f9444a..c9215fe6 100644 --- a/iron/tests/stream/kernel_layouts.py +++ b/iron/tests/stream/kernel_layouts.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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. diff --git a/iron/tests/stream/names.py b/iron/tests/stream/names.py index f2358ee5..cdb58d57 100644 --- a/iron/tests/stream/names.py +++ b/iron/tests/stream/names.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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. @@ -41,6 +41,14 @@ def test_golden_weights_load_into_the_module(): 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)" ) @@ -50,25 +58,23 @@ def test_golden_weights_load_into_the_module(): DIMS = (256, 512, 2048) -@pytest.mark.parametrize("split_groups", [False, True]) -def test_group_ports_are_named_by_the_exported_graph(split_groups): +@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, split_groups): + 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, split_groups=True - ) + (_, 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, split_groups=True) + 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 = [ diff --git a/iron/tests/stream/placement.py b/iron/tests/stream/placement.py index f051617b..0615715a 100644 --- a/iron/tests/stream/placement.py +++ b/iron/tests/stream/placement.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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. -An accelerator description gives every compute tile a ``[column, row]`` -coordinate; :class:`~iron.common.stream.hardware.ComputeArray` is the only place -that reads it, so operators never spell out a core id. +: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 @@ -15,12 +15,19 @@ "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 +ARRAY = stream_design.array() -# The core ids this operator was measured with, on the whole-array Strix target. +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], @@ -30,7 +37,7 @@ } -def test_array_matches_the_accelerator(): +def test_array_matches_the_device(): assert (ARRAY.num_columns, ARRAY.num_rows) == (8, 4) assert ARRAY.all_columns == tuple(range(8)) @@ -46,6 +53,13 @@ def test_cores_follow_column_then_row_order(): 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)) @@ -58,12 +72,38 @@ def test_allocate_rejects_an_oversubscribed_array(): ARRAY.allocate([ARRAY.num_columns, 1]) -def test_emitted_allocation_is_the_measured_one(tmp_path): +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 - _, mapping_path = stream_design.build_inputs( - 256, 512, 2048, 32, 32, 64, tmp_path / "design" + 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"] @@ -71,15 +111,26 @@ def test_emitted_allocation_is_the_measured_one(tmp_path): assert emitted == MEASURED_ALLOCATION -@pytest.mark.parametrize("accelerator", ["whole_array.yaml", "single_col.yaml"]) -def test_other_accelerators_resolve(accelerator): - import os - - import stream +def test_layer_by_layer_gives_every_layer_the_whole_array(tmp_path): + import yaml - path = os.path.join( - os.path.dirname(stream.__file__), "inputs", "aie", "hardware", accelerator + _, mapping_path = stream_design.build_inputs( + *DIMS, tmp_path / "design_k5", k=stream_design.LAYER_BY_LAYER ) - array = ComputeArray.from_accelerator(path) - assert array.num_rows == 4 + 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 index be597fbb..e1ecd4c0 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -1,39 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# 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). # -# It is 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/run the operator and its test: +# 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) that cannot be PyPI -# # dependencies. As of stream-dse 1.13.7 it does NOT install -# # mlir_aie/llvm-aie: IRON already pins those in -# # requirements.txt, and codegen only emits text MLIR via -# # xdsl (it never imports the aie bindings). Earlier releases -# # reinstalled an older mlir_aie here, clobbering IRON's wheel. +# # dialects (xdsl-aie, snax-mlir), which cannot be PyPI +# # dependencies. It leaves mlir_aie/llvm-aie alone; IRON +# # pins those in requirements.txt. # -# Notes: -# - stream-dse generates the fused MLIR design at build time (license-free -# OR-Tools GSCIP solver; no Gurobi needed) and writes its generated workload/ -# mapping files into its own installed package directory, so that environment -# must be writable. -# - >=1.13.8 is required: stream_design.py feeds IRON-authored operand layouts -# into code generation via optimize_allocation_co(kernels=...), the override -# hook added in stream-dse 1.13.4; the k=2 variant (op_k2.py) additionally needs -# the two-fusion-group support (make_swiglu_mapping(split_groups=...) + the -# multi-group AIE codegen pipeline) added in stream-dse 1.13.5; 1.13.6 makes -# the debug workload-graph visualization non-fatal, so code generation no longer -# requires graphviz (`dot`); 1.13.7 makes `stream-setup-aie` stop reinstalling -# mlir_aie/llvm-aie, which previously downgraded IRON's pinned mlir_aie (1.3.5 -> -# 0.0.1) and numpy (2.x -> 1.26.4) and broke the whole operator test suite; and -# 1.13.8 orders the generated design's runtime arguments by the exported graph -# instead of by tensor name, which is what lets the weights below carry their -# reference names. +# 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.8 From 58b36e3c909024831389df32f4dacdbf420c9e2e Mon Sep 17 00:00:00 2001 From: asyms Date: Tue, 28 Jul 2026 20:12:22 +0200 Subject: [PATCH 12/21] ci: require stream-dse>=1.13.9 The layer-by-layer (k=5) design needs the fusion-group and AIE codegen fixes released in 1.13.9; on 1.13.8 its groups fail to generate. --- requirements_stream.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements_stream.txt b/requirements_stream.txt index e1ecd4c0..513ee7f9 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -19,4 +19,4 @@ # package directory, so that environment must be writable. onnxscript>=0.7 -stream-dse>=1.13.8 +stream-dse>=1.13.9 From 63b721528d1fd046eccfc745e9694e0dd01dd639 Mon Sep 17 00:00:00 2001 From: asyms Date: Tue, 28 Jul 2026 22:59:57 +0200 Subject: [PATCH 13/21] swiglu_prefill_stream: read whole rows in the layer-by-layer elementwise groups The k=5 elementwise layers read from and write to memory rather than sitting behind a GEMM, so give them a contiguous operand layout and a tile of whole rows. Their transfers then run the length of a row instead of one MAC tile at a time. --- .../swiglu_prefill_stream/stream_design.py | 39 ++++++++++++------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index 0f4c7ac7..ae6d56ff 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -71,6 +71,12 @@ GATE_UP_TILES = (SEQ_TILE, EMBEDDING_TILE, HIDDEN_TILE) DOWN_TILES = (SEQ_TILE, HIDDEN_TILE, EMBEDDING_TILE) +# 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 to the # next is derived from the exported graph, not named here. k=5 is layer by layer, @@ -91,7 +97,7 @@ def array() -> ComputeArray: return ComputeArray.from_device(aie_utils.get_current_device()) -def _placements(k): +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 @@ -101,47 +107,52 @@ def _placements(k): 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. + 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() def gemm(tiles): return dict(zip("mkn", tiles), utilization=61.8, layout="default") - elementwise = {"utilization": 50.0, "layout": "default"} + def elementwise(rows, columns, layout): + return {"utilization": 50.0, "layout": layout, "m": rows, "n": columns} 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(GATE_UP_TILES)), UP: Placement(wide, gemm_split, gemm(GATE_UP_TILES)), - SILU: Placement(wide, elementwise_split, elementwise, rows=[0]), - MUL: Placement(wide, elementwise_split, elementwise, rows=[0]), + SILU: Placement(wide, elementwise_split, rows_wide, rows=[0]), + MUL: Placement(wide, elementwise_split, rows_wide, rows=[0]), DOWN: Placement(wide, gemm_split, gemm(DOWN_TILES)), } 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 keep the layout the GEMM writes. + fused = elementwise(SEQ_TILE, HIDDEN_TILE, "default") return { GATE: Placement(columns[GATE], gemm_split, gemm(GATE_UP_TILES)), UP: Placement(columns[UP], gemm_split, gemm(GATE_UP_TILES)), - SILU: Placement(columns[SILU], elementwise_split, elementwise), - MUL: Placement(columns[MUL], elementwise_split, elementwise), + SILU: Placement(columns[SILU], elementwise_split, fused), + MUL: Placement(columns[MUL], elementwise_split, fused), DOWN: Placement(columns[DOWN], gemm_split, gemm(DOWN_TILES)), } -def _layer_tiling(layer): +def _layer_tiling(layer, hidden_dim): """Intra-core tiling of one layer, over the dimensions it iterates.""" if layer is DOWN: contraction, output = HIDDEN_TILE, EMBEDDING_TILE elif layer in (GATE, UP): contraction, output = EMBEDDING_TILE, HIDDEN_TILE else: - return [(layer, "D1", HIDDEN_TILE), (layer, "D0", SEQ_TILE)] + return [(layer, "D1", hidden_dim), (layer, "D0", ELEMENTWISE_ROWS)] return [ (layer, "D1", contraction), (layer, "D2", output), @@ -149,7 +160,7 @@ def _layer_tiling(layer): ] -def _groups(k): +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 @@ -157,9 +168,9 @@ def _groups(k): the output dimension. """ if k == LAYER_BY_LAYER: - tiling = [_layer_tiling(layers[0]) for layers in GROUP_LAYERS[k]] + tiling = [_layer_tiling(layers[0], hidden_dim) for layers in GROUP_LAYERS[k]] elif k == 2: - tiling = [_layer_tiling(GATE), _layer_tiling(DOWN)] + tiling = [_layer_tiling(GATE, hidden_dim), _layer_tiling(DOWN, hidden_dim)] else: tiling = [ [ @@ -227,8 +238,8 @@ def build_inputs(seq_len, embedding_dim, hidden_dim, output_dir, k=1): workload.write(output_dir / "workload.onnx"), emit_mapping( workload, - _placements(k), - _groups(k), + _placements(k, hidden_dim), + _groups(k, hidden_dim), array(), output_dir / "mapping.yaml", ), From 964c5a52bfd569fbdf574bf9fd3254c5377d1076 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 29 Jul 2026 01:18:02 +0200 Subject: [PATCH 14/21] sequence: configure an empty device to close the re-entrancy gap aiecc --expand-load-pdis marks each configure point by loading one of two otherwise empty PDIs, alternating 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 empty device makes the count even, which costs one PDI load and no register writes. Removes the need to rebuild the callable per dispatch for the affected sequences. --- iron/common/compilation/sequence.py | 40 +++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index c941a857..03b3f2b4 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, as the fused sequence below skips reconfiguring for them. + """ + 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,19 @@ def fuse_mlir(artifact: SequenceMLIRArtifact) -> None: dev_op.sym_name = ir.StringAttr.get(op_name) ctx.module.body.append(dev_op) + # An empty device, configured at the end of the sequence to leave the array + # in a state the next dispatch can start from. + 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 +310,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)) From 4a10562da908533c3ab81504e66e4dc233f58dd4 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 29 Jul 2026 02:20:36 +0200 Subject: [PATCH 15/21] ci: require stream-dse>=1.13.10 Picks up the reuse-gated memory-tile staging, the vectorized elementwise multiply and the graph-derived transfer topology the k=5 designs need. --- requirements_stream.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements_stream.txt b/requirements_stream.txt index 513ee7f9..80e7cd59 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -19,4 +19,4 @@ # package directory, so that environment must be writable. onnxscript>=0.7 -stream-dse>=1.13.9 +stream-dse>=1.13.10 From bacaeb3d1a9159f964bc9f73e4107d9f87f2cb42 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 30 Jul 2026 00:16:03 +0200 Subject: [PATCH 16/21] iron/common: build and configure a design once when operators share it An operator can report a design_key naming the design it compiles to. A sequence that opts in with share_designs groups its operators by that key, so a design several of them compile to is emitted once, prefixed once, and configured once. Sharing is off by default, and a key of None never shares, so a sequence that does not ask for it is unaffected. Operators reporting the same key must take the same runtime arguments, which is checked rather than assumed. --- iron/common/base.py | 9 ++++++++ iron/common/sequence.py | 46 +++++++++++++++++++++++++++++++++++------ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/iron/common/base.py b/iron/common/base.py index e8dea253..701e90df 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/sequence.py b/iron/common/sequence.py index a7a33cef..020ce49d 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -98,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", @@ -123,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}" @@ -264,6 +265,7 @@ def __init__( buffer_sizes=None, dispatch="auto", extra_flags=None, + share_designs=False, *args, **kwargs, ): @@ -278,7 +280,9 @@ 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 = ( @@ -288,6 +292,9 @@ def __init__( # 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 [] + # Build and configure a design once when several operators compile to it. + # Off by default so sequences that do not opt in are unaffected. + self.share_designs = share_designs self._dispatch = dispatch @staticmethod @@ -306,6 +313,33 @@ 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. + + Operators are de-duplicated by identity as above, and additionally by + ``design_key`` when ``share_designs`` is set, so a design that several + operators compile to is built 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 = ( From 49e536f22102ac96af183adca3fe6037a2e2b521 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 30 Jul 2026 00:16:03 +0200 Subject: [PATCH 17/21] iron/common/stream: take the MAC tile the kernel object is compiled for mm.cc emulates bf16 matmuls on the bfp16 MACs when built with AIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16, and that shape takes an 8 row MAC tile rather than a 4 row one. The registry 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. The cross-check against stream-dse's own layouts now covers both settings. --- iron/common/stream/ops.py | 25 +++++++++++++++++---- iron/tests/stream/kernel_layouts.py | 34 +++++++++++++++++++---------- 2 files changed, 44 insertions(+), 15 deletions(-) diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index 242fae92..4571ea97 100644 --- a/iron/common/stream/ops.py +++ b/iron/common/stream/ops.py @@ -32,7 +32,10 @@ # 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) @@ -56,14 +59,24 @@ def custom_op(name: str, arity: int = 1) -> Op: return Op(CUSTOM_DOMAIN, name, schema) -def gemm_layouts(m: int, k: int, n: int) -> tuple[TiledStridedLayout, ...]: +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.""" - return (tiled_2d(m, k, R, S), tiled_2d(k, n, S, T), tiled_2d(m, n, R, T)) + 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) -> tuple[TiledStridedLayout, ...]: +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, R, T),) * nb_operands + 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): @@ -87,6 +100,10 @@ def _gemm_artifacts(base_dir, kernel_dir, m: int, k: int, n: int): f"-DDIM_K={k}", f"-DDIM_N={n}", "-Dbf16_bf16_ONLY", + # The tile operators build mm.cc this way by default: bf16 matmuls + # emulated on the bfp16 MACs, with even rounding on the conversion. + "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", + "-DROUND_CONV_EVEN", ], rename_symbols={ "matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}", diff --git a/iron/tests/stream/kernel_layouts.py b/iron/tests/stream/kernel_layouts.py index c9215fe6..d6efdad5 100644 --- a/iron/tests/stream/kernel_layouts.py +++ b/iron/tests/stream/kernel_layouts.py @@ -35,17 +35,29 @@ def _assert_same(iron_layouts, stream_kernel): assert [str(layout) for layout in expected] == [str(layout) for layout in actual] -@pytest.mark.parametrize("m,k,n", [(32, 32, 64), (32, 64, 32)]) -def test_gemm_layouts_match_stream(m, k, n): - _assert_same(gemm_layouts(m, k, n), AIEKernels[GEMM.key](61.8, m, k, n, "default")) - - -def test_silu_layouts_match_stream(): - _assert_same(elementwise_layouts(2), AIEKernels[SILU.key](50.0, "default")) - - -def test_eltwise_mul_layouts_match_stream(): - _assert_same(elementwise_layouts(3), AIEKernels[ELTWISE_MUL.key](50.0, "default")) +@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]) From 26149f8bc7d303988244053f53f0377034d43a83 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 30 Jul 2026 00:16:03 +0200 Subject: [PATCH 18/21] swiglu_prefill_stream: choose the kernel tile per fusion granularity 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. Layer by layer takes 64x64x64 where the fused granularities take 32x32x64. The GEMMs are built for the bfp16 MACs, and the elementwise layers that keep the layout a neighbouring GEMM writes are tiled to match. Groups whose generated design is byte-identical report the same design_key, so gate and up share one design. --- iron/operators/swiglu_prefill_stream/op.py | 23 +++- .../swiglu_prefill_stream/stream_design.py | 125 ++++++++++++------ 2 files changed, 105 insertions(+), 43 deletions(-) diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 9e23b194..c20aef9b 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -65,10 +65,11 @@ def get_kernel_artifacts(self): # op registry, so they stay in step with the design stream-dse generates # against IRON's aie_kernels library. Only this group's layers are built. design = self._design + gemm_tiles = design.gemm_tiles(self.k) per_layer = { - design.GATE: (GEMM, design.GATE_UP_TILES), - design.UP: (GEMM, design.GATE_UP_TILES), - design.DOWN: (GEMM, design.DOWN_TILES), + 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), } @@ -82,6 +83,17 @@ def get_kernel_artifacts(self): ) ] + 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. @@ -134,7 +146,9 @@ class SwiGLUPrefillStream(OperatorSequence): importing this module does not. """ - def __init__(self, seq_len, embedding_dim, hidden_dim, k=1, context=None): + 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( @@ -155,5 +169,6 @@ def __init__(self, seq_len, embedding_dim, hidden_dim, k=1, context=None): input_args=inputs, output_args=outputs, extra_flags=["--dynamic-objFifos"], + share_designs=share_designs, context=context, ) diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index ae6d56ff..3da68b4f 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -18,6 +18,7 @@ building it does. """ +import hashlib import os import re from functools import lru_cache @@ -64,12 +65,13 @@ MUL: reference.HIDDEN, } -# The kernel tile every layer is compiled and mapped for. The mapping carries these -# and no absolute dimension, so it holds across problem sizes; _check_shapes rejects -# a workload they cannot tile evenly. -SEQ_TILE, EMBEDDING_TILE, HIDDEN_TILE = 32, 32, 64 -GATE_UP_TILES = (SEQ_TILE, EMBEDDING_TILE, HIDDEN_TILE) -DOWN_TILES = (SEQ_TILE, HIDDEN_TILE, EMBEDDING_TILE) +# 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. The mapping carries the tile and no +# absolute dimension, so it holds across problem sizes; _check_shapes rejects a +# workload it cannot tile evenly. +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 @@ -89,6 +91,21 @@ } +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.""" @@ -111,12 +128,22 @@ def _placements(k, hidden_dim): 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") + return dict( + zip("mkn", tiles), utilization=61.8, layout="default", bfp16_mmul=True + ) - def elementwise(rows, columns, layout): - return {"utilization": 50.0, "layout": layout, "m": rows, "n": columns} + 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 @@ -124,39 +151,36 @@ def elementwise(rows, columns, layout): elementwise_split = (("D0", grid.num_columns),) rows_wide = elementwise(ELEMENTWISE_ROWS, hidden_dim, "contiguous") return { - GATE: Placement(wide, gemm_split, gemm(GATE_UP_TILES)), - UP: Placement(wide, gemm_split, gemm(GATE_UP_TILES)), + 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(DOWN_TILES)), + 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 keep the layout the GEMM writes. - fused = elementwise(SEQ_TILE, HIDDEN_TILE, "default") + fused = elementwise(sequence_tile, hidden_tile, "default", bfp16_mmul=True) return { - GATE: Placement(columns[GATE], gemm_split, gemm(GATE_UP_TILES)), - UP: Placement(columns[UP], gemm_split, gemm(GATE_UP_TILES)), + 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(DOWN_TILES)), + DOWN: Placement(columns[DOWN], gemm_split, gemm(tiles[DOWN])), } -def _layer_tiling(layer, hidden_dim): +def _layer_tiling(layer, hidden_dim, k): """Intra-core tiling of one layer, over the dimensions it iterates.""" - if layer is DOWN: - contraction, output = HIDDEN_TILE, EMBEDDING_TILE - elif layer in (GATE, UP): - contraction, output = EMBEDDING_TILE, HIDDEN_TILE - else: + 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", SEQ_TILE), + (layer, "D0", sequence), ] @@ -167,17 +191,21 @@ def _groups(k, hidden_dim): 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) for layers in GROUP_LAYERS[k]] + tiling = [_layer_tiling(layers[0], hidden_dim, k) for layers in GROUP_LAYERS[k]] elif k == 2: - tiling = [_layer_tiling(GATE, hidden_dim), _layer_tiling(DOWN, hidden_dim)] + 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", SEQ_TILE), + (GATE, "D1", embedding_tile), + (DOWN, "D2", embedding_tile), + (GATE, "D2", hidden_tile), + (GATE, "D0", sequence_tile), ] ] return [ @@ -189,21 +217,22 @@ def _groups(k, hidden_dim): 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 < SEQ_TILE * grid.num_rows: + 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 {SEQ_TILE * grid.num_rows}" + f"least {sequence_tile * grid.num_rows}" ) - if embedding_dim % (EMBEDDING_TILE * gemm_split): + if embedding_dim % (embedding_tile * gemm_split): raise ValueError( f"embedding_dim ({embedding_dim}) must be a multiple of " - f"{EMBEDDING_TILE * gemm_split}" + f"{embedding_tile * gemm_split}" ) - if hidden_dim % (HIDDEN_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}" + f"{hidden_tile * gemm_split}" ) @@ -340,6 +369,19 @@ def region_module(mlir_text: str, func_prefix: str = ""): 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 ): @@ -349,7 +391,12 @@ def load_group( ``func_prefix`` is injected by ``OperatorSequence``. Every group loader calls this; the first generates the design and the rest reuse the files on disk. """ - 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 region_module(Path(finals[group_index]).read_text(), func_prefix) + 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) From 51092df83b5b9af6f71fd3a9a7dd43503ab97843 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 30 Jul 2026 00:16:03 +0200 Subject: [PATCH 19/21] swiglu_prefill_stream: document the flow, the k modes and the numbers --- .../operators/swiglu_prefill_stream/README.md | 145 +++++++++++------- 1 file changed, 90 insertions(+), 55 deletions(-) diff --git a/iron/operators/swiglu_prefill_stream/README.md b/iron/operators/swiglu_prefill_stream/README.md index cfa71445..8167f977 100644 --- a/iron/operators/swiglu_prefill_stream/README.md +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -5,82 +5,117 @@ SPDX-License-Identifier: Apache-2.0 # SwiGLU prefill (stream-dse codegen) -This operator is **fused**: the whole SwiGLU-prefill block (both GEMMs + SiLU + -elementwise-mul) is emitted as a **single MLIR design generated by -[`stream-dse`](https://github.com/KULeuven-MICAS/stream)**, then compiled by IRON's normal -flow into one full ELF. Unlike the other operators, its MLIR is not written by hand — it is -produced at build time by [`stream_design.py`](./stream_design.py), which calls the installed -`stream` package (`stream.api.optimize_allocation_co(..., enable_codegen=True)`). +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. -## Where the design comes from +## What IRON provides and what stream-dse returns -Both inputs stream-dse needs are generated **by IRON**, from one source: +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` → [`iron/common/stream/workload.py`](../../common/stream/workload.py) | +| --- | --- | --- | +| 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 | [`iron/common/stream/ops.py`](../../common/stream/ops.py) registry | +| Kernels (`.cc`) | IRON's `aie_kernels` library | the registry in [`iron/common/stream/ops.py`](../../common/stream/ops.py) | -The module that generates the design is the same one the test checks the result -against, so the two cannot drift. Mapping node names are taken from the exported -graph rather than restated, so workload and mapping cannot disagree either. Both -files are written into the experiment's output directory at build time — nothing -is committed to the tree. +`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. -Changing the problem size means re-exporting with different example shapes: the -placement carries tile sizes and core sets but no absolute dimensions, so it is -unaffected (within the divisibility limits `stream_design._check_shapes` enforces). +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. -`k` selects how many fused groups the block is split into, and so how many designs -the ELF holds: `1` keeps the whole block on the array, `2` splits after the -elementwise multiply, and `5` runs layer by layer like `swiglu_prefill` does. The -groups themselves are `stream_design.GROUP_LAYERS`; the kernel tiles every design -is built for are `SEQ_TILE`, `EMBEDDING_TILE` and `HIDDEN_TILE` in the same module. +Nothing crosses the boundary except those files, which is why stream-dse stays an +optional dependency. A checkout without it still imports and tests everything else. -Supporting another op (say a softmax block) is one `StreamKernel` plus one -`TORCH_OPS` entry in `iron/common/stream/ops.py`, pointing at the existing -`aie_kernels//softmax.cc`, plus that operator's own placement. +## Fusion granularity: the `k` modes -## Enabling stream codegen +`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`. -`stream-dse` is an **optional, separately-installed** dependency (it is *not* in IRON's -`requirements.txt`). Install it into the **same environment** as IRON via the extra -requirements file: +| `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. + +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. + +## 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 deps +stream-setup-aie # required: installs stream-dse's AIE codegen dialects ``` -Notes: -- MLIR generation uses the open-source **OR-Tools GSCIP** solver (`backend="ortools_gscip"`), - so **no Gurobi license** is required. -- `stream-setup-aie` is **required**: it installs the pure-Python AIE codegen dialects - stream-dse needs that cannot be plain PyPI dependencies (`xdsl-aie`, `snax-mlir`), since they - are direct git installs. It does not install `mlir_aie` / `llvm-aie`: IRON's - `requirements.txt` already pins those, and stream-dse's codegen only emits text MLIR via - xdsl (it never imports the `aie` bindings). -- Importing the operator does **not** require `stream-dse` (the launcher is imported lazily); - only **building** (`operator.compile()` / running the test) does. +`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 & run +## Build and run ```bash -# build + run on an NPU2 (Strix) device -source /opt/xilinx/xrt/setup.sh # XRT on PATH (provides pyxrt + xclbinutil) +source /opt/xilinx/xrt/setup.sh pytest iron/operators/swiglu_prefill_stream/test.py ``` -The feasible/verified shape is **seq 256 / embedding 512 / hidden 2048**, tiles -**32 / 32 / 64**, target **npu2**. +## Adding another operator + +One `StreamKernel` plus one `TORCH_OPS` entry in `iron/common/stream/ops.py`, pointing +at the existing `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 -## Caveats (stream-dse packaging) +The hardware description (`whole_array_strix.yaml`) is resolved from the installed +`stream` package, where it ships as package data; nothing is vendored here. -- The hardware-description YAML (`whole_array_strix.yaml` + `hardware/cores/*.yaml`) is - resolved from the **installed `stream` package**, where it ships as package data; - nothing is vendored in this operator. -- Nodes and their results are renamed from the ATen names `torch.export` produces - (`matmul`, `matmul_1`, …). Those are prefixes of one another, which stream-dse's - tensor bookkeeping does not distinguish reliably (it raises a `KeyError` while - building the memory-capacity constraints). +Node names are set explicitly rather than taken from the ATen names `torch.export` +produces (`matmul`, `matmul_1`, ...), which are prefixes of one another. From 5e07f426e9233be25162f521552db572ff9d574a Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 30 Jul 2026 00:16:03 +0200 Subject: [PATCH 20/21] ci: require stream-dse>=1.13.11 Picks up the MAC tile the kernel object is compiled for, the offchip traffic term in the allocation objective, and draining only the transfers that carry data out. --- requirements_stream.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements_stream.txt b/requirements_stream.txt index 80e7cd59..092dc823 100644 --- a/requirements_stream.txt +++ b/requirements_stream.txt @@ -19,4 +19,4 @@ # package directory, so that environment must be writable. onnxscript>=0.7 -stream-dse>=1.13.10 +stream-dse>=1.13.11 From 14e890cfdb20c904970008b47687298a8afcd10a Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 30 Jul 2026 00:37:42 +0200 Subject: [PATCH 21/21] swiglu_prefill_stream: time the dispatches that reuse the callable The timed runs each built a callable of their own, so every one of them paid for a fresh hardware context and the reported latency was roughly twice what a dispatch costs. Also trims comments that restated the code beside them. --- iron/common/compilation/sequence.py | 4 +--- iron/common/sequence.py | 7 ++----- iron/common/stream/ops.py | 4 ++-- .../operators/swiglu_prefill_stream/README.md | 19 ++++++++++++++----- iron/operators/swiglu_prefill_stream/op.py | 5 ++--- .../swiglu_prefill_stream/stream_design.py | 12 +++++------- iron/operators/swiglu_prefill_stream/test.py | 18 +++++++----------- 7 files changed, 33 insertions(+), 36 deletions(-) diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index 03b3f2b4..52382b80 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -101,7 +101,7 @@ def needs_additional_reset(runlist: list[Any]) -> bool: 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, as the fused sequence below skips reconfiguring for them. + configure point. """ points = 0 previous = None @@ -193,8 +193,6 @@ def fuse_mlir(artifact: SequenceMLIRArtifact) -> None: dev_op.sym_name = ir.StringAttr.get(op_name) ctx.module.body.append(dev_op) - # An empty device, configured at the end of the sequence to leave the array - # in a state the next dispatch can start from. needs_reset = needs_additional_reset(artifact.runlist) if needs_reset: diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 020ce49d..c4034cad 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -292,8 +292,6 @@ def __init__( # 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 [] - # Build and configure a design once when several operators compile to it. - # Off by default so sequences that do not opt in are unaffected. self.share_designs = share_designs self._dispatch = dispatch @@ -316,9 +314,8 @@ def unique_operators(self): def unique_designs(self): """The designs to build, and which design each operator uses. - Operators are de-duplicated by identity as above, and additionally by - ``design_key`` when ``share_designs`` is set, so a design that several - operators compile to is built and configured once. + 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 = {} diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index 4571ea97..825527ef 100644 --- a/iron/common/stream/ops.py +++ b/iron/common/stream/ops.py @@ -100,8 +100,8 @@ def _gemm_artifacts(base_dir, kernel_dir, m: int, k: int, n: int): f"-DDIM_K={k}", f"-DDIM_N={n}", "-Dbf16_bf16_ONLY", - # The tile operators build mm.cc this way by default: bf16 matmuls - # emulated on the bfp16 MACs, with even rounding on the conversion. + # 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", ], diff --git a/iron/operators/swiglu_prefill_stream/README.md b/iron/operators/swiglu_prefill_stream/README.md index 8167f977..d39b6de1 100644 --- a/iron/operators/swiglu_prefill_stream/README.md +++ b/iron/operators/swiglu_prefill_stream/README.md @@ -30,8 +30,8 @@ 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 stays an -optional dependency. A checkout without it still imports and tests everything else. +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 @@ -52,6 +52,10 @@ afford shrinks as more layers fuse onto it. That is why the tile is chosen per ` (`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. @@ -71,6 +75,10 @@ an idle NPU2. `swiglu_prefill` is the hand-written operator at the same shape. 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`. @@ -108,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 the existing `aie_kernels//.cc`, plus that operator's own placement. The +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. @@ -117,5 +125,6 @@ generated DMAs produce and the layout the compiled object expects come from one 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 rather than taken from the ATen names `torch.export` -produces (`matmul`, `matmul_1`, ...), which are prefixes of one another. +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 index c20aef9b..68de6187 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/iron/operators/swiglu_prefill_stream/op.py @@ -61,9 +61,8 @@ def get_mlir_artifact(self): ) def get_kernel_artifacts(self): - # Each kernel's source, compile flags and symbol names come from the stream - # op registry, so they stay in step with the design stream-dse generates - # against IRON's aie_kernels library. Only this group's layers are built. + # 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 = { diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index 3da68b4f..b4edae17 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/iron/operators/swiglu_prefill_stream/stream_design.py @@ -67,9 +67,8 @@ # 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. The mapping carries the tile and no -# absolute dimension, so it holds across problem sizes; _check_shapes rejects a -# workload it cannot tile evenly. +# 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 @@ -80,9 +79,8 @@ 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 to the -# next is derived from the exported graph, not named here. k=5 is layer by layer, -# every layer its own design, as in :mod:`iron.operators.swiglu_prefill`. +# 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]], @@ -161,7 +159,7 @@ def elementwise(rows, columns, layout, bfp16_mmul=False): 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 keep the layout the GEMM writes. + # 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])), diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index c3a762a9..c0a9deaa 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -33,14 +33,11 @@ TIMED_RUNS = 3 -def _dispatch(operator, golden_ref): - """A callable on a freshly configured array, with its inputs staged. +def _staged(operator, golden_ref): + """A callable with its inputs staged. - Creating the hw_context applies the design's init CDO. The full-ELF path does - not reset array state between dispatches, so every dispatch takes a new - callable: reusing one reads stale state, and a design of several fused groups - deadlocks on the second dispatch. Inputs are named by the golden reference, - and the design consumes the weights in their natural (K, N) layout. + 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): @@ -71,7 +68,7 @@ def test_swiglu_prefill_stream(k, aie_context): # 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 = _dispatch(operator, golden_ref) + run = _staged(operator, golden_ref) run() output = run.get_buffer(OUTPUT).to_torch().reshape((SEQ_LEN, EMBEDDING_DIM)) errors = verify_buffer( @@ -84,11 +81,10 @@ def test_swiglu_prefill_stream(k, aie_context): ) assert not errors, f"Test failed with errors: {errors}" - # Performance: time dispatches of their own, each on a freshly configured - # array, so every timed run is one that computed the verified result. + # The first dispatch on a callable pays for its hardware context, so time the + # ones after it. latencies = [] for _ in range(TIMED_RUNS): - run = _dispatch(operator, golden_ref) start = time.perf_counter() run() latencies.append((time.perf_counter() - start) * 1e6)