operators: emit trace configuration when trace_size is set

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

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.

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()

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()

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

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.

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

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])

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.

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: