operators: emit trace configuration when trace_size is set#142
Open
atassis wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR A -- amd/IRON, branch
feat/iron-trace-parity-devel@4d9f60c(basedevel)Title: operators: emit trace configuration when trace_size is set
Most designs in
iron/operatorsaccept atrace_sizeargument but never callRuntime.enable_trace, so they run fine and return atrace.txtof 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 explicittrace_sizewhen one ispassed and otherwise falls back to
IRON_TRACE_SIZE, so the designs whoseop.pydoesn't plumbtrace_sizethrough can still be traced;IRON_TRACE_NTILEScaps how many workers areinstrumented (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.
DesignGeneratorloads these files bypath via
spec_from_file_location+exec_modulewith no package context(
iron/common/compilation/base.py), so a relative import raises "attempted relative import with noknown parent package" at generation time. It also passes
py_compile, so it won't show up in alint pass. Worth knowing before anyone tidies it. This is also the first dependency from a
design.pyonto the surroundingironpackage -- if you'd rather designs stayed standalone, sayso and I'll inline the helper instead.
dequantis deliberately untouched. It already wires trace, but in a simpler form (noworkers, nocoretile_events), so folding it onto the helper would change its behaviour ratherthan 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.tracewhentrace_sizeis set, tell me the shape you'd prefer and I'll add it.Testing
Generated
softmaxthrough that same package-less loader: 0aie.traceops before this changewith tracing requested, 18 after, and still 0 when tracing is off. The 18 are a full
aie.trace @trace_core_1block (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 --checkandreuse lintboth clean.