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

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/compilation/sequence.py

@@ -27,6 +27,9 @@
     MLIRArtifact,
 )
 
+RESET_DEVICE = "reset_device"
+
+
 # Compilation Artifacts
 # ##########################################################################
 
@@ -89,6 +92,26 @@ def get_child_mlir_module(mlir_artifact: PythonGeneratedMLIRArtifact) -> Any:
     return callback_function(*gen.args, **gen.kwargs)
 
 
+def needs_additional_reset(runlist: list[Any]) -> bool:
+    """Whether the sequence must configure one more device than the runlist asks for.
+
+    ``aiecc --expand-load-pdis`` marks each configure point by loading one of two
+    otherwise empty PDIs, alternating between them from a fixed start. A load of the
+    PDI already loaded has no effect, so a sequence with an odd number of configure
+    points ends on the one the next dispatch starts with, and that dispatch
+    reconfigures over the state the last design left. Configuring one more device
+    makes the count even. Consecutive entries running the same operator share a
+    configure point.
+    """
+    points = 0
+    previous = None
+    for op_name, *_ in runlist:
+        if op_name != previous:
+            points += 1
+            previous = op_name
+    return points % 2 == 1
+
+
 def fuse_mlir(artifact: SequenceMLIRArtifact) -> None:
     """Fuse multiple MLIR modules by inlining their device operations and adding a new main device and runtime sequence that call into sequence of operations based on a runlist."""
 
@@ -170,6 +193,17 @@ def fuse_mlir(artifact: SequenceMLIRArtifact) -> None:
             dev_op.sym_name = ir.StringAttr.get(op_name)
             ctx.module.body.append(dev_op)
 
+        needs_reset = needs_additional_reset(artifact.runlist)
+        if needs_reset:
+
+            @aie.device(device_ty)
+            def reset():
+                @aiex.runtime_sequence()
+                def sequence():
+                    pass
+
+            reset.operation.attributes["sym_name"] = ir.StringAttr.get(RESET_DEVICE)
+
         # Create the main device -- this contains the runtime sequence calling into the other devices
         @aie.device(device_ty)
         def main():
@@ -274,6 +308,10 @@ def sequence(input_buf, output_buf, scratch_buf):
                         sequence_sym_ref_attr = ir.FlatSymbolRefAttr.get("sequence")
                         run_op = aiex.RunOp(sequence_sym_ref_attr, buffer_ssa_values)
 
+                if needs_reset:
+                    reset_op = aiex.ConfigureOp(ir.FlatSymbolRefAttr.get(RESET_DEVICE))
+                    reset_op.body.blocks.append()
+
         # Write the fused MLIR to file
         with open(artifact.filename, "w") as f:
             f.write(str(ctx.module))

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])
 
@@ -97,18 +98,19 @@ def build_fused_mlir(self, seq):
         """
         operator_mlir_map = {}
         comp_runlist = []
-        op_names = {}  # id(op) -> op_name
+        designs, design_of = seq.unique_designs()
+        design_names = []
 
-        for idx, op in enumerate(seq.unique_operators()):
+        for idx, op in enumerate(designs):
             mlir_artifact = op.get_mlir_artifact()
             if len(op.get_kernel_artifacts()) > 0:
                 mlir_artifact.generator.kwargs["func_prefix"] = f"op{idx}_"
             op_name = f"op{idx}_{op.__class__.__name__}"
-            op_names[id(op)] = op_name
+            design_names.append(op_name)
             operator_mlir_map[op_name] = mlir_artifact
 
         for op, *bufs in seq.runlist:
-            comp_runlist.append((op_names[id(op)], *bufs))
+            comp_runlist.append((design_names[design_of[id(op)]], *bufs))
 
         return comp.SequenceMLIRArtifact(
             seq.name + "_fused.mlir",
@@ -122,7 +124,7 @@ def build_fused_mlir(self, seq):
     def _collect_kernel_artifacts(self, seq):
         """Kernel artifacts from all child operators, prefixed per operator index."""
         kernel_artifacts = []
-        for idx, op in enumerate(seq.unique_operators()):
+        for idx, op in enumerate(seq.unique_designs()[0]):
             objs = op.get_kernel_artifacts()
             for obj in objs:
                 obj.filename = f"op{idx}_{obj.filename}"
@@ -262,6 +264,8 @@ def __init__(
         output_args,
         buffer_sizes=None,
         dispatch="auto",
+        extra_flags=None,
+        share_designs=False,
         *args,
         **kwargs,
     ):
@@ -276,12 +280,19 @@ def __init__(
             )
         super().__init__(*args, **kwargs)
         self.runlist = runlist
-        self.name = name
+        # Sharing changes which designs are built, so it belongs in the name that
+        # keys the build artifacts.
+        self.name = name + "_shared" if share_designs else name
         self.input_args = input_args
         self.output_args = output_args
         self.explicit_buffer_sizes = (
             buffer_sizes or {}
         )  # Optional dict: buffer_name -> size_in_bytes
+        # Extra aiecc flags forwarded to the full-ELF build (e.g. --dynamic-objFifos
+        # for placed/routed whole-array designs that would otherwise overflow AIE2p
+        # program memory). Empty by default, so other sequences are unaffected.
+        self.extra_flags = extra_flags or []
+        self.share_designs = share_designs
         self._dispatch = dispatch
 
     @staticmethod
@@ -300,6 +311,32 @@ def unique_operators(self):
             seen.setdefault(id(op), op)
         return list(seen.values())
 
+    def unique_designs(self):
+        """The designs to build, and which design each operator uses.
+
+        With ``share_designs`` set, operators reporting the same ``design_key``
+        collapse onto one design, so it is built, prefixed and configured once.
+        """
+        designs = []
+        design_of = {}
+        first_with_key = {}
+        for op in self.unique_operators():
+            key = op.design_key() if self.share_designs else None
+            if key is not None and key in first_with_key:
+                shared = designs[first_with_key[key]]
+                if op.get_arg_spec() != shared.get_arg_spec():
+                    raise ValueError(
+                        f"{op.name} and {shared.name} report the same design_key but "
+                        "different runtime arguments, so the design cannot be shared"
+                    )
+                design_of[id(op)] = first_with_key[key]
+                continue
+            if key is not None:
+                first_with_key[key] = len(designs)
+            design_of[id(op)] = len(designs)
+            designs.append(op)
+        return designs, design_of
+
     def calculate_buffer_layout(self):
         args = {}  # base_buffer_name -> args_spec
         sliced_buffers = (

iron/common/stream/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Building blocks for stream-dse-backed operators.
+
+An operator supplies a reference ``nn.Module`` and a placement; these modules turn
+that into everything stream-dse needs:
+
+* :mod:`~iron.common.stream.ops` -- the registry binding a torch ATen op to its ONNX
+  form, its stream-dse kernel and IRON's ``aie_kernels`` source.
+* :mod:`~iron.common.stream.workload` -- ``torch.export`` of the module into the ONNX
+  workload stream-dse optimizes.
+* :mod:`~iron.common.stream.mapping` -- the mapping YAML, named from that same graph.
+
+The submodules are not re-exported here: they need ``onnx``/``pyyaml`` (installed
+with stream-dse, see ``requirements_stream.txt``), so importing an operator must not
+pull them in. Import them directly from the module that builds the design.
+"""