Skip to content

operators: emit trace configuration when trace_size is set - #142

Open
atassis wants to merge 3 commits into
amd:develfrom
atassis:feat/iron-trace-parity-devel
Open

operators: emit trace configuration when trace_size is set#142
atassis wants to merge 3 commits into
amd:develfrom
atassis:feat/iron-trace-parity-devel

Conversation

@atassis

@atassis atassis commented Jul 25, 2026

Copy link
Copy Markdown
Contributor

Most designs in iron/operators accept a trace_size argument but never call
Runtime.enable_trace, so they run fine and return a trace.txt of all zeros. That reads as
"this op can't be traced" rather than "the wiring is missing", which is what sent me looking.

Affected: axpy, binary_elementwise, channeled_unary (so gelu/silu/relu/sigmoid/tanh),
gemm, leaky_relu, mem_copy, mha, rope, softmax.

Rather than repeat the same event setup per design, the wiring goes in
iron/operators/_trace.py. maybe_enable_trace() uses an explicit trace_size when one is
passed and otherwise falls back to IRON_TRACE_SIZE, so the designs whose op.py doesn't plumb
trace_size through can still be traced; IRON_TRACE_NTILES caps how many workers are
instrumented (0 traces none). It's a no-op when neither is set, so default behaviour is unchanged.

Three things I'd rather flag than have you find

The import is absolute, not relative, and that matters. DesignGenerator loads these files by
path via spec_from_file_location + exec_module with no package context
(iron/common/compilation/base.py), so a relative import raises "attempted relative import with no
known parent package" at generation time. It also passes py_compile, so it won't show up in a
lint pass. Worth knowing before anyone tidies it. This is also the first dependency from a
design.py onto the surrounding iron package -- if you'd rather designs stayed standalone, say
so and I'll inline the helper instead.

dequant is deliberately untouched. It already wires trace, but in a simpler form (no
workers, no coretile_events), so folding it onto the helper would change its behaviour rather
than just deduplicate. Happy to do it as a follow-up if you want one path.

No automated coverage for the fix itself. The verification below is a manual check; nothing in
CI would catch this regressing again. If you want a test asserting a design emits aie.trace when
trace_size is set, tell me the shape you'd prefer and I'll add it.

Testing

Generated softmax through that same package-less loader: 0 aie.trace ops before this change
with tracing requested, 18 after
, and still 0 when tracing is off. The 18 are a full
aie.trace @trace_core_1 block (8 core events, 3 port configs, start/stop broadcasts,
host_config {buffer_size = 8192}).

On device (Strix, NPU2): 935 tests pass across gelu, softmax, axpy, mem_copy, rope.
black --check and reuse lint both clean.

Every design in iron/operators takes a trace_size argument, but a design only
emits trace configuration if it calls Runtime.enable_trace. Most designs never
make that call, so they run fine and hand back a trace.txt full of zeros, which
reads as "this op cannot be traced" rather than "the wiring is missing".

Affected: axpy, binary_elementwise, channeled_unary (so gelu/silu/relu/sigmoid/
tanh), gemm, leaky_relu, mem_copy, mha, rope, softmax. dequant already wires
trace and is left alone here.

Rather than repeat the same event setup in each design, this puts the wiring in
iron/operators/_trace.py. maybe_enable_trace() takes an explicit trace_size when
one is passed and otherwise falls back to IRON_TRACE_SIZE, so the designs whose
op.py does not plumb trace_size through can still be traced; IRON_TRACE_NTILES
caps how many workers are instrumented (0 traces none). It is a no-op when
neither is set, so default behaviour is unchanged.

The import is absolute rather than relative because DesignGenerator loads these
files by path via spec_from_file_location + exec_module, with no package context
(iron/common/compilation/base.py), and a relative import raises "attempted
relative import with no known parent package" at generation time. Note this is
the first dependency from a design.py onto the surrounding iron package.

Verified by generating softmax through that same loader: 0 aie.trace ops before
this change with tracing requested, 18 after, and still 0 when tracing is off.
On device, 935 tests pass across gelu, softmax, axpy, mem_copy and rope.

@andrej andrej left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for cleaning up our messes. :)

Comment thread iron/operators/_trace.py
Comment thread iron/operators/_trace.py Outdated
atassis and others added 2 commits July 28, 2026 19:25
Co-authored-by: André Rösti <androsti@amd.com>
Co-authored-by: André Rösti <androsti@amd.com>
@atassis
atassis requested a review from andrej July 29, 2026 22:23
atassis added a commit to atassis/xdna-engine that referenced this pull request Jul 30, 2026
…was already plumbed, never wired)

residual_add_iron.py, glu_iron.py, silu_iron.py accepted -t/--trace_size on the CLI and threaded it
into the design function, but never called Program.enable_trace -- the same dead-plumbed pattern
already fixed upstream for the IRON operators library (amd/IRON#142). Wire it the way
ln_affine_cast_iron.py already does: trace one worker (egress_shim_col=1), and add -n/--cores so a
traced build can drop to cores=1 (production's 8 cores saturate the shim south ports for trace
egress). trace_size=0 (production) is unaffected -- verified byte-identical MLIR output against the
pristine generators, and 0 aie.trace ops with tracing off vs 17 with it on (cores=1) for all three.

Add traced Makefile variants (resaddtr/glutr/silutr) mirroring Makefile.lnaffcasttr, wire them into
sync_kernels.sh, and promote trace_brick.py (previously an untracked sandbox script) into
route_b_kernels/ctx_ln/ with resadd/glu ABI siblings, so a device lane has a ready-to-run capture
command instead of hand-writing one.

Built (CPU only, no device) final_resaddtr_512x1024_s050.xclbin and final_glutr_512x1024.xclbin
end to end through aiecc --dump-intermediates, confirming the shim-egress recipe transfers from
ln_affine_cast to these bricks. Device capture not run (npu contended by another lane tonight).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants