Conversation
`get_arg_spec()` is about to move out of each `op.py` and onto a shape function declared with the operator itself. That is meant to be a pure refactor -- the specs must come out identical for every operator and every configuration -- but "identical across 22 classes" is not something a reviewer can check by eye, so record it instead. `arg_spec_snapshot.json` is the pre-refactor truth, generated on devel. The test re-derives specs and diffs against it, so a shape or dtype that changes during the refactor fails loudly and names the case. Both device generations are covered. An operator reads the ShimDMA column limit at construction, so a spec can depend on the device, and one that silently differed between npu1 and npu2 would otherwise land as a correctness bug on whichever generation CI does not run. Setting a device *description* via `from_name` rather than a live device keeps this runnable anywhere -- no XRT, no NPU. Three findings while building the matrix, each now pinned by a test: - `num_aie_columns` defaults (AXPY's is 8) exceed the ShimDMA limit of narrow devices, so cases pin it rather than inherit a value that varies by width. - Presence of `get_arg_spec` proves nothing, since every class inherits the attribute. The three SwiGLU composites are `OperatorSequence` subclasses and raise from it on purpose; `test_every_operator_with_a_spec_is_covered` excludes those and fails on anything else missing from the matrix, so the gate cannot quietly start checking less. - `_SwiGLUStreamGroup` needs the optional `stream` package. It is skipped on both sides of the comparison rather than dropped, so a snapshot generated without it does not read as "case added" on a machine that has it. Verified the gate discriminates: swapping GEMM's first spec from (M, K) to (K, M) fails 20 cases across both devices with a readable diff, and reverting returns it to green. Co-Authored-By: Claude <noreply@anthropic.com>
Clustering the operators by their recorded spec pattern, to find which ones could share a base, turned up two cases in this matrix that prove less than they appear to. StridedCopy read as "(in, out), same shape", which would have grouped it with the elementwise family. It is not: its input and output sizes are independent parameters, and the case simply set both to 1024. With output_buffer_size=256 it reports in (1024,), out (256,). It belongs with Dequant and Repeat instead. Transpose was only exercised square (64x64), which cannot distinguish a flat (M*N,) buffer from a shape that tracks (M, N) -- a refactor emitting (N, M) would have passed. Non-square confirms both buffers stay flat, since a transpose changes layout rather than size. Both are the same failure: a case whose parameters coincide cannot pin the relationship between them. Added a differing-size StridedCopy and a non-square Transpose. Co-Authored-By: Claude <noreply@anthropic.com>
angle_rows is an independent parameter constrained to divide rows, and it merely defaults to rows -- so with rows=32, angle_rows=8 the spec is in (32,64), in (8,64), out (32,64). Every case here left it defaulted, which made RoPE read as three buffers of one shape and grouped it with the elementwise binaries it is not. Third case in this matrix where coinciding parameters hid a relationship, after StridedCopy's equal buffer sizes and Transpose's square dimensions. Co-Authored-By: Claude <noreply@anthropic.com>
ChanneledUnaryOperator bundles two things that do not have to travel
together: the arg spec, and the design that generates the MLIR. Bundling them
is why Softmax, MemCopy and Transpose each hand-wrote
[AIERuntimeArgSpec("in", (size,)), AIERuntimeArgSpec("out", (size,))]
instead of inheriting it -- they share the shape rule with the elementwise
activations but emit completely different MLIR, so the base was not available
to them.
Split the two. same_shape_unary() and same_shape_binary() are plain functions,
so an operator reuses a shape rule by calling it rather than by inheriting a
design it does not want. Thirteen operators now route through them: the seven
on ChanneledUnaryOperator, the three on BinaryElementwiseOperator, plus
Softmax, MemCopy and Transpose.
Also add reads/writes predicates and nbytes() to AIERuntimeArgSpec. Callers
that want to know whether a step touches a buffer compare `direction` against
a set inline, which is fine until "inout" appears -- it must answer yes to
both, and a caller that partitions arguments into inputs and outputs counts it
once and places it wrong. Liveness analysis for memory planning is exactly
such a caller.
Four operators that look like they belong in these clusters do not, each
verified against the code rather than the spec shape alone:
- RoPE's angles broadcast (angle_rows defaults to rows but is independent, so
rows=32/angle_rows=8 gives in (32,64), in (8,64), out (32,64)).
- StridedCopy's input and output sizes are independent parameters.
- RMSNorm carries an optional weight between its input and output.
- Dequant, Repeat, GEMM, GEMV and MHA have shapes that genuinely differ.
Specs are unchanged: the snapshot gate passes untouched across all 22
operators on both device generations. Full suite is 35 failed / 40 errors
before and after -- the XRT-dependent tests, which cannot run on a host
without a device.
Co-Authored-By: Claude <noreply@anthropic.com>
An operator's spec is a function of its parameters, so declare it as one. `arg_spec` is a staticmethod taking the fields it needs by name, and the base binds them via inspect.signature -- so `get_arg_spec()` stops being a method every operator reimplements and becomes something derived. `bind()` matches parameters to attributes by name, filling from fields and properties alike, and raises naming both the operator and the parameter it could not supply. That failure is the point: the hand-written kwargs dicts it replaces restated field names with nothing checking the two sides, so a rename on one side surfaced as a TypeError from inside the callee or, when the parameter had a default, as a silently wrong value compiled into a design. Converted softmax, gemm and mha. The latter two are the ones worth having: - GEMM's b_col_maj/c_col_maj transpose a declared shape rather than resize it. - MHA's shape needs a helper call (seq_len rounds up to a pipeline multiple) and a branch (num_KV_heads == 0 means K/V are as wide as Q). Neither is expressible in a declarative shape notation without that notation becoming Python, which is why these stay plain functions -- the same conclusion JAX reaches with abstract_eval and PyTorch with register_fake. Making a shape rule a function of parameters rather than of an instance also means a caller can ask what shape an operator *would* produce before building it, which is what graph capture needs to place a value it has not constructed. MHA._calculate_seq_padding becomes a staticmethod; both remaining callers go through self and are unaffected. Specs unchanged -- the snapshot gate passes across all 22 operators on both device generations. Full suite 35 failed / 40 errors before and after, the XRT-dependent tests. Deliberately not done here: binding the *design* kwargs the same way. Several design parameters are spelled differently from the field feeding them (softmax's num_elements <- size, tile_size <- cols), so that change renames design parameters, and nothing in this gate covers design output. It needs the per-operator hardware tests. Co-Authored-By: Claude <noreply@anthropic.com>
Every operator whose spec is a function of its parameters now says so. The eight remaining conversions plus both shared bases, leaving one override. Several were not mechanical, and the shape function is where the reason now lives rather than in an instance attribute computed at construction: - Dequant derived input_size in __post_init__; the packing rule (two 4-bit values per byte, plus a bf16 scale and zero point per group) is now stated where the shape is. - RoPE's angles broadcast, so angle_rows defaults to rows inside the function rather than only in __post_init__ -- otherwise the rule would be correct when bound from an operator and wrong when called directly. - GEMV and Transpose carry no batch dimension at all when num_batches == 1, rather than one of extent 1. - RMSNorm's optional weight sits between input and output, which is why it is not a same-shape unary despite both ends matching. - StridedCopy's two sizes are independent: it may gather from a large buffer into a small one. _SwiGLUStreamGroup keeps a get_arg_spec() override. Its spec comes from the exported workload graph reached through an instance attribute, so it is not a function of the operator's fields -- exactly the case the override exists for. That leaves 21 of 22 operators declaring a shape rule callable without an instance, which is what graph capture needs to place a value before building the operator that produces it. Specs unchanged: the snapshot gate passes across all 22 operators on both device generations. Full suite 35 failed / 40 errors, matching baseline. Co-Authored-By: Claude <noreply@anthropic.com>
…alog Two problems, one surfaced by my own change. The check read this process's sys.modules, so it measured session history rather than the import graph. Adding a test that imports MHA at module scope (operator_binding.py) broke it, even though the catalog was still perfectly lazy -- and the converse is worse: a session that happened not to touch MHA would pass even if the catalog had turned eager. Running the import in a subprocess reads the graph itself and is order-independent. The witnesses were also wrong in spirit. The docstring I first wrote claimed MHA and swiglu_decode were "expensive to import"; measured, MHA imports in 180ms against ReLU's 198ms and pulls in no top-level module ReLU does not. They were never expensive -- they were arbitrary catalog members standing in for "the rest of the catalog", which is what PEP 562 re-export (iron/operators/__init__.py) actually saves. So drive the test from _OPERATOR_MODULES instead: all fourteen operators are covered, and one added to the table is covered without touching a test. Doing that immediately found something the two witnesses could not: composite operators legitimately import their parts. SwiGLUDecode pulls in elementwise_mul, gemv and silu because it is an OperatorSequence built from them. So leaves are held to "imports nothing but itself" and composites to "does not import the entire catalog", which is the regression that matters. Added a guard-the-guard test as well. Every other assertion here is about a module being absent, so a probe that silently imported nothing -- or a typo in a module name -- would make them all vacuously true. Co-Authored-By: Claude <noreply@anthropic.com>
Each operator kept a hand-written map from its own fields to its design's signature -- GEMM's ran to eighteen entries, StridedCopy passed twelve positionally -- with nothing checking the two sides against each other. A field renamed on one side and not the other surfaced as a TypeError from inside the design, or, when the parameter had a default, as a silently wrong value compiled into a kernel. DesignGenerator now takes bind_from and fills the signature from the operator at call time. Binding late matters: design modules are imported lazily because they pull in the MLIR dialects, and reading a signature any earlier would defeat that. Explicit kwargs still win, so an operator can override or pass something it does not store. Converted StridedCopy (12 positional + 3 keys -> 0), Softmax (9 -> 0), GEMV (7 positional + 3 keys -> 0), MHA (12 -> 3) and GEMM (18 -> 7), aligning design parameter names to the operator's vocabulary where the rename was unambiguous. Two stopped short deliberately. GEMM keeps m/k/n and friends explicit: those are single letters appearing 30-odd times across a 400-line design, and a substitution could silently merge a parameter with an unrelated loop variable in a way the hardware tests would not reliably catch. MHA keeps S_q/S_kv because they are distinct design parameters that merely happen to be equal here, so neither can bind from seq_len. Both are better settled when op.py and design.py merge and the naming can be decided in one place. dev, trace_size and verbose move to the base, since every design takes them and no operator stored them. trace_size is a plain class attribute rather than a property: OperatorSequence and LayerNorm both assign self.trace_size, and a property without a setter cannot be shadowed by an instance attribute -- as a property it took the fusion suite from 80 passed to 80 failed. It is also unannotated so dataclass subclasses do not adopt it as a field. GEMM's kernel object name was written out twice, once for the design to link against and once for the artifact to build; they are now one property, so the object built and the object linked cannot drift apart. Verified on a Strix npu2: iron/tests 470 passed, converted operators 325 passed, fusion suite 80/80, arg-spec snapshot unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
ChanneledUnaryOperator and BinaryElementwiseOperator passed their designs a bare list matched by position -- [dev, size, num_aie_columns, num_channels, tile_size, 0] plus three more appended at the call site. Position is worse than the dicts already replaced: inserting a parameter into a design signature shifts every argument after it with nothing to notice. Both now bind by name, which converts ten operators at once: elementwise_add, elementwise_mul, gelu, layer_norm, relu, sigmoid, silu, tanh, and the two swiglu composites. The designs' parameter names move to the operators' vocabulary (num_columns -> num_aie_columns, num_elements -> size), and _kernel_link_file becomes kernel_obj_file to match what the design calls it. _mlir_callback_args survives for axpy and leaky_relu, which append an extra parameter and build their own artifact. layer_norm's override, which existed only to pass self.trace_size instead of a hardcoded 0, is now redundant since trace_size binds like anything else. Fixes a bug this surfaced: get_child_mlir_module repeated DesignGenerator's import-and-call rather than using it, so the fusion pass never saw bind_from and every fused dispatch died with "missing 7 required positional arguments". Both paths now share DesignGenerator.resolve(); the fusion pass needs the module object rather than its string form, which is the whole reason the duplicate existed. Verified on a Strix npu2: the ten converted operators 1015 passed, iron/tests 470 passed, fusion suite included. Co-Authored-By: Claude <noreply@anthropic.com>
_link_build_outputs_into fills an aiecc work dir from two directories and skips any name already present, so whichever is linked first wins. It linked the flat build dir first, which meant a leftover flat object shadowed the arch-scoped one -- defeating the per-arch scoping move_artifacts exists to provide, and making a design link against code compiled for another era. That is not theoretical. It is why silu, rope, rms_norm and the fused elementwise-add sequence all failed today with "undefined symbol" for symbols that were present in build/<arch>/ and absent from a months-old copy in build/: the stale one was being linked. Swapping the order fixes it. The flat directory still supplies everything that is not a kernel object -- mlir, xclbin, insts -- because those have no arch-scoped copy to take precedence. Verified by planting a deliberately corrupt flat silu.o dated August and forcing a full rebuild: flat-first fails all 75 silu tests, arch-first passes all 75 with the same corrupt file in place. Full iron/tests 470 passed. Not fixed here: something still writes both build/x.o and build/<arch>/x.o for every kernel, byte-identical and with the same mtime to the nanosecond. I could not identify the writer -- it is not shutil.copy/copy2/copyfile/move, not os.link/symlink/replace/rename, not compile_cxx_core_function (traced: one call, one arch-scoped path), and an strace of openat/creat/linkat/rename over a full rebuild shows no syscall naming the flat path. Whatever creates it, this change makes it harmless rather than hazardous. Co-Authored-By: Claude <noreply@anthropic.com>
op.py, design.py and reference.py were three files describing one operator, and the split cost more than it bought: the design was reached by path with a string function name, the reference by a function-local import, and a reader had to open all three to see what the operator was. They are now one module. DesignGenerator grows an `fn` field so a collapsed operator hands its design function over directly -- importing its own module by path would execute it a second time and build a duplicate of the class doing the asking. PythonGeneratedMLIRArtifact takes its staleness dependency from a new source_file property, which falls back to the function's own module when no path was given, so a design declared beside its operator is still tracked for rebuilds. The lazily-imported design was the one real argument for keeping them apart: loading MLIR dialects only when a design is actually generated. Measured, that costs 55 ms on top of a 225 ms operator import, and the invariant the catalog laziness test protects -- importing one operator must not import the others -- is untouched. Verified on a Strix npu2: softmax 15 passed, iron/tests 470 passed. Co-Authored-By: Claude <noreply@anthropic.com>
Same collapse as softmax, plus the bind conversion these three still needed: each passed its design a positional tuple, which matches by order alone, so inserting a parameter into a design signature shifted everything after it. transpose's design also spelled num_columns where the operator says num_aie_columns. rope keeps its README; the design and reference bodies move into op.py, and the tests and rope_reference_convention now import the reference from .op. Verified on a Strix npu2: these three plus iron/tests, 950 passed. Co-Authored-By: Claude <noreply@anthropic.com>
Four more collapses, all already bound so this is the merge alone. gemm and mha additionally carried an empty positional args tuple left over from the path-and-name form, which had to go with it -- an empty tuple after a keyword argument is a syntax error, and the merge refused to write rather than emit a broken file. gemm is the largest of these at 1127 lines in one module. That is not small, but it is the same code in one place instead of three, and the naming question it raises -- the design's m/k/n against the operator's tile_m/tile_k/tile_n -- can now be settled by reading a single file. Verified on a Strix npu2: gemm 130 passed; strided_copy, gemv, mha and iron/tests 650 passed. Co-Authored-By: Claude <noreply@anthropic.com>
Four more, each needing its bind conversion first: mem_copy and dequant still passed positional tuples, axpy and leaky_relu routed through _mlir_callback_args to append scalar_factor and alpha. Both of those are ordinary fields, so they bind by name like everything else and the override becomes unnecessary. dequant's and axpy's designs also spelled num_elements and num_columns where the operators say size and num_aie_columns. With the design local, the DesignGenerator no longer needs callback_fn -- a ClassVar holding the design's function name as a string -- since it can name the function directly. mem_copy still fails one config on this box (num_cores=16, num_channels=2, tile_size=64, size=1024, bypass=False; the five reported failures are that one case across iter0-4). That is unrelated: it reproduces identically on clean origin/devel in a worktree, compiles without error, and fails only at dispatch with ERT_CMD_STATE_TIMEOUT. Suspected driver or device difference, since CI is reportedly green. Verified on a Strix npu2: axpy and leaky_relu 325 passed, mem_copy and dequant 475 passed with the one known config failing, iron/tests 470 passed. Co-Authored-By: Claude <noreply@anthropic.com>
The last one, and the only one that needed more than a move: it had two designs, weighted and not, selected by building a different file path and function-name string. With both in the module that selection is a plain conditional over two functions, which is what it always meant. weight_length becomes a property -- the weighted design names it that, the operator calls the same quantity tile_size, and a property lets each keep its own vocabulary rather than renaming one to suit the other. Every operator with its own design is now a single file plus its kernel source. The two shared designs stay shared: channeled_unary_design.py serves seven operators and binary_elementwise_design.py three, so collapsing those would mean copying one design into ten files, which is the opposite of the point. Verified on a Strix npu2: rms_norm 295 passed, iron/tests 470 passed. Co-Authored-By: Claude <noreply@anthropic.com>
Ported from the graph-capture prototype, which had already got this right. calculate_buffer_layout assigns running offsets with no liveness analysis at all, so every intermediate in a sequence stays resident for the whole sequence and peak memory is the sum of all buffers rather than the maximum concurrently live. For a model that runs one block sixteen times that is sixteen copies of each scratch buffer. Two passes over a runlist: live_ranges gives each buffer the step interval it must stay resident for, and plan assigns byte offsets so buffers whose lifetimes do not overlap can share addresses. Greedy by size descending with best-fit placement -- Algorithm 3 of Pisarchyk & Lee (MLSys 2020), which is what TFLite ships and what TorchInductor approximates. peak_live_bytes gives the lower bound to check any plan against. Buffers the host addresses by name -- weights, caches, the sequence's own inputs and outputs -- are pinned and never pooled. The two tests covering OperatorSequence's buffer_offsets are marked xfail(strict): that parameter does not exist yet, and these pin the contract the wiring step has to satisfy. strict so they fail loudly once it lands rather than passing silently as xpass. Nothing is wired up yet, so no behaviour changes: iron/tests 470 passed plus 90 new allocator tests. Co-Authored-By: Claude <noreply@anthropic.com>
iron/tests/infrastructure/allocator.py shared a module basename with iron/common/allocator.py. Collection is nondeterministic under that: one run of the full suite came back with 30 failures and 40 errors, the next with the same tree came back clean. Renaming removes the ambiguity rather than relying on import mode to resolve it. iron/tests: 515 passed, 3 skipped, 10 xfailed. Co-Authored-By: Claude <noreply@anthropic.com>
calculate_buffer_layout assigned running offsets in declaration order, so every intermediate stayed resident for the whole sequence and an arena was as large as the sum of everything in it. OperatorSequence now takes buffer_offsets, and passing None keeps exactly the previous layout. Planned offsets are rebased past the unplanned buffers rather than applied from zero. A plan is relative to its own pool and starts at zero, so applying it directly drops the first planned intermediate on top of the weights -- and that aliasing is silent, because the arena simply does not grow. The test that caught it asserts a planned buffer starts at or after the end of every unplanned one. The two tests covering this were xfail(strict) pending the parameter. Removing the markers showed they had a second problem: they never set a device, so they died with "'NoneType' object has no attribute 'resolve'" -- which reads as a bug in the code under test rather than a missing fixture. Added the device fixture the other suites use. Nothing calls this with a plan yet, so behaviour is unchanged: iron/tests 525 passed. Co-Authored-By: Claude <noreply@anthropic.com>
OperatorSequence can now derive buffer_offsets from its own runlist: scratch_plan() walks the steps, takes each buffer's live range from first-write to last-read, and packs the ones whose lifetimes do not overlap into shared addresses. On a four-deep chain the scratch arena drops from 6144 to 4096 bytes, with the last intermediate reusing the first one's address. Buffers the host addresses -- the sequence's own inputs and outputs, and anything given an explicit size -- are pinned and never pooled: their contents outlive the sequence, so they need private, stable addresses. plan_scratch defaults to False. This is the first change here where a mistake is wrong numbers rather than a crash: two buffers aliased while both are live produce quietly incorrect results. Off by default means nothing moves until a caller asks, and the existing fusion tests keep exercising the old layout. Two tests state the contract: that planning shrinks the arena, and that no two buffers overlapping in time ever overlap in bytes. The second is the invariant a liveness bug would break, written as an assertion rather than left implicit in a numerical comparison. iron/tests: 545 passed. Co-Authored-By: Claude <noreply@anthropic.com>
Ported from the graph-capture prototype. A Traced value stands in for a
buffer, calling an operator records a step, and the buffer names that
OperatorSequence needs become generated internals rather than something a
model author types:
with capture() as g:
h = g(rms_norm, x, w)
q = g(q_proj, h, wq)
Graph.__call__ allocates outputs from the operator's own arg_spec, so a
recorded value knows its shape without a second rule -- which is what the
shape functions from layer 1 were for. infer_io derives inputs and outputs
from the recording: a buffer no step produced is an input, one no step
re-consumes is an output. scratch_plan reuses the layer 2 allocator, pinning
anything the caller named.
build() emits an OperatorSequence, so capture is a frontend over the existing
dispatch machinery rather than a replacement: fused and separate dispatch,
tracing and the ELF path are all reused untouched.
The prototype's mnist test is dropped rather than ported -- it imports an
application that does not exist here, and porting an application to satisfy a
test would be the wrong order.
iron/tests: 615 passed.
Co-Authored-By: Claude <noreply@anthropic.com>
Every capture test so far stopped at the recording -- runlist, inferred I/O, plan -- which are all statements about bookkeeping. None of them showed that a recorded graph computes anything. This adds the test that does: the same arithmetic expressed as dataflow and as a hand-written runlist must produce identical values, bit for bit. It found that only one of the two ways to get there worked. get_callable() went straight to the dispatch policy, but subbuffer_layout is populated during compile(), so dispatching without compiling first died with an AttributeError about a missing attribute rather than anything about compilation. Ahead-of-time worked because compile() was explicit; just-in-time did not work at all. get_callable() now compiles if that has not happened yet. compile() skips artifacts already on disk, so the ahead-of-time path is unchanged and arriving here twice costs nothing. Both are tested, parametrised aot/jit, and both must agree with the hand-written sequence. This is also the gate for buffer planning. build() pools scratch by default, so a captured graph already runs on a planned layout -- and two buffers aliased while both are live would show up here as wrong numbers and nowhere else, since nothing about it raises. iron/tests: 600 passed. Co-Authored-By: Claude <noreply@anthropic.com>
The numerical comparison ran only under dispatch="reference", the CPU path.
That shows the recorded wiring and the planned layout agree with a
hand-written runlist, but says nothing about the fused ELF -- which is the
path that actually runs on the device, and the one buffer planning affects.
Parametrised over both modes, so the four combinations of {aot, jit} x
{reference, fused} all have to produce identical values. Confirmed the fused
case is really the device path and not a silent fallback: the policy resolves
to FusedDispatch and the callable is SequenceFullELFCallable.
iron/tests: 600 passed.
Co-Authored-By: Claude <noreply@anthropic.com>
scratch_plan pinned the sequence's inputs, outputs and explicitly-sized buffers, but not slices. A step that writes a slice put it in the pool, and it came back with an offset of its own -- unrelated to its parent, which calculate_buffer_layout resolves it against. Nothing raises: the slice simply reads the wrong memory. Found by probing the written-slice case directly. The existing tests use whole buffers, so none of them could reach it, and Llama's decode path is full of slices -- it would have shown up there as wrong tokens. The capture prototype already pinned slices; porting scratch_plan onto OperatorSequence is where it was dropped. iron/tests: 615 passed. Co-Authored-By: Claude <noreply@anthropic.com>
Two things the old naming got wrong.
scratch_plan() read like a noun a caller supplies. It is not: the layout is
derived entirely from the runlist -- liveness from the recorded order, sizes
from each operator's arg_spec -- and nobody passes it in. Renamed to
infer_buffer_offsets(), which says what it does. buffer_offsets stays as the
escape hatch for a caller who wants to override the inference.
plan_scratch defaulted to False. That was right when planning was unproven:
leaving it off kept the fusion tests running on the old layout as a control.
It is no longer right. Planning is now checked bit-exact on the device across
{aot, jit} x {reference, fused}, and the one real hole -- pooling a sliced
buffer, which aliases silently -- is fixed and pinned by a test. Captured
graphs already planned by default, so hand-written sequences behaving
differently was an inconsistency rather than a safeguard.
So it defaults to True, and plan_scratch=False becomes the escape hatch back
to packing every buffer back to back.
iron/tests: 615 passed, the eighty fusion tests now running on inferred
layouts.
Co-Authored-By: Claude <noreply@anthropic.com>
…guish Retiring IRON's artifact graph onto CompilableDesign only works if its key tells two captured graphs apart. It does not, in the obvious encoding, and that had to be established before building on it. Two generators that close over different MLIR but share a code object get the SAME cache key: the recipe hash covers the code object and compile_kwargs, not closure contents. Handing captured graphs over as bare closures would give the second one the first one's artifacts, silently. I nearly concluded the opposite. Probing it with `lambda: a` and `lambda: b` shows different keys -- but those lambdas name different variables, so they have different code objects, and the difference had nothing to do with the graphs. Two captured graphs go through one call site and share a code object. The probe has to keep the code identical and vary only the closure, which is what these tests do. compile_kwargs IS part of the recipe hash, so that is where a graph's identity has to go. Also pinned: full_elf is in the key (fused and separate produce different artifacts from the same MLIR), and an unchanged graph keeps its key so the cache can hit at all. Feasibility itself is confirmed: a captured three-step graph produces fused MLIR with three aie.device blocks, and CompilableDesign accepts it. iron/tests: 620 passed. Co-Authored-By: Claude <noreply@anthropic.com>
The seam for retiring the artifact graph. It takes a sequence that has already produced its fused MLIR and compiles that half the upstream way, leaving the rest alone -- so the move can happen in steps instead of one deletion that has to land whole. Four things about the upstream API are not guessable from its signature, and each cost an iteration to find: - compile_kwargs keys must be in the generator's signature AND carry a CompileTime[T] annotation. Note that `from __future__ import annotations` breaks this: the annotation becomes a string and get_type_hints resolves it against module globals, so a function-local import of CompileTime leaves it unresolvable and the key is rejected as unexpected. - The generator must return an MLIR Module. _generate_uncached calls module.operation.verify() on whatever it gets, so text raises AttributeError. - object_files does NOT stage anything; it feeds the artifact hash only. Objects must be copied into the work dir under bare names, because the fused MLIR's link_with asks for "op0_add.o" with no directory. This corrects the plan, which had _link_build_outputs_into being deleted along with the DAG -- staging is load-bearing and has to survive. - The cache key does not see closure contents, so two graphs whose generators share a code object collide. The MLIR's digest rides in compile_kwargs to keep them distinct. Checked on hardware: a captured two-step graph compiles to a linked full ELF, verified by its magic bytes rather than by its existence. iron/tests: 670 passed. Co-Authored-By: Claude <noreply@anthropic.com>
The seam produced an ELF, which is not the same as producing the right one.
Compiled against the artifact rule it replaces, it came out 70,936 bytes to
the rule's 99,768 -- because the rule passes two flags the seam did not:
--expand-load-pdis switches PDIs between steps, which is what a
multi-device runlist is
--get-scratchpad-parameters emits the parameter table the host writes to
Neither is tuning. Without them the result links, loads and looks fine, and is
a different program. Nothing reports it.
Added a parity test asserting both paths build the same byte count, with those
two numbers in the docstring so a future flag regression reads as "the flags
diverged" rather than as an unexplained inequality. Byte-for-byte equality is
not available: aiecc embeds its working directory.
Still missing from the seam: --get-input-with-addresses, which the rule adds
when trace_size > 0. Tracing is not handled here yet.
iron/tests: 665 passed.
Co-Authored-By: Claude <noreply@anthropic.com>
The seam ignored trace_size, so switching FusedDispatch onto it would have broken traced builds without any sign. The rule adds --get-input-with-addresses when trace_size > 0 because the trace parser reads the lowered module for the buffer layout and each design's traced tiles; without it the ELF builds, loads and runs, and there is simply nothing to parse. That flag also has to reach the cache key. The MLIR is identical either way, so a traced and an untraced build of the same graph would otherwise share an entry, and the traced one would be handed an ELF with no trace in it. trace goes into compile_kwargs alongside the graph digest, which is the half of the key that sees them. extra_flags is threaded through too; the rule has always forwarded those. Parity is checked for a traced build as well, against the same rule. iron/tests: 680 passed. Co-Authored-By: Claude <noreply@anthropic.com>
SequenceFullELFCallable asserted artifacts[0] was a FullElfArtifact and read its filename, so dispatch was tied not just to an ELF existing but to the artifact graph having been the thing that built it. A sequence compiled through CompilableDesign has exactly the same ELF and no such artifact. full_elf_path(seq) returns an explicit elf_path when one is set and falls back to the artifact otherwise, so both producers work and the failure names what is actually wrong rather than tripping an isinstance assert. One of the three couplings that have to come apart before FusedDispatch can move. The other two are harder and are not addressed here: FullElfArtifact is what *causes* the fused MLIR and the kernel objects to be built -- they are its dependencies, and it is the only artifact registered -- so removing it removes the reason its own inputs exist. Those two have to become targets in their own right first. Behaviour is unchanged; nothing sets elf_path yet. iron/tests: 670 passed. Co-Authored-By: Claude <noreply@anthropic.com>
The switch. FullElfArtifact is no longer registered; the fused MLIR and the kernel objects become targets in their own right, and link_elf() produces the ELF through CompilableDesign once they exist. Untangling that required the artifact to stop being load-bearing in two ways at once. It was the only artifact registered, so it was both the output and the reason its own inputs got built -- its dependencies were the MLIR and the objects. And SequenceFullELFCallable asserted on its type to find the ELF path, which the previous commit replaced with full_elf_path(). The parity tests that gated this are removed, because they compared against a rule that no longer runs. One is replaced by a check that the trace flag still reaches aiecc -- and correcting it is worth recording: --get-input-with-addresses does not change the ELF, which comes out the same size either way. It emits a side file, input_with_addresses.mlir, and that file is what the trace parser reads. Asserting on ELF size passed for the wrong reason before the switch and failed for the right one after; the test now looks for the file. Verified on a Strix npu2: the eighty fusion tests pass on the new path, no FullElfArtifact is registered, elf_path points at the CompilableDesign output, and iron/tests is 665 passed. Co-Authored-By: Claude <noreply@anthropic.com>
# Conflicts: # iron/common/compilation/base.py Co-authored-by: hunhoffe <54562339+hunhoffe@users.noreply.github.com>
Co-authored-by: hunhoffe <54562339+hunhoffe@users.noreply.github.com>
…e ports Resolved onto the declaration layer, and carrying three commits from operator-model-rework that devel does not have: - edee6a4: a packed operand is declared in bfp16ebs8 blocks, not bytes. BoundBuffer grows host_shape/host_dtype (the host's run of uint8), sizes go through aie.utils.bfp.itemsize, and flm.GEMM's B is In((packed_blocks,), dtype=v8bfp16ebs8). This is the fix for every flm.GEMM shape timing out. - f884c74: flm/dequant as an overlay, adapted to the fork's _build(), the npu_runtime fixture and numpy-only packing. - acf5aaf: its pin is superseded by devel's dev43; nothing else to port. Against mlir-aie head (#3800 merged) the _EXTERN_CACHE workaround in build_design is gone: the cache it cleared no longer exists. Merge fallout fixed here: - flm l1_budget reads dev.arch, as the tuning does, instead of resolving the device through target_arch (devel's version), which a stub device cannot. - swiglu_prefill takes b_col_maj and passes it to all three GEMMs; its test traces the graph to check that, replacing devel's stale layout test. - The run_test timing test uses a real operator's buffers rather than SimpleNamespace stand-ins, which have no host_shape. Co-Authored-By: Claude <noreply@anthropic.com>
…alues" _write_values looked up the callable's params with getattr(instance, ..., None). params is a property, so an AttributeError raised inside it -- pyxrt.run.get_ctrl_scratchpad_bo missing on XRT 2.21 -- was swallowed and the call fell through to "SequenceFullELFCallable takes no per-call values". Look the property up on the class and let what it raises through. Co-Authored-By: Claude <noreply@anthropic.com>
mlir-aie head's LUT activation factories reject a tile below 1024 or not a multiple of 32; the declared lowering case used 256. 1024 is valid for the pinned wheel too. Co-Authored-By: Claude <noreply@anthropic.com>
Port of 3731ae0 from operator-model-rework. get_buffer() hands out writable views into the fused ELF's scratch buffer (weights, KV caches), but _sync_inputs only flushed the input buffer. The full-ELF callable calls run_handle.start() directly, skipping the host runtime's per-argument flush, and NPU access to scratch is not cache-coherent, so a graph's closed-over weights reached the device late or not at all: swiglu_decode returned wrong output on its first call. test_non_input_buffers_sync_without_explicit_flush is ported to numpy. Without the flush its fused case fails 5/5 iterations. Co-Authored-By: Claude <noreply@anthropic.com>
fused_design fuses the designs once outside compile() to digest them for
the cache key, and that run registers each design's ExternalFunctions.
compile() clears ExternalFunction._instances only when it generates, so on
a cache hit they stayed registered. The next fusion naming one of their
object files with other flags then raised a collision: GEMM's b_col_maj
changes its compile flags, not gemm_{m}x{k}x{n}.o, so swiglu_prefill
passed cold and failed on every warm rerun.
The key's call now owns the registry lifecycle, as upstream expects of
anything generating outside compile(). The toolchain test builds a
swiglu_prefill twice, checks the hit leaves nothing registered, then
builds the b_col_maj variant; it fails without the clear.
Co-Authored-By: Claude <noreply@anthropic.com>
The fork's harness and graph layer are numpy-only (43fd0da); these tests still called torch methods on numpy arrays, or closed torch bf16 weights over a graph, which np.asarray cannot take. swiglu_decode's reference stays torch, since a stream test checks a torch module against it bit for bit; as_numpy and bf16_matmul adapt it. swiglu_decode also read the up projection from step 2, which is SiLU; the up GEMV is now found by type. Co-Authored-By: Claude <noreply@anthropic.com>
mlir-aie head's LUT activations reject a line under 1024 elements, so the 256 default and the 128-512 tiles the case sweep split 2048 into all failed in the factory. The overlays default to 1024, and channeled_unary_cases takes a tile_floor that drops the splits below it. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
A fused build prefixed each design's kernels with its step index, named each aie.device for its step, and fused every design's MLIR again on a cache hit just to compute the key. So identical code compiled once per position, no device was reused across graphs or extents, and three quarters of a warm compile was regenerating text it threw away. - declare_kernel prefixes a kernel's symbol and object with a digest of its recipe (source bytes, flags, include dirs). Equal recipes share one object whether standalone, fused, or in another graph. Target loses func_prefix; stream's designs, which name their own kernels, opt out. - A fused device is named <Operator>_<design digest>, so one design is one device at any step of any sequence, and repeats merge. - fused_design keys on fused_identity (design code, parameters, source digests, the runlist and buffer layout) and generates MLIR only on a miss. Llama 1B, npu2: warm decode 0.52s -> 0.28s, warm prefill 0.89s -> 0.36s; cold decode 14.0s -> 9.5s (12 kernel compiles, was 18). With mlir-aie's referenced-ops device key, decode at L=1024 after L=2048 re-places 5 of 17 devices instead of all of them. Co-Authored-By: Claude <noreply@anthropic.com>
…he fork-side ports Stacks the operator-model rework on #220: kernels come from the aie.iron.kernels factories, each operator's tolerance comes from its kernel's contract, and verification and benchmarking use mlir-aie's Tolerance/compare, run_iters and get_trace_buffer. Ported onto the rework's architecture: - Overlay.tolerance(target) and Operator.reference_tolerance() read the kernel contract; the Testing catalog and elementwise ops use it. - SequenceCompareCallable judges each step by its operator's contract (typed tolerance / raise_on_mismatch); new Tanh test covers both verdicts. - The full-ELF callable binds the trace buffer from the lowered module's #aie.trace_buffer layout. - Operator.dev binds the device (ensure_current_device): the factories read only a bound device and otherwise fall back to aie2. - generator_for passes the operator's trace_size to build_design; it was never reaching the design, which the old mocked tracing test hid. The tracing test is now #220's real hardware one. - run_test takes a typed Operator and raises TypeError otherwise. - flm.GEMM keeps its hand-declared fused kernel (its -D flags are its own). Fallout fixed on the way: - Llama: the images return numpy but the harness samples in torch, so llama_forward_pass converts the logits at the boundary (_torch). - Llama test reads TTFT/TPS from stderr too, where the harness writes them; gains #220's accuracy (fp32 CPU reference) and determinism tests. - Dropped the relative_build_dir test, which targets devel's build tree. - AGENTS.md rewritten for the rework; MLIR_AIE_KERNEL_SOURCES rename. Verified on Strix Halo (npu2): catalog non-extensive 570 passed, 15 failed (DynamicSoftmax 2 columns x 2 channels, failing before this merge); infrastructure tests pass, including tracing. The Llama tests on the rework are not yet run. Co-Authored-By: Claude <noreply@anthropic.com>
When a B fill needs several shim descriptors (b_col_maj at 2048x8192x2048 on eight columns), issuing them column by column pushes the second row-block's B tasks onto a column whose cores still wait for A from the columns not yet issued. The shim task queue is only a few deep, so the push stalls the whole instruction stream and the dispatch hangs. Llama 3.2 1B prefill's down projection is that shape; add it to the extensive parameters. Co-Authored-By: Claude <noreply@anthropic.com>
CompiledGraph copied each weight in as `.astype(view.dtype)`, which builds a whole temporary first. For Llama's 501 MiB tied embedding that took 5-50 s per upload, most of it faulting in the temporary; assignment casts straight into the buffer view instead. The upload also ran lazily inside the first prefill and decode calls, so both counted towards time to first token and tokens per second, where devel writes its weights before timing starts. CompiledGraph.load() uploads now, and AIELlama calls it. Prompt 1024, 40 tokens: TTFT 14.56 -> 1.24 s, decode 0.64 -> 7.58 tok/s. Co-Authored-By: Claude <noreply@anthropic.com>
Prompt 1024, 40 tokens, teacher-forced: prefill KL 0.026, decode max 0.015, 2 top-1 mismatches; determinism 0/8 differing runs. #220's llama_npu.py on the same mlir-aie build measures 0.074 and 0.013, the figures the comment quoted before. Co-Authored-By: Claude <noreply@anthropic.com>
A resident (weight, state) is keyed by its storage and gets one offset, the same in every image placed in the arena; each image's transients are planned around the residents and reuse any other image's transient bytes, since only one image runs at a time. Nothing placed ever moves, so a later image puts new residents on top and the arena only grows; ScratchArena grows its XRT buffer to match and full-ELF callables rebind to it. Every offset, pooled or packed, is now a multiple of the coherence granule. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Calling a graph function at a new shape compiles a version for it, and every full-ELF version is placed in the function's one ArenaPlan: its weights and states are residents, uploaded once and read and written at the same offset by every version. Versions that cannot share (an xclbin chain) refuse a state rather than silently copying it. Bindings are typed (Binding: op, member, value, symbol) and callables take per-call values through write_values(), replacing the getattr/hasattr probing in CompiledGraph. A sequence settles its mode before placing anything in a shared arena. Co-Authored-By: Claude <noreply@anthropic.com>
npu.py, harness.py and graphs.py run on numpy: the weights are the mapped safetensors checkpoint, the embedding a numpy gather, the RoPE table the numpy one, sampling the seeded numpy Sampler. torch is left to the CPU reference (model.py, reference.py) and the accuracy entry point, which is its own module (accuracy.py) so the application never imports it. The NPU and the reference read one RoPE table. AIELlama takes its images as arguments, so the CPU tests stand the graph references in without a monkeypatch, and run the harness's accuracy and determinism checks. Co-Authored-By: Claude <noreply@anthropic.com>
lower() cleared ExternalFunction's registry before generating a design but left that design's kernels registered, so lowering_graph.py run before full_elf.py failed the cached-build test on GEMM's leftover kernels (a directory run orders them the other way, and passed). Co-Authored-By: Claude <noreply@anthropic.com>
LlamaGraph holds forward(x, angles, *, cache_offset, vector_size, last), which branches on the static shape of x: one row is a decode step (GEMV projections, the row written into the caches at cache_offset, the softmax masked to vector_size), many rows a prompt (GEMM projections, MHA, the caches written from row zero, the head over row `last` alone). The two versions it compiles, one per input shape, run in the function's one scratch arena: the weights and caches sit at one offset in both images and are uploaded once, so the caches a prompt writes are the ones the next decode step reads. prefill_to_decode and the second graph are gone. AIELlama passes the RoPE table as bf16, the dtype the versions are compiled for; a float32 table is another signature and another compile. On hardware (Strix Halo): test.py 6/6; accuracy identical to the two-graph version (prefill KL 0.035112, decode max KL 0.014146, 2 top-1 mismatches); the generated text for the same seed is byte-identical to it. Co-Authored-By: Claude <noreply@anthropic.com>
Llama 3.2 1B's checkpoint config has rope_scaling {factor 32,
low_freq_factor 1, high_freq_factor 4, original_max_position_embeddings
8192, rope_type llama3}; the table ignored it, and so did the CPU reference,
since both read config.angles. For head_dim 64 and base 500000 it keeps the
fastest 15 frequencies, interpolates 3 and divides the slowest 14 by 32, so
it changes the model at every position, not only past 8192.
Llama3RopeScaling applies it to the frequencies in float64, before their
one rounding to float32; tested against Meta's published apply_scaling.
Co-Authored-By: Claude <noreply@anthropic.com>
The config's table covers the model's 131072-position context; the images read at most max_seq_len rows of it, so the bf16 copy was 16 MiB where 0.25 MiB is used. Co-Authored-By: Claude <noreply@anthropic.com>
The prefill KL rose from 0.035 to 0.091 with the scaling. It is one position: over 140 positions of prompt.txt the NPU's KL against the fp32 reference has the same distribution with and without the scaling (median 0.006), and the bf16 RoPE table costs 5e-6 of it. The fp32 reference matches Hugging Face transformers' LlamaForCausalLM to KL 1e-8, scaled and unscaled. Co-Authored-By: Claude <noreply@anthropic.com>
--prompt-len slices prompt.txt by characters, but setup() added it to the generated-token count and checked the sum against MAX_SEQ_LEN's rows, so the defaults (2048 characters, about 580 tokens) failed an assert the run fits in. Check the tokenized prompt instead, and say what --prompt-len counts. Co-Authored-By: Claude <noreply@anthropic.com>
The mapped checkpoint stayed resident after upload, so a run held it and the buffers at once: VmHWM 5293 MiB, of which RssFile 2457 was the mapping and RssShmem 2605 the buffers. CompiledGraph.upload/load take a release callback, called with each weight as soon as it is in its buffer (in an arena it is never read again: a grown arena keeps the device's contents). SafetensorsFile.release madvises a view's pages away; the mapping is read-only and of the file, so a later read faults them back. AIELlama passes LlamaWeights.release, which skips arrays that are not the checkpoint's (the attention scale). VmHWM 5293 -> 3413 MiB, steady RSS 5293 -> 2975 MiB; the remaining ~440 MiB peak is the 501 MiB embedding uploaded as one weight. Accuracy bit-identical (prefill KL 0.090601, decode max 0.023189). Co-Authored-By: Claude <noreply@anthropic.com>
In a shared arena the first version loaded uploads every weight, so a later version's upload() found nothing to copy and never made its runtime; the first call did. For Llama that was the prompt's image: its hw context and ELF load landed in the first prefill, 88 ms of TTFT (1.280 s first prefill vs 1.192 s steady). Interleaved over 8 rounds, TTFT is now 1.214 s against the two-graph version's 1.234 s, decode unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
One step's KL is as much the position's as the NPU's: the prefill step's 0.091 is one of two positions in 140 above 0.05. The mean (0.0083) and p90 (0.018) over all 40 steps bound a drift; a loose max bounds a single broken step. Co-Authored-By: Claude <noreply@anthropic.com>
upload(release=...) now hands release every piece of a weight, a flat view of at most piece_bytes (64 MiB), as soon as it is in its buffer, rather than the whole weight after. A model whose weights view a mapped checkpoint then holds at most one piece of it beside the buffers. For Llama that was the 501 MiB embedding: peak RSS 3372 -> 2939-2976 MiB, now equal to the steady state; load time unchanged (1.7 s). Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
CI Test Results250c6e1 (2026_09_26_02_24_15) IRON - CI SummaryExamplesiron/applications/llama_3_2_1b
Smalliron/operators
iron/operators/flm/dequant
iron/operators/flm/gemm
iron/operators/gemm
iron/operators/gemv
iron/operators/mha
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
Krackan - SmallIRONTested on iron/operators
iron/operators/flm/dequant
iron/operators/flm/gemm
iron/operators/gemm
iron/operators/gemv
iron/operators/mha
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
Krackan - ExamplesIRONTested on iron/applications/llama_3_2_1b
Phoenix - SmallIRONTested on iron/operators
iron/operators/flm/dequant
iron/operators/flm/gemm
iron/operators/gemm
iron/operators/gemv
iron/operators/mha
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
Phoenix - ExamplesIRONTested on Trend tables omitted, the comment hit GitHub's size limit. Full report in the workflow run. |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe the intent of your PR here.
Added
Changed
Removed
PR Merge Checklist
develcommit and pointing todevel.