Composite GEMM: K-loop accumulator residency, pinned operands, sweep + deck

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/kernbench/sim_engine/op_log.py

@@ -44,11 +44,25 @@ class OpLogger:
         return self._records
 
     def record_start(self, t: float, component_id: str, msg: Any) -> None:
-        """Called by ComponentBase._on_process_start."""
+        """Called by ComponentBase._on_process_start.
+
+        Snapshots TileToken stage_type at start time so we can attribute the
+        record correctly even if the token advances stage_idx before
+        record_end fires.
+        """
+        snap: dict[str, Any] = {}
+        # TileToken (ADR-0021 pipeline) — capture which stage this is.
+        try:
+            stage = getattr(msg, "current_stage", None)
+            if stage is not None:
+                snap["stage_type"] = stage.stage_type.name
+        except Exception:
+            pass
         self._pending[id(msg)] = {
             "t_start": t,
             "component_id": component_id,
             "msg": msg,
+            "snap": snap,
         }
 
     def record_end(self, t: float, component_id: str, msg: Any) -> None:
@@ -57,6 +71,16 @@ class OpLogger:
         if pending is None:
             return
         op_kind, op_name, params = _extract_op_info(msg)
+        # Merge TileToken stage_type captured at record_start into params,
+        # and reflect it in op_name so reporting can disambiguate
+        # DMA_READ vs DMA_WRITE and FETCH vs STORE on the same component.
+        snap = pending.get("snap", {})
+        stage_type = snap.get("stage_type")
+        if stage_type is not None:
+            params = dict(params)
+            params["stage_type"] = stage_type
+            if op_name == "TileToken":
+                op_name = f"TileToken/{stage_type}"
         # Snapshot data at record time so Phase 2 replay sidesteps
         # downstream mutations of source addrs (e.g. a tl.store that
         # overwrites HBM after a load handle was sent, or a slot that