Enable trace in OperatorSequence

iron/common/compilation/__init__.py

@@ -31,4 +31,5 @@
 from .sequence import (
     SequenceMLIRArtifact,
     FusePythonGeneratedMLIRCompilationRule,
+    trace_buffer_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):
@@ -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,22 @@
 # ##########################################################################
 
 
+def trace_buffer_layout(mlir_text: str):
+    """Regions of the fused trace buffer, one per traced operator.
+
+    `-aie-fuse-trace-buffers` gives the dispatched sequence one trace buffer
+    covering every design it configures, and records the split on the sequence.
+    The host reads it for the buffer's size, the parser for which design wrote
+    which bytes.
+
+    Returns `(total_bytes, slices)`; `(0, [])` for an untraced build.
+    """
+    slices = get_trace_slices(mlir_text)
+    if not slices:
+        return 0, []
+    return max(s["offset"] + s["size"] for s in slices), slices
+
+
 class SequenceMLIRArtifact(MLIRArtifact):
     def __init__(
         self,
@@ -43,6 +60,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 +69,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, leaving those artifacts named as before."""
+    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,24 @@ 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. trace_size is per design and says
+        # nothing about how many channels or sub-designs claim a share, so the
+        # size comes from the lowered module.
+        self.trace_buffer = None
+        self.trace_slices = []
+        if self.op.trace_size:
+            total, self.trace_slices = comp.trace_buffer_layout(
+                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 +656,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()