cute/profiler: token-mode pairing, tid fix, bounds_check, raw event escape hatch

- Add region_start/region_end + RegionToken for explicit pairing across Python
  scopes or cutlass.range iterations. Token carries unit_id, event_idx,
  start_ns, target_warp so the start/end pair can't drift.
- Fix profile_region passing unit_id as tid to warp_stop; add a tid parameter
  (defaults to unit_id for back-compat).
- Add bounds_check (Constexpr) on static_start/stop/warp_start/warp_stop/
  profile_region to opt out of the inner scf.if when calling from deeply
  nested control flow.
- Add raw_event_stop primitive (no warp guard, lane-0-only stores) for call
  sites that already sit inside a warp election and trigger cf.br lowering
  failures with the standard profile_region path.
- Update example.py with a token-mode demo; README documents all three modes
  and notes the IKET analogue.

Author: drisspg

transformer_nuggets/cute/profiler/README.md

@@ -25,20 +25,20 @@ with profile_session(
     my_kernel(output, prof.tensor, prof.max_events_per_unit)
 ```
 
-## Two Modes
+## Three Modes
 
-### Atomic Mode (simple)
+### Atomic mode (default)
 
 Omit `event_idx` and indices are allocated via atomics at runtime:
 
 ```python
-with profile_region(prof_buf, max_events, TAG, tid):
+with profile_region(prof_buf, max_events, TAG, unit_id):
     do_work()
 ```
 
 Works in loops without manual bookkeeping. The tradeoff: **nested regions cause timing skew**. The inner region's atomic runs while the outer region's timer is still running, inflating the outer duration.
 
-### Manual Mode (accurate for nesting)
+### Static mode (accurate for nesting)
 
 Pass an explicit `event_idx` to avoid atomics entirely:
 
@@ -60,6 +60,25 @@ for i in cutlass.range(4):
 
 You set `max_events_per_unit` to something larger than you need; the decoder scans all slots and skips empty ones.
 
+### Token mode (explicit pairing across scopes)
+
+`with profile_region(...)` already pairs start and end structurally, but it can't bridge two Python scopes or open in iteration `i` and close in `i+1`. For those cases, use `region_start` / `region_end` and pass the returned `RegionToken` explicitly:
+
+```python
+from transformer_nuggets.cute.profiler import region_start, region_end, RegionToken
+
+outer = region_start(prof_buf, unit_id, max_events)
+for i in cutlass.range(N):
+    inner = region_start(prof_buf, unit_id, max_events)
+    do_work(i)
+    region_end(prof_buf, TAG_INNER, inner, max_events)
+region_end(prof_buf, TAG_OUTER, outer, max_events)
+```
+
+`RegionToken` captures `(unit_id, event_idx, start_ns, target_warp)`, so `region_end` only needs the token plus the tag, and the start/end pair can't drift on those fields. It's a `NamedTuple`, so the DSL can thread it through `cutlass.range` as a loop-carried value if you need to keep a region open across iterations.
+
+This is the closest analogue to NVIDIA IKET's `iket.range_start` / `iket.range_end` SSA-token pairing in `cutlass.cute.experimental.iket`. The shape is similar but the mechanism is different: IKET emits MLIR ops and lowers via the proprietary `iket` dialect, while this profiler emits inline PTX (`%globaltimer`, `st.global.cs.u64`) directly.
+
 ## API
 
 ### Host (`host.py`)
@@ -77,7 +96,9 @@ You set `max_events_per_unit` to something larger than you need; the decoder sca
 
 | Function | Description |
 |----------|-------------|
-| `profile_region(buf, max_events, tag, tid, event_idx=None)` | Context manager |
+| `profile_region(buf, max_events, tag, unit_id, target_warp=None, event_idx=None, tid=None)` | Context-manager API |
+| `region_start(buf, unit_id, max_events, target_warp=None, event_idx=None) -> RegionToken` | Open a region, return a pairing token |
+| `region_end(buf, tag, token, max_events_per_unit, tid=None)` | Close a region using its token |
 | `warp_start/warp_stop(...)` | Low-level start/stop (lane 0 of target_warp) |
 | `warp_atomic_alloc(...)` | Allocate event index atomically |
 

transformer_nuggets/cute/profiler/__init__.py

@@ -47,6 +47,10 @@ def my_kernel(output, prof_buf, max_events):
     warp_start,
     warp_stop,
     profile_region,
+    region_start,
+    region_end,
+    RegionToken,
+    raw_event_stop,
 )
 
 from transformer_nuggets.cute.profiler.postprocessors import (
@@ -78,6 +82,10 @@ def my_kernel(output, prof_buf, max_events):
     "warp_start",
     "warp_stop",
     "profile_region",
+    "region_start",
+    "region_end",
+    "RegionToken",
+    "raw_event_stop",
     # Post-processors
     "group_by_unit",
     "group_by_tag",

transformer_nuggets/cute/profiler/example.py

@@ -1,8 +1,10 @@
 """Example demonstrating NVIDIA intra-kernel profiling with CUTE DSL.
 
-This example shows BOTH profiling modes:
-1. Atomic mode: No event_idx needed, indices allocated via atomics (simple)
-2. Static mode: Explicit event_idx via runtime expressions (no atomics)
+This example shows THREE profiling shapes:
+1. Atomic mode (recommended default): ``with profile_region(...)``, no event_idx.
+2. Static mode: ``with profile_region(..., event_idx=...)`` for zero atomics.
+3. Token mode: ``region_start(...) -> token`` / ``region_end(..., token)`` for
+   explicit pairing across Python scopes or loop iterations.
 
 Run with:
     python -m transformer_nuggets.cute.profiler.example
@@ -19,7 +21,7 @@
 
 from transformer_nuggets.cute.base import CuteOp
 from transformer_nuggets.cute.profiler.host import profile_session
-from transformer_nuggets.cute.profiler.ops import profile_region
+from transformer_nuggets.cute.profiler.ops import profile_region, region_start, region_end
 from transformer_nuggets.cute.profiler.postprocessors import group_by_unit
 
 
@@ -214,10 +216,94 @@ def run_static_mode():
         print("✗ Output verification failed!")
 
 
+class ProfiledKernelToken(CuteOp):
+    """Kernel using TOKEN mode: ``region_start`` returns a token consumed by ``region_end``.
+
+    The token carries (event_idx, start_ns); pairing is data-flow, not scope.
+    Use this when a region naturally spans a loop iteration boundary or two
+    Python scopes that ``with`` can't bridge cleanly.
+    """
+
+    def __init__(self, num_iterations: int = 4):
+        super().__init__()
+        self.num_iterations = num_iterations
+
+    @cute.kernel
+    def kernel(
+        self,
+        output: cute.Tensor,
+        prof_buf: cute.Tensor,
+        max_events_per_unit: cutlass.Int32,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        bdim, _, _ = cute.arch.block_dim()
+
+        prof_tid = bidx
+        global_idx = bidx * bdim + tidx
+
+        outer_tok = region_start(prof_buf, prof_tid, max_events_per_unit)
+        for i in cutlass.range(self.num_iterations):
+            compute_tok = region_start(prof_buf, prof_tid, max_events_per_unit)
+            if global_idx < cute.size(output):
+                val = output[global_idx]
+                output[global_idx] = val + 1
+            region_end(prof_buf, TAG_COMPUTE, compute_tok, max_events_per_unit)
+        region_end(prof_buf, TAG_ITERATION, outer_tok, max_events_per_unit)
+
+    @cute.jit()
+    def __call__(
+        self,
+        output: cute.Tensor,
+        prof_buf: cute.Tensor,
+        max_events_per_unit: cutlass.Int32,
+    ):
+        self.kernel(output, prof_buf, max_events_per_unit).launch(
+            grid=(NUM_BLOCKS, 1, 1),
+            block=(THREADS_PER_BLOCK, 1, 1),
+        )
+
+    def interface(self, output: torch.Tensor, prof_buf: torch.Tensor, max_events: int):
+        self.__call__(from_dlpack(output), from_dlpack(prof_buf), Int32(max_events))
+
+
+def run_token_mode():
+    """Run the token mode example."""
+    print("\n" + "=" * 60)
+    print("TOKEN MODE: region_start returns a token, region_end consumes it")
+    print("=" * 60)
+
+    device = torch.device("cuda")
+    output = torch.zeros(256, dtype=torch.float32, device=device)
+
+    import transformer_nuggets
+
+    trace_path = transformer_nuggets.DATA_DIR / "profiler_token_trace.pftrace"
+
+    with profile_session(
+        max_events_per_unit=2 * NUM_ITERATIONS + 2,
+        num_units=(NUM_BLOCKS, "Block"),
+        tag_names=["iteration", "compute", "store"],
+        trace_path=str(trace_path),
+        device=device,
+    ) as (prof, tag_table):
+        print(f"Tags: {tag_table.names}")
+        kernel = ProfiledKernelToken(num_iterations=NUM_ITERATIONS)
+        kernel.interface(output, prof.tensor, prof.max_events_per_unit)
+
+    print(f"Trace: {trace_path}")
+
+    expected = torch.full_like(output, float(NUM_ITERATIONS))
+    if torch.allclose(output, expected):
+        print("✓ Output verification passed!")
+    else:
+        print("✗ Output verification failed!")
+
+
 def main():
     print("=" * 60)
     print("NVIDIA Intra-Kernel Profiling Example")
-    print("Demonstrating BOTH atomic and static modes")
+    print("Demonstrating atomic, static, and token modes")
     print("=" * 60)
 
     if not torch.cuda.is_available():
@@ -226,6 +312,7 @@ def main():
 
     run_atomic_mode()
     run_static_mode()
+    run_token_mode()
 
     print("\n" + "=" * 60)
     print("Done! View traces at https://ui.perfetto.dev/")