From 85ed572dbfc2c09b60bbcc12a741e15dc7488e58 Mon Sep 17 00:00:00 2001 From: asyms Date: Mon, 10 Aug 2026 19:19:57 +0200 Subject: [PATCH 1/4] 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 72e8f22368368e41a836de596db5c2b40a7e9419 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 20 Aug 2026 10:37:45 +0200 Subject: [PATCH 2/4] 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 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..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 6826ced28..2f3d12b70 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 @@ -562,6 +567,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 @@ -597,6 +604,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: @@ -624,6 +637,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 9cb83e21504b5a7e54cdcf8c81892083a07960ea Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 20 Aug 2026 11:51:08 +0200 Subject: [PATCH 3/4] 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 2f3d12b70..7da94e1b2 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 @@ -568,10 +569,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 0d8e10b30baadd0014fddeabe67f1e83d2d28ac3 Mon Sep 17 00:00:00 2001 From: asyms Date: Thu, 20 Aug 2026 14:32:27 +0200 Subject: [PATCH 4/4] 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 7da94e1b2..041632980 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 @@ -565,27 +564,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. @@ -619,10 +614,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): @@ -651,9 +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") - 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