Repository navigation
refactor(08-09): CuTeDSL cleanup, fast paths, and smem-aliasing fix - #20
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.
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.
Summary
Cleanup pass over examples 08 (
dynamic-mma) and 09 (pipelining), with adjacent touch-ups in 07 / 11–14 where the same patterns recurred. Covers both the C++ and the CuTeDSL ports.Correctness
sC/sDsmem aliasing under swizzle. The previous.inner/.outer+recast_ptrsplit corrupted the upper-half columns; aliased smem now usesmake_tensor(other.iterator, full_layout).Perf
cute.arch.launch_boundsto lift occupancy.__syncthreads()/cute.arch.sync_threads()that sat between the smem fill(0)/clear() and the subsequent predicated cp.async. Same-thread program order plus the existing post-cp_async_wait barrier suffice.Style / docs
acc_dtypeplumbing.AutoVectorizingCopyfault undercopy_if([BUG] Missing copy_if Implementation for AutoVectorizingCopyWithAssumedAlignment NVIDIA/cutlass#2354).Affected files