[Tool] Add example of gemm trace profile

Author: tridao

examples/example_gemm_trace.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+"""Trace an SM90 GEMM kernel and visualize in Perfetto.
+
+Run with:    QUACK_TRACE=1 python examples/example_gemm_trace.py
+Visualize:   Open /tmp/gemm_trace.json in https://ui.perfetto.dev
+"""
+
+import math
+
+import torch
+import cutlass
+from cutlass import Float32
+
+from quack.gemm import gemm
+from quack.gemm_default_epi import GemmDefaultSm90
+from quack.trace import TraceSession
+
+M, N, K = 4096, 4096, 4096
+TILE_M, TILE_N = 128, 192
+CLUSTER_M, CLUSTER_N = 2, 1
+OUT_PATH = "/tmp/gemm_trace.json"
+
+
+def main():
+    A = torch.randn(1, M, K, device="cuda", dtype=torch.float16)
+    B = torch.randn(1, N, K, device="cuda", dtype=torch.float16)
+    D = torch.empty(1, M, N, device="cuda", dtype=torch.float16)
+
+    # Query the GEMM config for block size (threads_per_cta).
+    g = GemmDefaultSm90(Float32, cutlass.Float16, (TILE_M, TILE_N), (CLUSTER_M, CLUSTER_N, 1))
+    # grid_size = math.ceil(M / TILE_M) * math.ceil(N / TILE_N)
+    grid_size = 132
+
+    with TraceSession(OUT_PATH, grid_size=grid_size, block_size=g.threads_per_cta,
+                      region_names=["tma_load", "mma", "epilogue"]) as sess:
+        gemm(A, B, D, C=None, tile_count_semaphore=None,
+             tile_M=TILE_M, tile_N=TILE_N,
+             cluster_M=CLUSTER_M, cluster_N=CLUSTER_N,
+             persistent=True, pingpong=True, trace_ptr=sess.ptr)
+
+    # Verify correctness.
+    ref = A[0] @ B[0].T
+    print(f"max error: {(D[0] - ref).abs().max().item():.4f}")
+    print(f"Open {OUT_PATH} in https://ui.perfetto.dev")
+
+
+if __name__ == "__main__":
+    main()

quack/gemm.py

@@ -55,6 +55,7 @@ def _compile_gemm(
     device_capacity,
     rounding_mode,
     sr_seed_mode,
+    has_trace_ptr,
 ):
     GemmCls = GemmDefaultSm100 if device_capacity[0] > 9 else GemmDefaultSm90
     mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors(
@@ -118,6 +119,7 @@ def fake_scalar(mode, dtype=Float32):
         epi_args,
         scheduler_args,
         varlen_args,
+        has_trace_ptr=has_trace_ptr,
     )
 
 
@@ -147,6 +149,7 @@ def gemm(
     add_to_output: bool = False,
     rounding_mode: int = RoundingMode.RN,
     sr_seed: int | Tensor = 0,
+    trace_ptr=None,  # Optional Int64 from TraceSession.ptr
 ) -> None:
     varlen_m = cu_seqlens_m is not None
     varlen_k = cu_seqlens_k is not None
@@ -216,6 +219,7 @@ def gemm(
         device_capacity,
         rounding_mode,
         sr_seed_mode,
+        trace_ptr is not None,
     )
 
     from quack.cache_utils import COMPILE_ONLY
@@ -251,6 +255,8 @@ def scalar_arg(scalar, mode, dtype=Float32):
     varlen_args = make_varlen_args(cu_seqlens_m, cu_seqlens_k, A_idx)
 
     if device_capacity[0] > 9:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None)
+        compiled_fn(
+            A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, trace_ptr
+        )
     else:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args)
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, trace_ptr)

quack/gemm_sm90.py

@@ -375,6 +375,7 @@ def __call__(
         scheduler_args: TileSchedulerOptions,
         varlen_args: Optional[VarlenArguments],
         stream: cuda.CUstream,
+        trace_ptr: Optional[cutlass.Int64] = None,
     ):
         """Execute the GEMM operation in steps:
         - Setup static attributes
@@ -542,6 +543,7 @@ class SharedStorage:
             self.epi_c_smem_layout_staged,
             tile_sched_params,
             TileSchedulerCls,
+            trace_ptr,
         ).launch(
             grid=grid,
             block=[self.threads_per_cta, 1, 1],
@@ -573,6 +575,7 @@ def kernel(
         epi_c_smem_layout: cute.ComposedLayout,
         tile_sched_params,
         TileSchedulerCls: cutlass.Constexpr[Callable],
+        trace_ptr: Optional[cutlass.Int64] = None,
     ):
         """
         GPU device kernel performing the batched GEMM computation.
