Add stream-dse fused SwiGLU-prefill operator

.github/actions/prereqs/action.yaml

@@ -21,4 +21,11 @@ runs:
         source ${{ inputs.env_name }}/bin/activate
         pip install --upgrade pip
         pip install -r requirements.txt
+        pip install -r requirements_stream.txt
+        # stream-dse's pure-Python AIE codegen dialects (xdsl-aie, snax-mlir) are
+        # git installs, not PyPI dependencies; this console script installs them so
+        # the stream-backed operators can generate their MLIR at build time. It does
+        # NOT install mlir_aie/llvm-aie -- those stay pinned by requirements.txt above
+        # (stream-dse>=1.13.7; older releases reinstalled/downgraded them here).
+        stream-setup-aie
         echo "Prerequisites installed into ${{ inputs.env_name }}"

.gitignore

@@ -22,3 +22,4 @@ id_ed25519.pub
 .cline_storage
 *.egg-info
 **/*.prj/**
+/outputs/

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/*

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}"

iron/common/__init__.py

@@ -18,3 +18,4 @@
     PythonGeneratedMLIRArtifact,
     DesignGenerator,
 )
+from .layout import Stride, TiledStride, TiledStridedLayout, tiled_2d

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(

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))),
+        )
+    )

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

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

iron/operators/swiglu_prefill_stream/README.md

@@ -0,0 +1,55 @@
+<!--
+SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
+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 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 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.
+
+## 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.