Fix Gluon tracing: runtime guard, 2D reduction, test cleanup

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: mawad-amd

iris/experimental/iris_gluon.py

@@ -151,21 +151,27 @@ def record_event_start(
         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.
+            address: Memory address(es) - 1D or 2D block of pointers.
             pid_m: Program ID in M dimension
             pid_n: Program ID in N dimension
-            mask: Optional mask tensor indicating valid elements.
+            mask: Optional mask tensor indicating valid elements (1D or 2D).
         """
         if not self.enabled:
             return tl.cast(0, tl.int32)
 
+        # Guard against runtime-disabled tracing: when the kernel is compiled
+        # with tracing=True but the host context has tracing disabled, the
+        # counter pointers are null and max_events is 0. Skip all work.
+        if self.max_events <= 0:
+            return tl.cast(0, 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:
             mask_i32 = tl.cast(mask, tl.int32)
-            num_elements = gl.sum(mask_i32, axis=0)
+            num_elements = gl.sum(mask_i32)
             elem_type = address.dtype.element_ty
             bitwidth = elem_type.primitive_bitwidth
             elem_size_bytes = bitwidth // 8
@@ -184,7 +190,7 @@ def record_event_start(
             tl.store(self.buf_cu_id + event_idx, device_utils.get_cu_id())
             tl.store(self.buf_timestamp + event_idx, device_utils.read_realtime())
             addr_i64 = tl.cast(address, tl.int64)
-            tl.store(self.buf_address + event_idx, gl.min(addr_i64, axis=0))
+            tl.store(self.buf_address + event_idx, gl.min(addr_i64))
             tl.store(self.buf_duration_cycles + event_idx, tl.cast(0, tl.int64))
             tl.store(self.buf_op_index + event_idx, op_index)
             tl.store(self.buf_payload_size + event_idx, tl.cast(payload_size, tl.int32))
@@ -263,9 +269,16 @@ def initialize(context_tensor, tracing: gl.constexpr = False):
         heap_bases = context_tensor + 2  # Offset pointer to start at heap bases
 
         if tracing:
-            # Extract tracing info: starts after heap_bases, then skip trace_enabled flag
+            # Extract tracing info: starts after heap_bases, then skip trace_enabled flag.
             # Layout: [cur_rank, num_ranks, heap_base_0..N-1, trace_enabled, max_events,
             #          trace_counter_ptr, op_index_counter_ptr, buf_event_id, ...(13 buffers)]
+            #
+            # When tracing is disabled at the host, the context tensor is padded with
+            # zeros in the same positions (max_events=0, null pointers). On device,
+            # the tracing helpers (e.g., record_event_start) first early-return when
+            # max_events <= 0 and then guard all writes with a bounds check
+            # (event_idx < max_events), so decoding potentially null pointers here is
+            # safe as long as those invariants are preserved.
             trace_info_base = 2 + num_ranks + 1  # skip cur_rank, num_ranks, heap_bases, trace_enabled
             max_events = tl.cast(gl.load(context_tensor + trace_info_base + 0), tl.int32)
             trace_counter_ptr = gl.load(context_tensor + trace_info_base + 1)
@@ -1001,7 +1014,13 @@ def _build_device_context(self):
                 self.tracing.op_index_counter.data_ptr(),
             ] + trace_buffer_ptrs
         else:
-            context_data += [0]  # trace_enabled = 0 (false)
+            # trace_enabled = 0, then pad with zeros so a kernel compiled with
+            # tracing=True can safely decode the same layout without reading
+            # out of bounds. The zeros produce max_events=0 and null pointers,
+            # so the bounds check (event_idx < max_events) in record_event_start
+            # prevents any actual writes.
+            # Padding: 1 (trace_enabled) + 1 (max_events) + 2 (counters) + 13 (buffers) = 17
+            context_data += [0] * 17
 
         self._device_context = torch.tensor(context_data, dtype=torch.int64, device=self.device)
 

tests/unittests/test_device_context_gluon.py

@@ -67,7 +67,7 @@ def test_device_context_gluon_tracing_1d_address():
 
     # Verify event data fields for the first recorded event
     bufs = shmem.tracing.trace_buffers
-    assert bufs["event_id"][0].item() == 3  # TraceEvent().put == 3
+    assert bufs["event_id"][0].item() == int(TraceEvent().put)
     assert bufs["cur_rank"][0].item() == source_rank
     assert bufs["target_rank"][0].item() == (source_rank + 1) % num_ranks
     assert bufs["timestamp"][0].item() > 0
@@ -104,26 +104,67 @@ def test_device_context_gluon_initialize():
     gc.collect()
 
 
+def test_device_context_gluon_tracing_compiled_but_disabled():
+    """Test kernel compiled with tracing=True against a context tensor with tracing disabled.
+
+    The context tensor is zero-padded so the kernel can safely decode the tracing
+    layout. max_events=0 ensures record_event_start never writes to the null buffers.
+    """
+    shmem = iris_gl.iris(1 << 20)
+    # Do NOT enable tracing on host — context tensor has trace_enabled=0
+    context_tensor = shmem.get_device_context()
+    source_rank = shmem.get_rank()
+    num_ranks = shmem.get_num_ranks()
+
+    BLOCK_SIZE = 64
+    dummy_buffer = shmem.zeros((BLOCK_SIZE,), dtype=torch.int64)
+
+    shmem.barrier()
+
+    # Launch kernel compiled with tracing=True against non-tracing context
+    device_context_tracing_1d_address_kernel[(1,)](
+        iris_gl.IrisDeviceCtx,
+        context_tensor,
+        dummy_buffer,
+        source_rank,
+        num_ranks,
+        BLOCK_SIZE,
+        num_warps=1,
+    )
+    shmem.barrier()
+
+    # Verify the padded layout still reports tracing disabled and that the
+    # dummy buffer remains untouched (no writes to null pointers).
+    trace_enabled_idx = 2 + num_ranks
+    assert context_tensor[trace_enabled_idx].item() == 0
+    assert torch.all(dummy_buffer == 0)
+
+    shmem.barrier()
+    del shmem
+    import gc
+
+    gc.collect()
+
+
 def test_device_context_gluon_tracing_enable():
     """Test that shmem.tracing.enable() rebuilds context tensor with tracing fields."""
     shmem = iris_gl.iris(1 << 20)
     num_ranks = shmem.get_num_ranks()
 
-    # Without tracing: [cur_rank, num_ranks, heap_base_0..N-1, trace_enabled=0]
+    # Without tracing: tensor is padded to same size as tracing-enabled layout
+    # (zeros for max_events, counter ptrs, buffer ptrs) so kernels compiled with
+    # tracing=True can safely decode the tensor without reading out of bounds.
     ctx_no_trace = shmem.get_device_context()
-    size_no_trace = ctx_no_trace.shape[0]
+    trace_enabled_idx = 2 + num_ranks
+    assert ctx_no_trace[trace_enabled_idx].item() == 0
 
     # Enable tracing and rebuild
     shmem.tracing.enable(max_events=1000)
     ctx_with_trace = shmem.get_device_context()
-    size_with_trace = ctx_with_trace.shape[0]
 
-    # With tracing the context tensor is larger:
-    # [cur_rank, num_ranks, heap_base_0..N-1, trace_enabled=1, max_events,
-    #  trace_counter_ptr, op_index_counter_ptr, 13 buffer ptrs]
-    assert size_with_trace > size_no_trace
+    # Both tensors should be the same size (no-trace is zero-padded)
+    assert ctx_with_trace.shape[0] == ctx_no_trace.shape[0]
     # trace_enabled flag should be 1
-    trace_enabled_idx = 2 + num_ranks
     assert ctx_with_trace[trace_enabled_idx].item() == 1
 
     shmem.barrier()