diff --git a/01-minimal-gemm/cutedsl_minimal_gemm.py b/01-minimal-gemm/cutedsl_minimal_gemm.py index ad65cd8..0726829 100644 --- a/01-minimal-gemm/cutedsl_minimal_gemm.py +++ b/01-minimal-gemm/cutedsl_minimal_gemm.py @@ -198,7 +198,7 @@ def main() -> None: make_cute_tensor(c), make_fake_stream(use_tvm_ffi_env_stream=True), True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) gemm_accum = cute.compile( minimal_gemm, @@ -207,7 +207,7 @@ def main() -> None: make_cute_tensor(c), make_fake_stream(use_tvm_ffi_env_stream=True), False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) for exp in exps: diff --git a/01-minimal-gemm/profile.sh b/01-minimal-gemm/profile.sh index e81e876..d15071f 100644 --- a/01-minimal-gemm/profile.sh +++ b/01-minimal-gemm/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file 01.csv --metrics gpu__time_duration.sum --kernel-name "minimal_gemm" python minimal_gemm.py # ncu-rep output -ncu -o ncu_prof_1 --import-source 1 --set full --kernel-name "minimal_gemm" -f python minimal_gemm.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_1 --import-source 1 --set full --kernel-name "minimal_gemm" -f python minimal_gemm.py +# CuTe DSL impl: +ncu -o ncu_prof_1 --import-source 1 --set full --kernel-name "regex:.*minimal_gemm.*" -f python cutedsl_minimal_gemm.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_1 python minimal_gemm.py diff --git a/02-mixed-precision-gemm/cutedsl_mixed_precision_gemm.py b/02-mixed-precision-gemm/cutedsl_mixed_precision_gemm.py index 0a6dcb0..7104142 100644 --- a/02-mixed-precision-gemm/cutedsl_mixed_precision_gemm.py +++ b/02-mixed-precision-gemm/cutedsl_mixed_precision_gemm.py @@ -242,7 +242,7 @@ def compile_for(out_torch_dtype: torch.dtype, out_cutlass_dtype): out_cutlass_dtype, True, is_cvt, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) gemm_accum = cute.compile( mixed_precision_gemm, @@ -254,7 +254,7 @@ def compile_for(out_torch_dtype: torch.dtype, out_cutlass_dtype): out_cutlass_dtype, False, is_cvt, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return gemm_clear, gemm_accum diff --git a/02-mixed-precision-gemm/profile.sh b/02-mixed-precision-gemm/profile.sh index 89b7e66..a00a324 100644 --- a/02-mixed-precision-gemm/profile.sh +++ b/02-mixed-precision-gemm/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file 02.csv --metrics gpu__time_duration.sum --kernel-name "mixed_precision_gemm" python mixed_precision_gemm.py # ncu-rep output -ncu -o ncu_prof_2 --import-source 1 --set full --kernel-name "mixed_precision_gemm" -f python mixed_precision_gemm.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_2 --import-source 1 --set full --kernel-name "mixed_precision_gemm" -f python mixed_precision_gemm.py +# CuTe DSL impl: +ncu -o ncu_prof_2 --import-source 1 --set full --kernel-name "regex:.*mixed_precision_gemm.*" -f python cutedsl_mixed_precision_gemm.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_2 python mixed_precision_gemm.py diff --git a/03-tiled-mma/cutedsl_tiled_mma.py b/03-tiled-mma/cutedsl_tiled_mma.py index 48fa76e..09966ad 100644 --- a/03-tiled-mma/cutedsl_tiled_mma.py +++ b/03-tiled-mma/cutedsl_tiled_mma.py @@ -195,7 +195,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) gemm_accum = cute.compile( tiled_mma, @@ -206,7 +206,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) print(f" M={M}, N={N}, K={K} ".center(PRINT_LENGTH, "-")) diff --git a/03-tiled-mma/profile.sh b/03-tiled-mma/profile.sh index 84e8795..411e2ee 100644 --- a/03-tiled-mma/profile.sh +++ b/03-tiled-mma/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file 03.csv --metrics gpu__time_duration.sum --kernel-name "tiled_mma" python tiled_mma.py # ncu-rep output -ncu -o ncu_prof_3 --import-source 1 --set full --kernel-name "tiled_mma" -f python tiled_mma.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_3 --import-source 1 --set full --kernel-name "tiled_mma" -f python tiled_mma.py +# CuTe DSL impl: +ncu -o ncu_prof_3 --import-source 1 --set full --kernel-name "regex:.*tiled_mma.*" -f python cutedsl_tiled_mma.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_3 python tiled_mma.py diff --git a/04-tiled-copy/cutedsl_tiled_copy.py b/04-tiled-copy/cutedsl_tiled_copy.py index 6e49ebe..f345ffe 100644 --- a/04-tiled-copy/cutedsl_tiled_copy.py +++ b/04-tiled-copy/cutedsl_tiled_copy.py @@ -235,7 +235,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) gemm_accum = cute.compile( tiled_copy_gemm, @@ -246,7 +246,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) print(f" M={M}, N={N}, K={K} ".center(PRINT_LENGTH, "-")) diff --git a/04-tiled-copy/profile.sh b/04-tiled-copy/profile.sh index d298ee9..2b55fe3 100644 --- a/04-tiled-copy/profile.sh +++ b/04-tiled-copy/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file tiled_copy.csv --metrics gpu__time_duration.sum --kernel-name "tiled_copy" python tiled_copy.py # ncu-rep output -ncu -o ncu_prof_4 --import-source 1 --set full --kernel-name "tiled_copy" -f python tiled_copy.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_4 --import-source 1 --set full --kernel-name "tiled_copy" -f python tiled_copy.py +# CuTe DSL impl: +ncu -o ncu_prof_4 --import-source 1 --set full --kernel-name "regex:.*tiled_copy.*" -f python cutedsl_tiled_copy.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_4 python tiled_copy.py diff --git a/05-block-mma/cutedsl_block_mma.py b/05-block-mma/cutedsl_block_mma.py index 2dc9e75..b302d60 100644 --- a/05-block-mma/cutedsl_block_mma.py +++ b/05-block-mma/cutedsl_block_mma.py @@ -232,7 +232,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) gemm_accum = cute.compile( block_mma_gemm, @@ -243,7 +243,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) print(f" M={M}, N={N}, K={K} ".center(PRINT_LENGTH, "-")) diff --git a/05-block-mma/profile.sh b/05-block-mma/profile.sh index 3efe7c8..b604e44 100644 --- a/05-block-mma/profile.sh +++ b/05-block-mma/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file block_mma.csv --metrics gpu__time_duration.sum --kernel-name "block_mma" python block_mma.py # ncu-rep output -ncu -o ncu_prof_5 --import-source 1 --set full --kernel-name "block_mma" -f python block_mma.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_5 --import-source 1 --set full --kernel-name "block_mma" -f python block_mma.py +# CuTe DSL impl: +ncu -o ncu_prof_5 --import-source 1 --set full --kernel-name "regex:.*block_mma.*" -f python cutedsl_block_mma.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_5 python block_mma.py diff --git a/06-block-copy/cutedsl_block_copy.py b/06-block-copy/cutedsl_block_copy.py index 0b94929..0d0e85e 100644 --- a/06-block-copy/cutedsl_block_copy.py +++ b/06-block-copy/cutedsl_block_copy.py @@ -380,7 +380,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) gemm_accum = cute.compile( block_copy_gemm, @@ -391,7 +391,7 @@ def main() -> None: make_fake_stream(use_tvm_ffi_env_stream=True), cutlass.BFloat16, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) print(f" M={M}, N={N}, K={K} ".center(PRINT_LENGTH, "-")) diff --git a/06-block-copy/profile.sh b/06-block-copy/profile.sh index 5ea3e49..7f61f31 100644 --- a/06-block-copy/profile.sh +++ b/06-block-copy/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file block_copy.csv --metrics gpu__time_duration.sum --kernel-name "block_copy" python block_copy.py # ncu-rep output -ncu -o ncu_prof_6 --import-source 1 --set full --kernel-name "block_copy" -f python block_copy.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_6 --import-source 1 --set full --kernel-name "block_copy" -f python block_copy.py +# CuTe DSL impl: +ncu -o ncu_prof_6 --import-source 1 --set full --kernel-name "regex:.*block_copy.*" -f python cutedsl_block_copy.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_6 python block_copy.py diff --git a/07-swizzling/cutedsl_swizzling.py b/07-swizzling/cutedsl_swizzling.py index a4d91c2..0c92e42 100644 --- a/07-swizzling/cutedsl_swizzling.py +++ b/07-swizzling/cutedsl_swizzling.py @@ -396,7 +396,7 @@ def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out acc_dtype, out_dtype, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( swizzling_gemm, @@ -408,7 +408,7 @@ def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out acc_dtype, out_dtype, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/07-swizzling/profile.sh b/07-swizzling/profile.sh index a867ffa..c4ba53e 100644 --- a/07-swizzling/profile.sh +++ b/07-swizzling/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file swizzling.csv --metrics gpu__time_duration.sum --kernel-name "swizzling" python swizzling.py # ncu-rep output -ncu -o ncu_prof_7 --import-source 1 --set full --kernel-name "swizzling" -f python swizzling.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_7 --import-source 1 --set full --kernel-name "swizzling" -f python swizzling.py +# CuTe DSL impl: +ncu -o ncu_prof_7 --import-source 1 --set full --kernel-name "regex:.*swizzling.*" -f python cutedsl_swizzling.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_7 python swizzling.py diff --git a/08-dynamic-mma/cutedsl_dynamic_mma.py b/08-dynamic-mma/cutedsl_dynamic_mma.py index a3c129d..b3899eb 100644 --- a/08-dynamic-mma/cutedsl_dynamic_mma.py +++ b/08-dynamic-mma/cutedsl_dynamic_mma.py @@ -651,7 +651,7 @@ def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out acc_dtype, out_dtype, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( dynamic_mma_gemm, @@ -663,7 +663,7 @@ def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out acc_dtype, out_dtype, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/08-dynamic-mma/profile.sh b/08-dynamic-mma/profile.sh index a43904d..e9ad66b 100644 --- a/08-dynamic-mma/profile.sh +++ b/08-dynamic-mma/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file dynamic_mma.csv --metrics gpu__time_duration.sum --kernel-name "dynamic_mma" python dynamic_mma.py # ncu-rep output -ncu -o ncu_prof_8 --import-source 1 --set full --kernel-name "dynamic_mma" -f python dynamic_mma.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_8 --import-source 1 --set full --kernel-name "dynamic_mma" -f python dynamic_mma.py +# CuTe DSL impl: +ncu -o ncu_prof_8 --import-source 1 --set full --kernel-name "regex:.*dynamic_mma.*" -f python cutedsl_dynamic_mma.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_8 python dynamic_mma.py diff --git a/09-pipelining/cutedsl_pipelining.py b/09-pipelining/cutedsl_pipelining.py index 2968baf..251c4a5 100644 --- a/09-pipelining/cutedsl_pipelining.py +++ b/09-pipelining/cutedsl_pipelining.py @@ -839,7 +839,7 @@ def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out acc_dtype, out_dtype, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( pipelining_gemm, @@ -851,7 +851,7 @@ def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out acc_dtype, out_dtype, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/09-pipelining/cutedsl_pipelining_profiled.py b/09-pipelining/cutedsl_pipelining_profiled.py new file mode 100644 index 0000000..5cdcc43 --- /dev/null +++ b/09-pipelining/cutedsl_pipelining_profiled.py @@ -0,0 +1,1243 @@ +"""Multi-stage cp.async pipelined GEMM in CuTe DSL. + +A near-verbatim copy of ``cutedsl_pipelining.py`` plus an in-kernel timeline +profiler (see ``cutedsl_sm_profiler.py``): the only kernel changes are +``%globaltimer`` stamps at the phase boundaries, gated by a ``profile_enabled`` +constexpr so the kernel is unchanged when profiling is off. + +CuTe DSL counterpart of ``pipelining.cu`` / ``pipelining.py``. Extends +``08-dynamic-mma`` with two pipelines stacked on top: + + 1. **G2S (smem) pipeline** with ``NUM_STAGES`` ring slots (default 3). + Stages ``0 .. NUM_STAGES-2`` are prefetched in the prologue. The + mainloop's last k-block triggers ``cp_async_wait_group(NUM_STAGES-2)`` + so at most one outstanding cp.async group is in flight while the + register pipeline consumes the freshly arrived stage. + 2. **S2R (register) pipeline**, gated by the ``REG_STAGES`` constant. + With ``REG_STAGES >= 2`` (default) the smem -> rmem load for + k_block + 1 is hoisted ahead of the MMA for k_block, so the LDS / + LDSM latency is hidden behind the tensor-core math. Setting + ``REG_STAGES == 1`` disables the register pipeline (``pipelining_no_ + reg_prefetch.cu``): each k-tile loads its whole A/B fragment, then + runs every MMA, leaving only the G2S pipeline to overlap work. + +The prologue mirrors ``pipelining.cu`` lines 156-214: + * predicated K-residue copy for stage 0 (shifted by ``k_residue``); + * full-tile copies for stages ``1 .. NUM_STAGES-2`` with an overshoot + guard that clears the M/N predicates when ``ik == k_tile_count``; + * ``cp_async_wait_group(NUM_STAGES-2)`` and a single ``ldmatrix`` of + k_block=0 from smem stage 0 before entering the mainloop. + +The mainloop walks all ``num_k_block`` MMA k-iterations per k-tile, +prefetches the next rmem at the top, fires the next stage's cp.async +on ``k_block == 0``, then issues the tensor-core gemm. Three counters +(``smem_pipe_read``, ``smem_pipe_write``, ``k_tile_index``) track the +ring-buffer state. + +Three dtype specs are exercised, matching ``pipelining.py``: + + * fp16 in, fp16 acc, fp16 out (exercise only) + * fp16 in, fp32 acc, fp16 out (validated) + * bf16 in, fp32 acc, bf16 out (validated) + +Run with ``python cutedsl_pipelining.py``. + +Smem footprint / occupancy: + * Only the multi-stage A/B pipeline gets real smem (``A_pipe + B_pipe``; + 96KB for the default 128x128x64 tile at 3 stages). The C addend and the + O store buffer are never live during the mainloop, so the epilogue + aliases them over the (drained) A/B smem, mirroring 13-/14-*. This keeps + the block under the ~113KB needed for two resident CTAs per SM. + * To realise that second block the launch sets ``min_blocks_per_mp=2`` + (mirrors 08-dynamic-mma): without it CuTeDSL's allocator runs + ILP-greedy to ~160 reg/thread and caps occupancy at one block/SM. The + hint tightens the budget to 128 reg/thread, letting two blocks + co-reside (~24% achieved occupancy vs ~12.5%). At that cap the + no-prefetch path (``REG_STAGES == 1``) fits spill-free, while the + register-prefetch path (``REG_STAGES >= 2``) keeps an extra k_block + fragment live and spills lightly to (L1-resident) local memory — which + is why the no-prefetch path is the faster configuration on H200 here + (448us vs 527us at M=N=K=4096, bf16). + * C is folded into the accumulator in the epilogue (acc starts at zero, + acc += C after the matmul) rather than preloaded, so sC need not coexist + with the pipeline. The epilogue then narrows to out_dtype and drains + R2S -> S2G via a swizzled sO (Swizzle<3,3,3> o (8 x min(64, BLK_N))), + with a 2-D (CCPY_M, CCPY_N) S2G predicate for M/N residue handling. +""" + +from pathlib import Path + +import cutedsl_sm_profiler as smprof +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import CUstream +from cutlass.cute.runtime import from_dlpack, make_fake_stream + + +# Block tile (matches pipelining.cu's default KernelSpec defaults) +BLK_M = 128 +BLK_N = 128 +BLK_K = 64 +# G2S (smem) ring-buffer depth. +NUM_STAGES = 3 +# Register-pipeline (S2R) depth, mirroring ``prefetch_s2r_tiles`` in the C++: +# * ``REG_STAGES >= 2`` double-buffers the smem -> rmem load, hoisting the +# k_block + 1 fetch ahead of the k_block MMA (the default fast path). +# * ``REG_STAGES == 1`` disables register prefetch entirely: each k-tile +# loads its whole A/B fragment, then runs every MMA — only the +# GMEM -> SMEM -> RF copy pipeline overlaps work. Mirrors +# ``pipelining_no_reg_prefetch.cu``. +REG_STAGES = 1 + +MMA_INST_MNK = (16, 8, 16) +ATOM_LAYOUT_MNK = (2, 4, 1) +VAL_EXPAND_MNK = (1, 1, 2) +MMA_TILE_MNK = ( + ATOM_LAYOUT_MNK[0] * VAL_EXPAND_MNK[0] * MMA_INST_MNK[0], # 32 + ATOM_LAYOUT_MNK[1] * VAL_EXPAND_MNK[1] * MMA_INST_MNK[1], # 32 + ATOM_LAYOUT_MNK[2] * VAL_EXPAND_MNK[2] * MMA_INST_MNK[2], # 32 +) +NUM_THREADS = ( + ATOM_LAYOUT_MNK[0] * ATOM_LAYOUT_MNK[1] * ATOM_LAYOUT_MNK[2] * 32 # 256 +) + +# Profiler event types for THIS kernel (the profiler itself is generic and knows +# nothing about these -- they are registered with names on the host side). The +# four phases are recorded once each; compute/wait are recorded once per +# mainloop k-tile, so a warp emits 4 + 2*k_tile_count events. +PROF_TOTAL = 0 +PROF_PROLOGUE = 1 +PROF_MAINLOOP = 2 +PROF_EPILOGUE = 3 +PROF_COMPUTE = 4 +PROF_WAIT = 5 +PROF_EVENT_NAMES = { + PROF_TOTAL: "total", + PROF_PROLOGUE: "prologue", + PROF_MAINLOOP: "mainloop", + PROF_EPILOGUE: "epilogue", + PROF_COMPUTE: "compute", + PROF_WAIT: "wait", +} + + +# ----------------------------------------------------------------------------- +# Device kernel +# ----------------------------------------------------------------------------- + + +@cute.kernel +def pipelining_kernel( + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + mO: cute.Tensor, + tiled_mma: cute.TiledMma, + g2s_tiled_copy_a: cute.TiledCopy, + g2s_tiled_copy_b: cute.TiledCopy, + g2s_tiled_copy_c: cute.TiledCopy, + s2r_tiled_copy_a: cute.TiledCopy, + s2r_tiled_copy_b: cute.TiledCopy, + s2r_tiled_copy_c: cute.TiledCopy, + r2s_tiled_copy_o: cute.TiledCopy, + s2g_tiled_copy_o: cute.TiledCopy, + sA_layout: cute.ComposedLayout, + sB_layout: cute.ComposedLayout, + sC_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + out_dtype: cutlass.Constexpr, + is_gemm: cutlass.Constexpr[bool], + # Profiling is opt-in via this single bundled handle. A disabled ``prof`` + # (its buffer is None) lowers to no device argument and folds every + # profiler call to nothing, so the kernel is then identical to + # ``cutedsl_pipelining.py``. + prof: smprof.ProfilerArgs = smprof.ProfilerArgs(None, False), +): + tid, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + M, K = mA.shape + N, _ = mB.shape + + # Open a profiling session once (resolves this warp's ids + buffer sizes a + # single time, from the grid dims and the buffer header); call sites then + # just name an event type. + prof_ctx = prof.begin() + prof_ctx.event_start(PROF_TOTAL) + prof_ctx.event_start(PROF_PROLOGUE) + + # 3-D block tiler — gA / gB carry a K-tile iteration mode. + tiler = (BLK_M, BLK_N, BLK_K) + gA = cute.local_tile(mA, tiler=tiler, coord=(bidy, bidx, None), proj=(1, None, 1)) + gB = cute.local_tile(mB, tiler=tiler, coord=(bidy, bidx, None), proj=(None, 1, 1)) + gC = cute.local_tile(mC, tiler=tiler, coord=(bidy, bidx, 0), proj=(1, 1, None)) + gO = cute.local_tile(mO, tiler=tiler, coord=(bidy, bidx, 0), proj=(1, 1, None)) + + # Per-block residue clamps. + m_max = M - BLK_M * bidy + n_max = N - BLK_N * bidx + k_tile_count = cute.size(gA, mode=[2]) + # k_residue <= 0: how far the first K-tile's origin sits before the + # true K=0. Shifting gA/gB along K by k_residue makes every K-tile + # after the first trivially in-bounds. + k_residue = K - BLK_K * k_tile_count + gA = cute.domain_offset((0, k_residue, 0), gA) + gB = cute.domain_offset((0, k_residue, 0), gB) + + # ----- Smem allocation ----- + # Only the multi-stage A/B pipeline gets real storage. sC (accumulate + # addend) and sO (epilogue store buffer) are never live during the + # mainloop, so the epilogue aliases them over this same A/B smem (see + # below). Footprint stays at sA + sB so two CTAs fit per SM. + smem = cutlass.utils.SmemAllocator() + sA = smem.allocate_tensor(mA.element_type, sA_layout, byte_alignment=16) + sB = smem.allocate_tensor(mB.element_type, sB_layout, byte_alignment=16) + + # ----- G2S partitions (predicated) ----- + thr_g2s_a = g2s_tiled_copy_a.get_slice(tid) + tAgA = thr_g2s_a.partition_S(gA) # (CPY, CPY_M, CPY_K, k_tiles) + tAsA = thr_g2s_a.partition_D(sA) # (CPY, CPY_M, CPY_K, PIPE) + + thr_g2s_b = g2s_tiled_copy_b.get_slice(tid) + tBgB = thr_g2s_b.partition_S(gB) # (CPY, CPY_N, CPY_K, k_tiles) + tBsB = thr_g2s_b.partition_D(sB) # (CPY, CPY_N, CPY_K, PIPE) + + # C is loaded in the epilogue (folded into the accumulator there), so its + # smem partition is built then; only the gmem-side partition is needed up + # front for the identity-coord predicate setup. + thr_g2s_c = g2s_tiled_copy_c.get_slice(tid) + tCgC = thr_g2s_c.partition_S(gC) # (CPY, CPY_M, CPY_N) + + # ----- Identity tensors for predication ----- + cA = cute.make_identity_tensor((BLK_M, BLK_K)) + cB = cute.make_identity_tensor((BLK_N, BLK_K)) + cC = cute.make_identity_tensor((BLK_M, BLK_N)) + tAcA = thr_g2s_a.partition_S(cA) + tBcB = thr_g2s_b.partition_S(cB) + tCcC = thr_g2s_c.partition_S(cC) + + # ----- M / N predicate (used for prologue stages 1.. and the mainloop) ----- + # Three-mode layout (rest_v, CPY_M, CPY_K) with the K mode broadcast at + # stride 0 so the same M/N predicate is replayed across every K iter. + # K-bound is not needed here because the domain_offset shift guarantees + # every K-tile from stage 1 onward is in-bounds; only stage 0 (handled + # by ``tApA_first`` below) needs the per-element K check. + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAgA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBgB.shape[0][1], + cute.size(tBsB, mode=[1]), + cute.size(tBsB, mode=[2]), + ), + stride=(cute.size(tBsB, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tApA.shape[0]): + for m in cutlass.range_constexpr(tApA.shape[1]): + tApA[rest_v, m, 0] = cute.elem_less( + tAcA[(0, rest_v), m, 0][0], + m_max, + ) + for rest_v in cutlass.range_constexpr(tBpB.shape[0]): + for n in cutlass.range_constexpr(tBpB.shape[1]): + tBpB[rest_v, n, 0] = cute.elem_less( + tBcB[(0, rest_v), n, 0][0], + n_max, + ) + + # ----- Fragments ----- + thr_mma = tiled_mma.get_slice(tid) + tCrA = tiled_mma.make_fragment_A(thr_mma.partition_A(sA)[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(thr_mma.partition_B(sB)[None, None, None, 0]) + tCrC = tiled_mma.make_fragment_C(thr_mma.partition_C(gC)) + + thr_s2r_a = s2r_tiled_copy_a.get_slice(tid) + thr_s2r_b = s2r_tiled_copy_b.get_slice(tid) + thr_s2r_c = s2r_tiled_copy_c.get_slice(tid) + + tAsA_s2r = thr_s2r_a.partition_S(sA) # (CPY, CPY_M, CPY_K, PIPE) + tBsB_s2r = thr_s2r_b.partition_S(sB) # (CPY, CPY_N, CPY_K, PIPE) + tArA_s2r = thr_s2r_a.retile(tCrA) # (CPY, CPY_M, CPY_K) + tBrB_s2r = thr_s2r_b.retile(tCrB) # (CPY, CPY_N, CPY_K) + + # ========================================================================= + # G2S prologue: prefetch NUM_STAGES-1 stages. + # ========================================================================= + + # Stage 0 (residue tile) needs a per-element predicate combining M-bound + # AND K-bound. With CPY_K = 1 a single iter carries lanes spread across + # the whole BLK_K range — some at valid in-bounds K-positions, others + # not after the domain_offset shift — so the gate must fire per lane / + # per val element rather than at the C++ iteration level. + tApA_first = cute.make_rmem_tensor( + cute.make_layout( + ( + tAgA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=( + cute.size(tAsA, mode=[1]) * cute.size(tAsA, mode=[2]), + cute.size(tAsA, mode=[2]), + 1, + ), + ), + cutlass.Boolean, + ) + tBpB_first = cute.make_rmem_tensor( + cute.make_layout( + ( + tBgB.shape[0][1], + cute.size(tBsB, mode=[1]), + cute.size(tBsB, mode=[2]), + ), + stride=( + cute.size(tBsB, mode=[1]) * cute.size(tBsB, mode=[2]), + cute.size(tBsB, mode=[2]), + 1, + ), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tApA_first.shape[0]): + for m in cutlass.range_constexpr(tApA_first.shape[1]): + for k in cutlass.range_constexpr(tApA_first.shape[2]): + tApA_first[rest_v, m, k] = cute.elem_less( + tAcA[(0, rest_v), m, k][0], + m_max, + ) and cute.elem_less( + cutlass.Int32(-1), + tAcA[(0, rest_v), m, k][1] + k_residue, + ) + for rest_v in cutlass.range_constexpr(tBpB_first.shape[0]): + for n in cutlass.range_constexpr(tBpB_first.shape[1]): + for k in cutlass.range_constexpr(tBpB_first.shape[2]): + tBpB_first[rest_v, n, k] = cute.elem_less( + tBcB[(0, rest_v), n, k][0], + n_max, + ) and cute.elem_less( + cutlass.Int32(-1), + tBcB[(0, rest_v), n, k][1] + k_residue, + ) + + # Pre-zero stage 0 so predicate-off slots read as 0. Each thread fills its + # own g2s partition slots and the following cp.async writes that same + # partition, so program order suffices — the post-wait_group sync_threads + # below publishes both the zeros and the loads to the s2r readers. + tAsA[None, None, None, 0].fill(0) + tBsB[None, None, None, 0].fill(0) + cute.copy( + g2s_tiled_copy_a, + tAgA[None, None, None, 0], + tAsA[None, None, None, 0], + pred=tApA_first, + ) + cute.copy( + g2s_tiled_copy_b, + tBgB[None, None, None, 0], + tBsB[None, None, None, 0], + pred=tBpB_first, + ) + cute.arch.cp_async_commit_group() + + # Stages 1 .. NUM_STAGES-2: full-tile copies, with overshoot guard. + k_tile_index = cutlass.Int32(1) + for ik in cutlass.range_constexpr(1, NUM_STAGES - 1): + # Once we'd read past the end of K, mask everything off. + if k_tile_index >= k_tile_count: + tApA.fill(False) + tBpB.fill(False) + cute.copy( + g2s_tiled_copy_a, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, ik], + pred=tApA, + ) + cute.copy( + g2s_tiled_copy_b, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, ik], + pred=tBpB, + ) + cute.arch.cp_async_commit_group() + k_tile_index = k_tile_index + 1 + + # ========================================================================= + # Wait for the prologue's stage-0 cp.async, then prime the accumulator. + # ========================================================================= + cute.arch.cp_async_wait_group(NUM_STAGES - 2) + cute.arch.sync_threads() + # Prologue prefetch done; close the prologue range and open the mainloop in + # one stamp (the two slices abut). + prof_ctx.event_switch(PROF_PROLOGUE, PROF_MAINLOOP) + + num_k_block = cute.size(tCrA, mode=[2]) + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(NUM_STAGES - 1) + + # ----- Initialise accumulator to zero ----- + # The C addend (accumulate path) is folded in during the epilogue, not + # preloaded here, so sC stays out of the mainloop and can alias the A/B + # smem. Both mainloops therefore start from a zeroed accumulator. + tCrC.fill(0.0) + + if cutlass.const_expr(REG_STAGES >= 2): + # ===================================================================== + # Register-prefetch (S2R) pipeline. The smem -> rmem load for + # k_block + 1 is hoisted ahead of the MMA for k_block so the LDS / + # LDSM latency hides behind the tensor-core math. Double-buffered at + # the register level. The outer loop's iteration variable is unused — + # gmem progress is tracked by ``k_tile_index`` (the cp.async read + # pointer, NUM_STAGES-1 ahead of the smem read pointer because the + # prologue prefetched those stages) and smem ring positions by + # ``smem_pipe_{read,write}``. Mirrors ``pipelining.cu``. + # ===================================================================== + # Prefetch first k_block from stage 0 into registers. + cute.copy( + s2r_tiled_copy_a, + tAsA_s2r[None, None, 0, smem_pipe_read], + tArA_s2r[None, None, 0], + ) + cute.copy( + s2r_tiled_copy_b, + tBsB_s2r[None, None, 0, smem_pipe_read], + tBrB_s2r[None, None, 0], + ) + + for _ in cutlass.range(k_tile_count, unroll_full=False): + for k_block in cutlass.range_constexpr(num_k_block): + # When we're on the last k_block of this tile, wait for the + # next smem stage to arrive and bump smem_pipe_read. + if k_block == num_k_block - 1: + cute.arch.cp_async_wait_group(NUM_STAGES - 2) + cute.arch.sync_threads() + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == NUM_STAGES: + smem_pipe_read = cutlass.Int32(0) + + # Prefetch next k_block from smem to rmem. After the bump + # above (which only fires on the last k_block of the tile), + # smem_pipe_read already points at the freshly-synced stage, + # so the wrap to k_block=0 reads from the new stage. + k_block_next = (k_block + 1) % num_k_block + cute.copy( + s2r_tiled_copy_a, + tAsA_s2r[None, None, k_block_next, smem_pipe_read], + tArA_s2r[None, None, k_block_next], + ) + cute.copy( + s2r_tiled_copy_b, + tBsB_s2r[None, None, k_block_next, smem_pipe_read], + tBrB_s2r[None, None, k_block_next], + ) + + # On the first k_block of the tile, fire cp.async for the + # next smem stage (writes ahead of the current read pointer). + if k_block == 0: + # Overshoot guard: when k_tile_index has stepped past + # the end of K, mask everything off so we don't issue + # OOB cp.async. + if k_tile_index >= k_tile_count: + tApA.fill(False) + tBpB.fill(False) + cute.copy( + g2s_tiled_copy_a, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.copy( + g2s_tiled_copy_b, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, smem_pipe_write], + pred=tBpB, + ) + cute.arch.cp_async_commit_group() + k_tile_index = k_tile_index + 1 + smem_pipe_write = smem_pipe_write + 1 + if smem_pipe_write == NUM_STAGES: + smem_pipe_write = cutlass.Int32(0) + + # Tensor-core gemm for this k_block. + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + else: + # ===================================================================== + # No register prefetch: only the GMEM -> SMEM -> RF copy pipeline + # overlaps work. Each k-tile loads its whole A/B fragment from smem in + # one copy, fires the next smem stage's cp.async, then runs the MMA + # over every k_block. Because all S2R loads finish before the first + # MMA, the smem -> rmem latency is *not* hidden behind the math — this + # is the baseline the register pipeline above improves on. Mirrors + # ``pipelining_no_reg_prefetch.cu``'s mainloop. + # ===================================================================== + for _ in cutlass.range(k_tile_count, unroll_full=False): + # Per-tile "compute" event opens here (S2R load + MMAs). The append + # counter gives it the next event_id, so compute events land in + # k-tile order; folds to nothing when profiling is off. + prof_ctx.event_start(PROF_COMPUTE) + + # Whole-tile smem -> rmem load (every k_block at once). + cute.copy( + s2r_tiled_copy_a, + tAsA_s2r[None, None, None, smem_pipe_read], + tArA_s2r, + ) + cute.copy( + s2r_tiled_copy_b, + tBsB_s2r[None, None, None, smem_pipe_read], + tBrB_s2r, + ) + + # Fire cp.async for the next smem stage (writes ahead of read ptr). + if k_tile_index >= k_tile_count: + tApA.fill(False) + tBpB.fill(False) + cute.copy( + g2s_tiled_copy_a, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.copy( + g2s_tiled_copy_b, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, smem_pipe_write], + pred=tBpB, + ) + cute.arch.cp_async_commit_group() + k_tile_index = k_tile_index + 1 + smem_pipe_write = smem_pipe_write + 1 + if smem_pipe_write == NUM_STAGES: + smem_pipe_write = cutlass.Int32(0) + + # MMA over every k_block of the freshly-loaded tile. + for k_block in cutlass.range_constexpr(num_k_block): + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + + # Close "compute" and open "wait" with a single stamp (after this + # tile's MMAs, before the cp.async wait). This boundary is the one + # mark not anchored by a __syncthreads; the compiler barrier inside + # the timestamp read keeps it from drifting. Fusing the two halves the + # boundary's timer reads and makes the slices abut exactly. + prof_ctx.event_switch(PROF_COMPUTE, PROF_WAIT) + + # Wait for the next stage to land, advance the smem read pointer. + cute.arch.cp_async_wait_group(NUM_STAGES - 2) + cute.arch.sync_threads() + # Close "wait" (barrier-anchored by the sync above). + prof_ctx.event_end(PROF_WAIT) + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == NUM_STAGES: + smem_pipe_read = cutlass.Int32(0) + + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + # Mainloop drained; close the mainloop range and open the epilogue (one stamp). + prof_ctx.event_switch(PROF_MAINLOOP, PROF_EPILOGUE) + + # ========================================================================= + # Epilogue. C and O reuse the now-drained A/B smem: the mainloop's final + # cp_async_wait_group(0) + sync above guarantees every thread is done + # reading sA / sB, so the region is free to overwrite. + # ========================================================================= + + # ----- Fold the C addend into the accumulator (accumulate path only) ----- + if cutlass.const_expr(not is_gemm): + # sC (compute-C dtype) aliases sA's storage. An fp32 BLK_M x BLK_N + # tile spans sA plus part of the contiguous sB, both dead here. The + # swizzle must stay in the make_tensor layout (where partition_D + # composes it correctly); recast_ptr only adjusts the element dtype + # (and strips to a bare pointer). Splitting the swizzle into + # recast_ptr instead mis-maps partition_D under Swizzle<3,3,3> on + # SM80-class ldmatrix/cp.async copies — that form is for SM90 TMA. + sC = cute.make_tensor( + cute.recast_ptr(sA.iterator, dtype=mC.element_type), + sC_layout, + ) + tCsC = thr_g2s_c.partition_D(sC) + + # Coalesced gmem -> smem load of C with the M / N residue predicate + # (C has no K mode, so no K bound). Predicate-off slots stay zero and + # add nothing. + tCsC.fill(0) + for m in cutlass.range_constexpr(cute.size(tCgC, mode=[1])): + for n in cutlass.range_constexpr(cute.size(tCgC, mode=[2])): + if cute.elem_less(tCcC[0, m, n][0], m_max) and cute.elem_less( + tCcC[0, m, n][1], + n_max, + ): + cute.copy( + g2s_tiled_copy_c, + tCgC[None, m, n], + tCsC[None, m, n], + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # smem -> rmem into a C-shaped fragment, then acc += C (in acc dtype). + tCrC_add = cute.make_fragment_like(tCrC, mC.element_type) + cute.copy( + s2r_tiled_copy_c, + thr_s2r_c.partition_S(sC), + thr_s2r_c.retile(tCrC_add), + ) + tCrC.store(tCrC.load() + tCrC_add.load().to(tCrC.element_type)) + + # All threads must finish reading sC before sO overwrites the region. + cute.arch.sync_threads() + + # ----- Narrow to out_dtype, R2S -> S2G via a swizzled smem buffer ----- + # sO (out dtype, which equals A's dtype for every spec) aliases sA's + # storage via a plain full-ComposedLayout make_tensor. The split + # .outer / .inner form mis-maps swizzled upper-half columns in the r2s + # partition_D store, so the plain form is mandatory here. + sO = cute.make_tensor(sA.iterator, sO_layout) + + tCrO = cute.make_fragment_like(tCrC, out_dtype) + tCrO.store(tCrC.load().to(out_dtype)) + + # R2S: register fragment -> swizzled smem buffer sO. + thr_r2s_o = r2s_tiled_copy_o.get_slice(tid) + tOrO_r2s = thr_r2s_o.retile(tCrO) + tOsO_r2s = thr_r2s_o.partition_D(sO[None, None, 0]) + cute.copy(r2s_tiled_copy_o, tOrO_r2s, tOsO_r2s) + + cute.arch.sync_threads() + + # S2G: smem -> gmem with a TV layout that packs threads contiguously + # along the N dim, matching pipelining.cu's TiledCopyO_S2G. The + # partitioned tensors are 3-mode ``((CPY_INNER, REST_V), CPY_M, CPY_N)``; + # the pred uses a stride-0 rest_v slot so the same iter-level scalar + # is replayed across the (trivial) CPY mode. + thr_s2g_o = s2g_tiled_copy_o.get_slice(tid) + tOsO_s2g = thr_s2g_o.partition_S(sO[None, None, 0]) + tOgO_s2g = thr_s2g_o.partition_D(gO) + tOcO_s2g = thr_s2g_o.partition_S(cC) + + rest_v_size = tOgO_s2g.shape[0][1] + ccpy_m_size = tOgO_s2g.shape[1] + ccpy_n_size = tOgO_s2g.shape[2] + tOpO_s2g = cute.make_rmem_tensor( + cute.make_layout( + (rest_v_size, ccpy_m_size, ccpy_n_size), + stride=(0, 1, ccpy_m_size), + ), + cutlass.Boolean, + ) + for m in cutlass.range_constexpr(ccpy_m_size): + for n in cutlass.range_constexpr(ccpy_n_size): + tOpO_s2g[0, m, n] = cute.elem_less( + tOcO_s2g[(0, 0), m, n][0], + m_max, + ) and cute.elem_less( + tOcO_s2g[(0, 0), m, n][1], + n_max, + ) + cute.copy(s2g_tiled_copy_o, tOsO_s2g, tOgO_s2g, pred=tOpO_s2g) + + # Close the epilogue and total ranges. The compiler barrier inside event_end + # already keeps the final clock read from being hoisted above the S2G store. + prof_ctx.event_end(PROF_EPILOGUE) + prof_ctx.event_end(PROF_TOTAL) + + +@cute.jit +def pipelining_gemm( + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + mO: cute.Tensor, + stream: CUstream, + acc_dtype: cutlass.Constexpr, + out_dtype: cutlass.Constexpr, + is_gemm: cutlass.Constexpr[bool], + prof_buffer=None, + profile_enabled: cutlass.Constexpr[bool] = False, + profile_show_overhead: cutlass.Constexpr[bool] = False, +): + # ----- Tiled MMA ----- + op = cute.nvgpu.warp.MmaF16BF16Op( + mA.element_type, + acc_dtype, + MMA_INST_MNK, + ) + tm = cute.make_tiled_mma( + op, + atom_layout_mnk=ATOM_LAYOUT_MNK, + permutation_mnk=MMA_TILE_MNK, + ) + + # ----- Swizzled smem layouts (Swizzle<3,3,3>) ; A/B carry a PIPE mode ----- + swz = cute.make_swizzle(3, 3, 3) + inner_AB = min(64, BLK_K) + atom_AB = cute.make_composed_layout( + swz, + 0, + cute.make_layout((8, inner_AB), stride=(inner_AB, 1)), + ) + sA_layout = cute.tile_to_shape( + atom_AB, + (BLK_M, BLK_K, NUM_STAGES), + order=(0, 1, 2), + ) + sB_layout = cute.tile_to_shape( + atom_AB, + (BLK_N, BLK_K, NUM_STAGES), + order=(0, 1, 2), + ) + # sC mirrors pipelining.cu's SmemLayoutC: Swizzle<3,3,3> o (8 x min(64, BLK_N)). + inner_C = min(64, BLK_N) + atom_C = cute.make_composed_layout( + swz, + 0, + cute.make_layout((8, inner_C), stride=(inner_C, 1)), + ) + sC_layout = cute.tile_to_shape(atom_C, (BLK_M, BLK_N), order=(0, 1)) + # sO mirrors pipelining.cu's SmemLayoutO: Swizzle<3,3,3> o (8 x min(64, BLK_N)). + # Single-stage — the epilogue runs after the mainloop finishes draining, + # so no PIPE mode is needed. + inner_O = min(64, BLK_N) + atom_O = cute.make_composed_layout( + swz, + 0, + cute.make_layout((8, inner_O), stride=(inner_O, 1)), + ) + sO_layout = cute.tile_to_shape(atom_O, (BLK_M, BLK_N, 1), order=(0, 1, 2)) + + # ----- G2S copies (cp.async) ----- + g2s_op = cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL, + ) + elt_bytes_ab = mA.element_type.width // 8 + k_chunk = min(64, BLK_K) // (16 // elt_bytes_ab) + tlAB_thr = cute.make_layout( + (NUM_THREADS // k_chunk, k_chunk), + stride=(k_chunk, 1), + ) + tlAB_val = cute.make_layout((1, 16 // elt_bytes_ab)) + g2s_atom_a = cute.make_copy_atom(g2s_op, mA.element_type, num_bits_per_copy=128) + g2s_atom_b = cute.make_copy_atom(g2s_op, mB.element_type, num_bits_per_copy=128) + g2s_tiled_copy_a = cute.make_tiled_copy_tv(g2s_atom_a, tlAB_thr, tlAB_val) + g2s_tiled_copy_b = cute.make_tiled_copy_tv(g2s_atom_b, tlAB_thr, tlAB_val) + + elt_bytes_c = mC.element_type.width // 8 + n_chunk = min(64, BLK_N) // (16 // elt_bytes_c) + tlC_thr = cute.make_layout( + (NUM_THREADS // n_chunk, n_chunk), + stride=(n_chunk, 1), + ) + tlC_val = cute.make_layout((1, 16 // elt_bytes_c)) + g2s_atom_c = cute.make_copy_atom(g2s_op, mC.element_type, num_bits_per_copy=128) + g2s_tiled_copy_c = cute.make_tiled_copy_tv(g2s_atom_c, tlC_thr, tlC_val) + + universal = cute.nvgpu.CopyUniversalOp() + + # ----- S2R copies (ldmatrix for 16-bit operands) ----- + # A/B own 4 32-bit packets per thread (VAL_EXPAND_K=2), so the x4 ldmatrix + # variant fits exactly. C has no K val-expand, so each thread only owns + # 2 32-bit packets — use the x2 ldmatrix variant when C is 16-bit. For + # wider C (e.g. fp32) fall back to a universal copy. + ldm_op_ab = cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4) + ldm_op_c = cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2) + s2r_atom_a = cute.make_copy_atom(ldm_op_ab, mA.element_type) + s2r_atom_b = cute.make_copy_atom(ldm_op_ab, mB.element_type) + s2r_tiled_copy_a = cute.make_tiled_copy_A(s2r_atom_a, tm) + s2r_tiled_copy_b = cute.make_tiled_copy_B(s2r_atom_b, tm) + if cutlass.const_expr(mC.element_type.width == 16): + s2r_atom_c = cute.make_copy_atom(ldm_op_c, mC.element_type) + else: + s2r_atom_c = cute.make_copy_atom(universal, mC.element_type) + s2r_tiled_copy_c = cute.make_tiled_copy_C(s2r_atom_c, tm) + + # ----- R2S + S2G copies for the smem-staged epilogue ----- + # R2S: MMA-derived TiledCopy so each thread's accumulator fragment lands + # at its natural smem position under the swizzled sO layout. Pick the op + # the same way pipelining.cu does: SM90+ with a 16-bit output uses + # ``stmatrix`` (x2 because the C/O fragment has no K val-expand and so + # owns 2 32-bit packets per thread); otherwise fall back to a universal + # STS lowering. + sm_major, _ = torch.cuda.get_device_capability() + if cutlass.const_expr(sm_major >= 9 and mO.element_type.width == 16): + stm_op = cute.nvgpu.warp.StMatrix8x8x16bOp(False, 2) + r2s_atom_o = cute.make_copy_atom(stm_op, mO.element_type) + else: + r2s_atom_o = cute.make_copy_atom(universal, mO.element_type) + r2s_tiled_copy_o = cute.make_tiled_copy_C(r2s_atom_o, tm) + + # S2G: explicit TV layout — threads packed contiguously along N, + # each thread emits 16 / sizeof(out) elements in a 128-bit store. + elt_bytes_o = mO.element_type.width // 8 + n_chunk_o = min(64, BLK_N) // (16 // elt_bytes_o) + tlO_thr = cute.make_layout( + (NUM_THREADS // n_chunk_o, n_chunk_o), + stride=(n_chunk_o, 1), + ) + tlO_val = cute.make_layout((1, 16 // elt_bytes_o)) + s2g_atom_o = cute.make_copy_atom(universal, mO.element_type, num_bits_per_copy=128) + s2g_tiled_copy_o = cute.make_tiled_copy_tv(s2g_atom_o, tlO_thr, tlO_val) + + # ----- Launch ----- + M, _ = mA.shape + N, _ = mB.shape + grid_n = (N + BLK_N - 1) // BLK_N + grid_m = (M + BLK_M - 1) // BLK_M + + # Bundle into the single handle the kernel takes. With profiling off the + # default args make this ProfilerArgs(None, False) -> no buffer arg. + prof = smprof.ProfilerArgs(prof_buffer, profile_enabled, profile_show_overhead) + + pipelining_kernel( + mA, + mB, + mC, + mO, + tm, + g2s_tiled_copy_a, + g2s_tiled_copy_b, + g2s_tiled_copy_c, + s2r_tiled_copy_a, + s2r_tiled_copy_b, + s2r_tiled_copy_c, + r2s_tiled_copy_o, + s2g_tiled_copy_o, + sA_layout, + sB_layout, + sC_layout, + sO_layout, + out_dtype, + is_gemm, + prof, + ).launch( + grid=(grid_n, grid_m, 1), + block=(NUM_THREADS, 1, 1), + stream=stream, + # Require >= 2 resident blocks per SM (nvvm.minctasm). Unconstrained, + # CuTeDSL's register allocator runs ILP-greedy to ~160 reg/thread (not + # a spill, just loose) and caps occupancy at one block/SM on H100/H200. + # The hint tightens the budget to <=128 reg/thread so two blocks share + # an SM; combined with the sC/sO smem aliasing the footprint (96KB for + # the default tile) leaves room for that second block. + min_blocks_per_mp=2, + ) + + +# ----------------------------------------------------------------------------- +# Host-side test harness (mirrors pipelining.py / cutedsl_dynamic_mma.py) +# ----------------------------------------------------------------------------- + + +PRINT_LENGTH = 100 + + +def relative_error(target: torch.Tensor, ref: torch.Tensor, eps: float = 1e-8) -> float: + diff = target - ref + norm_diff = torch.norm(diff, p=2) + norm_diff_ref = torch.norm(ref, p=2) + return (norm_diff / (norm_diff_ref + eps)).item() + + +def compare_matrix( + kernel_output: torch.Tensor, + torch_output: torch.Tensor, + counters: dict, +) -> None: + kernel_output = kernel_output.float() + torch_output = torch_output.float() + max_diff = torch.max(torch.abs(torch_output - kernel_output)) + mean_diff = torch.mean(torch.abs(torch_output - kernel_output)) + re = relative_error(kernel_output, torch_output) + is_correct = re < 0.001 + + if not is_correct: + counters["failed"] += 1 + print(f" Kernel Output: {tuple(kernel_output.shape)} ".center(PRINT_LENGTH, "-")) + print(kernel_output[:8, :8]) + print(f" Torch Output: {tuple(torch_output.shape)} ".center(PRINT_LENGTH, "-")) + print(torch_output[:8, :8]) + else: + counters["succeed"] += 1 + status = "Success" if is_correct else "Failed" + print( + f" Result: {status}, Max diff = {max_diff:.5f}, Mean diff = {mean_diff:.5f}, RE = {(re * 100):.2f}% ".center( + PRINT_LENGTH, "-" + ) + ) + + +def make_cute_tensor(t: torch.Tensor) -> cute.Tensor: + divisibility = max(1, 16 // t.element_size()) + return ( + from_dlpack(t, assumed_align=16, enable_tvm_ffi=True) + .mark_layout_dynamic(leading_dim=1) + .mark_compact_shape_dynamic(mode=1, divisibility=divisibility) + ) + + +def _compile_pair(a_template, b_template, c_template, o_template, acc_dtype, out_dtype): + g_clear = cute.compile( + pipelining_gemm, + make_cute_tensor(a_template), + make_cute_tensor(b_template), + make_cute_tensor(c_template), + make_cute_tensor(o_template), + make_fake_stream(use_tvm_ffi_env_stream=True), + acc_dtype, + out_dtype, + True, + options="--enable-tvm-ffi --generate-line-info", + ) + g_accum = cute.compile( + pipelining_gemm, + make_cute_tensor(a_template), + make_cute_tensor(b_template), + make_cute_tensor(c_template), + make_cute_tensor(o_template), + make_fake_stream(use_tvm_ffi_env_stream=True), + acc_dtype, + out_dtype, + False, + options="--enable-tvm-ffi --generate-line-info", + ) + return g_clear, g_accum + + +def _events_per_warp(K: int) -> int: + """How many range events one warp emits for a given K. + + 4 phase ranges (total / prologue / mainloop / epilogue) plus, when the + register pipeline is disabled (``REG_STAGES == 1``), 2 ranges (compute, wait) + per k-tile. Each ``event_start`` bumps the per-warp append counter once, so + this is exactly the largest ``event_id`` the warp reaches; the profiler's + ``max_events_per_warp`` must be at least this or trailing events are dropped. + """ + k_tiles = (K + BLK_K - 1) // BLK_K + return 4 + (2 * k_tiles if REG_STAGES < 2 else 0) + + +def _profile_once(M: int, N: int, K: int, out_dir: Path, show_overhead: bool) -> None: + """Compile + run ONE profiled bf16 GEMM and dump its summary + Perfetto trace. + + ``show_overhead`` is a compile-time constant: with it on, the kernel records + the extra start/end stamps and the export draws the two overhead bands; with + it off, neither the device nor the trace carries any overhead data (the + minimal-overhead reference). The two settings write to distinct files. + """ + num_warps = NUM_THREADS // 32 + num_blocks = ((M + BLK_M - 1) // BLK_M) * ((N + BLK_N - 1) // BLK_N) + + torch.cuda.manual_seed_all(9527) + a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + b = torch.randn(N, K, device="cuda", dtype=torch.bfloat16) + c = torch.randn(M, N, device="cuda", dtype=torch.float32) + + # Per-warp event capacity, sized from K with a small margin so nothing is + # dropped (events beyond it are silently dropped — the counter keeps + # advancing but no record is written). + max_events = _events_per_warp(K) + 8 + profiler = smprof.SmProfiler(num_blocks, num_warps, max_events) + for event_no, name in PROF_EVENT_NAMES.items(): + profiler.register_event(event_no, name) + g = cute.compile( + pipelining_gemm, + make_cute_tensor(a), + make_cute_tensor(b), + make_cute_tensor(c), + make_cute_tensor(torch.empty(M, N, device="cuda", dtype=torch.bfloat16)), + make_fake_stream(use_tvm_ffi_env_stream=True), + cutlass.Float32, + cutlass.BFloat16, + False, + profiler.cute_buffer(), + True, + show_overhead, + options="--enable-tvm-ffi --generate-line-info", + ) + out = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + profiler.reset() + g(a, b, c.clone(), out, profiler.buffer) + torch.cuda.synchronize() + + ref = torch.addmm(c, a.float(), b.T.float()).bfloat16() + tag = "overhead-on" if show_overhead else "overhead-off" + print( + f" profiled bf16 {M}x{N}x{K} [{tag}]: RE = {relative_error(out.float(), ref.float()) * 100:.2f}% ".center( + PRINT_LENGTH, "=" + ) + ) + profiler.summarize(f"bf16 {M}x{N}x{K} [{tag}]") + + # Discriminating dump: with the bands recorded, show the raw start/end + # overhead widths for a handful of events so a zero band (CSE'd reads / + # sub-resolution timer) is distinguishable from a real one. Also persist the + # raw buffer for offline re-analysis. + if show_overhead: + print(" raw overhead widths (ns) for first events ".center(PRINT_LENGTH, "-")) + print(f"{'event':<12}{'start_ovh':>12}{'end_ovh':>12}{'interval':>12}") + for _ in range(8): + for record in profiler._records(): + name = PROF_EVENT_NAMES.get(record.event_no, f"event_{record.event_no}") + # A band is only defined when its inner stamp was recorded (> 0). An + # event closed by event_switch shares the boundary stamp and has no + # end band, so en_early stays 0 -> report 0 rather than en - 0. + start_ovh = record.st_late - record.st if record.st_late > 0 else 0 + end_ovh = record.en - record.en_early if record.en_early > 0 else 0 + print(f"{name:<12}{start_ovh:>12d}{end_ovh:>12d}{record.en - record.st:>12d}") + buf_path = out_dir / f"pipelining_bf16_M{M}_N{N}_K{K}_buffer.pt" + torch.save(profiler.buffer.cpu(), str(buf_path)) + print(f" raw buffer -> {buf_path} ".center(PRINT_LENGTH, "-")) + + suffix = "_overhead" if show_overhead else "" + json_path = out_dir / f"pipelining_bf16_M{M}_N{N}_K{K}{suffix}.json" + n_events = profiler.export_perfetto(str(json_path), show_overhead=show_overhead) + print(f" Perfetto: {n_events} slices -> {json_path} ".center(PRINT_LENGTH, "-")) + + +def profile_demo() -> None: + """Dump a profiled 4096x4096x4096 bf16 GEMM trace (overhead bands OFF). + + This is the only path that passes a profiler buffer (profiling on); the + sweep in ``main`` compiles with profiling off, so no buffer is threaded + through and the kernel matches ``cutedsl_pipelining.py``. + + Exports the full all-block trace without overhead bands (the overhead-on + variant at this size would be a ~3x larger JSON; use a smaller M/N if you + want to inspect the per-event overhead bands). ``max_events`` is derived from + K so the per-warp capacity always covers the kernel's event count. + """ + out_dir = Path(__file__).resolve().parent.parent / ".claude" / "perf" / "inkernel_profile" + out_dir.mkdir(parents=True, exist_ok=True) + + M, N, K = 4096, 4096, 4096 + _profile_once(M, N, K, out_dir, show_overhead=False) + + +def bench_overhead(M: int = 512, N: int = 512, K: int = 4096, iters: int = 100, repeats: int = 10) -> None: + """Profiling-off vs profiling-on wall-clock comparison + overhead breakdown. + + Compiles three variants of the SAME bf16 GEMM and times each over + ``repeats`` batches of ``iters`` launches (CUDA events, warmup discarded): + + * off -- profiling disabled (the production kernel) + * on (phases) -- only the 4 coarse phase events (outside the hot loop) + * on (full) -- phases + the 2 per-k-tile events inside the mainloop + + ``on(phases) - off`` is the profiler's FIXED cost; ``on(full) - on(phases)`` + is the per-tile (hot-loop) cost -- i.e. where the overhead lives. + """ + num_warps = NUM_THREADS // 32 + num_blocks = ((M + BLK_M - 1) // BLK_M) * ((N + BLK_N - 1) // BLK_N) + k_tiles = (K + BLK_K - 1) // BLK_K + + torch.cuda.manual_seed_all(9527) + a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + b = torch.randn(N, K, device="cuda", dtype=torch.bfloat16) + c = torch.randn(M, N, device="cuda", dtype=torch.float32).clone() + out = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + o_tmpl = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + + def _compile(profile_enabled, prof_buffer): + return cute.compile( + pipelining_gemm, + make_cute_tensor(a), + make_cute_tensor(b), + make_cute_tensor(c), + make_cute_tensor(o_tmpl), + make_fake_stream(use_tvm_ffi_env_stream=True), + cutlass.Float32, + cutlass.BFloat16, + False, + prof_buffer, + profile_enabled, + options="--enable-tvm-ffi --generate-line-info", + ) + + g_off = _compile(False, None) + profiler = smprof.SmProfiler(num_blocks, num_warps, _events_per_warp(K) + 8) + g_on = _compile(True, profiler.cute_buffer()) + + def _time(run): + for _ in range(10): # warmup + run() + torch.cuda.synchronize() + per_iter = [] + for _ in range(repeats): + start, stop = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + run() + stop.record() + torch.cuda.synchronize() + per_iter.append(start.elapsed_time(stop) / iters * 1000.0) # us + t = torch.tensor(per_iter) + return t.mean().item(), t.std().item() + + off = _time(lambda: g_off(a, b, c, out)) + profiler.reset() + on = _time(lambda: g_on(a, b, c, out, profiler.buffer)) + + events_per_warp = 4 + 2 * k_tiles # 4 phases + 2 (compute, wait) per k-tile + overhead = on[0] - off[0] + print( + f" Profiler overhead: bf16 {M}x{N}x{K}, {num_blocks} blocks x {num_warps} warps, " + f"{k_tiles} k-tiles, {iters}x{repeats} ".center(PRINT_LENGTH, "=") + ) + print(f"{'variant':<18}{'us/iter':>12}{'std':>10}{'vs off':>12}") + print(f"{'off (no profile)':<18}{off[0]:>12.2f}{off[1]:>10.2f}{'--':>12}") + print(f"{'on (profiled)':<18}{on[0]:>12.2f}{on[1]:>10.2f}{overhead:>+12.2f}") + print("-" * PRINT_LENGTH) + print( + f"total overhead: {overhead:+.2f} us ({overhead / off[0] * 100:+.1f}% of off), " + f"{events_per_warp} events/warp -> ~{overhead / events_per_warp * 1000:.0f} ns/event, " + f"{overhead / k_tiles * 1000:.0f} ns/k-tile" + ) + + +def main() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("This example requires a CUDA-capable GPU.") + + # Sweep matches pipelining.py exactly. + Ms = [16, 64, 128, 192, 256, 1024, 4096, 8192] + Ns = [16, 64, 128, 192, 256, 1024, 4096, 8192] + Ks = [16, 64, 128, 192, 256, 1024, 4096, 8192] + exps = [(m, n, k) for m in Ms for n in Ns for k in Ks] + + counters = {"succeed": 0, "failed": 0} + torch.cuda.manual_seed_all(9527) + + M0, N0, K0 = 128, 128, 64 + + # ----- Spec 1: fp16 in, fp16 acc, fp16 out (exercise only) ----- + print(" Compiling fp16 in / fp16 acc / fp16 out ... ".center(PRINT_LENGTH, "-")) + a_t = torch.empty(M0, K0, device="cuda", dtype=torch.float16) + b_t = torch.empty(N0, K0, device="cuda", dtype=torch.float16) + c_t = torch.empty(M0, N0, device="cuda", dtype=torch.float16) + o_t = torch.empty(M0, N0, device="cuda", dtype=torch.float16) + fp16_clear, fp16_accum = _compile_pair( + a_t, + b_t, + c_t, + o_t, + cutlass.Float16, + cutlass.Float16, + ) + + # ----- Spec 2: fp16 in, fp32 acc, fp16 out ----- + print(" Compiling fp16 in / fp32 acc / fp16 out ... ".center(PRINT_LENGTH, "-")) + a_t = torch.empty(M0, K0, device="cuda", dtype=torch.float16) + b_t = torch.empty(N0, K0, device="cuda", dtype=torch.float16) + c_t = torch.empty(M0, N0, device="cuda", dtype=torch.float32) + o_t = torch.empty(M0, N0, device="cuda", dtype=torch.float16) + fp16f32_clear, fp16f32_accum = _compile_pair( + a_t, + b_t, + c_t, + o_t, + cutlass.Float32, + cutlass.Float16, + ) + + # ----- Spec 3: bf16 in, fp32 acc, bf16 out ----- + print(" Compiling bf16 in / fp32 acc / bf16 out ... ".center(PRINT_LENGTH, "-")) + a_t = torch.empty(M0, K0, device="cuda", dtype=torch.bfloat16) + b_t = torch.empty(N0, K0, device="cuda", dtype=torch.bfloat16) + c_t = torch.empty(M0, N0, device="cuda", dtype=torch.float32) + o_t = torch.empty(M0, N0, device="cuda", dtype=torch.bfloat16) + bf16_clear, bf16_accum = _compile_pair( + a_t, + b_t, + c_t, + o_t, + cutlass.Float32, + cutlass.BFloat16, + ) + + # ----- Sweep: fp16 = fp16 * fp16 + fp16 (Spec 1, exercise only) ----- + print(" fp16 = fp16 * fp16 + fp16 (exercise only) ".center(PRINT_LENGTH, "=")) + torch.cuda.manual_seed_all(9527) + for m, n, k in exps: + print(f" M={m}, N={n}, K={k} ".center(PRINT_LENGTH, "-")) + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(n, k, device="cuda", dtype=torch.float16) + c = torch.randn(m, n, device="cuda", dtype=torch.float16) + + out = torch.empty(m, n, device="cuda", dtype=torch.float16) + fp16_clear(a, b, c.clone(), out) + out = torch.empty(m, n, device="cuda", dtype=torch.float16) + fp16_accum(a, b, c.clone(), out) + torch.cuda.synchronize() + + # ----- Sweep: fp16 in, fp32 acc, fp16 out (Spec 2) ----- + print(" fp16 = fp32_acc(fp16 * fp16) + fp32 ".center(PRINT_LENGTH, "=")) + torch.cuda.manual_seed_all(9527) + for m, n, k in exps: + print(f" M={m}, N={n}, K={k} ".center(PRINT_LENGTH, "-")) + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(n, k, device="cuda", dtype=torch.float16) + c = torch.randn(m, n, device="cuda", dtype=torch.float32) + + out = torch.empty(m, n, device="cuda", dtype=torch.float16) + fp16f32_clear(a, b, c.clone(), out) + torch.cuda.synchronize() + compare_matrix(out, torch.matmul(a, b.T), counters) + + out = torch.empty(m, n, device="cuda", dtype=torch.float16) + fp16f32_accum(a, b, c.clone(), out) + torch.cuda.synchronize() + compare_matrix(out, torch.addmm(c, a.float(), b.T.float()).half(), counters) + + # ----- Sweep: bf16 in, fp32 acc, bf16 out (Spec 3) ----- + print(" bf16 = fp32_acc(bf16 * bf16) + fp32 ".center(PRINT_LENGTH, "=")) + torch.cuda.manual_seed_all(9527) + for m, n, k in exps: + print(f" M={m}, N={n}, K={k} ".center(PRINT_LENGTH, "-")) + a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + c = torch.randn(m, n, device="cuda", dtype=torch.float32) + + out = torch.empty(m, n, device="cuda", dtype=torch.bfloat16) + bf16_clear(a, b, c.clone(), out) + torch.cuda.synchronize() + compare_matrix(out, torch.matmul(a.float(), b.T.float()).bfloat16(), counters) + + out = torch.empty(m, n, device="cuda", dtype=torch.bfloat16) + bf16_accum(a, b, c.clone(), out) + torch.cuda.synchronize() + compare_matrix( + out, + torch.addmm(c, a.float(), b.T.float()).bfloat16(), + counters, + ) + + print(f" Summary: {counters['succeed']} Succeed, {counters['failed']} Failed ".center(PRINT_LENGTH, "-")) + + +if __name__ == "__main__": + profile_demo() + main() diff --git a/09-pipelining/cutedsl_sm_profiler.py b/09-pipelining/cutedsl_sm_profiler.py new file mode 100644 index 0000000..163ceb2 --- /dev/null +++ b/09-pipelining/cutedsl_sm_profiler.py @@ -0,0 +1,600 @@ +"""Reusable, application-agnostic in-kernel SM profiler for CuTe DSL kernels. + +A close CuTe DSL port of ``magnus_a2a``'s ``sm_profiler`` (CUDA C++). The +profiler is kernel-agnostic: callers register their own event *types* +(``event_no``) via :meth:`SmProfiler.register_event` and instrument code with +:func:`event_start` / :func:`event_end`. + +Model +----- +Each (block, warp) appends events into one global ``int64`` buffer: a per-warp +counter hands out a monotonically increasing ``event_id``, so events land in +order. A range event stores a type (``event_no``) plus start/end timestamps +from ``%globaltimer``, each preceded by a compiler barrier so the read is not +reordered across the surrounding work. No atomics -- only lane 0 writes a given +slot. Overflow drops: past ``max_events_per_warp`` the counter keeps advancing +but no records are written. + +Two deviations from the original cut per-event memory traffic: a :class:`Session` +(opened once via :func:`begin`) caches the header + this warp's ids in registers +instead of re-parsing per call; and ``event_start`` returns the ``event_id`` in a +register for ``event_end`` to reuse, replacing the original's global ``active`` +table round-trip. + +Buffer layout (single 1-D int64 tensor, element offsets; +``total_warps = num_blocks * num_warps``, ``E = max_events_per_warp``, +``F = NUM_EV_FIELDS``):: + + header : [0, HEADER) num_blocks, num_warps, max_events + smid : [HEADER, +total_warps) sm_id per (block, warp), written once + counters : [.., +total_warps) one per (block, warp) + events : [.., +total_warps*E*F) F fields per (block, warp, event_id) + +(block, warp, event_id) are implicit in the index and ``sm_id`` is per-warp, so +each record stores only (event_no, st_ts, en_ts). +""" + +from __future__ import annotations + +import json +from collections import namedtuple +from dataclasses import dataclass, field + +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, nvvm +from cutlass.cute.runtime import from_dlpack + + +# --------------------------------------------------------------------------- +# Special-register reads + compiler barrier. +# --------------------------------------------------------------------------- +def _globaltimer_ns(): + """Read ``%globaltimer`` (ns) via a SIDE-EFFECTING inline asm. + + Not the pure ``nvvm.read_ptx_sreg_globaltimer`` intrinsic: two pure SREG + reads in straight-line code get CSE'd into one, making an event's two stamps + (e.g. ``EV_ST`` / ``EV_ST_LATE``) identical and collapsing the overhead band + to zero. A ``volatile`` ``mov.u64`` cannot be CSE'd, so each call reads at its + own program point. + """ + return cutlass.Int64( + llvm.inline_asm( + ir.IntegerType.get_signless(64), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def _smid(): + return cutlass.Int64(nvvm.read_ptx_sreg_smid(ir.IntegerType.get_signless(32))) + + +def _compiler_barrier(): + """Empty ``asm volatile`` with a memory clobber -- a pure compiler barrier. + + Emits no instruction but stops the compiler from scheduling the adjacent + ``%globaltimer`` read across the work it brackets; only constrains ordering, + unlike a hardware ``membar``. + """ + llvm.inline_asm( + None, + [], + "", + "~{memory}", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +# --------------------------------------------------------------------------- +# Buffer schema (shared by device writers and host readers). No application +# event types live here -- callers register their own. +# --------------------------------------------------------------------------- +HDR_NUM_BLOCKS = 0 +HDR_NUM_WARPS = 1 +HDR_MAX_EVENTS = 2 +HEADER = 3 + +# Per-event record fields (int64). ``EV_ST`` / ``EV_EN`` bound the event and are +# stretched to swallow the profiler's own bookkeeping: ``EV_ST`` is the FIRST +# instruction of ``event_start``, ``EV_EN`` the LAST of ``event_end``, so +# adjacent events abut instead of leaving an overhead gap. +# +# ``EV_ST_LATE`` / ``EV_EN_EARLY`` are written ONLY under ``show_overhead``; they +# mark the inner edges of the two overhead bands (``[EV_ST, EV_ST_LATE]`` = +# start-side cost, ``[EV_EN_EARLY, EV_EN]`` = end-side cost, with the inner +# ``[EV_ST_LATE, EV_EN_EARLY]`` the pure work). Zero (and ignored on readout) +# when off or for dropped events. +EV_EVENT_NO = 0 +EV_ST = 1 +EV_EN = 2 +EV_ST_LATE = 3 +EV_EN_EARLY = 4 +NUM_EV_FIELDS = 5 + +# Each warp's append counter gets its own 128-byte L2 line (16 int64) instead of +# sitting in a packed ``counters[total_warps]`` array. Reason: the warps of one +# block run on the same SM and their counter slots would otherwise share a cache +# line, so the per-event counter read-modify-write false-shares and serializes +# across warps. Padding to a line-per-counter removes that contention. Only slot +# 0 of each stride is used; the padding bytes are never read. +COUNTER_STRIDE = 16 + +# One parsed range event. ``st_late`` / ``en_early`` are 0 unless the run had +# ``show_overhead`` on; the readout treats a 0 as "no band". +Record = namedtuple( + "Record", + "block warp event_id event_no sm_id st en st_late en_early", +) + + +# --------------------------------------------------------------------------- +# Device-side API (call from inside a @cute.kernel): open a session once with +# ``prof.begin()``, then ``session.event_start(no)`` / ``session.event_end(no)``. +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class Session: + """Per-invocation profiler context (open with :meth:`ProfilerArgs.begin`). + + Holds this warp's coordinates, cached once by ``begin``. It is a + DynamicExpression (implements the flatten protocol below) so it can be + referenced inside a dynamic ``cutlass.range`` loop -- the DSL must flatten + every value live in the loop body, which a plain Python object cannot do. + Its only mutable state is the ``_open`` map, which lives purely at trace time + (no runtime cost). When profiling is disabled the coordinate fields are + ``None`` and every method folds to nothing. + """ + + prof: ProfilerArgs + warp_global: object + max_events: object + counters_base: object + events_base: object + # Compile-time map ``event_no -> open event_id`` so callers can close an event + # by its (constexpr) TYPE rather than threading the returned id. Resolved + # entirely at trace time: it parks a register ``Value``, materializes nothing, + # and adds no runtime traffic. Mutated in place (allowed on a frozen instance); + # excluded from the flatten protocol and from eq/hash. The DSL hands the loop + # body a freshly reconstructed Session (empty map), which is why same-type + # events must open and close within the same trace region -- true for this + # profiler's call pattern, and the same single-live-event-per-type limit the + # original's ``active`` table had. + _open: dict = field(default_factory=dict, compare=False) + + # --- DynamicExpression protocol: flatten to the buffer's MLIR values plus + # the four cached coordinates; ``enabled`` stays a compile-time constant. + def __extract_mlir_values__(self): + if not self.prof.enabled: + return [] + vals = list(self.prof.__extract_mlir_values__()) + for coord in (self.warp_global, self.max_events, self.counters_base, self.events_base): + vals += list(coord.__extract_mlir_values__()) + return vals + + def __new_from_mlir_values__(self, values): + if not self.prof.enabled: + return self + n = len(self.prof.__extract_mlir_values__()) + new_prof = self.prof.__new_from_mlir_values__(values[:n]) + rest = values[n:] + coords = [] + for coord in (self.warp_global, self.max_events, self.counters_base, self.events_base): + k = len(coord.__extract_mlir_values__()) + coords.append(coord.__new_from_mlir_values__(rest[:k])) + rest = rest[k:] + # Share the SAME _open dict across reconstructions: the DSL flattens + + # reconstructs the session around any region (e.g. the prologue scf.for) + # it stays live across, and a fresh default_factory dict would lose the + # open-id parked before that region -- the type-keyed close would then + # KeyError. Sharing keeps one map for the whole trace. + return Session(new_prof, *coords, _open=self._open) + + def _rec(self, event_id): + return self.events_base + (self.warp_global * self.max_events + event_id) * NUM_EV_FIELDS + + @cute.jit + def _alloc(self, event_no: cutlass.Constexpr, start_ts): + """Lane-0 helper: bump the counter, write the start record, return the id. + + Drops the record (counter still advances) once the warp is full. ``sm_id`` + is per-warp (written by begin) and ``EV_EN`` is left zeroed for the close. + """ + counter_off = self.counters_base + self.warp_global * COUNTER_STRIDE + event_id = self.prof.buffer[counter_off] + self.prof.buffer[counter_off] = event_id + 1 + if event_id < self.max_events: + rec = self._rec(event_id) + self.prof.buffer[rec + EV_EVENT_NO] = cutlass.Int64(event_no) + self.prof.buffer[rec + EV_ST] = start_ts + if cutlass.const_expr(self.prof.show_overhead): + # Latest start stamp: [EV_ST, EV_ST_LATE] = start-side overhead. + _compiler_barrier() + self.prof.buffer[rec + EV_ST_LATE] = _globaltimer_ns() + return event_id + + @cute.jit + def event_start(self, event_no: cutlass.Constexpr): + """Open a range event of type ``event_no`` (close it with the SAME type). + + No-op when disabled. Lane 0 allocates the next event_id (counter RMW, no + atomics) and writes the start record; the id is parked in the session's + compile-time ``_open`` map so :meth:`event_end` can recover it from the + type alone. The start stamp is the FIRST instruction so the interval + swallows this call's bookkeeping; under ``show_overhead`` a second stamp + (``EV_ST_LATE``) bounds the start-side overhead band. + + ``@cute.jit`` runs the AST preprocessor so the dynamic ``if`` below is + plain Python (lowered to ``scf.if``); inlined into the kernel trace. + """ + if cutlass.const_expr(self.prof.enabled): + # Lane 0 is the sole writer; other lanes keep the sentinel -1. Seeded + # before the branch because the DSL requires a value live past a + # dynamic ``if`` to exist beforehand (the branch only updates it). + event_id = cutlass.Int64(-1) + if cute.arch.lane_idx() == 0: + # Earliest start stamp: before the counter RMW and record writes. + # The barrier stops the compiler from sinking it past them. + _compiler_barrier() + event_id = self._alloc(event_no, _globaltimer_ns()) + # Park the warp-wide id (the scf.if result) under its type, at trace + # time -- nothing is emitted; event_end reads it back. + self._open[event_no] = event_id + return event_id + return None + + @cute.jit + def event_end(self, event_no: cutlass.Constexpr): + """Close the open range of type ``event_no`` (opened by :meth:`event_start`). + + Takes the event TYPE, not a runtime id: the id is recovered from the + compile-time ``_open`` map (a trace-time lookup, no global round-trip). + No-op when disabled; skips an id that overflowed the per-warp capacity. + ``EV_EN`` is the LAST stamp taken, so the interval reaches as far as it + can toward the next event; under ``show_overhead`` an earlier stamp + (``EV_EN_EARLY``) bounds the end-side band ``[EV_EN_EARLY, EV_EN]``. + """ + if cutlass.const_expr(self.prof.enabled): + event_id = self._open[event_no] + # Lane 0 (the lane that opened the event) is the sole writer. + if cute.arch.lane_idx() == 0: + if event_id < self.max_events: + rec = self._rec(event_id) + if cutlass.const_expr(self.prof.show_overhead): + # Earliest end stamp; the later EV_EN below covers this + # store, so [EV_EN_EARLY, EV_EN] = end-side overhead. + _compiler_barrier() + self.prof.buffer[rec + EV_EN_EARLY] = _globaltimer_ns() + _compiler_barrier() + self.prof.buffer[rec + EV_EN] = _globaltimer_ns() + + @cute.jit + def event_switch(self, close_no: cutlass.Constexpr, open_no: cutlass.Constexpr): + """Close ``close_no`` and open ``open_no`` sharing ONE ``%globaltimer`` read. + + For a back-to-back boundary (``event_end`` immediately followed by + ``event_start``) this is both cheaper -- one timer read instead of two -- + and more faithful: the closing ``EV_EN`` and opening ``EV_ST`` are the + SAME stamp, so the two slices abut exactly with no profiler-overhead gap. + Under ``show_overhead`` the switch's bookkeeping shows up as the opening + event's start band (the closed event gets no end band). + """ + if cutlass.const_expr(self.prof.enabled): + open_id = cutlass.Int64(-1) + if cute.arch.lane_idx() == 0: + _compiler_barrier() + boundary_ts = _globaltimer_ns() + close_id = self._open[close_no] + if close_id < self.max_events: + self.prof.buffer[self._rec(close_id) + EV_EN] = boundary_ts + open_id = self._alloc(open_no, boundary_ts) + self._open[open_no] = open_id + return open_id + return None + + +@dataclass(frozen=True) +class ProfilerArgs: + """Device-side profiler handle bundling the buffer and the enable flag. + + Passed as a SINGLE kernel argument. The jit-arg protocol below sends only + ``buffer`` across the kernel boundary as an MLIR value; ``enabled`` stays a + compile-time constant so ``const_expr(prof.enabled)`` folds every profiler + call to nothing when off, and a ``None`` buffer adds no device argument. + Open a session inside the kernel with ``prof.begin()``. + """ + + buffer: cute.Tensor + enabled: cutlass.Constexpr[bool] + # When on, each event records a second timestamp (``EV_ST_LATE`` / + # ``EV_EN_EARLY``) for the overhead bands. Compile-time constant, folds away + # when off. + show_overhead: cutlass.Constexpr[bool] = False + + def __extract_mlir_values__(self): + return [] if self.buffer is None else self.buffer.__extract_mlir_values__() + + def __new_from_mlir_values__(self, values): + if self.buffer is None: + return ProfilerArgs(self.buffer, self.enabled, self.show_overhead) + return ProfilerArgs(self.buffer.__new_from_mlir_values__(values), self.enabled, self.show_overhead) + + @cute.jit + def begin(self) -> Session: + """Open a profiling session: read the header + this warp's ids ONCE. + + Also records this warp's ``sm_id`` once (constant for the warp's + lifetime), so per-event records need not carry it. No-op when disabled. + """ + if cutlass.const_expr(self.enabled): + num_blocks = self.buffer[HDR_NUM_BLOCKS] + num_warps = self.buffer[HDR_NUM_WARPS] + max_events = self.buffer[HDR_MAX_EVENTS] + # Row-major linear ids, general for 1-/2-/3-D grids and blocks: the + # formula uses only the X/Y extents plus the Z coordinate (unused dims + # are 1 and fall out). + bidx, bidy, bidz = cute.arch.block_idx() + gdx, gdy, _ = cute.arch.grid_dim() + block_linear = (bidz * gdy + bidy) * gdx + bidx + tidx, tidy, tidz = cute.arch.thread_idx() + bdx, bdy, _ = cute.arch.block_dim() + warp = ((tidz * bdy + tidy) * bdx + tidx) // 32 + warp_global = block_linear * num_warps + warp + total_warps = num_blocks * num_warps + # Layout: header | smid[tw] | counters[tw*STRIDE] | events[tw*E*F]. + # Counters are line-padded (COUNTER_STRIDE per warp) to avoid cross-warp + # false sharing; only slot 0 of each stride is written. + smid_off = HEADER + warp_global + counters_base = HEADER + total_warps + events_base = HEADER + total_warps + total_warps * COUNTER_STRIDE + if cute.arch.lane_idx() == 0: # one sm_id write per warp, not per event + self.buffer[smid_off] = _smid() + return Session(self, warp_global, max_events, counters_base, events_base) + return Session(self, None, None, None, None) + + +# --------------------------------------------------------------------------- +# Host-side management, parsing and export +# --------------------------------------------------------------------------- +def make_cute_buffer(buffer: torch.Tensor) -> cute.Tensor: + """Wrap the 1-D int64 profiler buffer as a (static-layout) CuTe tensor.""" + return from_dlpack(buffer, assumed_align=16) + + +class SmProfiler: + """Host-side owner of the profiler buffer and its readout. + + Mirrors the original ``sm_profiler_create_buffer`` / ``register_event`` / + ``export_to_file`` API. + + Parameters + ---------- + num_blocks: + Number of CTAs in the launch grid. + num_warps: + Warps per block (the profiler's per-block "warp" count). + max_events_per_warp: + Per-warp event capacity; events beyond it are dropped. + """ + + PRINT_WIDTH = 100 + + def __init__(self, num_blocks: int, num_warps: int, max_events_per_warp: int): + self.num_blocks = int(num_blocks) + self.num_warps = int(num_warps) + self.max_events = int(max_events_per_warp) + self.total_warps = self.num_blocks * self.num_warps + self.names: dict[int, str] = {} + # header + smid[total_warps] + counters[total_warps*COUNTER_STRIDE] + + # events[total_warps*max_events*F]. Counters are line-padded to avoid + # cross-warp false sharing on the per-event RMW. + size = HEADER + self.total_warps * (1 + COUNTER_STRIDE + self.max_events * NUM_EV_FIELDS) + self.buffer = torch.zeros(size, dtype=torch.int64, device="cuda") + self._write_header() + + def _write_header(self) -> None: + self.buffer[HDR_NUM_BLOCKS] = self.num_blocks + self.buffer[HDR_NUM_WARPS] = self.num_warps + self.buffer[HDR_MAX_EVENTS] = self.max_events + + def reset(self) -> None: + self.buffer.zero_() + self._write_header() + + def register_event(self, event_no: int, name: str) -> None: + """Name an event type for the readout (host-side registry).""" + self.names[int(event_no)] = str(name) + + def cute_buffer(self) -> cute.Tensor: + """CuTe view of the buffer, for compile templates / kernel launch.""" + return make_cute_buffer(self.buffer) + + # ----- parsing ----- + def _records(self): + """Yield a :class:`Record` per written event (skips unclosed/malformed).""" + buf = self.buffer.cpu() + # Layout: header | smid[total_warps] | counters[total_warps*COUNTER_STRIDE] | events[..]. + smid = buf[HEADER : HEADER + self.total_warps].reshape(self.num_blocks, self.num_warps) + counters_end = HEADER + self.total_warps + self.total_warps * COUNTER_STRIDE + # Each warp's counter occupies COUNTER_STRIDE slots; only slot 0 is used. + counters = buf[HEADER + self.total_warps : counters_end].reshape( + self.num_blocks, self.num_warps, COUNTER_STRIDE + )[:, :, 0] + ev_base = counters_end + events = buf[ev_base:].reshape(self.num_blocks, self.num_warps, self.max_events, NUM_EV_FIELDS) + for block in range(self.num_blocks): + for warp in range(self.num_warps): + sm_id = int(smid[block, warp].item()) # per-warp, recorded once + n = min(int(counters[block, warp].item()), self.max_events) + for event_id in range(n): + rec = events[block, warp, event_id] + st, en = int(rec[EV_ST].item()), int(rec[EV_EN].item()) + if en <= 0 or en < st: # unclosed / malformed + continue + yield Record( + block=block, + warp=warp, + event_id=event_id, + event_no=int(rec[EV_EVENT_NO].item()), + sm_id=sm_id, + st=st, + en=en, + st_late=int(rec[EV_ST_LATE].item()), + en_early=int(rec[EV_EN_EARLY].item()), + ) + + # ----- summary (generic: per event type) ----- + def summarize(self, tag: str) -> None: + """Per-event-type duration stats (ns, globaltimer) over all (block, warp). + + ``en - st`` is the interval (work + instrumentation), since the stamps + are stretched to cover bookkeeping. With ``show_overhead`` recorded, the + inner ``[EV_ST_LATE, EV_EN_EARLY]`` span (work-only) is added as a column. + """ + intervals: dict[int, list] = {} + works: dict[int, list] = {} + for record in self._records(): + intervals.setdefault(record.event_no, []).append(record.en - record.st) + if record.st_late > 0 and record.en_early > 0 and record.en_early >= record.st_late: + works.setdefault(record.event_no, []).append(record.en_early - record.st_late) + have_work = bool(works) + print( + f" Profile ({tag}): {self.num_blocks} blocks x {self.num_warps} warps, units = ns (globaltimer) ".center( + self.PRINT_WIDTH, "-" + ) + ) + header = f"{'event':<14}{'min':>13}{'mean':>13}{'max':>13}{'samples':>10}" + if have_work: + header += f"{'work_mean':>13}" + print(header) + for event_no in sorted(intervals): + vals = intervals[event_no] + name = self.names.get(event_no, f"event_{event_no}") + t = torch.tensor(vals, dtype=torch.float64) + line = ( + f"{name:<14}{int(t.min().item()):>13d}{int(t.mean().item()):>13d}" + f"{int(t.max().item()):>13d}{len(vals):>10d}" + ) + if have_work: + work_vals = works.get(event_no) + if work_vals: + work_mean = int(torch.tensor(work_vals, dtype=torch.float64).mean().item()) + line += f"{work_mean:>13d}" + else: + line += f"{'--':>13}" + print(line) + + # ----- export ----- + def export_perfetto(self, path: str, show_overhead: bool = False) -> int: + """Write a Chrome-Trace JSON (open in https://ui.perfetto.dev). + + Each range event becomes one ``ph:"X"`` slice over the full ``[EV_ST, + EV_EN]`` interval; ``ts`` / ``dur`` are microseconds rebased to the + earliest start. + + One Chrome "process" per block, one "thread" per warp. ``tid`` is kept + disjoint from every ``pid`` because Perfetto's importer folds a thread + whose ``tid`` equals its ``pid`` into the process's "main thread": hence + ``pid = block + 1`` and ``tid = tid_base + block*num_warps + warp`` with + ``tid_base`` above the pid range. + + Under ``show_overhead``, each event also emits ``Δstart`` /``Δend`` child + slices over the two bands (``overhead`` category), each only when its + inner stamp is present and the band has positive width. + """ + records = list(self._records()) + base = min((record.st for record in records), default=0) + tid_base = self.num_blocks + 1 + + def pid_of(block): + return block + 1 + + def tid_of(block, warp): + return tid_base + block * self.num_warps + warp + + def us(ts_ns): + return (ts_ns - base) / 1000.0 + + sm_of = {} + for record in records: + sm_of.setdefault((record.block, record.warp), record.sm_id) # first event's SM + + meta = [] + for block in range(self.num_blocks): + pid = pid_of(block) + sm_id = sm_of.get((block, 0), -1) + meta.append( + {"name": "process_name", "ph": "M", "pid": pid, "args": {"name": f"block {block} (sm {sm_id})"}} + ) + meta.append({"name": "process_sort_index", "ph": "M", "pid": pid, "args": {"sort_index": block}}) + for warp in range(self.num_warps): + tid = tid_of(block, warp) + meta.append( + {"name": "thread_name", "ph": "M", "pid": pid, "tid": tid, "args": {"name": f"warp {warp}"}} + ) + meta.append( + {"name": "thread_sort_index", "ph": "M", "pid": pid, "tid": tid, "args": {"sort_index": warp}} + ) + + slices = [] + for record in records: + name = self.names.get(record.event_no, f"event_{record.event_no}") + pid, tid = pid_of(record.block), tid_of(record.block, record.warp) + slices.append( + { + "name": name, + "cat": name, + "ph": "X", + "ts": us(record.st), + "dur": (record.en - record.st) / 1000.0, + "pid": pid, + "tid": tid, + "args": {"sm_id": record.sm_id}, + } + ) + if not show_overhead: + continue + # Start-side band [EV_ST, EV_ST_LATE]: profiler bookkeeping cost. + if record.st_late > record.st: + slices.append( + { + "name": f"{name} Δstart", + "cat": "overhead", + "ph": "X", + "ts": us(record.st), + "dur": (record.st_late - record.st) / 1000.0, + "pid": pid, + "tid": tid, + "args": {"sm_id": record.sm_id}, + } + ) + # End-side band [EV_EN_EARLY, EV_EN]: trailing store cost. + if 0 < record.en_early < record.en: + slices.append( + { + "name": f"{name} Δend", + "cat": "overhead", + "ph": "X", + "ts": us(record.en_early), + "dur": (record.en - record.en_early) / 1000.0, + "pid": pid, + "tid": tid, + "args": {"sm_id": record.sm_id}, + } + ) + trace = { + "traceEvents": meta + slices, + "displayTimeUnit": "ns", + "metadata": {"unit": "globaltimer ns -> us, grid-relative"}, + } + with open(path, "w") as handle: + json.dump(trace, handle) + return len(slices) diff --git a/09-pipelining/profile.sh b/09-pipelining/profile.sh index 00d9444..e7e13ab 100644 --- a/09-pipelining/profile.sh +++ b/09-pipelining/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file pipelining.csv --metrics gpu__time_duration.sum --kernel-name "pipelining" python pipelining.py # ncu-rep output -ncu -o ncu_prof_9 --import-source 1 --set full --kernel-name "pipelining" -f python pipelining.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_9 --import-source 1 --set full --kernel-name "pipelining" -f python pipelining.py +# CuTe DSL impl: +ncu -o ncu_prof_9 --import-source 1 --set full --kernel-name "regex:.*pipelining.*" -f python cutedsl_pipelining.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_9 python pipelining.py diff --git a/11-tma-load-store/cutedsl_tma_load_store.py b/11-tma-load-store/cutedsl_tma_load_store.py index 0de98a8..55bd7af 100644 --- a/11-tma-load-store/cutedsl_tma_load_store.py +++ b/11-tma-load-store/cutedsl_tma_load_store.py @@ -569,7 +569,7 @@ def _compile_pair(a_template, b_template, c_template, d_template, acc_dtype, out acc_dtype, out_dtype, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( tma_load_store, @@ -581,7 +581,7 @@ def _compile_pair(a_template, b_template, c_template, d_template, acc_dtype, out acc_dtype, out_dtype, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/11-tma-load-store/profile.sh b/11-tma-load-store/profile.sh index af3c072..975e55b 100644 --- a/11-tma-load-store/profile.sh +++ b/11-tma-load-store/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file tma_load_store.csv --metrics gpu__time_duration.sum --kernel-name "tma_load_store" python tma_load_store.py # ncu-rep output -ncu -o ncu_prof_11 --import-source 1 --set full --kernel-name "tma_load_store" -f python tma_load_store.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_11 --import-source 1 --set full --kernel-name "tma_load_store" -f python tma_load_store.py +# CuTe DSL impl: +ncu -o ncu_prof_11 --import-source 1 --set full --kernel-name "regex:.*tma_load_store.*" -f python cutedsl_tma_load_store.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_11 python tma_load_store.py diff --git a/12-tma-multicast-reduce/cutedsl_tma_multicast_reduce.py b/12-tma-multicast-reduce/cutedsl_tma_multicast_reduce.py index c3f2986..1c0c287 100644 --- a/12-tma-multicast-reduce/cutedsl_tma_multicast_reduce.py +++ b/12-tma-multicast-reduce/cutedsl_tma_multicast_reduce.py @@ -746,7 +746,7 @@ def _compile_variants(a_template, b_template, c_template, d_template, acc_dtype, out_dtype, True, use_reduce_add, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( tma_multicast_reduce, @@ -759,7 +759,7 @@ def _compile_variants(a_template, b_template, c_template, d_template, acc_dtype, out_dtype, False, use_reduce_add, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/12-tma-multicast-reduce/profile.sh b/12-tma-multicast-reduce/profile.sh index f381ebd..12320b3 100644 --- a/12-tma-multicast-reduce/profile.sh +++ b/12-tma-multicast-reduce/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file tma_multicast_reduce.csv --metrics gpu__time_duration.sum --kernel-name "tma_multicast_reduce" python tma_multicast_reduce.py # ncu-rep output -ncu -o ncu_prof_12 --import-source 1 --set full --kernel-name "tma_multicast_reduce" -f python tma_multicast_reduce.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_12 --import-source 1 --set full --kernel-name "tma_multicast_reduce" -f python tma_multicast_reduce.py +# CuTe DSL impl: +ncu -o ncu_prof_12 --import-source 1 --set full --kernel-name "regex:.*tma_multicast_reduce.*" -f python cutedsl_tma_multicast_reduce.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_12 python tma_multicast_reduce.py diff --git a/13-warpgroup-mma/cutedsl_warpgroup_mma.py b/13-warpgroup-mma/cutedsl_warpgroup_mma.py index be1ccee..7b9e5ae 100644 --- a/13-warpgroup-mma/cutedsl_warpgroup_mma.py +++ b/13-warpgroup-mma/cutedsl_warpgroup_mma.py @@ -591,7 +591,7 @@ def _compile_pair(a_template, b_template, c_template, d_template, acc_dtype, out acc_dtype, out_dtype, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( warpgroup_mma_host, @@ -603,7 +603,7 @@ def _compile_pair(a_template, b_template, c_template, d_template, acc_dtype, out acc_dtype, out_dtype, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/13-warpgroup-mma/profile.sh b/13-warpgroup-mma/profile.sh index f7887d3..d5687cb 100644 --- a/13-warpgroup-mma/profile.sh +++ b/13-warpgroup-mma/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file warpgroup_mma.csv --metrics gpu__time_duration.sum --kernel-name "warpgroup_mma" python warpgroup_mma.py # ncu-rep output -ncu -o ncu_prof_13 --import-source 1 --set full --kernel-name "warpgroup_mma" -f python warpgroup_mma.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_13 --import-source 1 --set full --kernel-name "warpgroup_mma" -f python warpgroup_mma.py +# CuTe DSL impl: +ncu -o ncu_prof_13 --import-source 1 --set full --kernel-name "regex:.*warpgroup_mma.*" -f python cutedsl_warpgroup_mma.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_13 python warpgroup_mma.py diff --git a/14-warp-specialization/cutedsl_warp_specialization.py b/14-warp-specialization/cutedsl_warp_specialization.py index 579945f..20e0948 100644 --- a/14-warp-specialization/cutedsl_warp_specialization.py +++ b/14-warp-specialization/cutedsl_warp_specialization.py @@ -632,7 +632,7 @@ def _compile_pair(a_template, b_template, c_template, d_template, acc_dtype, out acc_dtype, out_dtype, True, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) g_accum = cute.compile( warp_specialization_host, @@ -644,7 +644,7 @@ def _compile_pair(a_template, b_template, c_template, d_template, acc_dtype, out acc_dtype, out_dtype, False, - options="--enable-tvm-ffi", + options="--enable-tvm-ffi --generate-line-info", ) return g_clear, g_accum diff --git a/14-warp-specialization/profile.sh b/14-warp-specialization/profile.sh index d35e851..a503505 100644 --- a/14-warp-specialization/profile.sh +++ b/14-warp-specialization/profile.sh @@ -2,7 +2,10 @@ # ncu --csv --log-file warp_specialization.csv --metrics gpu__time_duration.sum --kernel-name "warp_specialization" python warp_specialization.py # ncu-rep output -ncu -o ncu_prof_14 --import-source 1 --set full --kernel-name "warp_specialization" -f python warp_specialization.py +# C++ / reference impl (was the default; commented out): +# ncu -o ncu_prof_14 --import-source 1 --set full --kernel-name "warp_specialization" -f python warp_specialization.py +# CuTe DSL impl: +ncu -o ncu_prof_14 --import-source 1 --set full --kernel-name "regex:.*warp_specialization.*" -f python cutedsl_warp_specialization.py # nsys-rep output # nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas,cublas-verbose,mpi,ucx,oshmem,python-gil,syscall --backtrace=dwarf --output=nsys_prof_14 python warp_specialization.py