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>
Dead once FusedDispatch stopped registering the artifact: nothing constructs one, so AieccFullElfCompilationRule can never match. 57 lines of the artifact graph, and the first of it to actually go rather than be routed around. full_elf_path() loses its fallback with them. It now fails with what is actually wrong -- link_elf() has not run -- instead of reporting that no FullElfArtifact is registered, which would have been true of every sequence and told the reader nothing. iron/tests: 665 passed, the eighty fusion tests among them. Co-Authored-By: Claude <noreply@anthropic.com>
Two things, both found by the xclbin path failing to link. Staging has to happen inside the generator. A cache miss runs _cleanup_failed_compilation on the work directory before compiling, so objects put there beforehand are wiped. The generator runs after that clear and before aiecc, which is the only window where they survive. The fused path worked by luck of ordering; the xclbin path did not, and its work directory was empty at link time. compile_xclbin_insts is the separate-dispatch counterpart. Chaining looked like it needed something CompilableDesign lacks -- each operator's xclbin links onto the previous one's via --xclbin-input -- but that and the kernel name are both aiecc flags, which it already forwards. No local subclass is needed. The predecessor goes into the cache key: two operators with identical MLIR chained onto different xclbins are different artifacts. It produces a 25,081-byte xclbin and a 3,248-byte insts stream for a standalone ElementwiseAdd. Worth recording why this looked broken first. The initial probe reused an operator whose .mlir had been written by a fused build, which mutates generator.kwargs["func_prefix"] without changing the artifact's filename. IRON's filename+mtime cache handed back the prefixed MLIR, so a standalone operator asked for op0_add.o and failed to link. That is a real defect in the artifact graph, not just a bad probe: a fused build poisons the standalone cache entry for the same operator, and content addressing makes it unrepresentable. iron/tests: 665 passed. Co-Authored-By: Claude <noreply@anthropic.com>
build_fused_mlir mutates each operator's MLIR generator to add a func_prefix without changing the artifact's filename. Those artifacts are dependencies of the SequenceMLIRArtifact, so they are compiled to disk -- writing symbol- prefixed MLIR to the exact path a standalone build of the same operator reads. The cache keys on filename and mtime, so a later standalone build trusts it and asks the linker for op0_add.o, which no standalone build produces. The failure lands a long way from the cause: an undefined symbol at link time, in a build that did nothing wrong, in a different process or session from the fused build that poisoned the slot. Reproduced deliberately -- fused build in one process, standalone in another, same operator config -- after an initial attempt failed to show it. That attempt used a size the fused build never touched, so the standalone read a file no fused build had written. Worth recording: I had already asserted this defect in a commit message off a single contaminated observation, then retracted it when the bad probe came back clean. Both the claim and the retraction were made on evidence that could not support them. The prefix changes what the MLIR is, so it now changes where it is written. iron/tests: 670 passed, with the new case failing before the fix and passing after. Co-Authored-By: Claude <noreply@anthropic.com>
The DAG's staleness check only compared filename and mtime, so a fused build mutating a shared operator's generator.kwargs (func_prefix) in place -- without touching the artifact's path -- left a standalone build trusting a stale, differently-prefixed file with a newer mtime than its source. 7335940 patched that one instance by also renaming the fused artifact's filename by hand. PythonGeneratedMLIRArtifact now stamps a recipe hash (generator code identity + kwargs, via upstream's _compute_recipe_hash) alongside its MLIR on generation, and is_available_in_filesystem() checks it. That makes the whole class of collision unrepresentable instead of relying on every mutation site remembering to rename around it, so the manual rename in FusedDispatch.build_fused_mlir is now redundant and removed -- the on-hardware regression test still passes without it. device kwargs go through _device_identity_key rather than str(device): the default object repr embeds a memory address, which would invalidate on every fresh from_name() call even when the device hasn't changed. iron/tests: 700 passed, 3 skipped (up from 670 baseline; the increase is the new recipe-hash unit tests). Co-Authored-By: Claude <noreply@anthropic.com>
Mirrors what FusedDispatch already does for the full ELF: kernel objects still go through the artifact graph (Peano/chess compile isn't on CompilableDesign yet), but the per-operator xclbin+insts chain is now built by jit_compile.compile_xclbin_insts() at make_callable() time instead of XclbinArtifact/InstsBinArtifact/AieccXclbinInstsCompilationRule. CompareDispatch inherits set_up_artifacts from SeparateDispatch and needed the same link_xclbins() call added to its make_callable(). iron/tests/infrastructure/sequence.py: 80 passed, including "separate" and "compare" dispatch and their bit-identical parity check against "fused" -- run on hardware, not mocked. Co-Authored-By: Claude <noreply@anthropic.com>
SeparateDispatch was the only caller that ever set xclbin_input on an
XclbinArtifact, and it now chains through jit_compile.compile_xclbin_insts()'s
own plain-path parameter instead (see the prior commit). Nothing else
constructs an XclbinArtifact with it, so the field and the
--xclbin-input branch in AieccXclbinInstsCompilationRule were dead.
iron/tests/infrastructure/{sequence,jit_compile_path,compilable_design_contract}.py:
180 passed. iron/operators/flm/gemm/test.py (the standalone XclbinArtifact
caller): 185 passed.
Co-Authored-By: Claude <noreply@anthropic.com>
SequenceMLIRArtifact + FusePythonGeneratedMLIRCompilationRule wrote the fused MLIR to disk purely as DAG bookkeeping -- fuse_mlir() never read that file back; it always called each operator's generator in-memory via get_child_mlir_module(). Removing the artifact+rule pair and making fuse_mlir() a plain function that returns text means a fused build no longer writes any per-operator .mlir file to disk at all, which makes the cache-poisoning bug class from the previous commits structurally impossible here rather than just detected. compile_sequence() (jit_compile.py) used to locate the fused MLIR by scanning the artifact graph for a "_fused.mlir"-suffixed filename; it now calls FusedDispatch.build_fused_mlir() directly, since there's no artifact left to scan for. iron/tests: 700 passed, 3 skipped (unchanged baseline). iron/operators + iron/applications: 3165 passed, 5 failed (the one known-unrelated mem_copy 16-core hardware timeout), 21 skipped. Co-Authored-By: Claude <noreply@anthropic.com>
CompilableDesign.compile() bypasses its own on-disk cache entirely whenever explicit xclbin_path/inst_path (or full_elf_path) are given -- confirmed in its source, not just the docstring: the cache-hit branch is gated on `not explicit_paths`. compile_fused_elf/compile_xclbin_insts always pass explicit paths, so every call recompiled through aiecc even with a byte-identical recipe -- measured directly: two independently constructed OperatorSequence instances with the same config each rebuilt the ELF (mtime changed both times). _compile_if_changed() reuses CompilableDesign's own content hash (already relied on by compilable_design_contract.py) rather than inventing a second one, and stamps it in a sidecar next to the first output, mirroring PythonGeneratedMLIRArtifact.recipe_hash()'s existing pattern. Both compile functions skip the rebuild (and the object-staging copy) on a hit. New tests in jit_compile_path.py assert ELF/xclbin mtime is unchanged across two independently-built, identical-recipe compiles. Full iron/tests: 710 passed, 3 skipped. Co-Authored-By: Claude <noreply@anthropic.com>
MLIROperator.set_up_artifacts() was the last dispatch mode still on the old XclbinArtifact/InstsBinArtifact/AieccXclbinInstsCompilationRule path; FusedDispatch and SeparateDispatch moved onto jit_compile.py earlier this session. Kernel objects stay on the artifact graph (Peano/ chess compile isn't on CompilableDesign yet -- its own kernel auto-compile only fires for upstream's ExternalFunction, which no IRON design uses, all 14 use plain Kernel(name, prebuilt_object)). link_xclbin() is new: lazy and idempotent, mirroring FusedDispatch.link_elf/SeparateDispatch.link_xclbins, called from get_callable() the first time a standalone operator actually needs a compiled binary. It's zero new compile logic -- compile_xclbin_insts() already supports the no-chaining single-operator case. get_artifacts() is deleted (confirmed zero callers left after SeparateDispatch moved off it earlier this session). Two operators are deliberately left untouched: flm.GEMM overrides set_up_artifacts() with a runtime-parameter scheme (one xclbin reused across every shape sharing a config, verified by its own test_one_xclbin_serves_every_shape/every_clamp_bound tests asserting the xclbin's path *and mtime* stay identical) and used to inherit get_callable() silently -- it now gets an explicit override, a verbatim copy of the old base implementation, so the base class changing under it can't break it silently. flm.MMPrebuilt already overrides get_callable() explicitly (hardcoded kernel name for its downloaded xclbin) and is unaffected. Neither operator's set_up_artifacts() calls super(), so neither is touched by this change. mlir_cache_poisoning.py's _linked_objects() helper called operator.compile() and read a .mlir artifact off disk; standalone operators no longer write one (same reasoning as the fused-path change earlier this session), so it now calls the generator directly -- and no longer needs to compile at all for this check. iron/tests: 710 passed, 3 skipped (unchanged). mlir_cache_poisoning.py + kernel_object_arch_isolation.py: 40/40 passed, exercising the new standalone path (ElementwiseAdd) and kernel-object registration directly. Full iron/operators + iron/applications regression still in progress. Co-Authored-By: Claude <noreply@anthropic.com>
The previous pin (dev4) predated two upstream changes IRON needs together: #3584, which adds symbol_prefix plumbing and ships llvm-nm in mlir_aie/bin, and the llvm-tool-discovery change, which teaches aie.utils.config to search the peano tree as well as its own. dev18 was the newest nightly when this was last checked and was literally the commit before #3584; dev26 has both. Verified on the 8-col Strix with XRT 2.26: iron/tests is 710 passed / 3 skipped before and after, i.e. every previously-passing test still passes. The bump also regenerates every dialect binding for a new LLVM, so the null result is the point. Co-Authored-By: Claude <noreply@anthropic.com>
mlir-aie now provides prefix_symbols_in_object() and resolves llvm-objcopy, llvm-nm and llvm-ar through aie.utils.config, so the copies here can go. _prefix_symbols() built the nm -> rename-map -> objcopy pipeline by hand and carried a separate implementation per platform: a Windows branch that shelled out to an inline python -c script, and a POSIX branch that ran nm and awk under sh. Both are replaced by one PythonCallbackCompilationCommand, which also drops the .symbol_map and .symbol_map.syms files this left beside every prefixed object. The upstream parser reads the symbol name as the last field rather than awk's positional $3. _find_tool/_find_working_tool/_tool_runs searched peano_dir, mlir_aie_dir and PATH, because upstream's resolvers looked only in the mlir-aie bin directory and would not find llvm-nm or llvm-ar, which ship with peano instead. That was a workaround, not duplication -- and it is still load-bearing, just upstream's now: on this box objcopy and nm resolve into mlir_aie/bin while ar resolves into llvm-aie/bin. The execute-it-first guard _find_working_tool added lives there too. peano_dir is no longer threaded into the rules, and ArchiveCompilationRule needs no constructor at all. Verified three ways on the 8-col Strix, since a cached object would make this vacuous: driving KernelCompilationRule over a real kernel with prefix_symbols set emits two commands instead of three and llvm-nm reads back op0_add_one and op0_helper_fn; a cold fused build from an empty build dir produces op0_add.o with op0_eltwise_add_bf16_* and zero .symbol_map sidecars; iron/tests is 765 passed / 13 skipped, unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
A checkpoint is a state_dict, so the thing that reads one should be an nn.Module. Declaring the tree once buys load_state_dict to fill it, named_parameters() to walk it, and -- the point here -- one name per weight that is the same string on the checkpoint, in the module tree, and on the device buffer. llama_npu.py currently spells out nine hand-typed HF keys per layer on the decode side and a matching list on the prefill side, with a .T on some and not others; llama_cpu.py keeps a third copy. This is where that becomes one row of FROM_HF. The tree holds parameters and nothing else -- no forward. What llama computes stays in llama_npu.py and llama_cpu.py; a third opinion on the same arithmetic would be the duplication this is meant to remove. Nothing consumes it yet, so this commit is additive: the uploads move over next. Three deviations from the obvious spelling, each measured rather than assumed: nn.RMSNorm(dim).eps is None, which silently means finfo(bfloat16).eps ~= 0.0078 instead of llama's 1e-5, so eps is passed explicitly; from_hf builds on the meta device and loads with assign=True, because otherwise the constructor kaiming-initialises 1.236 B parameters (~2.5 GB) purely to overwrite them and the filled tree holds a second 2.5 GB; and requires_grad_(False), since nothing here trains and a consumer feeds a weight straight into a host F.linear. Tested against the real 2.47 GB Llama-3.2-1B checkpoint, not just a stand-in: all 146 checkpoint keys map, none go unused, the tree reports 1.2358 B parameters, and every parameter shares storage with the checkpoint tensor it came from. The shaped stand-in tier cannot show that FROM_HF matches the real key spelling, so both tiers exist. Co-Authored-By: Claude <noreply@anthropic.com>
The decode side spelled out nine Hugging Face keys per layer, the prefill side kept a matching list with a .T on six of them, and llama_cpu.py kept a third copy -- 35 hand-typed key strings across three files, each one a chance for the checkpoint, the buffer and the reference to disagree silently. The device buffers are renamed to the module tree's parameter names, so decode's whole upload becomes a loop over named_parameters(). That is not just shorter: get_buffer() raises on a name only one side knows, so a rename now fails at startup instead of leaving a weight zeroed. The W_* strings were local to llama_npu.py and never reach MLIR or a filename -- the only parsing done on a buffer name is calculate_buffer_layout's "[" slice test -- so dots are safe. Prefill keeps its own layout, because its GEMM wants each projection K-major while decode's GEMV wants it as shipped. That disagreement is one keyword on _upload() rather than a second key list, so prefill now names weights by attribute access and a typo is an AttributeError at construction. out_head is the exception that rules out deriving the layout from the module type: it is an nn.Linear but must not be transposed. Verified byte-identical rather than argued: for all 146 parameters the tensor this uploads is torch.equal to the one the old key strings fetched, and the prefill transposes match too. Llama now runs end to end on the NPU -- 4/4 of iron/applications/llama_3.2_1b/test.py, first time, since XRT here is 2.26 and the hw_context path that needed >=2.21 is no longer blocked. Two fixes this uncovered, both pre-existing and only reachable once llama could actually run: SequenceFullELFCallable.params and .lowered_mlir_text still read self.op.artifacts[0].mlir_input, but 02f78d8 made the fused MLIR stop being an artifact, so artifacts[0] is now always a KernelObjectArtifact and both raised AttributeError. The work dir is derived from the ELF path instead, through a named fused_work_dir() so the convention lives in one place the way _aiecc_work_dir's docstring asks. The attention-scores GEMV used K == head_dim == the default vector size, and mv.cc in mlir_aie 1.4.4.dev26 added static_assert(k >= 2*r) to protect a pipelining pragma that assumes two iterations. Pinned to 32, checked against the golden reference at llama's exact shape (M=1024, K=64, 32 batches) with no errors. Upstream's assert is arguably too strong -- k == r is a well-defined matvec -- and that is worth raising there. Not resolved here: NPU decode output degrades after a few tokens relative to llama_cpu.py on the same prompt and seed (prefill reproduces the prompt exactly and the first generated tokens agree). It is not this change -- the bytes are identical -- but it predates any observation, because llama_npu.py could not run on this host until today. test.py asserts only returncode == 0, so it does not catch it. Co-Authored-By: Claude <noreply@anthropic.com>
351747c moved standalone operators onto CompilableDesign: set_up_artifacts() registers only kernel objects, and link_xclbin() builds the xclbin/insts lazily the first time get_callable() asks. For an operator whose design has no C++ kernel that leaves the artifact graph *empty*, so compile() did nothing at all -- it generated no MLIR, and returned success for configurations whose MLIR cannot be generated. That is not just a missing error. Operators validate their configuration while building the design, so "compile() succeeded" stopped meaning the operator is buildable, and the diagnosis surfaced later from get_callable(), or never. Repeat(cols=513) is the clearest case: cols has no divisor giving both a word-aligned chunk and a chunk count inside the 10-bit wrap field, the generator says so, and compile() reported success anyway. compile() now drives link_xclbin() after the artifact-graph pass, skipping it under dry_run. get_callable() still calls it too, so an operator that was never explicitly compiled keeps working. flm.GEMM and flm.MMPrebuilt override link_xclbin() to do nothing, symmetric with the get_callable() overrides they already carry: their xclbins genuinely are artifacts that the graph pass builds -- one for the config/shape RTP split, one downloaded prebuilt -- and the base implementation would compile a second one and defeat the point. Caught by iron/operators/{repeat,strided_copy} rejection tests, which had turned into 20 "DID NOT RAISE" failures. They pass again, flm is 47/47, and iron/tests is 765 passed / 13 skipped. Co-Authored-By: Claude <noreply@anthropic.com>
mv.cc requires DIM_K % VEC_SIZE == 0 *and* DIM_K >= 2*VEC_SIZE -- the inner
loop carries a pipelining pragma that assumes at least two iterations, and both
are static_asserts. __post_init__ only checked the first, one factor too weak,
so K == kernel_vector_size passed validation and then failed as a C++ error
from inside a Peano build: "static assertion failed due to requirement
'64U >= 2 * 64U'", with nothing pointing at the operator argument responsible.
kernel_vector_size now defaults to None and resolves to the widest legal width
for K. Passed explicitly it is checked, and the message names the rule and
lists what would work. K >= 128 still resolves to 64, so no existing
configuration changes width; K=64 gets 32 and K=32 gets 16.
This fixes the three gemv_batched shapes at K=64 that 1.4.4.dev26 broke, and
supersedes the explicit kernel_vector_size=32 that llama's attention-scores
GEMV carried -- the rule now lives in one place instead of at the call site.
Vector size is repr=False, so it is absent from the operator, MLIR and xclbin
names while appearing in the kernel object name (gemv_{K}k_{vs}vs.o). That is
the shape of a cache-poisoning bug, so it was tested rather than reasoned
about: building K=128 at the default, then explicitly at 32, then at the
default again shows the MLIR text differing between the two configurations,
link_with naming the matching object each time, the xclbin rebuilding on the
change, and the default build reproducing byte-identically.
Checked on hardware: gemv is 95/95; the vs=16 path that no test parameter
reaches compiles and matches the golden reference at K=32; llama is 4/4. A cold
llama build from an empty directory links op7_gemv_64k_32vs.o, with no
gemv_64k_64vs object, no flat duplicate objects and no .symbol_map files
anywhere in the tree.
Co-Authored-By: Claude <noreply@anthropic.com>
IRON's use of CompilableDesign was a text passthrough: link_xclbin() and SeparateDispatch called str(op.get_mlir_artifact().generator()) themselves, then wrapped the resulting text in a synthetic generator so upstream would accept it. Generation therefore happened outside compile(), and everything awkward about the seam followed from that one fact. The cache could not see the real generator, so identity had to be faked by hashing the emitted MLIR and smuggling the digest through compile_kwargs. Kernels could not be declared by the design, because ExternalFunction registers into a global set that compile() clears at the start of its own generation -- anything constructed earlier is wiped -- so objects had to be built beforehand by a separate rule and staged by hand. compile_xclbin_insts now takes the DesignGenerator and resolves but does not call it; the design runs inside compile(), under its lock, in the window where ExternalFunction._instances is collected. Identity comes from what it actually is: the design function, hashed by code identity, plus its bound parameters. The parameters need care, and the failure mode is silent. compile_kwargs values that are not callables are hashed by str(), and dev stringifies to "<abc.NPU2 object at 0x7f...>" -- an address. Passing it through would give every process a different key, which is not an error, just an aiecc run on every call forever. _params_key drops dev, since device identity already reaches the key via _compute_artifact_hash as (type, arch, cols, rows), and rejects any other parameter carrying an address rather than quietly degrading. Two processes now compute the same cache hash for the same operator, checked directly. Staging and object_files stay for operators that still declare prebuilt kernels; they go when the last one declares ExternalFunctions instead. The fused path keeps building its text up front -- fusing several designs into one module is a real transformation, not a passthrough. iron/tests 780 passed / 13 skipped; iron/operators 3165 passed with only the five known mem_copy 16-core timeouts, which is the pre-existing baseline exactly. Co-Authored-By: Claude <noreply@anthropic.com>
The first cut dropped the parameter called "dev" from the cache key, because a
device stringifies to "<abc.NPU2 object at 0x7f...>" and hashing that by str()
would re-key every process. Excluding it was right in spirit and wrong in two
ways.
Wrong to key on the name: what makes a value a device is its API, not what the
design happens to call it. A design spelling it "target" would have leaked an
address into the key, and a non-device parameter named "dev" would have been
silently dropped from it. Recognition is now duck-typed on exactly the
attributes _device_identity_key reads.
Wrong to drop it: a key that ignores the device is stable and incorrect -- two
designs differing only in target share an entry, so an NPU1 build can be handed
to NPU2. Upstream splits identity into a recipe (generator, parameters, flags)
and an artifact (sources, objects, tools, device), and a device belongs to the
second half; it is now spelled there the same way, reusing upstream's own
_device_identity_key rather than inventing a second spelling. It reads
('abc.NPU2', 'AIE2p', '8', '6'): stable across processes, and still telling NPU1
from NPU2.
The same reasoning applies to rebinding. Any device-valued parameter is re-read
from the bound device when the generator runs, rather than reusing what the
operator resolved earlier: compile() calls ensure_current_device() in between,
which can bind a device that was previously only inferred, and generating
against a different device than the key names is how a design silently ends up
built for the wrong target.
Tests now assert what matters rather than that the parameter is absent: no
address in the key, NPU1 and NPU2 keys differ, and a device is recognised when
the design calls it something else.
iron/tests 785 passed / 13 skipped; gemv, gemm and softmax 240 passed. The key
spelling changed, so this invalidates existing cache entries once.
Co-Authored-By: Claude <noreply@anthropic.com>
…dren
gemv named its kernel object twice -- once in the design's Kernel(), once in
get_kernel_artifacts() -- with the gemv_{K}k_{vs}vs.o formula written out
independently in both places and nothing keeping them in step. It now declares
one ExternalFunction per kernel, carrying the source, the -D flags and the
fusion prefix, and reports no kernel artifacts at all; upstream compiles them
and names the object by content. The gelu epilogue stops being an archive and
becomes a second ExternalFunction: each func.func carries its own link_with and
aie-assign-core-link-files aggregates them onto the core.
Two bugs had to be fixed to get there, both of which built and linked cleanly.
fuse_mlir inlines each child's device, runtime sequence included, and drives PDI
switching itself -- alternating two PDIs per configure point under
--expand-load-pdis, with needs_additional_reset keeping the count even. Moving
fused generation inside compile() put the children under _iron_full_elf, which
makes a runtime sequence load its own PDI because on that path no xclbin
configures the device. Both schemes then ran at once. Nothing failed to build:
the ELF linked and the device hung at dispatch with ERT_CMD_STATE_TIMEOUT.
Exactly one program in a fused build is a full ELF and it is not the children,
so _fuse_as_children shadows the flag for them -- which the old code got for
free by generating outside compile() entirely. The cache key is computed through
the same helper, because keying under one value and building under the other
describes a different program and nothing reports that either.
build_fused_mlir decided whether to prefix an operator by asking whether it had
kernel artifacts. That was a proxy for "has kernels", and ExternalFunction
breaks it: a migrated operator reports none, so it silently went unprefixed.
Every gemv shape in llama then defined matvec_vectorized_bf16_bf16, kept apart
only by each core linking its own object. It now asks the design whether it
takes func_prefix, which is the actual contract. Verified by reading the objects
back: op1_, op11_, op12_, op17_ prefixes on both names and symbols.
Note ExternalFunction joins its prefix with an underscore of its own, for the
symbol name and the rename pass alike, so IRON's "op0_" is handed over stripped.
iron/tests 785 passed / 13 skipped; iron/operators 3165 passed with only the five
known mem_copy 16-core timeouts; llama 4/4, and its generated text is
byte-identical to before this change, so the kernel move altered no numerics.
Co-Authored-By: Claude <noreply@anthropic.com>
ChanneledUnaryOperator and BinaryElementwiseOperator between them back eleven operators -- relu, gelu, silu, sigmoid, tanh, leaky_relu, layer_norm, elementwise_add, elementwise_mul, swiglu_prefill, swiglu_decode -- and each of them named its kernel object twice: once as the design's Kernel(), once as the operator's KernelObjectArtifact. Both bases now declare an ExternalFunction in the design and report no kernel artifacts at all. The two shared designs, and gemv before them, had the same five lines of declaration with the same trap in it, so it is one helper now: declare_kernel() in iron/operators/_kernels.py, alongside _trace as the other design-side shared piece. The trap is that IRON spells its fusion prefix "op0_" while ExternalFunction joins with an underscore of its own, for the symbol name and the rename pass alike, so the prefix has to be handed over stripped. The aie2 lut_based_ops case stays exactly as it was, and the reason is worth stating where it is now load-bearing: lut_based_ops.cpp defines tables the kernel references transitively from C++, with no MLIR call site, so aie-assign-core-link-files cannot discover that object by tracing func.call edges. It has to be archived and named by an ordinary link_with, which means a prebuilt Kernel. kernel_obj_file returning None is what selects the ExternalFunction path; needs_lut_archive names the condition. That branch is aie2-only and this box is aie2p, so it is untestable here and was not touched. kernel_object_arch_isolation used ElementwiseMul as a vehicle for testing that two arches cannot collide on one object path. That operator no longer produces an artifact to test, so the tests move to AXPY, which still does. They are guarding the operators left on the artifact path and should retire with it: upstream keys an ExternalFunction's object on content and on device identity, so the collision they describe is unrepresentable there. iron/tests 785 passed / 13 skipped; the eleven operators 1195 passed. Co-Authored-By: Claude <noreply@anthropic.com>
…, gemv Four more operators stop naming their kernel object twice, and gemv folds onto the declare_kernel helper the shared bases already use, so all five now read the same way. kernels_dir moves to MLIROperator beside dev: it is a parameter every design needs and no operator stores, which is exactly what those properties are for, and it keeps IRON_AIE_KERNELS_DIR redirecting the source while still reaching the cache key. mem_copy keeps its bypass branch, which means something different from the others: no kernel at all, rather than one declared elsewhere. kernel_object_arch_isolation no longer borrows an operator. It was written against ElementwiseMul, moved to AXPY when that migrated, and would have moved again -- the vehicle keeps disappearing because a design that declares an ExternalFunction produces no artifact to test. It builds the artifact directly now, which is honest about what it covers: a property of move_artifacts, for the operators still on the artifact path, to retire with it. Upstream keys such an object on content and device identity, so the collision it describes cannot occur there. iron/tests 785 passed / 13 skipped. axpy, transpose and dequant 96 passed; mem_copy 63 passed with only the known 16-core timeout. Co-Authored-By: Claude <noreply@anthropic.com>
rope and rms_norm are the straightforward shape -- one ExternalFunction per kernel, from one source each. rope's object stops being named for the method id and is named for the symbol instead, which is what actually distinguishes the two variants rope.cc defines. softmax is the first operator where declaring per kernel would have been wrong. softmax_bf16 and mask_bf16 both live in softmax.cc, so two defaulted declarations compile that translation unit twice into two objects, each defining *both* symbols -- a duplicate definition once a core links them. They share one object_file_name instead: identical source and flags give an identical content digest, so upstream neither reports a collision nor compiles twice. Checked rather than assumed, by reading the build back: one softmax.o per design, carrying mask_bf16 and softmax_bf16 together. declare_kernel grew object_file_name for that, and applies the fusion prefix to it. Upstream names a defaulted object after the prefixed symbol but takes an explicit one as given, so without that two fused operators would share one object. softmax keeps its aie2 archive, for the reason that keeps recurring: lut_based_ops.cpp's tables are reached transitively from C++ with no MLIR call site, so nothing discovers that object by tracing calls. kernel_obj_file returning None is what selects the ExternalFunction path. rope and rms_norm 115 passed; softmax 15 passed; iron/tests 785 passed / 13 skipped; llama still 2/2 on hardware, which is what exercises all three fused. Co-Authored-By: Claude <noreply@anthropic.com>
Both are cases where one translation unit backs several entry points, so both use the shared object name softmax needed: gemm's zero and matmul come out of mm.cc, and all six of mha's come out of mha.cc. Declared per kernel they would compile that unit once each and every copy would redefine every symbol in it. gemm keeps its aie2 quirk where it was, in the operator: that arch sources a patched mm.cc from the tree and needs an -I so the in-tree file's includes resolve against the package copies. It is now expressed as kernel_source and kernel_flags properties, which is all an ExternalFunction needs. The kernel_object property is gone -- the object name is derived in the design, where it is used. mha stops listing mm.cc and softmax.cc as dependencies. mha.cc #includes them, and upstream reads Peano's depfile and validates the manifest against it, which covers the transitive headers that hand-written list never did. Checked cold, from an empty build dir, because a warm run here proves nothing: mha produces exactly two objects, mha.o carrying all fifteen symbols and mha_passThrough.o, from six declarations and two compiles. gemm and mha 155 passed; iron/tests 785 passed / 13 skipped. Co-Authored-By: Claude <noreply@anthropic.com>
…e hashing This operator was the last one still on the artifact graph, kept there for its config/shape split: the xclbin is emitted at a reference shape so every shape sharing a configuration reuses it, and only the instruction stream is per shape. That turns out to be two compile_xclbin_insts calls, each discarding the half it did not want, so XclbinArtifact, InstsBinArtifact and the verbatim get_callable override all go. Only then can its kernels move: mm_fused.cc's three entry points are collected from ExternalFunction._instances, which is populated during generation inside compile(), and this operator never went through compile() at all. They share one object name, like softmax's and mha's, because they share a translation unit. The object it produces is byte-identical to the one the artifact graph built -- a9fcbabd, before and after -- which is a stronger statement than a benchmark on the most performance-tuned operator here: identical machine code cannot run at a different speed. Getting there exposed a caching bug that is not this operator's, and not new. _compute_artifact_hash reads get_current_device(probe_runtime=False), which is None until something binds a device -- and CompilableDesign.compile() binds it moments later, from inside. So _compile_if_changed hashed a "no device" identity for the first build in a process, stamped that, and the next identical build could never match it: every process silently rebuilt once, for every operator. Nothing failed, so nothing reported it. It surfaced here only because test_one_xclbin_serves_every_clamp_bound is the one test that asserts a second instance does not rebuild. Binding the device before hashing makes both sides agree. flm 235 passed, including the two one-xclbin tests; iron/tests 785 passed / 13 skipped. Those two tests now read _xclbin_path, since the artifact they used to reach through no longer exists. Co-Authored-By: Claude <noreply@anthropic.com>
…the bug Two gaps in the previous commit's fix, both found by checking it rather than trusting it. The bind was unguarded. CompilableDesign._bind_generation_device wraps the same call in try/except because binding probes the runtime, which a compile-only host without one cannot do; an unguarded call would turn a working offline build into a crash. Guarded identically now -- failing to bind leaves the device unset on both sides, which still agrees with itself. The test I added for it passed with the fix disabled, which made it worthless. It bound the device itself before calling _compile_if_changed, so the binding inside was never needed. It now leaves the device unset across the call, which is the only state where the fix does anything, and fails without it. Verified across processes as well as within one, on both paths, since a single process was never the interesting case: three separate runs of an operator and two of a fused sequence each reuse the first build rather than recompiling. iron/tests 790 passed / 13 skipped. Co-Authored-By: Claude <noreply@anthropic.com>
The tables in lut_based_ops.cpp are what aie2's exp/log kernels reference from C++, with no MLIR call site. aie-assign-core-link-files finds objects by tracing func.call edges, so it can never discover that one, and IRON stapled it on with an llvm-ar archive -- a whole artifact class, a compilation rule and a binutil, kept alive for one orphan object on one architecture. Compiling it into the kernel's own translation unit removes the orphan instead of working around it. declare_kernel grows bundled_sources, which generates a source that includes the bundle and then the kernel, and hands that to ExternalFunction as source_string -- no file on disk, and included by bare name against the search path so the digest upstream takes does not move with the checkout. A generated source rather than clang's -include, which was the first thing I tried: -include is processed before the arch macros are established and aie_api rejects it with "'__AIE_ARCH__' macro is required". So KernelArchiveArtifact, ArchiveCompilationRule, lut_based_ops_artifacts and the llvm-ar dependency are all gone, and this no longer waits on the upstream header-only change (the prompt for which stays valid -- it is still the better fix, it just is not a blocker any more). The path is aie2 and this box is aie2p, so it cannot be run here. It can be compiled here, and that is the property the archive supplied: softmax built for NPU1 produces one self-contained softmax.o, no undefined lut symbols, tables defined, both entry points present, and no .a anywhere in the build. Running it still needs NPU1 hardware. iron/tests 790 passed / 13 skipped; the lut users plus flm 745 passed on aie2p. Co-Authored-By: Claude <noreply@anthropic.com>
…hat that freed The last operator still building an artifact through the graph. Its xclbin is downloaded rather than compiled -- which is what the operator exists for, and the one thing RemoteFileArtifact expresses that the compile path cannot -- but its instruction stream was an InstsBinArtifact fed by a PythonGeneratedMLIRArtifact, which kept four more classes and two rules reachable for one caller. It is a compile_xclbin_insts call now, keeping the instructions and discarding the xclbin written beside them, the same way flm.GEMM discards the half each of its two builds did not want. Its design returns MLIR text where the others return a Module, which upstream reports as "AttributeError: 'str' object has no attribute 'operation'" -- naming neither the design nor the cause. _design_generator parses a string return rather than making every design agree on which to hand back. With that, nothing constructs an XclbinArtifact or an InstsBinArtifact, and no graph contains a PythonGeneratedMLIRArtifact, so the rules that compiled them are unreachable. Deleted: XclbinArtifact, InstsBinArtifact, _MLIRInputMixin, AieccCompilationRule, AieccXclbinInstsCompilationRule and GenerateMLIRFromPythonCompilationRule. PythonGeneratedMLIRArtifact itself stays -- every operator still uses it to carry its DesignGenerator, which is now all it does. compilation/base.py is 708 lines, from 977 on devel and 891 this morning. iron/tests 790 passed / 13 skipped; flm, softmax and gemv 345 passed; llama 2/2. Co-Authored-By: Claude <noreply@anthropic.com>
UNTESTED ON HARDWARE. stream-dse is not installed here, so all four stream tests skip and nothing below has been executed end to end. CI has to be the judge. What could be checked locally was, and is described at the end. stream was the last user of the artifact graph, and of rename_symbols, which upstream has no equivalent for. It needed one because stream-dse suffixes a GEMM's symbols with its tile shape so several shapes coexist in one design, while mm.cc defines them unsuffixed, and ExternalFunction can only prefix. The reconciliation goes the other way instead: the object is prefixed, and the generated MLIR is rewritten to agree. That works because IRON already rewrites this text -- _prefixed has been applying the fusion prefix to link_with and to every declared symbol and call site all along. _renamed is the same operation, applied first so the fusion prefix lands on top: mm.cc's matmul_bf16_bf16 becomes mm128_64_64_matmul_bf16_bf16, and op3_mm128_64_64_matmul_bf16_bf16 in a fused group. declare_kernel grew symbol_prefix for this. It composes with the fusion prefix rather than replacing it, and deliberately does not reach the object name: a generated design names the object it links, so renaming the file out from under it would break the link. That distinction was wrong in the first draft and the object came out as mm128_64_64_mm_128_64_64.o. Two other things this surfaced: group_index was positional, which the compile seam rejects -- a design's parameters reach the cache key by name. It is keyword-only now, and kernels_dir is threaded in rather than reached for through the default context, so pointing IRON at another kernel tree re-keys the build. lut_sources sat in iron/operators/_kernels and was imported by iron/common/operator_bases, so importing iron.operators._kernels first raised ImportError on a partially initialized module. It only ever worked because every test imported iron.common first. Moved to iron/common/device_utils. What was verified locally: both halves of the naming agree, checked by building the ExternalFunction and comparing against the rewritten text, fused and unfused; the rewrite composes correctly, checked by lifting _renamed and _prefixed out of the module (which cannot be imported without stream-dse) and running them over representative MLIR. iron/tests 790 passed / 13 skipped and 265 operator tests pass, so the paths that do run here are unaffected. With rename_symbols gone, compilation/base.py is 692 lines, from 977 on devel. Co-Authored-By: Claude <noreply@anthropic.com>
Every operator's kernels are ExternalFunctions now, so nothing constructs a KernelObjectArtifact, KernelCompilationRule never fires, and the objects land in the work dir because upstream compiles them there rather than because IRON copies them. Deleted: KernelObjectArtifact, KernelCompilationRule, _link_build_outputs_into, stage_objects, the object_files parameter threaded through both compile entry points, and the get_kernel_artifacts hook itself. move_artifacts loses its arch-scoping segment with them: it existed because two arches could collide on one build_dir path, and upstream keys an object on its content and on device identity instead. PythonGeneratedMLIRArtifact keeps only what it is still for -- carrying a DesignGenerator. Its recipe_hash and availability override detected a stale written .mlir, and nothing writes one; the collision they guarded against, a fused build mutating a shared operator's func_prefix under a standalone build's path, is caught by the compile cache key, which hashes the design and its parameters including that prefix. Three tests go with it. kernel_object_arch_isolation tested move_artifacts' arch scoping and could not even import once that was gone. mlir_recipe_hash and its fixture tested recipe_hash directly. mlir_cache_poisoning stays: it is the end-to-end check that the property still holds, and it now says where. One test was passing for a reason that had stopped being true -- "object_files does not stage, so this is the one thing the retirement cannot delete" -- when that step is exactly what was deleted. Renamed and rewritten to check the requirement rather than IRON's former way of meeting it. compilation/base.py is 518 lines, from 977 on devel. What is left of the artifact graph serves one thing: flm.MMPrebuilt's downloaded xclbin, which is fetched rather than built and so has nothing to compile. iron/tests 745 passed / 13 skipped; iron/operators 3165 passed with only the five known mem_copy 16-core timeouts. Co-Authored-By: Claude <noreply@anthropic.com>
Plan A proposes shape annotations on the design signature. Plan B moves the declaration into an interface() method body, where the operator's own parameters are already in scope -- which deletes the deferred-annotation machinery Plan A needs and keeps real dataclass fields for pyright. Plan B then splits what a compiled operator is. The toolchain already separates the overlay (per-core ELFs + PDI, a function of the design) from the runtime sequence (insts.bin, a function of the steps and the buffer ABI); aiecc emits them from one dependency graph that only diverges at the tail, and upstream's runtime already caches the two halves independently. IRON collapses that into dispatch="fused"|"separate", which bakes in four decisions and makes partial fusion unrepresentable. Plan B replaces the string with four constructors -- Overlay, StaticSequence/GeneratedSequence, Elf/Xclbin -- so the ELF/no-ELF question is which object you build, and Elf's signature rejects a generated sequence statically. Plan B also gives the overlay an interface of its own: shim bindings, resident symbols, buffer sizes. flm/mm_prebuilt already performs that agreement by hand against a downloaded xclbin, and records flm/gemm's failure to match it in a comment; all three fields are recoverable from files IRON already opens. Both are drafts for a plan-refining session, not approved work. Plan B gates everything behind two spikes (§9 step 0) because the one unverified claim -- that an expanded-load-PDI fused sequence dispatches correctly via the opcode-3 ABI -- fails as a device hang, not a build error. Co-Authored-By: Claude <noreply@anthropic.com>
Plan B stood on Plan A for its priorities, diagnosis, measurements, dismissal table and carried risk, so neither file was readable alone. This folds the load-bearing parts of A into B and deletes both. What came over from A, because B referenced it and could not stand without it: the 12 priorities; the llama_npu.py line-count diagnosis that motivates centering on the operator model rather than the graph recorder; the duplication-vs-verification argument and the live GEMV arg_spec disagreement that proves it (three restatements, all green, one of them wrong); the pyright measurements that rule out synthesised dataclass fields; the dismissal table; and the decode-drift risk a rewritten llama will inherit and look guilty for. Restored in the process: per-operator tuning got a section of its own again. B had reduced it to a code snippet, which lost priorities 5 and 6 and the MAX_WRAP FIXME it retires -- now with a table of which target model numbers are actually verified and which are still assumed. Renumbered throughout; §-refs, O1-O15 and E1-E31 all resolve. Still a draft, still gated on §19 step 0. Co-Authored-By: Claude <noreply@anthropic.com>
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
1 participant
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.