Repository navigation
feat(09): gate register prefetch via REG_STAGES, fit two CTAs per SM - #21
Merged
Merged
Conversation
- 07: remove the explicit ``acc_dtype`` parameter from the kernel signature; the same value is reachable via ``tCrC.element_type`` after the fragment is allocated, matching the convention already used in 06/08/09. - 08 / 09: replace the universal S2R atoms with explicit ``LdMatrix8x8x16bOp(False, 4)`` for A/B and ``LdMatrix8x8x16bOp(False, 2)`` for C when C is 16-bit (else ``CopyUniversalOp``). This mirrors 07's choice and the SM75_U32x4_LDSM_N / SM75_U32x2_LDSM_N selection in the C++ kernels, surfacing the ldmatrix path at the source level instead of relying on the DSL backend to lower a universal atom. - 08 / 09: keep the existing two-tensor predicate structure (``tApA`` M-only + ``tApA_first`` per-element M+K) rather than merging into one tensor; merging requires an extra refill between the residue tile and the mainloop that introduces a WAW dependency on the same predicate registers and blocks scheduling.
Replace the speculative comment on ``Copy_G2S_op`` selection with an evidence-based explanation derived from PTX inspection of the failing configuration. The previous comment claimed AutoVec lowered to a specific cp.async variant and blamed alignment; the actual mechanism is that ``AutoVectorizingCopyWithAssumedAlignment<128>`` is a 128-bit atom applied to a per-element-strided source, which under ``copy_if`` (with no recast specialization, unlike ``copy()``) unrolls to 8 wide loads at 2-byte offsets — 7 of which are misaligned and fault. The comment now points at NVIDIA/cutlass#2354 for the upstream tracker and references ``cutlass/examples/cute/tutorial/tiled_copy_if.cu`` as the official workaround pattern.
For blocks that lie strictly inside both M and N tiles AND K is tile-aligned, the per-element predicate fill + ``copy_if`` on the residue prologue tile does redundant work — every M-bound and K-bound check evaluates to true. Add a runtime check ``is_interior`` and branch to a plain ``copy()`` for the prologue's stage 0 in that case, falling back to the existing predicated path for boundary blocks and K-residue shapes. The fast path is intentionally limited to stage 0: later prologue stages and the mainloop already carry a per-iter overshoot guard (``if (ik == NTilesK) clear(tApA_g2s)``) which complicates short- circuiting and offers little benefit on those code paths.
The single-stage dynamic-mma kernel has no cp.async pipeline to hide memory latency, so it depends on occupancy for latency hiding. Left unconstrained, the CuTeDSL register allocator maximises per-thread ILP and lets the kernel grow to ~150 registers/thread (not a spill, just a loose allocation), which caps it at one resident block per SM on H100/H200. Pass ``min_blocks_per_mp=2`` to the kernel launch (lowered to the NVVM ``minctasm`` launch-bound attribute). This tightens the per-thread register budget so two blocks fit per SM; the kernel lands at 128 registers/thread with zero spill, and the higher occupancy gives the warp scheduler enough eligible warps to cover the serial load/compute latency. 2 is the largest spill-free setting here -- a higher value would force register spills to local memory. 09-pipelining is intentionally left untouched: its multi-stage cp.async pipeline already hides latency within each block, and forcing a second resident block there only adds issue-slot contention.
The mainloop cleared the A/B shared-memory tiles (.fill(0) / clear()) on every K-tile iteration so that predicated-off cp.async slots read as zero. That clear only needs to run once per block: the M-boundary rows masked off by the M-only predicate are never written by any K-tile, and the ik=0 K-residue columns are overwritten by every later full K-tile. Hoisting it out of the loop drops 63 of 64 redundant full-tile clears, trimming the per-iteration instruction count in the hot loop. Applied to both the CuTeDSL and C++ kernels. 09-pipelining already clears its smem once in the prologue, so it is unaffected.
The CuTeDSL test harnesses built the C (addend) tensor at the input dtype -- fp16 or bf16 -- for the fp32-accumulator dtype specs. The C++ bindings these tests mirror fix ComputeTypeC to float for the same specs, so the CuTeDSL tests were compiling and validating a different kernel configuration (16-bit C) than the C++ side (fp32 C). Build the C template and per-shape C tensors as fp32 so the CuTeDSL tests exercise the same operation as C++. Also reseed the RNG before each per-spec sweep, matching the C++ drivers' per-block reseed, and update the sweep banners accordingly.
Build the C addend as fp32 for the fp32-accumulator specs of the warpgroup-MMA and warp-specialization tests, matching the C++ ComputeTypeC=float bindings, as already done for the other examples. That correction exposed a latent kernel bug: both kernels aliased the epilogue sC/sD buffers onto the sB staged buffer alone, which only has room for a 16-bit C tile. An fp32 C tile overflows sB and triggers an illegal memory access. Re-base the alias from sB to sA -- sA and sB are contiguous in SharedStorage, so the alias now spans the combined sA+sB region, which holds the fp32 C tile. No extra shared memory is used. Also reseed the RNG before each per-spec sweep, matching the C++ drivers' per-block reseed.
The __syncthreads() / cute.arch.sync_threads() that sat between the A/B smem fill(0)/clear() and the subsequent predicated cp.async is unnecessary. Each thread's fill and its following cp.async target the same g2s partition slots, so single-thread program order already orders the two writes to the same addresses. Cross-thread visibility for the s2r readers is provided by the existing post-cp_async_wait __syncthreads(), which also publishes the zero-fill stores from every other thread. Affected sites: - 08-dynamic-mma/cutedsl_dynamic_mma.py (pre-K-loop A/B clear) - 08-dynamic-mma/dynamic_mma.cu (pre-K-loop A/B clear) - 09-pipelining/cutedsl_pipelining.py (stage-0 A/B clear) 09-pipelining/pipelining.cu and pipelining_no_reg_prefetch.cu already omitted this sync.
09-pipelining: - Add a REG_STAGES constant gating the S2R register pipeline. With REG_STAGES >= 2 the smem -> rmem load for k_block + 1 is hoisted ahead of the k_block MMA (double-buffered). REG_STAGES == 1 disables register prefetch: each k-tile loads its whole A/B fragment, then runs every MMA, leaving only the GMEM -> SMEM -> RF pipeline to overlap work. - Fold the C addend into the accumulator in the epilogue (acc starts at zero, acc += C after the matmul) instead of preloading it, so sC is never live during the mainloop. The epilogue then aliases sC and sO over the drained A/B pipeline smem, dropping the footprint to A_pipe + B_pipe. - Set min_blocks_per_mp=2 so two CTAs co-reside per SM at <=128 reg/thread. 08-dynamic-mma: - Reuse the precomputed tCsC partition for the smem clear.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
09-pipelining
REG_STAGESconstant gating the S2R register pipeline.REG_STAGES >= 2hoists the smem→rmem load fork_block + 1ahead of thek_blockMMA (double-buffered register pipeline);REG_STAGES == 1disables register prefetch so each k-tile loads its whole A/B fragment, then runs every MMA, leaving only the GMEM→SMEM→RF pipeline to overlap work.accstarts at zero,acc += Cafter the matmul) instead of preloading it, sosCis never live during the mainloop.sCandsOover the drained A/B pipeline smem, dropping the footprint toA_pipe + B_pipe.min_blocks_per_mp=2so two CTAs co-reside per SM at ≤128 reg/thread.08-dynamic-mma
tCsCpartition for the smem clear.