From 1fc0d9ef9a04fac381df5bb4e00d0ca811162ddb Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 1/5] 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 b66f3d7697f76d9a9633ccf8022cddb4cfc90014 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 2/5] 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 6fb2417035d3a8e9f6644d1607b0369dda9f2733 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 3/5] 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 85e52084aed8b30f3d52d065395a21914e09abfb Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 15 Jul 2026 23:38:39 +0200 Subject: [PATCH 4/5] 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 0a2151bbd770428753a35436bad1dc354f9310a9 Mon Sep 17 00:00:00 2001 From: asyms Date: Fri, 24 Jul 2026 19:38:32 +0200 Subject: [PATCH 5/5] 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