Wire the AIE trace buffer through the fused full ELF flow

iron/common/compilation/__init__.py

@@ -31,4 +31,5 @@
 from .sequence import (
     SequenceMLIRArtifact,
     FusePythonGeneratedMLIRCompilationRule,
+    trace_argument_layout,
 )

iron/common/compilation/base.py

@@ -320,11 +320,14 @@ def __init__(
         mlir_input: CompilationArtifact,
         dependencies: list[CompilationArtifact],
         extra_flags: list[str] | None = None,
+        trace_size: int = 0,
     ) -> 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 []
+        # Bytes of trace buffer per runlist step, 0 for an untraced build.
+        self.trace_size = trace_size
 
 
 class XclbinArtifact(_MLIRInputMixin, CompilationArtifact):

iron/common/compilation/sequence.py

@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+from itertools import count, islice
+
 import numpy as np
 import importlib.util
 from functools import partial
@@ -34,6 +36,27 @@
 # ##########################################################################
 
 
+def trace_argument_layout(arg_counts: dict[str, int], trace_size: int):
+    """Buffer slots for the fused runtime sequence, as (consolidated, trace, count).
+
+    Lowering patches a trace address against the dispatched kernel, not the callee, so
+    each operator needs its buffer at the index it uses. The rest take what is left.
+    """
+    if not trace_size:
+        return [0, 1, 2], {}, 3
+    trace_slots = dict(arg_counts)
+    counts = list(trace_slots.values())
+    shared = sorted({n for n in counts if counts.count(n) > 1})
+    if shared:
+        raise NotImplementedError(
+            "operators taking the same number of arguments would share one trace "
+            f"buffer (slots {shared}); trace them in separate dispatches"
+        )
+    trace_indices = sorted(set(counts))
+    consolidated_idx = list(islice((i for i in count() if i not in trace_indices), 3))
+    return consolidated_idx, trace_slots, max(trace_indices + consolidated_idx) + 1
+
+
 class SequenceMLIRArtifact(MLIRArtifact):
     def __init__(
         self,
@@ -43,6 +66,7 @@ def __init__(
         subbuffer_layout: dict[str, tuple[str, int, int]],
         buffer_sizes: tuple[int, int, int],
         slice_info: dict[str, tuple[str, int, int]] | None = None,
+        trace_size: int = 0,
     ) -> None:
         dependencies = list(operator_mlir_map.values())
         super().__init__(filename, dependencies)
@@ -51,6 +75,8 @@ def __init__(
         self.subbuffer_layout = subbuffer_layout
         self.buffer_sizes = buffer_sizes
         self.slice_info = slice_info or {}
+        # Bytes of trace buffer per runlist step, 0 for an untraced build.
+        self.trace_size = trace_size
 
 
 # Helper Functions
@@ -213,12 +239,33 @@ def main():
             itemsize = np.dtype(ml_dtypes.bfloat16).itemsize
 
             # RuntimeSequenceOp
-            @aiex.runtime_sequence(
-                np.ndarray[(input_buffer_size // itemsize,), buf_dtype],
-                np.ndarray[(output_buffer_size // itemsize,), buf_dtype],
-                np.ndarray[(scratch_buffer_size // itemsize,), buf_dtype],
+            trace_size = artifact.trace_size
+            consolidated_idx, trace_slots, n_args = trace_argument_layout(
+                {name: len(sequence_arg_types[name]) for name, *_ in artifact.runlist},
+                trace_size,
             )
-            def sequence(input_buf, output_buf, scratch_buf):
+            trace_indices = sorted(set(trace_slots.values()))
+
+            sizes = dict(
+                zip(
+                    consolidated_idx,
+                    (input_buffer_size, output_buffer_size, scratch_buffer_size),
+                )
+            )
+            arg_types = [
+                (
+                    np.ndarray[(max(1, trace_size),), np.dtype[np.int8]]
+                    if i in trace_indices
+                    else np.ndarray[(max(1, sizes.get(i, 0) // itemsize),), buf_dtype]
+                )
+                for i in range(n_args)
+            ]
+
+            @aiex.runtime_sequence(*arg_types)
+            def sequence(*all_bufs):
+                input_buf, output_buf, scratch_buf = (
+                    all_bufs[i] for i in consolidated_idx
+                )
                 consolidated_buffers = {
                     "input": input_buf,
                     "output": output_buf,
@@ -304,6 +351,10 @@ def sequence(input_buf, output_buf, scratch_buf):
                             )
                             buffer_ssa_values.append(reinterpreted)
 
+                        # Trace lowering appends a buffer to the callee's signature.
+                        if trace_size:
+                            buffer_ssa_values.append(all_bufs[trace_slots[op_name]])
+
                         # Run Op
                         sequence_sym_ref_attr = ir.FlatSymbolRefAttr.get("sequence")
                         run_op = aiex.RunOp(sequence_sym_ref_attr, buffer_ssa_values)

iron/common/sequence.py

@@ -85,6 +85,7 @@ def set_up_artifacts(self, seq):
             mlir_input=mlir_artifact,
             dependencies=[mlir_artifact] + kernel_objects,
             extra_flags=seq.extra_flags,
+            trace_size=seq.trace_size,
         )
         seq.add_artifacts([full_elf_artifact])
 
@@ -118,6 +119,7 @@ def build_fused_mlir(self, seq):
             subbuffer_layout=seq.subbuffer_layout,
             buffer_sizes=seq.buffer_sizes,
             slice_info=seq.slice_info,
+            trace_size=seq.trace_size,
         )
 
     def _collect_kernel_artifacts(self, seq):
@@ -264,6 +266,7 @@ def __init__(
         buffer_sizes=None,
         dispatch="auto",
         extra_flags=None,
+        trace_size=0,
         share_designs=False,
         *args,
         **kwargs,
@@ -290,6 +293,8 @@ def __init__(
         # Extra aiecc flags forwarded to the full-ELF build. Empty by default, so
         # other sequences are unaffected.
         self.extra_flags = extra_flags or []
+        # Bytes of hardware trace buffer per runlist step; 0 leaves the design untraced.
+        self.trace_size = trace_size
         self.share_designs = share_designs
         self._dispatch = dispatch
 
@@ -559,9 +564,20 @@ def __init__(self, op, device_name="main", sequence_name="sequence"):
         # ctrl-scratchpad backing buffer (and any ParameterScratchpad state
         # built on top of it) stays valid across calls.
         self.run_handle = pyxrt.run(self.xrt_kernel)
-        self.run_handle.set_arg(0, self.input_buffer.buffer_object())
-        self.run_handle.set_arg(1, self.output_buffer.buffer_object())
-        self.run_handle.set_arg(2, self.scratch_buffer.buffer_object())
+        consolidated_idx, trace_slots, _ = comp.trace_argument_layout(
+            {
+                f"op{i}_{o.__class__.__name__}": len(o.get_arg_spec())
+                for i, (o, *_) in enumerate(self.op.runlist)
+            },
+            self.op.trace_size,
+        )
+        for idx, buf in zip(
+            consolidated_idx,
+            (self.input_buffer, self.output_buffer, self.scratch_buffer),
+        ):
+            self.run_handle.set_arg(idx, buf.buffer_object())
+        for name, idx in trace_slots.items():
+            self.run_handle.set_arg(idx, self.trace_buffers[name].buffer_object())
 
         self._params = None
 
@@ -597,6 +613,15 @@ def _allocate_buffers(self):
         self.scratch_buffer = XRTTensor(
             (_n_elements(scratch_sz),), dtype=ml_dtypes.bfloat16
         )
+        trace_size = self.op.trace_size
+        self.trace_buffers = (
+            {
+                f"op{i}_{o.__class__.__name__}": XRTTensor((trace_size,), dtype=np.int8)
+                for i, (o, *_) in enumerate(self.op.runlist)
+            }
+            if trace_size
+            else {}
+        )
 
     def get_buffer(self, buffer_name):
         if buffer_name in self._buffer_cache:
@@ -624,6 +649,9 @@ def _sync_outputs(self):
         # range "cpu" (otherwise a looped dispatch would read stale output).
         self.output_buffer.device = "npu"
         self.output_buffer.to("cpu")
+        for buf in self.trace_buffers.values():
+            buf.device = "npu"
+            buf.to("cpu")
 
     def _run(self):
         self.run_handle.start()

iron/operators/swiglu_prefill_stream/op.py

@@ -148,6 +148,8 @@ class SwiGLUPrefillStream(OperatorSequence):
     def __init__(
         self, seq_len, embedding_dim, hidden_dim, k=1, context=None, share_designs=True
     ):
+        from iron.operators.swiglu_prefill_stream.stream_design import trace_size
+
         ports, inputs, outputs = _wiring(seq_len, embedding_dim, hidden_dim, k)
         groups = [
             _SwiGLUStreamGroup(
@@ -167,6 +169,7 @@ def __init__(
             ],
             input_args=inputs,
             output_args=outputs,
+            trace_size=trace_size(),
             share_designs=share_designs,
             context=context,
         )

iron/operators/swiglu_prefill_stream/stream_design.py

@@ -277,12 +277,27 @@ 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 ""
+    if trace_size():
+        suffix += "_traced"
     return (
         f"{hardware}-swiglu{suffix}_{seq_len}_{embedding_dim}_{hidden_dim}"
         f"-{grid.num_rows}_row_{grid.num_columns}_col"
     )
 
 
+def trace_size():
+    """DDR trace buffer in bytes, 0 for an untraced build.
+
+    Opt-in: tracing adds a runtime-sequence argument, so it changes the ABI.
+    """
+    return int(os.environ.get("IRON_TRACE_SIZE", "0"))
+
+
+def trace_tiles():
+    """How many tiles to trace. Routing, not the packet id space, is the real limit."""
+    return int(os.environ.get("IRON_TRACE_NTILES", "4"))
+
+
 def _design_paths(seq_len, embedding_dim, hidden_dim, k):
     """Where stream-dse writes each group's MLIR.
 
@@ -319,7 +334,8 @@ def _run_codegen(seq_len, embedding_dim, hidden_dim, npu, k):
         output_path=OUTPUT_ROOT,
         skip_if_exists=False,
         enable_codegen=True,
-        trace_size=0,
+        trace_size=trace_size(),
+        trace_max_tiles=trace_tiles(),
         nb_cols_to_use=grid.num_columns,
         npu=npu,
         backend=BACKEND,

iron/tests/infrastructure/trace_layout.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Buffer slots for a traced fused sequence. Wrong indices hang the device silently."""
+
+import pytest
+
+from iron.common.compilation import trace_argument_layout
+
+
+def test_untraced_keeps_the_three_consolidated_buffers_first():
+    assert trace_argument_layout({"a": 5}, 0) == ([0, 1, 2], {}, 3)
+
+
+def test_trace_buffer_lands_at_the_operator_argument_count():
+    consolidated, slots, n_args = trace_argument_layout({"a": 5}, 65536)
+    assert slots == {"a": 5}
+    assert consolidated == [0, 1, 2]
+    assert n_args == 6
+
+
+def test_consolidated_buffers_move_aside_for_a_low_trace_slot():
+    # An operator taking two arguments wants slot 2, which the scratch buffer would
+    # otherwise hold.
+    consolidated, slots, n_args = trace_argument_layout({"a": 2, "b": 4}, 65536)
+    assert slots == {"a": 2, "b": 4}
+    assert not set(consolidated) & set(slots.values())
+    assert consolidated == [0, 1, 3]
+    assert n_args == 5
+
+
+def test_every_operator_gets_its_own_slot():
+    _, slots, _ = trace_argument_layout({"a": 3, "b": 4, "c": 5}, 65536)
+    assert sorted(slots.values()) == [3, 4, 5]
+
+
+def test_operators_sharing_an_argument_count_are_refused():
+    with pytest.raises(NotImplementedError, match=r"slots \[3\]"):
+        trace_argument_layout({"a": 3, "b": 3, "c": 2}, 65536)

requirements_stream.txt

@@ -4,9 +4,10 @@
 # Optional dependencies for the stream-dse-backed fused SwiGLU-prefill operator
 # (iron/operators/swiglu_prefill_stream).
 #
-# Not installed by the default CI (requirements.txt); the operator's test skips
-# itself (pytest.importorskip) when stream-dse is absent. Install this file to
-# build and run the operator and its test:
+# Kept out of requirements.txt so an install without stream-dse still works: the
+# operator's test skips itself (pytest.importorskip) when it is absent. CI does
+# install this file (.github/actions/prereqs), so the operator runs there. 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
@@ -19,4 +20,4 @@
 # package directory, so that environment must be writable.
 
 onnxscript>=0.7
-stream-dse>=1.13.11
+stream-dse>=1.13.14