Enable trace in OperatorSequence

iron/common/compilation/__init__.py

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

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):
@@ -544,6 +547,10 @@ def compile(self, graph):
                 "--expand-load-pdis",
                 "--get-scratchpad-parameters",
             ] + artifact.extra_flags
+            if artifact.trace_size:
+                # The trace parser reads the lowered module for the buffer layout
+                # and each design's traced tiles and events.
+                options.append("--get-input-with-addresses")
 
             def _compile(
                 artifact=artifact,

iron/common/compilation/sequence.py

@@ -14,6 +14,7 @@
 from aie import ir
 from aie.dialects import aie, aiex, memref
 from aie.extras.context import mlir_mod_ctx
+from aie.utils.trace import get_trace_slices
 import ml_dtypes
 
 from typing import Any
@@ -34,6 +35,17 @@
 # ##########################################################################
 
 
+def trace_buffer_size(mlir_text: str) -> int:
+    """Bytes of the fused trace buffer the dispatched sequence takes.
+
+    `-aie-fuse-trace-buffers` gives the sequence one buffer covering every design
+    it configures, and records the split on the sequence. Returns 0 for an
+    untraced build.
+    """
+    slices = get_trace_slices(mlir_text)
+    return max((s["offset"] + s["size"] for s in slices), default=0)
+
+
 class SequenceMLIRArtifact(MLIRArtifact):
     def __init__(
         self,
@@ -43,6 +55,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 +64,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

iron/common/sequence.py

@@ -65,6 +65,12 @@ def resolve(self, device):
         return SeparateDispatch()
 
 
+def _trace_tag(seq):
+    """Tracing adds a runtime-sequence argument, so a traced build cannot reuse an
+    untraced one's ELF. Empty when untraced."""
+    return f"_traced{seq.trace_size}" if seq.trace_size else ""
+
+
 class FusedDispatch(SequenceDispatch):
     """Single-ELF dispatch (NPU2 only): all operators fused into one ELF."""
 
@@ -81,10 +87,11 @@ def set_up_artifacts(self, seq):
         mlir_artifact = self.build_fused_mlir(seq)
         kernel_objects = self._collect_kernel_artifacts(seq)
         full_elf_artifact = comp.FullElfArtifact(
-            f"{seq.name}.elf",
+            f"{seq.name}{_trace_tag(seq)}.elf",
             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])
 
@@ -112,12 +119,13 @@ def build_fused_mlir(self, seq):
             comp_runlist.append((design_names[design_of[id(op)]], *bufs))
 
         return comp.SequenceMLIRArtifact(
-            seq.name + "_fused.mlir",
+            f"{seq.name}{_trace_tag(seq)}_fused.mlir",
             operator_mlir_map=operator_mlir_map,
             runlist=comp_runlist,
             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 +272,7 @@ def __init__(
         buffer_sizes=None,
         dispatch="auto",
         extra_flags=None,
+        trace_size=0,
         share_designs=False,
         *args,
         **kwargs,
@@ -289,6 +298,8 @@ def __init__(
         )  # Optional dict: buffer_name -> size_in_bytes
         # Extra aiecc flags forwarded to the full-ELF build.
         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
 
@@ -561,6 +572,8 @@ def __init__(self, op, device_name="main", sequence_name="sequence"):
         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())
+        if self.trace_buffer is not None:
+            self.run_handle.set_arg(3, self.trace_buffer.buffer_object())
 
         self._params = None
 
@@ -598,6 +611,20 @@ def _allocate_buffers(self):
         self.scratch_buffer = XRTTensor(
             (_n_elements(scratch_sz),), dtype=ml_dtypes.bfloat16
         )
+        # Trace lowering appends one buffer covering every configured design, after
+        # the consolidated three. Its size depends on how many channels and
+        # sub-designs claim a share, so read it from the lowered module.
+        self.trace_buffer = None
+        if self.op.trace_size:
+            total = comp.trace_buffer_size(self.lowered_mlir_text())
+            if total:
+                self.trace_buffer = XRTTensor((total,), dtype=np.int8)
+
+    def lowered_mlir_text(self) -> str:
+        """aiecc's post-lowering module, which carries the trace buffer layout."""
+        mlir_filename = self.op.artifacts[0].mlir_input.filename
+        path = comp._aiecc_work_dir(mlir_filename) / "input_with_addresses.mlir"
+        return path.read_text()
 
     def get_buffer(self, buffer_name):
         if buffer_name in self._buffer_cache:
@@ -625,6 +652,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")
+        if self.trace_buffer is not None:
+            self.trace_buffer.device = "npu"
+            self.trace_buffer.to("cpu")
 
     def _run(self):
         self.run_handle.start()

iron/common/tracing_utils.py

@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Write a traced run's hardware trace buffer as Perfetto JSON.
+
+Tracing is configured at build time (``IRON_TRACE_SIZE`` / ``IRON_TRACE_NTILES``,
+read by the operator's design), and the runtime syncs the buffer device->host after
+every dispatch. Call :func:`dump_traces` after ``run()`` to write it out:
+
+    from iron.common.tracing_utils import dump_traces
+
+    run = operator.get_callable()
+    run()
+    dump_traces(run, "my_operator")
+
+On an untraced build the call returns an empty list, so a test can call it
+unconditionally.
+
+A dump writes the raw 32-bit words as hex text, plus one JSON file per traced
+design for https://ui.perfetto.dev. Keep the text: :func:`parse_trace_buffer`
+reparses it with a different column shift for the price of no further dispatch.
+
+:func:`dump_traces` also prints mlir-aie's per-tile cycles summary for each file it
+writes.
+
+Environment:
+  * ``IRON_TRACE_DIR``      where to write (default ``outputs/traces``)
+  * ``IRON_TRACE_MLIR``     override the MLIR the parser reads
+  * ``IRON_TRACE_COLSHIFT`` force the column shift; unset means auto-detect
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+
+from aie.utils.trace import parse_trace_slices, print_cycles_summary
+
+from . import compilation as comp
+
+__all__ = [
+    "dump_traces",
+    "parse_trace_buffer",
+    "lowered_mlir",
+]
+
+DEFAULT_TRACE_DIR = "outputs/traces"
+
+
+def lowered_mlir(run) -> tuple[Path, str]:
+    """The post-lowering MLIR for a callable, as ``(path, text)``.
+
+    mlir-aie's trace parser matches ``aiex.npu.write32`` ops against the trace unit's
+    config addresses. ``aie-insert-trace-flows`` emits those writes inside aiecc, so
+    the parser needs aiecc's lowered module. A traced build requests it with
+    ``--get-input-with-addresses``, which lands it in the work dir beside the source
+    (``<source>.mlir.d/``).
+    """
+    override = os.environ.get("IRON_TRACE_MLIR")
+    if override:
+        path = Path(override)
+        return path, path.read_text()
+
+    source = Path(run.op.artifacts[0].mlir_input.filename)
+    path = comp._aiecc_work_dir(str(source)) / "input_with_addresses.mlir"
+    if not path.exists():
+        raise FileNotFoundError(
+            f"{path} is missing; a traced build passes --get-input-with-addresses "
+            "to aiecc. Point IRON_TRACE_MLIR at a lowered module to override."
+        )
+    return path, path.read_text()
+
+
+def parse_trace_buffer(words, mlir_text: str, colshift: int | None = None):
+    """A trace buffer's words as ``(slice_info, events)`` per traced design.
+
+    The parser splits the buffer by the layout the compiler recorded on the
+    dispatched sequence, and decodes each region against the device that wrote it.
+
+    ``colshift`` of None lets the parser align the columns itself, which is what you
+    want by default: a design configured for one column may be loaded into another.
+    Override it when that alignment picks the wrong columns.
+
+    The parser calls ``sys.exit`` on some malformed input, so SystemExit becomes a
+    RuntimeError here: a visualisation failure must not fail a test.
+    """
+    try:
+        return parse_trace_slices(
+            np.asarray(words, dtype=np.uint32), mlir_text, colshift
+        )
+    except SystemExit as exc:
+        raise RuntimeError(
+            "mlir-aie's trace parser exited; the usual cause is an MLIR without the "
+            "trace register writes, or a column shift that does not match the data. "
+            "Run with logging at DEBUG to see the tiles it found."
+        ) from exc
+
+
+def _slug(text: str) -> str:
+    keep = "-_."
+    return "".join(c if c.isalnum() or c in keep else "_" for c in text)
+
+
+def dump_traces(
+    run,
+    tag: str,
+    out_dir=None,
+    colshift: int | None = None,
+    summary: bool = True,
+) -> list[Path]:
+    """Write a completed run's trace buffer as hex text and Perfetto JSON.
+
+    Call it after ``run()``: the callable syncs its trace buffer device->host as part
+    of the dispatch, so this only reads host memory. Returns the JSON paths written,
+    empty on an untraced build.
+
+    ``tag`` distinguishes one dump from another - a test name or parameter id. The
+    layout the compiler recorded on the dispatched sequence splits the buffer, so a
+    fused sequence yields one JSON file per configured design.
+    """
+    buffer = getattr(run, "trace_buffer", None)
+    if buffer is None:
+        if getattr(getattr(run, "op", None), "trace_size", 0):
+            raise TypeError(
+                f"{type(run).__name__} was built with tracing enabled but exposes no "
+                "trace_buffer; only the full-ELF sequence callable allocates one."
+            )
+        return []
+
+    out_dir = Path(out_dir or os.environ.get("IRON_TRACE_DIR", DEFAULT_TRACE_DIR))
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    if colshift is None:
+        env = os.environ.get("IRON_TRACE_COLSHIFT")
+        colshift = int(env) if env else None
+
+    mlir_path, mlir_text = lowered_mlir(run)
+    print(f"[trace] parsing against {mlir_path}")
+
+    words = buffer.to_torch().numpy().astype(np.uint8).view(np.uint32)
+    tag = _slug(tag)
+    raw = (out_dir / tag).with_suffix(".txt")
+    raw.write_text("\n".join(f"{w:08x}" for w in words) + "\n")
+    if not words.any():
+        print("[trace] buffer is all zeros, no trace data captured")
+        return []
+
+    try:
+        parsed = parse_trace_buffer(words, mlir_text, colshift)
+    except Exception as exc:  # a visualisation failure must not fail a run
+        print(f"[trace] parse failed ({exc}); raw words kept at {raw}")
+        return []
+
+    written = []
+    for index, (entry, events) in enumerate(parsed):
+        # A device may hold several runtime sequences, so both names identify a slice.
+        name = f"{index}_{entry['device']}_{entry['sequence']}" if entry else "trace"
+        if entry and words[(entry["offset"] + entry["size"]) // 4 - 1]:
+            print(
+                f"[trace] {name}: slice full ({entry['size']} B), trace is likely "
+                "truncated - raise IRON_TRACE_SIZE"
+            )
+
+        target = (out_dir / f"{tag}_{_slug(name)}").with_suffix(".json")
+        target.write_text(json.dumps(events))
+        print(f"[trace] {target} ({len(events)} events)")
+        written.append(target)
+
+        if summary:
+            print_cycles_summary(target)
+    return written

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