From 1478cbe22c9e19e555b71278db7ff787643368a3 Mon Sep 17 00:00:00 2001 From: asyms Date: Mon, 10 Aug 2026 19:19:57 +0200 Subject: [PATCH 1/9] require the stream-dse release that generates mlir-aie 1.4.0 buffer descriptors --- requirements_stream.txt | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/requirements_stream.txt b/requirements_stream.txt index 092dc823c..b0a8edc76 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.12 From aaa07b9509eac1b2a38c03423bc567fb3db0fa16 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 20 Aug 2026 10:37:45 +0200 Subject: [PATCH 2/9] carry a hardware trace buffer through the fused sequence --- =0.7 | 0 iron/common/compilation/base.py | 3 ++ iron/common/compilation/sequence.py | 27 +++++++++++++++++- iron/common/sequence.py | 16 +++++++++++ .../swiglu_prefill_stream/.README.md.swp | Bin 0 -> 16384 bytes iron/operators/swiglu_prefill_stream/op.py | 7 +++++ .../swiglu_prefill_stream/stream_design.py | 19 +++++++++++- 7 files changed, 70 insertions(+), 2 deletions(-) create mode 100644 =0.7 create mode 100644 iron/operators/swiglu_prefill_stream/.README.md.swp diff --git a/=0.7 b/=0.7 new file mode 100644 index 000000000..e69de29bb diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 5ecc8746a..4905b9965 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..dcca05aff 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -43,6 +43,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 +52,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 +216,20 @@ def main(): itemsize = np.dtype(ml_dtypes.bfloat16).itemsize # RuntimeSequenceOp + trace_size = getattr(artifact, "trace_size", 0) + n_traced = len(artifact.runlist) if trace_size else 0 + @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], + *( + [np.ndarray[(max(1, n_traced * trace_size),), np.dtype[np.int8]]] + if trace_size + else [] + ), ) - def sequence(input_buf, output_buf, scratch_buf): + def sequence(input_buf, output_buf, scratch_buf, *trace_bufs): consolidated_buffers = { "input": input_buf, "output": output_buf, @@ -228,6 +239,7 @@ def sequence(input_buf, output_buf, scratch_buf): # Execute operations in runlist order configure_op = None last_op_name = None + run_index = 0 for op_name, *buffer_names in artifact.runlist: expected_arg_types = sequence_arg_types[op_name] @@ -304,9 +316,22 @@ def sequence(input_buf, output_buf, scratch_buf): ) buffer_ssa_values.append(reinterpreted) + # Trace lowering appends a buffer to the callee, so the call + # has to carry one too. Each op writes its own slice. + if trace_size: + buffer_ssa_values.append( + memref.subview( + trace_bufs[0], + [run_index * trace_size], + [trace_size], + [1], + ) + ) + # Run Op sequence_sym_ref_attr = ir.FlatSymbolRefAttr.get("sequence") run_op = aiex.RunOp(sequence_sym_ref_attr, buffer_ssa_values) + run_index += 1 if needs_reset: reset_op = aiex.ConfigureOp(ir.FlatSymbolRefAttr.get(RESET_DEVICE)) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 0d33b482c..92a13c4b7 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, @@ -289,6 +292,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 +566,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 +605,12 @@ def _allocate_buffers(self): self.scratch_buffer = XRTTensor( (_n_elements(scratch_sz),), dtype=ml_dtypes.bfloat16 ) + trace_size = self.op.trace_size + self.trace_buffer = ( + XRTTensor((max(1, len(self.op.runlist) * trace_size),), dtype=np.int8) + if trace_size + else None + ) def get_buffer(self, buffer_name): if buffer_name in self._buffer_cache: @@ -625,6 +638,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/operators/swiglu_prefill_stream/.README.md.swp b/iron/operators/swiglu_prefill_stream/.README.md.swp new file mode 100644 index 0000000000000000000000000000000000000000..1bbe0a8af804f15205c665ed41b7f475b6a1bb12 GIT binary patch literal 16384 zcmeI3NsJ^%6^6rOb}v|PfD2M!3%knh%ARg!hH=|7(9A~5-7}-^X>7|fmD!n5Rq4*k zOy|4JPxSt%)Hn`)LIwgh)VqvpED1mR|Fiu6A8%Lc>)6SMM5=m-cL&R;C!6C=Pf)XQ>WHajxT(=YvRR1Gv)u?zz*x zzNweq|M0`|<()jWuWz~w%l5*l)64F{+4+SfdtvF!X)c`f&Yo#59DVn~wR-BTd-$Zg zr`PLoeCV&bsV|x`9;Te$kfEl!QqK`)=2Uo_={ zF3q!g`ScmRsLgoCM3InwdTE*CQ5;2dwS4l_*#+8QU;I9OX+Jk%Wg^WhDTAyS2We=O z^69xS{m6=1p1T|*rSYX_8TWjsRi+V5##D_Mb~YGzs?K#@WY5rMoTa6vl>>2HbBr*v z=Ps>Z)+j=Rj%0PAc_U_I3TZ0U1BmqWmZ_9uTL`0vFoTq=uJq2dGQAjKD&E>jY%e=? zoKzZ0xhoyS!@5AI3~5$`Qq;!wIYhVQCK1uTl+(!Q^}~?~cl(y{($^-1j1V=fM$DG9 znld47-jLIzmewd9o3unwk13C|KPSA;72iJQJJQ{Wah+6go``l2WH>S6#^^~!JFe|8 zOX_i|%RGrIPu+lRdjtGudMB$>$dW05Vy1qq%lJ7{uApJ&*?wYDUyth&b`X(YuCJym zI4XjvF0x6vpp=jm*>0gyAH1HlAY#s}G|ZShkxP0kM_Ew~&^b?aj!}1cd+XB0{*ATu zOIwZ1>#I**+HB5rJ-ymUDN`6_NSbC9Q(TFAD2(U|Qx#My7(jYNS4Eue!g{I4GK3_U zngS)MOC~=>>7Yj6kj8j^GMd`)S(-QwM_GxEqs{mVU#+4IG~v`tLHG(p$kzfUSu@c7#X$Wmz%3jnnk3zQ?@Ic2kSf6MrF$*tR%! z-^{%I!N8>5D%hRPyz|daV<4ug3;Z2ziD&i2gch`ADNgDIQ8Zsar2PnS#ECL(!A`a( zV-zjc!2q{z=crN1scmcNHQHK0MO;Hd^5K<&2^V`_IG{xQ-ZL3IK#JH69OrRvk~l@7 zWWa|!xBJaG)Pk@pP;d22by) ztC>3ri>yR-ZD-c~1Q?9j4^1ws6qbe=#h_zUEgw6OtY{C$V#Dqm7 z2X&HYM1(ncOL7tKKIRKwRx^C#zVQ3y}3ZC<-7rr_Ec6xmpWFU*QkhNZF)k6%4NOLdh) zXA~E>uMAeu&I7Odh`2b*r?@2Kl0`o96itn-?M|{{H^~TVWaQ%@&xL!)63>WI2gAa! z0DhQ89K&{ro5c#20-Cbx!|Y++FvY8h4|CS-)~lU>5Njp{Gbb~TZJ)B_Gg(B(-86|N3{5I;xzd~T&m{AXD@4v_ zmOJGBW*$o-fDkieSuE;mfoNkbptLB}c!IfJr}Q{dB22zK4cMG$zx?EQ!wb zofl_OPvV_I6!@m?xiG^R6`6|bf9j|P+Q@*nbW71e>(KG(Gxa=tk*?@|qnW?XldBu+ z(z^(jCqZaLSzP;08E4UkYj*MYNqP2`Lun@YwTE+XgB#N0vpMNpyL$Br*X|Mik!tbq ziSB|OwZAoqudHwT+D}t&!y1=(OCJ>vN1oFkU+=}it;ac2h1VYFp*!42Z{?!#%1$3{ z>JQtnj-?8xe8`|vVu^U79cK1PE4RCX;f8t zxw5!8#6#6P9!j9(aYizzg7h@N;Sc-vZwR0r0^!a1~qz_k&lZP5}M{egR$v&w;KS zQ!jWPJO`cyId~6v8M$5p&w~em$ax!RKP>|-11$qB11$sp-wZVAiIgR6fGAm$1V%2I zwK zQ$v)_OH)ZkmBe9OCDS8iw?b;;wj$#aITDE}kxW2)O?6aSUcF(`c@fjs#7?`xklKig zPf5{Ais!_oAq zuzE@Xu9Ec8FrlpZ9((8cI-wkZKNeUdsf z_2`2!s3t);mO(B)O1gi3@ljEe^O9*&Zc1tizPW!^6_?rob-j7|XB9(TBM0%fY$D*2 zRi;cHmE6R?ac%R$NBY;UQ$BTqkdgkIgp^*T<|Vnh5WMP2{DD*)(2&_UP!8W$IfZi4 zd0StZU3snlI+Z&~+*Ff`lmAwW`!U6|{l)#FqApWWASp-E?oPnBbS~+cD`Joau^DAc zOOm(Gi6(PADEb!EcBq&fO0yn;f@JDSCr!RDC1Q_kFQbT}uhf*fV^Vh#1_}fyoTPpd z(c9;moQv$h)pe>7b-m$8??+|A1_GKXjY|2|j;Lv-H=w+tuWj~jWLZ+`D_a-Vu4^fk zU=a(d(O*nFiHklJ1fRVVKS}n+JvmsJ7Z)07B*`zmYQJJmd6teSjJbr=DWOe&l-t2@ zB_1mDCNEn(UzDr2_w0#wXBjd)j=hFwv*Ab6J&=t zztIYvi9yHeF~>mronzs#P+*Rf0{II-n5>7?!)r!WJW6o z-5q3on1N)$$NN*%XJ@zT9-BU{4NGD3-_OG2fsXly@316$hwA|70~N+u@e@-vc_2k6 zRBG;q4)}#o9JA;UWj5G%RU)g%YmQoLRAM-&vo|zVf2U@*r|8@}=}SI*Z`ZSxkp5n5 zCY|n6D^KfldSWC)Ej_%9Y~^*j@-XLT_Ga6eZP#&?<}piVgM<%H(yTuQE}HsSQ`KZ& z#}+qawQ=3MUv=)jn`GQ~wJ-|MVvw6Po`GF5H*Os7y!h8t&gQ+Kp{};ao+#U?nyM^6 zPMvC&gVivsiy&la 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, or 0 for an untraced build. + + Tracing adds a trailing runtime-sequence argument, so it changes the operator's + calling convention and has to be asked for rather than defaulted on. + """ + 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 +335,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, From 979f9b4e24bc33bda28a871b91f519e69194cb3b Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 20 Aug 2026 11:51:08 +0200 Subject: [PATCH 3/9] patch the trace buffer address against the dispatched kernel Trace lowering records the buffer's index in the sequence it configures, but that index is resolved against the kernel the host dispatches, which for a fused build is the wrapper rather than the operator. The address was patched from an argument the wrapper did not have, so the trace DMA wrote nowhere and the buffer came back empty. Give the wrapper the buffer at the same index. --- iron/common/compilation/sequence.py | 24 +++++++++++++++++++++++- iron/common/sequence.py | 16 +++++++++++++++- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index dcca05aff..1068791c5 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -218,18 +218,40 @@ def main(): # RuntimeSequenceOp trace_size = getattr(artifact, "trace_size", 0) n_traced = len(artifact.runlist) if trace_size else 0 + # Trace lowering patches the buffer's address by the index it holds in the + # sequence it configures, and that index is resolved against the dispatched + # kernel rather than the callee. Giving it the same index here is what makes + # the two agree. Each operator would need its own index, and they collide + # with the consolidated buffers once an operator takes three arguments or + # fewer, so this covers a single-operator sequence only. + trace_arg_idx = 0 + if trace_size: + indices = { + len(sequence_arg_types[name]) for name, *_ in artifact.runlist + } + if len(indices) > 1 or min(indices) <= 2: + raise NotImplementedError( + "tracing a sequence needs one trace buffer per operator at the " + "index that operator gives it, and these operators want " + f"{sorted(indices)}, which does not leave room for the " + "consolidated buffers. Trace one operator at a time." + ) + trace_arg_idx = indices.pop() + n_pad = max(0, trace_arg_idx - 3) if trace_size else 0 @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], + *([np.ndarray[(1,), buf_dtype]] * n_pad), *( [np.ndarray[(max(1, n_traced * trace_size),), np.dtype[np.int8]]] if trace_size else [] ), ) - def sequence(input_buf, output_buf, scratch_buf, *trace_bufs): + def sequence(input_buf, output_buf, scratch_buf, *rest): + trace_bufs = rest[n_pad:] consolidated_buffers = { "input": input_buf, "output": output_buf, diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 92a13c4b7..5b8c1bb41 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -3,6 +3,7 @@ import hashlib import logging +import json import time from pathlib import Path import numpy as np @@ -567,10 +568,23 @@ def __init__(self, op, device_name="main", sequence_name="sequence"): 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()) + # Trace lowering appends the trace buffer to the runtime sequence it + # configures, so the kernel takes it as its last argument rather than + # after the three consolidated ones. + self.run_handle.set_arg( + self._kernel_arg_count() - 1, self.trace_buffer.buffer_object() + ) self._params = None + def _kernel_arg_count(self): + """How many arguments the built ELF declares, from aiecc's own config.""" + config = ( + Path(self.op.artifacts[0].mlir_input.filename + ".prj") + / "full_elf_config.json" + ) + return len(json.loads(config.read_text())["xrt-kernels"][0]["arguments"]) + @property def params(self): """Lazy ParameterScratchpad bound to this ELF's ctrl scratchpad BO. From fc25b474e472b447f4ea18a5d1fcc1f9657cf1ba Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 20 Aug 2026 14:32:27 +0200 Subject: [PATCH 4/9] give every traced operator its own trace buffer slot --- =0.7 | 0 iron/common/compilation/__init__.py | 1 + iron/common/compilation/sequence.py | 98 +++++++++--------- iron/common/sequence.py | 48 ++++----- .../swiglu_prefill_stream/.README.md.swp | Bin 16384 -> 0 bytes iron/operators/swiglu_prefill_stream/op.py | 10 +- .../swiglu_prefill_stream/stream_design.py | 5 +- iron/tests/infrastructure/trace_layout.py | 39 +++++++ requirements_stream.txt | 2 +- 9 files changed, 120 insertions(+), 83 deletions(-) delete mode 100644 =0.7 delete mode 100644 iron/operators/swiglu_prefill_stream/.README.md.swp create mode 100644 iron/tests/infrastructure/trace_layout.py diff --git a/=0.7 b/=0.7 deleted file mode 100644 index e69de29bb..000000000 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/sequence.py b/iron/common/compilation/sequence.py index 1068791c5..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, @@ -216,42 +239,33 @@ def main(): itemsize = np.dtype(ml_dtypes.bfloat16).itemsize # RuntimeSequenceOp - trace_size = getattr(artifact, "trace_size", 0) - n_traced = len(artifact.runlist) if trace_size else 0 - # Trace lowering patches the buffer's address by the index it holds in the - # sequence it configures, and that index is resolved against the dispatched - # kernel rather than the callee. Giving it the same index here is what makes - # the two agree. Each operator would need its own index, and they collide - # with the consolidated buffers once an operator takes three arguments or - # fewer, so this covers a single-operator sequence only. - trace_arg_idx = 0 - if trace_size: - indices = { - len(sequence_arg_types[name]) for name, *_ in artifact.runlist - } - if len(indices) > 1 or min(indices) <= 2: - raise NotImplementedError( - "tracing a sequence needs one trace buffer per operator at the " - "index that operator gives it, and these operators want " - f"{sorted(indices)}, which does not leave room for the " - "consolidated buffers. Trace one operator at a time." - ) - trace_arg_idx = indices.pop() - n_pad = max(0, trace_arg_idx - 3) if trace_size else 0 - - @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], - *([np.ndarray[(1,), buf_dtype]] * n_pad), - *( - [np.ndarray[(max(1, n_traced * trace_size),), np.dtype[np.int8]]] - if trace_size - else [] - ), + 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, *rest): - trace_bufs = rest[n_pad:] + 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, @@ -261,7 +275,6 @@ def sequence(input_buf, output_buf, scratch_buf, *rest): # Execute operations in runlist order configure_op = None last_op_name = None - run_index = 0 for op_name, *buffer_names in artifact.runlist: expected_arg_types = sequence_arg_types[op_name] @@ -338,22 +351,13 @@ def sequence(input_buf, output_buf, scratch_buf, *rest): ) buffer_ssa_values.append(reinterpreted) - # Trace lowering appends a buffer to the callee, so the call - # has to carry one too. Each op writes its own slice. + # Trace lowering appends a buffer to the callee's signature. if trace_size: - buffer_ssa_values.append( - memref.subview( - trace_bufs[0], - [run_index * trace_size], - [trace_size], - [1], - ) - ) + 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) - run_index += 1 if needs_reset: reset_op = aiex.ConfigureOp(ir.FlatSymbolRefAttr.get(RESET_DEVICE)) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 5b8c1bb41..9dc5b6f5f 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -3,7 +3,6 @@ import hashlib import logging -import json import time from pathlib import Path import numpy as np @@ -564,27 +563,23 @@ 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()) - if self.trace_buffer is not None: - # Trace lowering appends the trace buffer to the runtime sequence it - # configures, so the kernel takes it as its last argument rather than - # after the three consolidated ones. - self.run_handle.set_arg( - self._kernel_arg_count() - 1, self.trace_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 - def _kernel_arg_count(self): - """How many arguments the built ELF declares, from aiecc's own config.""" - config = ( - Path(self.op.artifacts[0].mlir_input.filename + ".prj") - / "full_elf_config.json" - ) - return len(json.loads(config.read_text())["xrt-kernels"][0]["arguments"]) - @property def params(self): """Lazy ParameterScratchpad bound to this ELF's ctrl scratchpad BO. @@ -620,10 +615,13 @@ def _allocate_buffers(self): (_n_elements(scratch_sz),), dtype=ml_dtypes.bfloat16 ) trace_size = self.op.trace_size - self.trace_buffer = ( - XRTTensor((max(1, len(self.op.runlist) * trace_size),), dtype=np.int8) + 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 None + else {} ) def get_buffer(self, buffer_name): @@ -652,9 +650,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") + 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/.README.md.swp b/iron/operators/swiglu_prefill_stream/.README.md.swp deleted file mode 100644 index 1bbe0a8af804f15205c665ed41b7f475b6a1bb12..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 16384 zcmeI3NsJ^%6^6rOb}v|PfD2M!3%knh%ARg!hH=|7(9A~5-7}-^X>7|fmD!n5Rq4*k zOy|4JPxSt%)Hn`)LIwgh)VqvpED1mR|Fiu6A8%Lc>)6SMM5=m-cL&R;C!6C=Pf)XQ>WHajxT(=YvRR1Gv)u?zz*x zzNweq|M0`|<()jWuWz~w%l5*l)64F{+4+SfdtvF!X)c`f&Yo#59DVn~wR-BTd-$Zg zr`PLoeCV&bsV|x`9;Te$kfEl!QqK`)=2Uo_={ zF3q!g`ScmRsLgoCM3InwdTE*CQ5;2dwS4l_*#+8QU;I9OX+Jk%Wg^WhDTAyS2We=O z^69xS{m6=1p1T|*rSYX_8TWjsRi+V5##D_Mb~YGzs?K#@WY5rMoTa6vl>>2HbBr*v z=Ps>Z)+j=Rj%0PAc_U_I3TZ0U1BmqWmZ_9uTL`0vFoTq=uJq2dGQAjKD&E>jY%e=? zoKzZ0xhoyS!@5AI3~5$`Qq;!wIYhVQCK1uTl+(!Q^}~?~cl(y{($^-1j1V=fM$DG9 znld47-jLIzmewd9o3unwk13C|KPSA;72iJQJJQ{Wah+6go``l2WH>S6#^^~!JFe|8 zOX_i|%RGrIPu+lRdjtGudMB$>$dW05Vy1qq%lJ7{uApJ&*?wYDUyth&b`X(YuCJym zI4XjvF0x6vpp=jm*>0gyAH1HlAY#s}G|ZShkxP0kM_Ew~&^b?aj!}1cd+XB0{*ATu zOIwZ1>#I**+HB5rJ-ymUDN`6_NSbC9Q(TFAD2(U|Qx#My7(jYNS4Eue!g{I4GK3_U zngS)MOC~=>>7Yj6kj8j^GMd`)S(-QwM_GxEqs{mVU#+4IG~v`tLHG(p$kzfUSu@c7#X$Wmz%3jnnk3zQ?@Ic2kSf6MrF$*tR%! z-^{%I!N8>5D%hRPyz|daV<4ug3;Z2ziD&i2gch`ADNgDIQ8Zsar2PnS#ECL(!A`a( zV-zjc!2q{z=crN1scmcNHQHK0MO;Hd^5K<&2^V`_IG{xQ-ZL3IK#JH69OrRvk~l@7 zWWa|!xBJaG)Pk@pP;d22by) ztC>3ri>yR-ZD-c~1Q?9j4^1ws6qbe=#h_zUEgw6OtY{C$V#Dqm7 z2X&HYM1(ncOL7tKKIRKwRx^C#zVQ3y}3ZC<-7rr_Ec6xmpWFU*QkhNZF)k6%4NOLdh) zXA~E>uMAeu&I7Odh`2b*r?@2Kl0`o96itn-?M|{{H^~TVWaQ%@&xL!)63>WI2gAa! z0DhQ89K&{ro5c#20-Cbx!|Y++FvY8h4|CS-)~lU>5Njp{Gbb~TZJ)B_Gg(B(-86|N3{5I;xzd~T&m{AXD@4v_ zmOJGBW*$o-fDkieSuE;mfoNkbptLB}c!IfJr}Q{dB22zK4cMG$zx?EQ!wb zofl_OPvV_I6!@m?xiG^R6`6|bf9j|P+Q@*nbW71e>(KG(Gxa=tk*?@|qnW?XldBu+ z(z^(jCqZaLSzP;08E4UkYj*MYNqP2`Lun@YwTE+XgB#N0vpMNpyL$Br*X|Mik!tbq ziSB|OwZAoqudHwT+D}t&!y1=(OCJ>vN1oFkU+=}it;ac2h1VYFp*!42Z{?!#%1$3{ z>JQtnj-?8xe8`|vVu^U79cK1PE4RCX;f8t zxw5!8#6#6P9!j9(aYizzg7h@N;Sc-vZwR0r0^!a1~qz_k&lZP5}M{egR$v&w;KS zQ!jWPJO`cyId~6v8M$5p&w~em$ax!RKP>|-11$qB11$sp-wZVAiIgR6fGAm$1V%2I zwK zQ$v)_OH)ZkmBe9OCDS8iw?b;;wj$#aITDE}kxW2)O?6aSUcF(`c@fjs#7?`xklKig zPf5{Ais!_oAq zuzE@Xu9Ec8FrlpZ9((8cI-wkZKNeUdsf z_2`2!s3t);mO(B)O1gi3@ljEe^O9*&Zc1tizPW!^6_?rob-j7|XB9(TBM0%fY$D*2 zRi;cHmE6R?ac%R$NBY;UQ$BTqkdgkIgp^*T<|Vnh5WMP2{DD*)(2&_UP!8W$IfZi4 zd0StZU3snlI+Z&~+*Ff`lmAwW`!U6|{l)#FqApWWASp-E?oPnBbS~+cD`Joau^DAc zOOm(Gi6(PADEb!EcBq&fO0yn;f@JDSCr!RDC1Q_kFQbT}uhf*fV^Vh#1_}fyoTPpd z(c9;moQv$h)pe>7b-m$8??+|A1_GKXjY|2|j;Lv-H=w+tuWj~jWLZ+`D_a-Vu4^fk zU=a(d(O*nFiHklJ1fRVVKS}n+JvmsJ7Z)07B*`zmYQJJmd6teSjJbr=DWOe&l-t2@ zB_1mDCNEn(UzDr2_w0#wXBjd)j=hFwv*Ab6J&=t zztIYvi9yHeF~>mronzs#P+*Rf0{II-n5>7?!)r!WJW6o z-5q3on1N)$$NN*%XJ@zT9-BU{4NGD3-_OG2fsXly@316$hwA|70~N+u@e@-vc_2k6 zRBG;q4)}#o9JA;UWj5G%RU)g%YmQoLRAM-&vo|zVf2U@*r|8@}=}SI*Z`ZSxkp5n5 zCY|n6D^KfldSWC)Ej_%9Y~^*j@-XLT_Ga6eZP#&?<}piVgM<%H(yTuQE}HsSQ`KZ& z#}+qawQ=3MUv=)jn`GQ~wJ-|MVvw6Po`GF5H*Os7y!h8t&gQ+Kp{};ao+#U?nyM^6 zPMvC&gVivsiy&la=0.7 -stream-dse>=1.13.12 +stream-dse>=1.13.14 From dac0b4c968ca104a3b05f02e52c920c170463f75 Mon Sep 17 00:00:00 2001 From: asyms Date: Wed, 26 Aug 2026 13:48:34 +0200 Subject: [PATCH 5/9] key the fused sequence artifacts on the trace size Tracing adds a runtime-sequence argument, so a traced ELF cannot be reused by an untraced run or the other way round: the two now get different artifact names and toggling the flag no longer needs the build directory wiped. --- iron/common/sequence.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 9dc5b6f5f..bcc8da88f 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, 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,7 +87,7 @@ 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, @@ -113,7 +119,7 @@ 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, From 18b17d3160b84052abb2c9a01093eb83d83d55ed Mon Sep 17 00:00:00 2001 From: Arne Symons Date: Fri, 28 Aug 2026 13:05:40 +0200 Subject: [PATCH 6/9] separate out tracing_utils.py and call in swiglu_prefill_stream/test.py --- iron/common/tracing_utils.py | 367 +++++++++++++++++++ iron/operators/swiglu_prefill_stream/test.py | 3 + 2 files changed, 370 insertions(+) create mode 100644 iron/common/tracing_utils.py diff --git a/iron/common/tracing_utils.py b/iron/common/tracing_utils.py new file mode 100644 index 000000000..52342822f --- /dev/null +++ b/iron/common/tracing_utils.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read a traced run's hardware trace buffers back and write Perfetto JSON. + +Tracing is configured at build time (``IRON_TRACE_SIZE`` / ``IRON_TRACE_NTILES``, +consumed by the operator's design) and the runtime already syncs the resulting +buffers device->host after every dispatch. Nothing reads them, though, so a traced +run leaves its data sitting in host memory. This module is that last step: one call +after ``run()`` turns those buffers into files. + + from iron.common.tracing_utils import dump_traces + + run = operator.get_callable() + run() + dump_traces(run, "my_operator") + +No-op on an untraced build, so the call can stay in a test unconditionally. + +Two files land per traced dispatch: the raw 32-bit words as hex text, and the +parsed JSON for https://ui.perfetto.dev. The raw text is kept because reparsing is +free and re-dispatching is not - see :func:`parse_trace_words` to reparse it with a +different column shift without touching the device. + +``dump_traces`` also prints a per-tile summary, since the Perfetto timeline of a +few hundred short kernel calls is hard to read at a glance and the numbers a +designer wants - how much of the run a core spent computing, and how much waiting - +are a few sums away. :func:`print_trace_summary` does the same for a JSON file +written earlier. + +Environment: + * ``IRON_TRACE_DIR`` where to write (default ``outputs/traces``) + * ``IRON_TRACE_MLIR`` override the MLIR the parser reads (see below) + * ``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 . import compilation as comp + +__all__ = [ + "dump_traces", + "parse_trace_words", + "lowered_mlir", + "trace_words", + "summarize_trace", + "print_trace_summary", +] + +DEFAULT_TRACE_DIR = "outputs/traces" + +# The kernel brackets: aie_kernels sources wrap their body in event0()/event1(), +# so one pair is one kernel invocation. Everything between two pairs is the core +# waiting - on its input object FIFO, on a lock, on the next descriptor. +KERNEL_START, KERNEL_END = "INSTR_EVENT_0", "INSTR_EVENT_1" + + +def lowered_mlir(run) -> tuple[Path, str]: + """The post-lowering MLIR for a callable, as ``(path, text)``. + + mlir-aie's trace parser recovers which tiles and events were traced by scanning + for ``aiex.npu.write32`` ops and pattern-matching the trace-unit config + addresses. It does not understand the declarative ``aie.trace`` ops a design is + written with, so the module handed to aiecc is useless to it: those ops only + become register writes inside aiecc, in ``aie-insert-trace-flows``. What we want + is one of aiecc's own intermediates, which it leaves in the work dir beside the + source (``.mlir.d/``). + + Picks the file with the most ``write32`` ops rather than hardcoding a stage name, + since those names are aiecc's business and have changed before. + """ + 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) + work_dir = comp._aiecc_work_dir(str(source)) + candidates = [] + for path in sorted(work_dir.rglob("*.mlir")): + text = path.read_text() + if "write32" in text: + candidates.append((text.count("write32"), path, text)) + if not candidates: + raise FileNotFoundError( + f"no lowered MLIR with write32 ops under {work_dir}. aiecc may not be " + "retaining its intermediates for this build; point IRON_TRACE_MLIR at a " + "lowered module to override." + ) + _, path, text = max(candidates, key=lambda c: c[0]) + return path, text + + +def trace_words(buf) -> np.ndarray: + """A trace buffer's contents as uint32 words, with the unfilled tail dropped. + + The buffer is allocated at the full trace size and only partly written, so the + trailing zeros are absence of events rather than events. Trimming them keeps the + JSON small and stops the parser inventing a long idle tail. + """ + raw = buf.to_torch().numpy().astype(np.uint8) + raw = raw[: raw.size - raw.size % 4] + words = raw.view(np.uint32) # little-endian on x86, matching the DMA layout + if not words.any(): + return words[:0] + return words[: int(np.nonzero(words)[0][-1]) + 1] + + +def parse_trace_words(words, mlir_text: str, colshift: int | None = None): + """Trace words plus the lowered MLIR into Trace Event Format events. + + ``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 only when the tiles in the output do not match the placement. + + The parser calls ``sys.exit`` rather than raising on some malformed input, so + SystemExit is caught here - a visualisation failure should never take a test + down with it. + """ + from aie.utils.trace.parse import parse_trace + + try: + return parse_trace(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 _by_tile(events): + """Group Trace Event Format records by pid, resolving each pid's tile name. + + The parser emits one process per traced tile (``process_name`` metadata), and + one thread per monitored event slot. Metadata records carry no timestamp, so + they are separated out here rather than filtered at every use. + """ + names, records = {}, {} + for index, event in enumerate(events): + pid = event.get("pid") + if event.get("ph") == "M": + if event.get("name") == "process_name": + names[pid] = event.get("args", {}).get("name", str(pid)) + continue + if "ts" in event: + records.setdefault(pid, []).append((event["ts"], index, event)) + for pid in records: + records[pid].sort() # index breaks ts ties, keeping emission order + return names, records + + +def _state_cycles(records): + """Cycles each event name was asserted, summed over its begin/end intervals. + + A level event (a stall, vector activity) is emitted as ``B``/``E`` pairs on its + own thread, re-asserted at every trace command, so one logical stall arrives as + many short intervals. Summing them gives the time in that state. These overlap + each other and the kernel brackets - a core stalls *during* a kernel call - so + they are shares of the window, not a partition of it. + """ + open_at, totals = {}, {} + for ts, _, event in records: + key = (event.get("tid"), event.get("name")) + if event.get("ph") == "B": + open_at.setdefault(key, ts) + elif event.get("ph") == "E" and key in open_at: + totals[key[1]] = totals.get(key[1], 0) + ts - open_at.pop(key) + return totals + + +def _invocations(records): + """Kernel invocations as ``(start, end)`` cycle pairs. + + Pairs each ``event0`` with the next ``event1``, ignoring repeats of either - + the same rule mlir-aie's own summary uses, so the call counts agree. + """ + spans, start = [], None + for ts, _, event in records: + if event.get("ph") != "B": + continue + if event.get("name") == KERNEL_START and start is None: + start = ts + elif event.get("name") == KERNEL_END and start is not None: + spans.append((start, ts)) + start = None + return spans + + +def _stats(values): + if not values: + return None + ordered = sorted(values) + return { + "count": len(ordered), + "total": sum(ordered), + "min": ordered[0], + "max": ordered[-1], + "mean": sum(ordered) / len(ordered), + } + + +def summarize_trace(source) -> dict: + """Per-tile cycle accounting for a parsed trace. + + Accepts a JSON path or an already-parsed event list. Returns + ``{tile_name: {...}}`` with, per tile: the traced ``window`` in cycles, the + kernel invocations (``busy``), the gaps between them (``waiting``), and the + cycles spent in each monitored state. + + All figures are AIE core cycles, from the trace unit's own timer. + """ + if isinstance(source, (str, Path)): + source = json.loads(Path(source).read_text()) + + names, per_pid = _by_tile(source) + summary = {} + for pid, records in per_pid.items(): + window = records[-1][0] - records[0][0] + spans = _invocations(records) + busy = _stats([end - start for start, end in spans]) + waiting = _stats([nxt[0] - cur[1] for cur, nxt in zip(spans, spans[1:])]) + summary[names.get(pid, str(pid))] = { + "window": window, + "busy": busy, + "waiting": waiting, + "states": _state_cycles(records), + } + return summary + + +def _pct(part, whole): + return f"{100.0 * part / whole:5.1f}%" if whole else " -" + + +def print_trace_summary(source, title: str | None = None) -> dict: + """Print :func:`summarize_trace` as a short per-tile report, and return it. + + Reads as: how much of the traced window each core spent inside a kernel, how + much it spent between kernels, and what it was stalled on meanwhile. + """ + summary = summarize_trace(source) + if title is None and isinstance(source, (str, Path)): + title = Path(source).name + if title: + print(f"\n[trace] {title}") + + for tile, data in summary.items(): + window = data["window"] + busy, waiting = data["busy"], data["waiting"] + print(f" {tile} - {window} cycles traced") + + if busy: + print( + f" in kernel {busy['count']:>6} calls {busy['total']:>10} cyc " + f"{_pct(busy['total'], window)} " + f"min/mean/max {busy['min']}/{busy['mean']:.1f}/{busy['max']}" + ) + else: + print( + f" in kernel no {KERNEL_START}/{KERNEL_END} pairs - does " + "this kernel call event0()/event1()?" + ) + if waiting: + print( + f" between {waiting['count']:>6} gaps " + f"{waiting['total']:>10} cyc {_pct(waiting['total'], window)} " + f"min/mean/max {waiting['min']}/{waiting['mean']:.1f}/{waiting['max']}" + ) + + # Stalls and vector activity overlap the above, so they are listed apart. + states = { + name: cycles + for name, cycles in data["states"].items() + if name not in (KERNEL_START, KERNEL_END) and cycles + } + for name, cycles in sorted(states.items(), key=lambda s: -s[1]): + print(f" {name.lower():<12} {cycles:>22} cyc {_pct(cycles, window)}") + if summary: + print( + " (stall and vector shares overlap the kernel time above, " + "they are not a partition)" + ) + return summary + + +def dump_traces( + run, + tag: str, + out_dir=None, + colshift: int | None = None, + summary: bool = True, +) -> list[Path]: + """Write every trace buffer of a completed run out as hex text and Perfetto JSON. + + Call it after ``run()``: the callable syncs its trace buffers 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. Each + buffer is named for the runlist entry it belongs to, so a fused sequence yields + one pair of files per operator in the sequence. + """ + buffers = getattr(run, "trace_buffers", None) + if not buffers: + if getattr(getattr(run, "op", None), "trace_size", 0): + raise TypeError( + f"{type(run).__name__} was built with tracing enabled but exposes no " + "trace_buffers; only the full-ELF sequence callable allocates them." + ) + 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}") + + tag = _slug(tag) + written = [] + for name, buf in buffers.items(): + words = trace_words(buf) + if not words.size: + print(f"[trace] {name}: buffer is all zeros, no trace data captured") + continue + capacity = buf.to_torch().numel() // 4 + if words.size == capacity: + print( + f"[trace] {name}: buffer full ({capacity * 4} B), trace is likely " + "truncated - raise IRON_TRACE_SIZE" + ) + + stem = out_dir / f"{tag}_{_slug(name)}" + stem.with_suffix(".txt").write_text("\n".join(f"{w:08x}" for w in words) + "\n") + + try: + events = parse_trace_words(words, mlir_text, colshift) + except Exception as exc: # never let a visualisation failure fail a run + print(f"[trace] {name}: parse failed ({exc}); raw words kept at {stem}.txt") + continue + + target = stem.with_suffix(".json") + target.write_text(json.dumps(events)) + print(f"[trace] {target} ({len(events)} events)") + written.append(target) + + if summary: + try: + print_trace_summary(events, title=target.name) + except Exception as exc: # a summary is never worth failing a run over + print(f"[trace] {name}: summary failed ({exc})") + return written \ No newline at end of file diff --git a/iron/operators/swiglu_prefill_stream/test.py b/iron/operators/swiglu_prefill_stream/test.py index c0a9deaa8..a6d9032af 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. @@ -70,6 +72,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, From 19b8f780974f5aabefb5573a0db198a7d1595430 Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 28 Aug 2026 16:15:53 -0600 Subject: [PATCH 7/9] Let the compiler place the trace buffer for a fused sequence mlir-aie appends a single trace buffer to the dispatched runtime sequence and hands each configured design a slice of it. The fused sequence no longer has to place one buffer per operator at the index that operator patches against. Drop trace_argument_layout. The three consolidated buffers keep indices 0-2 and the trace buffer lands after them, so operators taking the same number of arguments no longer contend for one slot. Size the host buffer from aie.trace_slices in the lowered module. trace_size is per design and says nothing about how many channels or sub-designs claim a share. Bind the buffer at index 3 and sync it back with the outputs. Ask aiecc for --get-input-with-addresses on a traced build and read the lowered module from that path, rather than picking the work-dir file with the most write32 ops. Split the buffer by recorded slice when dumping traces, and parse each slice against the device that wrote it: sub-designs routinely occupy the same tiles, and a merged parse cannot tell their events apart. --- iron/common/compilation/__init__.py | 2 +- iron/common/compilation/base.py | 4 ++ iron/common/compilation/sequence.py | 68 ++++++-------------- iron/common/sequence.py | 52 ++++++++-------- iron/common/tracing_utils.py | 67 ++++++++++---------- iron/tests/infrastructure/trace_layout.py | 76 ++++++++++++----------- 6 files changed, 124 insertions(+), 145 deletions(-) diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index 64d710988..71b33b500 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -31,5 +31,5 @@ from .sequence import ( SequenceMLIRArtifact, FusePythonGeneratedMLIRCompilationRule, - trace_argument_layout, + trace_buffer_layout, ) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 4905b9965..d79dcd60d 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -547,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 to recover 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 52cedd929..500251b45 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -7,8 +7,6 @@ from __future__ import annotations -from itertools import count, islice - import numpy as np import importlib.util from functools import partial @@ -36,25 +34,22 @@ # ########################################################################## -def trace_argument_layout(arg_counts: dict[str, int], trace_size: int): - """Buffer slots for the fused runtime sequence, as (consolidated, trace, count). +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 a single trace buffer + covering every design it configures, and records the split on the sequence. + Reading it back is how the host learns the buffer's size and how the parser + learns which design wrote which bytes. - 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. + Returns `(total_bytes, slices)`; `(0, [])` for an untraced build. """ - 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 + from aie.utils.trace import get_trace_slices + + 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): @@ -239,33 +234,12 @@ def main(): itemsize = np.dtype(ml_dtypes.bfloat16).itemsize # RuntimeSequenceOp - 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, + @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_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 - ) + def sequence(input_buf, output_buf, scratch_buf): consolidated_buffers = { "input": input_buf, "output": output_buf, @@ -351,10 +325,6 @@ def sequence(*all_bufs): ) 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 bcc8da88f..ff7b7a081 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -569,20 +569,11 @@ 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) - 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.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 @@ -620,15 +611,24 @@ 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 {} - ) + # Trace lowering appends one buffer covering every configured design, + # after the consolidated three. Its size comes from the lowered module + # rather than from trace_size, which is per design and says nothing + # about how many channels or sub-designs claim a share. + 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: @@ -656,9 +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") - for buf in self.trace_buffers.values(): - buf.device = "npu" - buf.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 index 52342822f..eb0857d45 100644 --- a/iron/common/tracing_utils.py +++ b/iron/common/tracing_utils.py @@ -68,12 +68,9 @@ def lowered_mlir(run) -> tuple[Path, str]: for ``aiex.npu.write32`` ops and pattern-matching the trace-unit config addresses. It does not understand the declarative ``aie.trace`` ops a design is written with, so the module handed to aiecc is useless to it: those ops only - become register writes inside aiecc, in ``aie-insert-trace-flows``. What we want - is one of aiecc's own intermediates, which it leaves in the work dir beside the - source (``.mlir.d/``). - - Picks the file with the most ``write32`` ops rather than hardcoding a stage name, - since those names are aiecc's business and have changed before. + become register writes inside aiecc, in ``aie-insert-trace-flows``. A traced + build asks aiecc for ``--get-input-with-addresses``, which lands the lowered + module in the work dir beside the source (``.mlir.d/``). """ override = os.environ.get("IRON_TRACE_MLIR") if override: @@ -81,20 +78,13 @@ def lowered_mlir(run) -> tuple[Path, str]: return path, path.read_text() source = Path(run.op.artifacts[0].mlir_input.filename) - work_dir = comp._aiecc_work_dir(str(source)) - candidates = [] - for path in sorted(work_dir.rglob("*.mlir")): - text = path.read_text() - if "write32" in text: - candidates.append((text.count("write32"), path, text)) - if not candidates: + path = comp._aiecc_work_dir(str(source)) / "input_with_addresses.mlir" + if not path.exists(): raise FileNotFoundError( - f"no lowered MLIR with write32 ops under {work_dir}. aiecc may not be " - "retaining its intermediates for this build; point IRON_TRACE_MLIR at a " - "lowered module to override." + f"{path} is missing; a traced build passes --get-input-with-addresses " + "to aiecc. Point IRON_TRACE_MLIR at a lowered module to override." ) - _, path, text = max(candidates, key=lambda c: c[0]) - return path, text + return path, path.read_text() def trace_words(buf) -> np.ndarray: @@ -112,13 +102,19 @@ def trace_words(buf) -> np.ndarray: return words[: int(np.nonzero(words)[0][-1]) + 1] -def parse_trace_words(words, mlir_text: str, colshift: int | None = None): +def parse_trace_words( + words, mlir_text: str, colshift: int | None = None, device: str | None = None +): """Trace words plus the lowered MLIR into Trace Event Format events. ``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 only when the tiles in the output do not match the placement. + ``device`` names the ``aie.device`` whose trace configuration these words were + written by. A fused sequence holds one per sub-design, and they routinely share + tile coordinates, so leaving it unset merges their event assignments. + The parser calls ``sys.exit`` rather than raising on some malformed input, so SystemExit is caught here - a visualisation failure should never take a test down with it. @@ -126,7 +122,9 @@ def parse_trace_words(words, mlir_text: str, colshift: int | None = None): from aie.utils.trace.parse import parse_trace try: - return parse_trace(np.asarray(words, dtype=np.uint32), mlir_text, colshift) + return parse_trace( + np.asarray(words, dtype=np.uint32), mlir_text, colshift, device + ) except SystemExit as exc: raise RuntimeError( "mlir-aie's trace parser exited; the usual cause is an MLIR without the " @@ -308,16 +306,16 @@ def dump_traces( 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. Each - buffer is named for the runlist entry it belongs to, so a fused sequence yields - one pair of files per operator in the sequence. + ``tag`` distinguishes one dump from another - a test name or parameter id. The + buffer is split by the layout the compiler recorded on the dispatched sequence, + so a fused sequence yields one pair of files per configured design. """ - buffers = getattr(run, "trace_buffers", None) - if not buffers: + 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_buffers; only the full-ELF sequence callable allocates them." + "trace_buffer; only the full-ELF sequence callable allocates one." ) return [] @@ -331,17 +329,20 @@ def dump_traces( mlir_path, mlir_text = lowered_mlir(run) print(f"[trace] parsing against {mlir_path}") + all_words = buffer.to_torch().numpy().astype(np.uint8).view(np.uint32) tag = _slug(tag) written = [] - for name, buf in buffers.items(): - words = trace_words(buf) + for index, entry in enumerate(run.trace_slices): + name = f"{index}_{entry['device']}" + start = entry["offset"] // 4 + region = all_words[start : start + entry["size"] // 4] + words = region[: int(np.nonzero(region)[0][-1]) + 1] if region.any() else region if not words.size: print(f"[trace] {name}: buffer is all zeros, no trace data captured") continue - capacity = buf.to_torch().numel() // 4 - if words.size == capacity: + if words.size == region.size: print( - f"[trace] {name}: buffer full ({capacity * 4} B), trace is likely " + f"[trace] {name}: slice full ({entry['size']} B), trace is likely " "truncated - raise IRON_TRACE_SIZE" ) @@ -349,7 +350,7 @@ def dump_traces( stem.with_suffix(".txt").write_text("\n".join(f"{w:08x}" for w in words) + "\n") try: - events = parse_trace_words(words, mlir_text, colshift) + events = parse_trace_words(words, mlir_text, colshift, entry["device"]) except Exception as exc: # never let a visualisation failure fail a run print(f"[trace] {name}: parse failed ({exc}); raw words kept at {stem}.txt") continue @@ -364,4 +365,4 @@ def dump_traces( print_trace_summary(events, title=target.name) except Exception as exc: # a summary is never worth failing a run over print(f"[trace] {name}: summary failed ({exc})") - return written \ No newline at end of file + return written diff --git a/iron/tests/infrastructure/trace_layout.py b/iron/tests/infrastructure/trace_layout.py index 9b658c373..d16b2626c 100644 --- a/iron/tests/infrastructure/trace_layout.py +++ b/iron/tests/infrastructure/trace_layout.py @@ -1,39 +1,43 @@ # 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) +"""Reading back the trace buffer layout the compiler recorded on the sequence.""" + +from iron.common.compilation import trace_buffer_layout + +LOWERED = """ +module { + aie.device(npu1_1col) { + aie.runtime_sequence @sequence(%arg0: memref<4xi32>, %arg1: memref<12288xi8>) + attributes {aie.trace_slices = [ + {device = "dev_a", offset = 0 : i64, sequence = "seq", size = 8192 : i64}, + {device = "dev_b", offset = 8192 : i64, sequence = "seq", size = 4096 : i64}]} { + } + } +} +""" + +UNTRACED = """ +module { + aie.device(npu1_1col) { + aie.runtime_sequence @sequence(%arg0: memref<4xi32>) { + } + } +} +""" + + +def test_total_spans_every_slice(): + total, slices = trace_buffer_layout(LOWERED) + assert total == 12288 + assert [s["offset"] for s in slices] == [0, 8192] + assert [s["size"] for s in slices] == [8192, 4096] + + +def test_each_slice_names_the_design_that_wrote_it(): + _, slices = trace_buffer_layout(LOWERED) + assert [s["device"] for s in slices] == ["dev_a", "dev_b"] + + +def test_untraced_build_has_no_trace_buffer(): + assert trace_buffer_layout(UNTRACED) == (0, []) From c30c33231d960cda0616606c905ab9c2729a6ecb Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 28 Aug 2026 16:38:38 -0600 Subject: [PATCH 8/9] Simplify the trace comments and docstrings Rewrite in the plain style: active voice, one idea per sentence, and no attribution of understanding to the parser. Match the module docstring of tracing_utils to a single trace buffer. It described one buffer per operator, which the compiler no longer produces. --- iron/common/compilation/base.py | 4 ++-- iron/common/compilation/sequence.py | 6 ++--- iron/common/sequence.py | 6 ++--- iron/common/tracing_utils.py | 35 ++++++++++++++--------------- 4 files changed, 25 insertions(+), 26 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index d79dcd60d..32d1864f3 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -548,8 +548,8 @@ def compile(self, graph): "--get-scratchpad-parameters", ] + artifact.extra_flags if artifact.trace_size: - # The trace parser reads the lowered module to recover the - # buffer layout and each design's traced tiles and events. + # 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( diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index 500251b45..7ee706d19 100644 --- a/iron/common/compilation/sequence.py +++ b/iron/common/compilation/sequence.py @@ -37,10 +37,10 @@ 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 a single trace buffer + `-aie-fuse-trace-buffers` gives the dispatched sequence one trace buffer covering every design it configures, and records the split on the sequence. - Reading it back is how the host learns the buffer's size and how the parser - learns which design wrote which bytes. + 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. """ diff --git a/iron/common/sequence.py b/iron/common/sequence.py index ff7b7a081..125ee6001 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -612,9 +612,9 @@ def _allocate_buffers(self): (_n_elements(scratch_sz),), dtype=ml_dtypes.bfloat16 ) # Trace lowering appends one buffer covering every configured design, - # after the consolidated three. Its size comes from the lowered module - # rather than from trace_size, which is per design and says nothing - # about how many channels or sub-designs claim a share. + # 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: diff --git a/iron/common/tracing_utils.py b/iron/common/tracing_utils.py index eb0857d45..28127d04f 100644 --- a/iron/common/tracing_utils.py +++ b/iron/common/tracing_utils.py @@ -1,13 +1,13 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Read a traced run's hardware trace buffers back and write Perfetto JSON. +"""Read a traced run's hardware trace buffer back and write Perfetto JSON. Tracing is configured at build time (``IRON_TRACE_SIZE`` / ``IRON_TRACE_NTILES``, consumed by the operator's design) and the runtime already syncs the resulting -buffers device->host after every dispatch. Nothing reads them, though, so a traced +buffer device->host after every dispatch. Nothing reads it, though, so a traced run leaves its data sitting in host memory. This module is that last step: one call -after ``run()`` turns those buffers into files. +after ``run()`` turns it into files. from iron.common.tracing_utils import dump_traces @@ -17,7 +17,7 @@ No-op on an untraced build, so the call can stay in a test unconditionally. -Two files land per traced dispatch: the raw 32-bit words as hex text, and the +Two files land per traced design: the raw 32-bit words as hex text, and the parsed JSON for https://ui.perfetto.dev. The raw text is kept because reparsing is free and re-dispatching is not - see :func:`parse_trace_words` to reparse it with a different column shift without touching the device. @@ -64,13 +64,12 @@ def lowered_mlir(run) -> tuple[Path, str]: """The post-lowering MLIR for a callable, as ``(path, text)``. - mlir-aie's trace parser recovers which tiles and events were traced by scanning - for ``aiex.npu.write32`` ops and pattern-matching the trace-unit config - addresses. It does not understand the declarative ``aie.trace`` ops a design is - written with, so the module handed to aiecc is useless to it: those ops only - become register writes inside aiecc, in ``aie-insert-trace-flows``. A traced - build asks aiecc for ``--get-input-with-addresses``, which lands the lowered - module in the work dir beside the source (``.mlir.d/``). + mlir-aie's trace parser reads ``aiex.npu.write32`` ops and matches the + trace-unit config addresses. ``aie-insert-trace-flows`` produces those writes + from the declarative ``aie.trace`` ops inside aiecc, so the module handed to + aiecc carries none of them. A traced build passes + ``--get-input-with-addresses``, which lands the lowered module in the work dir + beside the source (``.mlir.d/``). """ override = os.environ.get("IRON_TRACE_MLIR") if override: @@ -111,9 +110,9 @@ def parse_trace_words( want by default: a design configured for one column may be loaded into another. Override it only when the tiles in the output do not match the placement. - ``device`` names the ``aie.device`` whose trace configuration these words were - written by. A fused sequence holds one per sub-design, and they routinely share - tile coordinates, so leaving it unset merges their event assignments. + ``device`` names the ``aie.device`` that wrote these words. A fused sequence + holds one per sub-design, and two often share tile coordinates, so an unset + ``device`` merges their event assignments. The parser calls ``sys.exit`` rather than raising on some malformed input, so SystemExit is caught here - a visualisation failure should never take a test @@ -300,15 +299,15 @@ def dump_traces( colshift: int | None = None, summary: bool = True, ) -> list[Path]: - """Write every trace buffer of a completed run out as hex text and Perfetto JSON. + """Write a completed run's trace buffer as hex text and Perfetto JSON. - Call it after ``run()``: the callable syncs its trace buffers device->host as + 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 - buffer is split by the layout the compiler recorded on the dispatched sequence, - so a fused sequence yields one pair of files per configured design. + layout the compiler recorded on the dispatched sequence splits the buffer, so a + fused sequence yields one pair of files per configured design. """ buffer = getattr(run, "trace_buffer", None) if buffer is None: From b1a618ff8be89abc0affb1cd4d3af86c9a8044f2 Mon Sep 17 00:00:00 2001 From: andrej Date: Mon, 31 Aug 2026 14:43:46 -0600 Subject: [PATCH 9/9] Import the trace helpers at module scope Both files already load the aie bindings at import time, so a deferred import saves nothing. --- iron/common/compilation/sequence.py | 3 +-- iron/common/tracing_utils.py | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/iron/common/compilation/sequence.py b/iron/common/compilation/sequence.py index 7ee706d19..19ac7d75e 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 @@ -44,8 +45,6 @@ def trace_buffer_layout(mlir_text: str): Returns `(total_bytes, slices)`; `(0, [])` for an untraced build. """ - from aie.utils.trace import get_trace_slices - slices = get_trace_slices(mlir_text) if not slices: return 0, [] diff --git a/iron/common/tracing_utils.py b/iron/common/tracing_utils.py index 28127d04f..b7368c982 100644 --- a/iron/common/tracing_utils.py +++ b/iron/common/tracing_utils.py @@ -42,6 +42,8 @@ import numpy as np +from aie.utils.trace.parse import parse_trace + from . import compilation as comp __all__ = [ @@ -118,8 +120,6 @@ def parse_trace_words( SystemExit is caught here - a visualisation failure should never take a test down with it. """ - from aie.utils.trace.parse import parse_trace - try: return parse_trace( np.asarray(words, dtype=np.uint32), mlir_text, colshift, device