diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index de748ffb1..64d710988 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -31,4 +31,5 @@ from .sequence import ( SequenceMLIRArtifact, FusePythonGeneratedMLIRCompilationRule, + trace_argument_layout, ) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index c6f83f65e..84368c869 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): diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index 52382b805..52cedd929 100644 --- a/iron/common/compilation/sequence.py +++ b/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) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 6826ced28..041632980 100644 --- a/iron/common/sequence.py +++ b/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() 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..3d48b04cb 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, 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, diff --git a/iron/tests/infrastructure/trace_layout.py b/iron/tests/infrastructure/trace_layout.py new file mode 100644 index 000000000..9b658c373 --- /dev/null +++ b/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) 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