[Tool] Implement a trace profiler

Author: tridao

examples/example_trace.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+"""Minimal example: intra-kernel trace profiler in CuTe-DSL.
+
+Run with:    QUACK_TRACE=1 python examples/example_trace.py
+Run without: python examples/example_trace.py
+"""
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cutlass_dsl import Int32, Int64
+
+from quack.trace import TraceContext, TraceSession
+
+ITERS = 1000
+BLOCK_THREADS = 128
+GRID = 4
+REGION_NAMES = ("loop_body", "kernel_total")
+
+
+@cute.kernel
+def trace_record_kernel(trace_ptr: Optional[Int64], iters: Int32):
+    ctx = TraceContext.create(trace_ptr, region_names=REGION_NAMES)
+    ctx.b("kernel_total")
+    for i in cutlass.range(iters):
+        ctx.b("loop_body")
+        ctx.e("loop_body")
+    ctx.e("kernel_total")
+    ctx.flush()
+
+
+@cute.jit
+def launch(trace_ptr: Optional[Int64], iters: Int32):
+    trace_record_kernel(trace_ptr, iters).launch(
+        grid=(GRID, 1, 1), block=(BLOCK_THREADS, 1, 1),
+    )
+
+
+def main():
+    out_path = "/tmp/cute_dsl_trace.json"
+
+    # sess.ptr is None when QUACK_TRACE != 1 → TraceContext.create becomes a no-op.
+    with TraceSession(out_path, grid_size=GRID, block_size=BLOCK_THREADS,
+                      region_names=list(REGION_NAMES)) as sess:
+        launch(sess.ptr, ITERS)
+
+    if sess.ptr is not None:
+        print(f"Open {out_path} in https://ui.perfetto.dev to visualize.")
+
+
+if __name__ == "__main__":
+    main()

quack/copy_utils.py

@@ -13,6 +13,8 @@
 import cutlass.pipeline
 from cutlass._mlir.dialects import llvm
 from cutlass._mlir import ir
+
+from quack.utils import make_vector
 from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
 
 
@@ -1029,3 +1031,87 @@ def copy_fn(src_idx, dst_idx, tma_bar_ptr: cute.Pointer):
                 tma_gather4_load_fn(smem_ptr, tma_bar_ptr, col_idx, row_indices)
 
     return copy_fn
+
+
+# ---------------------------------------------------------------------------
+# Store helpers
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+@cute.jit
+def store(
+    ptr: cute.Pointer,
+    val,
+    pred: Optional[Boolean] = None,
+    cop: cutlass.Constexpr = None,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Store a scalar value via cute.arch.store.
+
+    ptr:  cute.Pointer (any address space).
+    val:  DSL Numeric value.
+    pred: None → unconditional.  DSL Boolean → skipped when pred == 0.
+    cop:  Cache operator — "wb" (default), "cg", "cs" (streaming), "wt".
+    """
+    if const_expr(pred is None):
+        cute.arch.store(ptr.llvm_ptr, type(val)(val), cop=cop, loc=loc, ip=ip)
+    else:
+        if pred:
+            cute.arch.store(ptr.llvm_ptr, type(val)(val), cop=cop, loc=loc, ip=ip)
+
+
+@dsl_user_op
+@cute.jit
+def store_v2(
+    ptr: cute.Pointer,
+    v0,
+    v1,
+    pred: Optional[Boolean] = None,
+    cop: cutlass.Constexpr = None,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Vectorized store of 2 elements via cute.arch.store.
+
+    Packs v0, v1 into an MLIR <2 x T> vector.
+    ptr:  cute.Pointer (any address space, must be aligned for vector width).
+    cop:  Cache operator — "wb" (default), "cg", "cs" (streaming), "wt".
+    """
+    vec = make_vector(type(v0), v0, v1, loc=loc, ip=ip)
+    if const_expr(pred is None):
+        cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)
+    else:
+        if pred:
+            cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)
+
+
+@dsl_user_op
+@cute.jit
+def store_v4(
+    ptr: cute.Pointer,
+    v0,
+    v1,
+    v2,
+    v3,
+    pred: Optional[Boolean] = None,
+    cop: cutlass.Constexpr = None,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Vectorized store of 4 elements via cute.arch.store.
+
+    Packs v0–v3 into an MLIR <4 x T> vector.
+    ptr:  cute.Pointer (any address space, must be aligned for vector width).
+    cop:  Cache operator — "wb" (default), "cg", "cs" (streaming), "wt".
+    """
+    vec = make_vector(type(v0), v0, v1, v2, v3, loc=loc, ip=ip)
+    if const_expr(pred is None):
+        cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)
+    else:
+        if pred:
+            cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)