@@ -601,6 +604,11 @@ def kernel(
         :type epi_smem_layout: cute.ComposedLayout
         """
 
+        from quack.trace import TraceContext
+
+        GEMM_REGIONS = ("tma_load", "mma", "epilogue")
+        tctx = TraceContext.create(trace_ptr, region_names=GEMM_REGIONS)
+
         varlen_m = const_expr(varlen_params.cu_seqlens_m is not None)
         varlen_k = const_expr(varlen_params.cu_seqlens_k is not None)
         assert not (varlen_m and varlen_k)
@@ -703,6 +711,7 @@ def kernel(
                     pipeline.PipelineUserType.Producer, self.ab_stage
                 )
                 while work_tile.is_valid_tile:
+                    tctx.b("tma_load")
                     tile_coord_mnkl = work_tile.tile_idx
                     batch_idx = tile_coord_mnkl[3]
                     # Local_tile partition global tensors
@@ -804,6 +813,7 @@ def kernel(
                             k_tile_cnt,
                             varlen_m=varlen_m,
                         )
+                    tctx.e("tma_load")
                     tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp)
                     work_tile = tile_scheduler.get_current_work()
                     # End of persistent scheduler loop
@@ -882,16 +892,19 @@ def kernel(
                 batch_idx = tile_coord_mnkl[3]
                 len_k = varlen_manager.len_k(batch_idx)
                 k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2])
+                tctx.b("mma")
                 ab_read_state = self.mma(
                     ab_pipeline, ab_read_state, mma_fn, acc, acc_slow, k_tile_cnt, warp_group_idx
                 )
+                tctx.e("mma")
                 if const_expr(varlen_k):
                     if k_tile_cnt == 0:
                         acc.fill(0.0)
 
                 # EPILOGUE
                 if const_expr(self.pingpong):
                     self.pingpong_barrier_sync(warp_group_idx, "epi")
+                tctx.b("epilogue")
 
                 copy_D = None
                 if const_expr(has_D):
@@ -966,6 +979,8 @@ def kernel(
                         epi_store_pipeline.producer_tail()
                     self.pingpong_barrier_arrive(1 - warp_group_idx, stage="epi")
 
+                tctx.e("epilogue")
+
                 if const_expr(not self.pingpong):
                     tile_scheduler.advance_to_next_work()
                     work_tile = tile_scheduler.get_current_work()
@@ -994,6 +1009,8 @@ def kernel(
                 if is_tma_warp:
                     epi_store_pipeline.producer_tail()
 
+        tctx.flush()
+
     @cute.jit
     def load_AB(
         self,

quack/gemm_tvm_ffi_utils.py

@@ -5,7 +5,7 @@
 
 
 import cutlass.cute as cute
-from cutlass import Int32, Float32
+from cutlass import Int32, Int64, Float32
 from cutlass.cute.runtime import make_ptr
 
 from quack.compile_utils import make_fake_tensor as fake_tensor
@@ -185,6 +185,7 @@ def compile_gemm_kernel(
     post_init=None,
     mSFA=None,
     mSFB=None,
+    has_trace_ptr=False,
 ):
     """Build GemmCls instance, apply SM90 partial, and cute.compile with TVM-FFI."""
     if device_capacity[0] == 9:
@@ -202,6 +203,9 @@ def compile_gemm_kernel(
         post_init(gemm_obj)
     stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
     sf_args = () if device_capacity[0] == 9 else (mSFA, mSFB)
+    # Trace pointer: Optional[Int64]. Compile with Int64(0) when tracing is
+    # requested, None otherwise. TVM-FFI caches each variant separately.
+    trace_ptr = Int64(0) if has_trace_ptr else None
     return cute.compile(
         gemm_obj,
         mA,
@@ -213,5 +217,6 @@ def compile_gemm_kernel(
         varlen_args,
         stream,
         *sf_args,
+        trace_ptr,
         options="--enable-tvm-ffi",
     )