From 4d9f60c0469f92ca3eb0ba2f2b611cf8ba539109 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Sat, 25 Jul 2026 16:15:06 +0300 Subject: [PATCH] operators: emit trace configuration when trace_size is set 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. --- iron/operators/_trace.py | 94 +++++++++++++++++++++ iron/operators/axpy/design.py | 2 + iron/operators/binary_elementwise_design.py | 2 + iron/operators/channeled_unary_design.py | 2 + iron/operators/gemm/design.py | 2 + iron/operators/leaky_relu/design.py | 2 + iron/operators/mem_copy/design.py | 2 + iron/operators/mha/design.py | 7 +- iron/operators/rope/design.py | 2 + iron/operators/softmax/design.py | 2 + 10 files changed, 116 insertions(+), 1 deletion(-) create mode 100644 iron/operators/_trace.py diff --git a/iron/operators/_trace.py b/iron/operators/_trace.py new file mode 100644 index 00000000..3e3765d1 --- /dev/null +++ b/iron/operators/_trace.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared per-op NPU hardware-trace wiring for the iron/operators designs. + +Why this exists +--------------- +Every design in this package takes a ``trace_size`` argument, but a design only actually +emits trace configuration if it calls ``Runtime.enable_trace``. Designs that accept +``trace_size`` and never make that call compile and run fine and then hand back a +``trace.txt`` full of zeros, which reads as "this op is untraceable" rather than +"nobody wired it up". + +Before this helper the wiring had drifted three ways: ``channeled_unary_design`` honored +the ``trace_size`` parameter *or* the ``IRON_TRACE_SIZE`` env var and passed a full +core-tile event set; ``gemm/design`` honored *only* the env var, so passing +``trace_size=`` did nothing; ``dequant/design`` honored only the parameter and passed no +events at all. Nine other designs had no wiring whatsoever. Keeping one implementation +here is what stops that from happening again. + +Semantics (the ``channeled_unary`` behaviour, which was the most complete): + * explicit ``trace_size`` wins; otherwise fall back to ``IRON_TRACE_SIZE`` + * ``IRON_TRACE_NTILES`` (default 1) caps how many workers get traced; 0 traces none + * no-op when neither is set, so production paths are unaffected + +Note this is deliberately NOT the upstream mechanism being reimplemented: +``Runtime.enable_trace`` is the same API upstream ships and upstream's own designs +(``aie.iron.algorithms.reduce`` / ``transform``) already call it. This package is +fork-carried, so the gap was ours alone. +""" + +import os + +__all__ = ["maybe_enable_trace", "resolve_trace_size"] + + +def resolve_trace_size(trace_size=None): + """Effective trace size: explicit argument first, then ``IRON_TRACE_SIZE``, else 0.""" + if trace_size and trace_size > 0: + return int(trace_size) + # Deliberately unguarded: a malformed IRON_TRACE_SIZE should raise rather than + # silently disable tracing, which would reproduce the all-zeros trace.txt + # confusion this module exists to remove. + return int(os.environ.get("IRON_TRACE_SIZE", "0")) + + +def _default_coretile_events(): + import aie.utils.trace as trace_utils + + ev = trace_utils.events + return [ + ev.PortEvent(ev.CoreEvent.PORT_RUNNING_0, ev.WireBundle.DMA, 0, True), + ev.PortEvent(ev.CoreEvent.PORT_RUNNING_1, ev.WireBundle.DMA, 1, True), + ev.PortEvent(ev.CoreEvent.PORT_RUNNING_2, ev.WireBundle.DMA, 0, False), + ev.CoreEvent.INSTR_EVENT_0, + ev.CoreEvent.INSTR_EVENT_1, + ev.CoreEvent.MEMORY_STALL, + ev.CoreEvent.LOCK_STALL, + ev.CoreEvent.INSTR_VECTOR, + ] + + +def maybe_enable_trace(rt, trace_size, workers, coretile_events=None): + """Configure per-op hardware trace on ``rt`` if tracing is requested. + + Call inside the ``rt.sequence(...)`` block, before ``rt.start(...)``. + + Args: + rt: the ``Runtime`` being built. + trace_size: the design's ``trace_size`` argument (may be None/0). + workers: the design's workers; the first ``IRON_TRACE_NTILES`` are traced. + coretile_events: override the default core-tile event set. + + Returns: + The effective trace size (0 when tracing is off and nothing was configured). + """ + ts = resolve_trace_size(trace_size) + if ts <= 0: + return 0 + + # A count, so 0 legitimately means "trace no tiles"; only negatives are + # meaningless (a negative slice index would silently drop the LAST worker). + ntiles = max(0, int(os.environ.get("IRON_TRACE_NTILES", "1"))) + + rt.enable_trace( + ts, + workers=list(workers)[:ntiles], + coretile_events=( + coretile_events + if coretile_events is not None + else _default_coretile_events() + ), + ) + return ts diff --git a/iron/operators/axpy/design.py b/iron/operators/axpy/design.py index cde40a63..12685bd6 100644 --- a/iron/operators/axpy/design.py +++ b/iron/operators/axpy/design.py @@ -7,6 +7,7 @@ from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ +from iron.operators._trace import maybe_enable_trace def my_axpy( @@ -85,6 +86,7 @@ def core_body(of_in1, of_in2, of_out, axpy): # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C): + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. diff --git a/iron/operators/binary_elementwise_design.py b/iron/operators/binary_elementwise_design.py index 9155ffb6..5b75f815 100644 --- a/iron/operators/binary_elementwise_design.py +++ b/iron/operators/binary_elementwise_design.py @@ -7,6 +7,7 @@ from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ +from iron.operators._trace import maybe_enable_trace def binary_elementwise_design( @@ -84,6 +85,7 @@ def core_body(of_in1, of_in2, of_out, eltwise_fn): # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C): + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) tg = rt.task_group() diff --git a/iron/operators/channeled_unary_design.py b/iron/operators/channeled_unary_design.py index 34aad71a..ab8e0fcb 100644 --- a/iron/operators/channeled_unary_design.py +++ b/iron/operators/channeled_unary_design.py @@ -7,6 +7,7 @@ from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ +from iron.operators._trace import maybe_enable_trace def channeled_unary_design( @@ -98,6 +99,7 @@ def core_fn(of_in, of_out, kernel_line): # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(transfer_type, transfer_type) as (a_in, b_out): + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) tg = rt.task_group() diff --git a/iron/operators/gemm/design.py b/iron/operators/gemm/design.py index d376b770..bfdb8442 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -21,6 +21,7 @@ from aie.iron.device import NPU1Col1, NPU1Col2, NPU1, NPU2, Tile from aie.helpers.taplib import TensorAccessSequence, TensorTiler2D, TensorAccessPattern from aie.iron.controlflow import range_ +from iron.operators._trace import maybe_enable_trace microkernel_mac_dim_map = { "npu1": { @@ -552,6 +553,7 @@ def core_fn( # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(A_ty, B_ty, C_ty) as (A, B, C): + maybe_enable_trace(rt, trace_size, workers) rt.start(*workers) # Set runtime parameters diff --git a/iron/operators/leaky_relu/design.py b/iron/operators/leaky_relu/design.py index bb33230b..b3ea1fb4 100644 --- a/iron/operators/leaky_relu/design.py +++ b/iron/operators/leaky_relu/design.py @@ -7,6 +7,7 @@ from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ +from iron.operators._trace import maybe_enable_trace def my_leaky_relu( @@ -94,6 +95,7 @@ def core_fn(of_in, of_out, leaky_relu_line): # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(transfer_type, transfer_type) as (a_in, b_out): + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. diff --git a/iron/operators/mem_copy/design.py b/iron/operators/mem_copy/design.py index 4f7faeb4..1eb3685e 100644 --- a/iron/operators/mem_copy/design.py +++ b/iron/operators/mem_copy/design.py @@ -20,6 +20,7 @@ from aie.iron.controlflow import range_ from aie.iron.runtime.endpoint import RuntimeEndpoint from aie.iron.device import AnyShimTile +from iron.operators._trace import maybe_enable_trace # The maximum value the 4th dimension of DMA BD can be set TAP_REPEAT_MAX = 64 @@ -245,6 +246,7 @@ def core_fn(of_in, of_out, mem_copy_line): with rt.sequence(transfer_type, transfer_type) as (a_in, b_out): # Start the workers if not bypass if not bypass: + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) # Calculate how much of workload can be partitioned evenly and what's remaining diff --git a/iron/operators/mha/design.py b/iron/operators/mha/design.py index d9ee167c..d5dac245 100644 --- a/iron/operators/mha/design.py +++ b/iron/operators/mha/design.py @@ -23,6 +23,7 @@ from aie.iron.controlflow import range_ from aie.helpers.taplib import TensorTiler2D, TensorAccessSequence, TensorAccessPattern from aie.helpers.dialects.scf import if_, else_ +from iron.operators._trace import maybe_enable_trace, resolve_trace_size dtype_map = { "bf16": bfloat16, @@ -121,7 +122,7 @@ def fused_mha( of_depth = 2 vectorized = True - enable_tracing = trace_size > 0 + enable_tracing = resolve_trace_size(trace_size) > 0 dtype_str = "bf16" if number_of_pipelines > 6: @@ -791,6 +792,10 @@ def set_mha_rtps(): for i in range(number_of_pipelines): rt.set_barrier(worker_barrier_list[j][i], 1) + maybe_enable_trace( + rt, trace_size, matmul_workers + softmax_workers + matmul_pv_workers + ) + for i in range(number_of_pipelines): rt.start(matmul_workers[i]) rt.start(softmax_workers[i]) diff --git a/iron/operators/rope/design.py b/iron/operators/rope/design.py index 79de08e2..5d8e2ccf 100644 --- a/iron/operators/rope/design.py +++ b/iron/operators/rope/design.py @@ -22,6 +22,7 @@ from aie.helpers.taplib.tap import TensorAccessPattern from aie.helpers.dialects.scf import _for as range_ from ml_dtypes import bfloat16 +from iron.operators._trace import maybe_enable_trace def rope( @@ -129,6 +130,7 @@ def core_body(of_in, of_lut, of_out, rope_kernel): # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(tensor_ty, angle_ty, tensor_ty) as (A, B, C): + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. diff --git a/iron/operators/softmax/design.py b/iron/operators/softmax/design.py index 6492b9c7..1aca7379 100644 --- a/iron/operators/softmax/design.py +++ b/iron/operators/softmax/design.py @@ -18,6 +18,7 @@ from aie.helpers.taplib.tap import TensorAccessPattern from aie.helpers.dialects.scf import _for as range_ from ml_dtypes import bfloat16 +from iron.operators._trace import maybe_enable_trace def softmax( @@ -157,6 +158,7 @@ def worker_args(i, j): # Runtime operations to move data to/from the AIE-array rt = Runtime() with rt.sequence(tensor_ty, tensor_ty) as (A, C): + maybe_enable_trace(rt, trace_size, my_workers) rt.start(*my_workers) if use_scratchpad: