Add op_index and payload_size tracking to device tracing (#424)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: mawad-amd

examples/23_gemm_all_scatter_tracing/gemm_all_scatter.py

@@ -147,12 +147,15 @@ def persistent_gemm_all_scatter(
             else:
                 # Record duration event around remote store (compiles away if tracing=False)
                 # Pass 2D pointer tensor; record_event_start takes min as representative address
+                # op_index is automatically tracked internally (0, 1, 2, ...)
+                # payload_size is automatically calculated from mask
                 handle = ctx.tracing.record_event_start(
                     event_id=TraceEvent().put,
                     target_rank=remote_rank,
                     address=c_global + global_offset,
                     pid_m=pid_m,
                     pid_n=pid_n,
+                    mask=sub_mask,
                 )
 
                 # Use DeviceContext.put for remote stores

iris/iris.py

@@ -946,11 +946,14 @@ def get_device_context(self):
                 self.tracing.trace_buffers["timestamp"].data_ptr(),
                 self.tracing.trace_buffers["address"].data_ptr(),
                 self.tracing.trace_buffers["duration_cycles"].data_ptr(),
+                self.tracing.trace_buffers["op_index"].data_ptr(),
+                self.tracing.trace_buffers["payload_size"].data_ptr(),
             ]
             context_data += [
                 1,  # trace_enabled = 1 (true)
                 self.tracing.max_events,
                 self.tracing.trace_counter.data_ptr(),
+                self.tracing.op_index_counter.data_ptr(),
             ] + trace_buffer_ptrs
         else:
             context_data += [0]  # trace_enabled = 0 (false)
@@ -1384,7 +1387,8 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False):
             >>>
             >>>     # With tracing
             >>>     ctx = DeviceContext.initialize(context_tensor, rank, world_size, tracing=True)
-            >>>     ctx.tracing.record_event_start(event_id=TraceEvent().put, target_rank=1, address=ptr)
+            >>>     mask = tl.full([64], True, dtype=tl.int1)  # Example mask
+            >>>     ctx.tracing.record_event_start(event_id=TraceEvent().put, target_rank=1, address=ptr, pid_m=0, pid_n=0, mask=mask)
         """
         # Extract heap bases (from index 2 onwards)
         heap_bases = context_tensor + 2  # Offset pointer to start at heap bases
@@ -1394,12 +1398,14 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False):
             trace_info_idx = 2 + world_size + 1  # Skip: cur_rank, num_ranks, heap_bases, trace_enabled flag
             max_events = tl.load(context_tensor + trace_info_idx + 0)
             trace_counter_ptr = tl.load(context_tensor + trace_info_idx + 1)
+            op_index_counter_ptr = tl.load(context_tensor + trace_info_idx + 2)
 
-            # Cast trace_counter_ptr to pointer type
+            # Cast counter pointers to pointer type
             trace_counter = tl.cast(trace_counter_ptr, tl.pointer_type(tl.int32))
+            op_index_counter = tl.cast(op_index_counter_ptr, tl.pointer_type(tl.int32))
 
-            # Extract trace buffer pointers (11 buffers)
-            base_idx = trace_info_idx + 2
+            # Extract trace buffer pointers (13 buffers)
+            base_idx = trace_info_idx + 3  # Updated: +3 because we now have op_index_counter
             trace_buf_event_id = tl.cast(tl.load(context_tensor + base_idx + 0), tl.pointer_type(tl.int32))
             trace_buf_pid = tl.cast(tl.load(context_tensor + base_idx + 1), tl.pointer_type(tl.int32))
             trace_buf_pid_m = tl.cast(tl.load(context_tensor + base_idx + 2), tl.pointer_type(tl.int32))
@@ -1411,13 +1417,16 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False):
             trace_buf_timestamp = tl.cast(tl.load(context_tensor + base_idx + 8), tl.pointer_type(tl.int64))
             trace_buf_address = tl.cast(tl.load(context_tensor + base_idx + 9), tl.pointer_type(tl.int64))
             trace_buf_duration_cycles = tl.cast(tl.load(context_tensor + base_idx + 10), tl.pointer_type(tl.int64))
+            trace_buf_op_index = tl.cast(tl.load(context_tensor + base_idx + 11), tl.pointer_type(tl.int32))
+            trace_buf_payload_size = tl.cast(tl.load(context_tensor + base_idx + 12), tl.pointer_type(tl.int32))
 
             # Create DeviceTracing instance
             device_tracing = DeviceTracing(
                 enabled=tracing,
                 rank=rank,
                 max_events=max_events,
                 counter=trace_counter,
+                op_index_counter=op_index_counter,
                 buf_event_id=trace_buf_event_id,
                 buf_pid=trace_buf_pid,
                 buf_pid_m=trace_buf_pid_m,
@@ -1429,6 +1438,8 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False):
                 buf_timestamp=trace_buf_timestamp,
                 buf_address=trace_buf_address,
                 buf_duration_cycles=trace_buf_duration_cycles,
+                buf_op_index=trace_buf_op_index,
+                buf_payload_size=trace_buf_payload_size,
             )
 
             return DeviceContext(rank, world_size, heap_bases, device_tracing)
@@ -1442,6 +1453,7 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False):
                 rank=rank,
                 max_events=max_events_zero,
                 counter=dummy_ptr_i32,
+                op_index_counter=dummy_ptr_i32,
                 buf_event_id=dummy_ptr_i32,
                 buf_pid=dummy_ptr_i32,
                 buf_pid_m=dummy_ptr_i32,
@@ -1453,6 +1465,8 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False):
                 buf_timestamp=dummy_ptr_i64,
                 buf_address=dummy_ptr_i64,
                 buf_duration_cycles=dummy_ptr_i64,
+                buf_op_index=dummy_ptr_i32,
+                buf_payload_size=dummy_ptr_i32,
             )
 
             return DeviceContext(rank, world_size, heap_bases, device_tracing)

iris/tracing/core.py

@@ -60,24 +60,29 @@ def enable(self, max_events=1_000_000):
             "timestamp": torch.zeros(max_events, dtype=torch.int64, device=device),
             "address": torch.zeros(max_events, dtype=torch.int64, device=device),
             "duration_cycles": torch.zeros(max_events, dtype=torch.int64, device=device),
+            "op_index": torch.zeros(max_events, dtype=torch.int32, device=device),
+            "payload_size": torch.zeros(max_events, dtype=torch.int32, device=device),
         }
 
         # Atomic counter for event indexing
         self.trace_counter = torch.zeros(1, dtype=torch.int32, device=device)
+        # Atomic counter for operation indexing (tracks operation order)
+        self.op_index_counter = torch.zeros(1, dtype=torch.int32, device=device)
 
         self.iris.info(f"Device tracing enabled with max {max_events} events")
 
     def reset(self):
         """
         Reset trace counter to start a new trace capture.
 
-        Clears the event counter but keeps buffers allocated.
+        Clears the event counter and operation index counter but keeps buffers allocated.
         """
         if not self.enabled:
             self.iris.warning("Tracing not enabled. Call tracing.enable() first.")
             return
 
         self.trace_counter.zero_()
+        self.op_index_counter.zero_()
         self.iris.debug("Trace buffers reset")
 
     def _collect_system_metadata(self):
@@ -139,6 +144,8 @@ def _build_trace_events(self, num_events):
                     "address": hex(int(self.trace_buffers["address"][i].item())),
                     "xcc_id": xcc_id,
                     "cu_id": cu_id,
+                    "op_index": int(self.trace_buffers["op_index"][i].item()),
+                    "payload_size": int(self.trace_buffers["payload_size"][i].item()),
                 },
             }
 
@@ -198,11 +205,31 @@ def export(self, filename="trace.json", merge=False):
             "traceEvents": trace_events,
             "displayTimeUnit": "ns",
             "metadata": {
-                "schema_version": "1.0",
+                "schema_version": "1.1",
                 "num_events": num_events,
                 "rank": self.iris.cur_rank,
                 "world_size": self.iris.num_ranks,
                 "time_unit": "raw cycles (s_memrealtime @ 100MHz)",
+                "fields": {
+                    "name": "Event type name (e.g., 'put', 'get', 'load', 'store')",
+                    "cat": "Event category (always 'iris')",
+                    "ts": "Start timestamp in raw cycles",
+                    "pid": "Process ID (current rank)",
+                    "tid": "Thread ID (XCC{id}_CU{id})",
+                    "ph": "Phase: 'X' for complete events, 'i' for instant events",
+                    "dur": "Duration in cycles (only for complete events)",
+                    "args": {
+                        "program_id": "Triton program ID (block ID)",
+                        "pid_m": "Program ID in M dimension",
+                        "pid_n": "Program ID in N dimension",
+                        "target_rank": "Target rank for the operation",
+                        "address": "Memory address (hex) - min of address block",
+                        "xcc_id": "XCC (chiplet) ID where event occurred",
+                        "cu_id": "Compute Unit ID where event occurred",
+                        "op_index": "Operation index (0, 1, 2, ...) - automatically tracked",
+                        "payload_size": "Payload size in bytes - automatically calculated from mask and datatype",
+                    },
+                },
                 **system_metadata,
             },
         }
@@ -255,13 +282,33 @@ def export(self, filename="trace.json", merge=False):
                 "traceEvents": all_events,
                 "displayTimeUnit": "ns",
                 "metadata": {
-                    "schema_version": "1.0",
+                    "schema_version": "1.1",
                     "total_events": len(all_events),
                     "max_events": self.max_events,
                     "time_unit": "cycles (s_memrealtime @ 100MHz)",
                     "world_size": self.iris.num_ranks,
                     "timestamp_offset": min_ts if all_timestamps else 0,
                     "aligned": "minimum timestamp across all ranks",
+                    "fields": {
+                        "name": "Event type name (e.g., 'put', 'get', 'load', 'store')",
+                        "cat": "Event category (always 'iris')",
+                        "ts": "Start timestamp in raw cycles",
+                        "pid": "Process ID (current rank)",
+                        "tid": "Thread ID (XCC{id}_CU{id})",
+                        "ph": "Phase: 'X' for complete events, 'i' for instant events",
+                        "dur": "Duration in cycles (only for complete events)",
+                        "args": {
+                            "program_id": "Triton program ID (block ID)",
+                            "pid_m": "Program ID in M dimension",
+                            "pid_n": "Program ID in N dimension",
+                            "target_rank": "Target rank for the operation",
+                            "address": "Memory address (hex) - min of address block",
+                            "xcc_id": "XCC (chiplet) ID where event occurred",
+                            "cu_id": "Compute Unit ID where event occurred",
+                            "op_index": "Operation index (0, 1, 2, ...) - automatically tracked",
+                            "payload_size": "Payload size in bytes - automatically calculated from mask and datatype",
+                        },
+                    },
                     **system_metadata,
                 },
             }

iris/tracing/device.py

@@ -29,7 +29,8 @@ class _DeviceTracingCls:
     enabled: tl.constexpr
     rank: tl.constexpr  # current rank (from ctx)
     max_events: tl.tensor  # scalar (0-dim)
-    counter: tl.tensor  # pointer to int32
+    counter: tl.tensor  # pointer to int32 (event counter)
+    op_index_counter: tl.tensor  # pointer to int32 (operation index counter)
     buf_event_id: tl.tensor
     buf_pid: tl.tensor
     buf_pid_m: tl.tensor
@@ -41,13 +42,16 @@ class _DeviceTracingCls:
     buf_timestamp: tl.tensor
     buf_address: tl.tensor
     buf_duration_cycles: tl.tensor
+    buf_op_index: tl.tensor
+    buf_payload_size: tl.tensor
 
     def __init__(
         self,
         enabled,
         rank,
         max_events,
         counter,
+        op_index_counter,
         buf_event_id,
         buf_pid,
         buf_pid_m,
@@ -59,12 +63,15 @@ def __init__(
         buf_timestamp,
         buf_address,
         buf_duration_cycles,
+        buf_op_index,
+        buf_payload_size,
     ):
         """Construct DeviceTracing (called from DeviceContext.initialize)."""
         self.enabled = enabled
         self.rank = rank
         self.max_events = max_events
         self.counter = counter
+        self.op_index_counter = op_index_counter
         self.buf_event_id = buf_event_id
         self.buf_pid = buf_pid
         self.buf_pid_m = buf_pid_m
@@ -76,6 +83,8 @@ def __init__(
         self.buf_timestamp = buf_timestamp
         self.buf_address = buf_address
         self.buf_duration_cycles = buf_duration_cycles
+        self.buf_op_index = buf_op_index
+        self.buf_payload_size = buf_payload_size
 
     @triton.jit
     def record_event_start(
@@ -85,18 +94,59 @@ def record_event_start(
         address,
         pid_m,
         pid_n,
+        mask=None,
     ):
         """
         Record start of a traced operation. Returns a handle for record_event_end.
 
         Only stores when event_idx.item() < max_events (bounds check).
         cur_rank is taken from the tracing context (ctx.rank).
+        op_index is automatically tracked internally (0, 1, 2, ...).
+        payload_size is automatically calculated from mask and datatype:
+        - Counts True values in mask to get number of elements
+        - Infers datatype size from address pointer type
+        - Multiplies elements * bytes_per_element to get total bytes
+        If mask is None, payload_size is set to 0 (unknown size).
+
+        Args:
+            event_id: Event type ID (constexpr)
+            target_rank: Target rank for the operation
+            address: Memory address(es) - can be 1D or 2D block of pointers.
+                     The element type is inferred from address.type.element_ty
+            pid_m: Program ID in M dimension
+            pid_n: Program ID in N dimension
+            mask: Optional mask tensor (1D or 2D) indicating valid elements.
+                  If provided, payload_size is calculated as:
+                  (count of True values) * (bytes per element from address dtype).
+                  If None, payload_size is set to 0.
         """
         if not self.enabled:
             # Return dummy handle; record_event_end will no-op (0 < max_events is false when disabled)
             return tl.full((), 0, dtype=tl.int32)
 
         event_idx = tl.atomic_add(self.counter, 1)
+        op_index = tl.atomic_add(self.op_index_counter, 1)
+
+        # Calculate payload_size from mask and datatype
+        if mask is not None:
+            # Count True values in mask (True=1, False=0, so sum gives count of elements)
+            mask_i32 = tl.cast(mask, tl.int32)
+            num_elements = tl.sum(mask_i32)
+
+            # Get element type from address pointer and calculate size in bytes
+            # address can be 1D or 2D block of pointers, all with same element type
+            # For blocks, use .dtype instead of .type (like in test_atomic_xchg_triton.py)
+            # address.dtype is the pointer type, address.dtype.element_ty is the element dtype
+            elem_type = address.dtype.element_ty
+            # Get size in bytes using primitive_bitwidth (bits / 8 = bytes)
+            bitwidth = elem_type.primitive_bitwidth
+            elem_size_bytes = bitwidth // 8
+            # Calculate total payload size in bytes
+            payload_size = num_elements * elem_size_bytes
+        else:
+            # No mask provided, set to 0 to indicate unknown size
+            payload_size = tl.full((), 0, dtype=tl.int32)
+
         if event_idx.item() < self.max_events.item():
             tl.store(self.buf_event_id + event_idx, event_id)
             tl.store(self.buf_pid + event_idx, tl.program_id(0))
@@ -111,6 +161,8 @@ def record_event_start(
             addr_i64 = tl.cast(address, tl.int64)
             tl.store(self.buf_address + event_idx, tl.min(addr_i64))
             tl.store(self.buf_duration_cycles + event_idx, tl.full((), 0, dtype=tl.int64))
+            tl.store(self.buf_op_index + event_idx, op_index)
+            tl.store(self.buf_payload_size + event_idx, payload_size)
         return event_idx
 
     @triton.jit

tests/unittests/test_device_context.py

@@ -25,12 +25,15 @@ def device_context_tracing_1d_address_kernel(
     # 1D block of pointers: dummy_buffer + offsets
     offsets = tl.arange(0, BLOCK_SIZE)
     address_1d = dummy_buffer + offsets
+    # Create a simple mask (all True for this test)
+    mask = tl.full([BLOCK_SIZE], True, dtype=tl.int1)
     handle = ctx.tracing.record_event_start(
         event_id=TraceEvent().put,
         target_rank=(cur_rank + 1) % num_ranks,
         address=address_1d,
         pid_m=tl.program_id(0),
         pid_n=0,
+        mask=mask,
     )
     ctx.tracing.record_event_end(handle)