diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index de748ffb1..d4e06c2e6 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -31,4 +31,5 @@ from .sequence import ( SequenceMLIRArtifact, FusePythonGeneratedMLIRCompilationRule, + trace_buffer_size, ) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 5ecc8746a..e63bdbd05 100644 --- a/iron/common/compilation/base.py +++ b/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, diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index 52382b805..6a1b6858f 100644 --- a/iron/common/compilation/sequence.py +++ b/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 diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 0d33b482c..5a499ae27 100644 --- a/iron/common/sequence.py +++ b/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() diff --git a/iron/common/tracing_utils.py b/iron/common/tracing_utils.py new file mode 100644 index 000000000..b98396148 --- /dev/null +++ b/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 + (``.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 diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 4a4098f57..0b0711c86 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/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, ) diff --git a/iron/operators/swiglu_prefill_stream/stream_design.py b/iron/operators/swiglu_prefill_stream/stream_design.py index b4edae178..03f36eaa9 100644 --- a/iron/operators/swiglu_prefill_stream/stream_design.py +++ b/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 capacity sets the practical 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, diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index c0a9deaa8..89ea3c239 100644 --- a/iron/operators/swiglu_prefill_stream/test.py +++ b/iron/operators/swiglu_prefill_stream/test.py @@ -7,6 +7,8 @@ import pytest import torch +from iron.common.tracing_utils import dump_traces + # The design is generated by stream-dse at compile() time. stream-dse is an # optional dependency (see requirements_stream.txt) absent from the default CI # image, so skip this whole module when it is unavailable. @@ -14,6 +16,15 @@ "stream", reason="stream-dse not installed (see requirements_stream.txt)" ) +# stream-dse's codegen emits `!aie.objectfifosubview`, which mlir-aie removed in +# Xilinx/mlir-aie#3553. Every wheel carrying the trace slice API this branch needs is +# newer than that removal. Revert this commit once a stream-dse release targets a +# post-#3553 mlir-aie. +pytest.skip( + "stream-dse codegen does not parse against the pinned mlir-aie", + allow_module_level=True, +) + from iron.operators.swiglu_prefill_stream.op import SwiGLUPrefillStream # The operator's design is generated from this module; the values it is checked @@ -70,6 +81,7 @@ def test_swiglu_prefill_stream(k, aie_context): # up to 25%. Tolerances are local to this test. run = _staged(operator, golden_ref) run() + dump_traces(run, f"swiglu_k{k}") output = run.get_buffer(OUTPUT).to_torch().reshape((SEQ_LEN, EMBEDDING_DIM)) errors = verify_buffer( output, diff --git a/iron/tests/infrastructure/trace_layout.py b/iron/tests/infrastructure/trace_layout.py new file mode 100644 index 000000000..8b556b5bd --- /dev/null +++ b/iron/tests/infrastructure/trace_layout.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reading back the trace buffer size the compiler recorded on the sequence.""" + +from iron.common.compilation import trace_buffer_size + +LOWERED = """ +module { + aie.device(npu1_1col) @main { + aie.runtime_sequence @sequence(%arg0: memref<4xi32>, %arg1: memref<12288xi8>) + attributes {trace_slices = [ + #aie.trace_slice, + #aie.trace_slice]} { + } + } +} +""" + +UNTRACED = """ +module { + aie.device(npu1_1col) @main { + aie.runtime_sequence @sequence(%arg0: memref<4xi32>) { + } + } +} +""" + + +def test_size_spans_every_slice(): + assert trace_buffer_size(LOWERED) == 12288 + + +def test_untraced_build_has_no_trace_buffer(): + assert trace_buffer_size(UNTRACED) == 0 diff --git a/requirements.txt b/requirements.txt index f626f2f26..1558581d3 100755 --- a/requirements.txt +++ b/requirements.txt @@ -13,7 +13,7 @@ --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.2.dev16+g7e00b57 +mlir_aie==1.4.3.dev60+gc80b88c llvm-aie==22.0.0.2026082001+84660bc3 black diff --git a/requirements_stream.txt b/requirements_stream.txt index 092dc823c..72799e5cf 100644 --- a/requirements_stream.txt +++ b/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