Cutlass or Ptx Method To Copy Data From Tensor-Memory Of SM Unit To Registers Efficiently

Accelerated Computing CUDA CUDA Programming and Performance
1

When converting 32bit fp accumulator data from tensor memory to 16bit fp on registers using cute::copy, compiler always generates full 32bit fp accumulator in registers before converting and copying to target registers. This creates unnecessary register spills (256 more than required).

Tensor memory: 256 kB float

Register destination: 128kB bf16

cute::copy: creates extra 256 kB float in local memory (spill)

I tried:

  • re-using a small chunk to convey the data,
  • separating loop iterations with memory clobbering,
  • scoping loop iteration algorithm inside nanosleep based gating,
  • synchronizing by tcgen05.wait::ld,
  • replacing the ptx instruction that used registers as [=r] with [+r]
  • volatile chunk definition

but it always ended up copying full accumulator tile. I wanted to copy in chunks so that register requirement wouldn’t go beyond 150 maybe.

What I didn’t try yet:

  • using a non-inlined function to copy with small chunk, then re-call function multiple times
  • hacking ptx or sass to dynamically change the instructions somehow, so that compiler can not know how many of them will be called (and can not allocate that many registers): probably impossible because how gpu works.
  • runtime-indexing in the loop (pragma unroll 1): this would just make register tile spill automatically even before tensor memory accessed
  • creating non-inlined function and gating it behind time-query to make it unknown number of calls in compile time (at cost of nanoseconds or microsecond)

Chunk size for copy: 16 registers, but compiler clones this many times for each loop iteration so it doesn’t work.

Algorithm: get 16 floats from tmem, convert to 16 bf16, write to register tile, repeat until full tile is copied.

Environment: nvrtc, cuda 13, C++20, cutlass 4.5.2, ubuntu 24.04 LTS, B200 GPU

2

Could you share a very small example with C/C++, PTX and SASS code?

3

The cuh file to run:

#pragma once
// =====================================================================
// tmem_drain_repro.cuh  --  standalone SM100 (B200) repro:
//   TMEM (tcgen05) fp32 accumulator  ->  bf16 register tile drain.
//
// THE ISSUE (CHUNKED=0, "naive"): draining a (TILE_M x TILE_N) fp32
// accumulator from Tensor Memory and converting to bf16 with one
// whole-tile cute::copy materializes the FULL per-thread fp32 fragment
// in registers before any conversion happens. At TILE_M=256 that is
// 512 fp32/thread + the 256-element bf16 destination tile -> far past
// the 255-register cap -> ptxas spills to local memory (see -Xptxas -v
// "spill stores/loads" and STL/LDL in the SASS).
//
// THE WORKAROUND (CHUNKED=1): slice the tmem-side partition into
// 16-value chunks (one 32dp32b16x LDTM each), reuse ONE 16-fp32 staging
// fragment, and convert into the bf16 tile IMMEDIATELY after each chunk
// load, before the next chunk's load is issued. ptxas then keeps the
// staging fragment in the same 16 registers for every iteration
// (load -> convert -> load -> convert).
//
// The bf16 tile is kept alive after the drain by a few rounds of
// arbitrary math + a store, so nothing is dead-code-eliminated.
//
// Build variants (see build_tmem_repro.sh):
//   TILE_M  = 128 | 256      (accumulator rows; N fixed at 256)
//   CHUNKED = 0   | 1        (naive whole-tile copy | chunked copy+convert)
//
// Environment: CUDA 13, CUTLASS 4.x includes, sm_100a. Compile-only
// repro (no launch needed to inspect PTX/SASS). Same behavior under
// NVRTC and nvcc (ptxas is the common backend).
// =====================================================================
#include <cute/tensor.hpp>
#include <cute/atom/mma_traits_sm100.hpp>   // SM100 UMMA atoms + tmem fragments
#include <cute/atom/copy_traits_sm100.hpp>  // SM100_TMEM_LOAD_* + make_tmem_copy
#include <cute/arch/tmem_allocator_sm100.hpp>
#include <cutlass/numeric_types.h>
#include <cuda_bf16.h>

using namespace cute;

#ifndef TILE_M
#define TILE_M 256
#endif
#ifndef CHUNKED
#define CHUNKED 0
#endif
constexpr int TILE_N_R = 256;

extern "C" __global__ void __launch_bounds__(128, 1)
tmem_drain_repro(__nv_bfloat16* __restrict__ out, __nv_bfloat162 mulk, __nv_bfloat162 addk) {
    // ---- TMEM allocation (full 512 columns), one warp allocates ----
    __shared__ uint32_t tmem_base;
    cute::TMEM::Allocator1Sm ta{};
    if ((threadIdx.x >> 5) == 0u) ta.allocate(cute::TMEM::Allocator1Sm::Sm100TmemCapacityColumns, &tmem_base);
    __syncthreads();
    if ((threadIdx.x >> 5) == 0u) ta.release_allocation_lock();
    __syncthreads();

    // ---- a (TILE_M x TILE_N) fp32 accumulator in TMEM, shaped by the 1SM UMMA atom (M=128, N=256).
    //      TILE_M=256 -> MMA_M=2 reps. No MMA is ever issued; only the fragment SHAPES matter here. ----
    auto tiled_mma = make_tiled_mma(SM100_MMA_F16BF16_SS<cutlass::bfloat16_t, cutlass::bfloat16_t, float,
                                                         128, TILE_N_R, UMMA::Major::K, UMMA::Major::K>{});
    auto cta_mma = tiled_mma.get_slice(0);
    Tensor gC   = make_tensor(make_gmem_ptr(reinterpret_cast<cutlass::bfloat16_t*>(out)),
                              make_layout(make_shape(Int<TILE_M>{}, Int<TILE_N_R>{}),
                                          make_stride(Int<TILE_N_R>{}, Int<1>{})));
    Tensor tCgC = cta_mma.partition_C(gC);            // ((M,N), MMA_M, MMA_N) gmem view (also the final store target)
    Tensor tAcc = cta_mma.make_fragment_C(tCgC);      // fp32 TMEM accumulator tensor
    tAcc.data() = tmem_base;

    // ---- TMEM->RMEM tiled copy: 32dp32b16x atom = 16 fp32 per thread per LDTM ----
    TiledCopy t2r = make_tmem_copy(SM100_TMEM_LOAD_32dp32b16x{}, tAcc);
    ThrCopy   thr = t2r.get_slice((int)threadIdx.x);
    Tensor tS = thr.partition_S(tAcc);                // TMEM source   (V, iter modes...)
    Tensor tD = thr.partition_D(tCgC);                // gmem reference (V, iter modes...)

    constexpr int R = decltype(rank(tD))::value;
    auto tSg = group_modes<1, R>(tS);                 // (V, (all chunks))
    auto tDg = group_modes<1, R>(tD);
    constexpr int NCH = decltype(size<1>(tDg))::value;   // 16 chunks at TILE_M=128, 32 at TILE_M=256

    // per-thread bf16 destination tile: ((V)=16 bf16, NCH) -> 128/256 bf16 = 64/128 registers
    Tensor held = make_tensor<cutlass::bfloat16_t>(make_shape(shape(tDg(_, 0)), Int<NCH>{}));

#if CHUNKED == 0
    // ================= NAIVE (the problem): whole-tile copy, then convert =================
    Tensor rAll = make_tensor<float>(shape(tDg));     // FULL per-thread fp32 fragment (256/512 fp32!)
    copy(t2r, tSg, rAll);                             // all LDTMs land before any conversion
    {
        auto f2 = recast<float2>(rAll);
        auto b2 = recast<__nv_bfloat162>(held);
        CUTE_UNROLL
        for (int e = 0; e < (int)size(b2); ++e) b2(e) = __float22bfloat162_rn(f2(e));
    }
#else
    // ================= CHUNKED (the fix): one chunk load -> convert immediately =================
    Tensor rC = make_tensor<float>(shape(tDg(_, 0)));  // 16 fp32 staging, REUSED by every chunk
    CUTE_UNROLL
    for (int c = 0; c < NCH; ++c) {
        copy(t2r, tSg(_, c), rC);                      // one LDTM.x16
        auto f2 = recast<float2>(rC);
        auto b2 = recast<__nv_bfloat162>(held(_, c));
        CUTE_UNROLL
        for (int e = 0; e < (int)size(f2); ++e) b2(e) = __float22bfloat162_rn(f2(e));   // consume BEFORE the next chunk's load
    }
#endif

    // ---- keep the bf16 register tile ALIVE: a few rounds of arbitrary math ----
    {
        auto b2 = recast<__nv_bfloat162>(held);
        CUTE_UNROLL
        for (int it = 0; it < 5; ++it) {
            CUTE_UNROLL
            for (int e = 0; e < (int)size(b2); ++e) b2(e) = __hfma2(b2(e), mulk, addk);
        }
    }

    // ---- store the tile (prevents DCE; plain elementwise STG) ----
    CUTE_UNROLL
    for (int c = 0; c < NCH; ++c) copy(held(_, c), tDg(_, c));

    // ---- release TMEM ----
    __syncthreads();
    if ((threadIdx.x >> 5) == 0u) ta.free(tmem_base, cute::TMEM::Allocator1Sm::Sm100TmemCapacityColumns);
}

The script that iterates 4 scenarios to compare & produce sass/ptx:

#!/usr/bin/env bash
# =====================================================================
# build_tmem_repro.sh -- compile tmem_drain_repro.cuh in 4 variants and
# dump PTX + SASS + register/spill usage. Compile-only (no GPU needed).
#   usage:  CUTLASS=/path/to/cutlass ./build_tmem_repro.sh
# Outputs (per variant m{128|256}_{naive|chunked}):
#   repro_<tag>.ptx / .sass / .resusage / .buildlog   + a summary table.
# =====================================================================
set -e
cd "$(dirname "$0")"
CUTLASS="${CUTLASS:-$HOME/Desktop/cutlass}"
ARCH=sm_100a
FLAGS="-std=c++20 -O3 --expt-relaxed-constexpr -I$CUTLASS/include -I$CUTLASS/tools/util/include"

echo "== building against CUTLASS at: $CUTLASS =="
for TM in 128 256; do
  for CH in 0 1; do
    tag="m${TM}_$([ "$CH" -eq 1 ] && echo chunked || echo naive)"
    echo "--- $tag ---"
    nvcc -arch=$ARCH $FLAGS -DTILE_M=$TM -DCHUNKED=$CH -x cu tmem_drain_repro.cuh \
         -cubin -o repro_$tag.cubin -Xptxas -v 2> repro_$tag.buildlog || { cat repro_$tag.buildlog; exit 1; }
    nvcc -arch=compute_100a $FLAGS -DTILE_M=$TM -DCHUNKED=$CH -x cu tmem_drain_repro.cuh \
         -ptx -o repro_$tag.ptx
    cuobjdump -sass       repro_$tag.cubin > repro_$tag.sass
    cuobjdump -res-usage  repro_$tag.cubin > repro_$tag.resusage
  done
done

echo
echo "===================== SUMMARY ====================="
for f in repro_*.buildlog; do
  tag=${f#repro_}; tag=${tag%.buildlog}
  regs=$(grep -oE "Used [0-9]+ registers" "$f" | head -1)
  spill=$(grep -oE "[0-9]+ bytes spill stores, [0-9]+ bytes spill loads" "$f" | head -1)
  stack=$(grep -oE "[0-9]+ bytes stack frame" "$f" | head -1)
  sass=repro_$tag.sass
  stl=$(grep -c " STL" "$sass" || true); ldl=$(grep -c " LDL" "$sass" || true)
  ldtm=$(grep -c "LDTM" "$sass" || true)
  printf "%-14s %-20s %-14s spills:[%s]  SASS: STL=%s LDL=%s LDTM=%s\n" \
         "$tag" "${regs:-?}" "${stack:-0 stack}" "${spill:-none}" "$stl" "$ldl" "$ldtm"
done
echo
echo "The tmem->rmem loop: grep -n -A12 'LDTM' repro_m256_chunked.sass | head -80"

The result:

== building against CUTLASS at: /home/tugrul-b200/Desktop/cutlass ==
--- m128_naive ---
tmem_drain_repro.cuh:1:9: warning: #pragma once in main file
    1 | #pragma once
      |         ^~~~
--- m128_chunked ---
tmem_drain_repro.cuh:1:9: warning: #pragma once in main file
    1 | #pragma once
      |         ^~~~
--- m256_naive ---
tmem_drain_repro.cuh:1:9: warning: #pragma once in main file
    1 | #pragma once
      |         ^~~~
--- m256_chunked ---
tmem_drain_repro.cuh:1:9: warning: #pragma once in main file
    1 | #pragma once
      |         ^~~~

===================== SUMMARY =====================
m128_chunked   Used 148 registers   0 bytes stack frame spills:[0 bytes spill stores, 0 bytes spill loads]  SASS: STL=0 LDL=0 LDTM=16
m128_naive     Used 255 registers   64 bytes stack frame spills:[64 bytes spill stores, 64 bytes spill loads]  SASS: STL=8 LDL=16 LDTM=16
m256_chunked   Used 255 registers   72 bytes stack frame spills:[68 bytes spill stores, 68 bytes spill loads]  SASS: STL=17 LDL=17 LDTM=32
m256_naive     Used 255 registers   2280 bytes stack frame spills:[2828 bytes spill stores, 2828 bytes spill loads]  SASS: STL=571 LDL=707 LDTM=32

The tmem->rmem loop: grep -n -A12 'LDTM' repro_m256_chunked.sass | head -80

So, despite the chunking, 256x256 version spills registers (and this is without any metadata register usage is in a real kernel with more logic than this, I had more spill really). It’s even worse register pressure than naive conversion of 128x256 fp32 to 128x256 bf16 (all at once, maximum spill).

Ekran görüntüsü 2026-07-02 123424
Ekran görüntüsü 2026-07-02 123424781×657 36 KB

I could not add the generated sass/ptx as forum api doesn’t allow them to be attached as files here.

Naive-m256:
/*0580*/  LOP3.LUT P0, RZ, R0, 0x3, R17, 0x48, !PT ;
/*0590*/  LDTM.x16 R20, tmem[UR4] ;
/*05a0*/  STL.64 [R1+0x400], R20 ;
/*05b0*/  STL.64 [R1+0x408], R22 ;
/*05c0*/  STL.64 [R1+0x410], R24 ;
/*05d0*/  STL.64 [R1+0x418], R26 ;
/*05e0*/  STL.64 [R1+0x420], R28 ;
          ... (tmem -> registers -> local memory -> reloaded later for F2FP)

Chunked-m256:
/*2690*/  F2FP.BF16.F32.PACK_AB R5, R5, R4 ;
/*26a0*/  F2FP.BF16.F32.PACK_AB R4, R7, R6 ;
/*26b0*/  F2FP.BF16.F32.PACK_AB R3, R13, R12 ;
/*26c0*/  STL [R1+0x4], R5 ;                    <-- converted bf16x2 word spilled
/*26d0*/  F2FP.BF16.F32.PACK_AB R0, R15, R14 ;
/*26e0*/  F2FP.BF16.F32.PACK_AB R252, R9, R8 ;
/*26f0*/  STL [R1], R4 ;                        <-- spilled
/*2700*/  F2FP.BF16.F32.PACK_AB R251, R11, R10 ;
/*2710*/  F2FP.BF16.F32.PACK_AB R250, R17, R16 ;
/*2720*/  STL [R1+0xc], R3 ;                    <-- spilled
/*2730*/  F2FP.BF16.F32.PACK_AB R249, R19, R18 ;
/*2740*/  STL [R1+0x8], R0 ;                    <-- spilled
          ... (17 STL total, local slots [R1+0x0 .. R1+0x40] = 68 bytes)

In real kernel, it clones chunk registers many times, with higher total spill, defeating purpose of chunking.

Real kernel:

// source: one reused 16-fp32 staging fragment per chunk.
// but, it's assigning a fresh 16-register window to each:
/*74c0*/  LDTM.x16 R232, tmem[UR4+0x10] ;
/*7650*/  LDTM.x16 R216, tmem[UR4+0x20] ;
/*77e0*/  LDTM.x16 R200, tmem[UR4+0x30] ;
/*7970*/  LDTM.x16 R184, tmem[UR4+0x40] ;
/*7b00*/  LDTM.x16 R168, tmem[UR4+0x50] ;
/*7c90*/  LDTM.x16 R152, tmem[UR4+0x60] ;
/*7e20*/  LDTM.x16 R136, tmem[UR4+0x70] ;
/*7fb0*/  LDTM.x16 R120, tmem[UR4+0x80] ;
/*8140*/  LDTM.x16 R104, tmem[UR4+0x90] ;
/*82d0*/  LDTM.x16 R88,  tmem[UR4+0xa0] ;
/*84a0*/  LDTM.x16 R72,  tmem[UR4+0xb0] ;
/*85f0*/  LDTM.x16 R56,  tmem[UR4+0xc0] ;   // ... register file exhausted along the way:
/*72c0*/  STL.64 [R1],     R24 ;             // some chunks spilled straight after their load
/*72d0*/  STL.64 [R1+0x8], R26 ;

// the conversions only appear after the whole load ladder (~1.5 KB later):
/*8ca0*/  F2FP.BF16.F32.PACK_AB R232, R233, R232 ;
/*8cb0*/  F2FP.BF16.F32.PACK_AB R233, R235, R234 ;
/*8cc0*/  F2FP.BF16.F32.PACK_AB R234, R237, R236 ;

Despite different ways to separate LDTM from F2FP, it always groups all LDTM together with different registers and F2FP together again with different registers.

Is there a way to stop ptxas from batching the LDTMs (a scheduling barrier that survives optimization), or better, a drain path that converts on load — e.g. a tcgen05.ld variant or CUTLASS copy atom emitting packed 16-bit values — so the fp32 staging never occupies registers? At 256×256 even perfect chunking can’t help, since the bf16 destination tile alone is 128 registers/thread.

4

Can you reuse registers in ptx (inline assembly) code? Perhaps it transfers to SASS?
(normally ptx has a virtual machine with infinite registers, so by default new registers are used, but one can manually reuse)

Probably ptxas wants to hide the latency by giving the memory access more time.

5

I tried that with and without [+r] and it didnt affect sass, produced same thing (cloned chunks).

6

Did you try for a single conversion per inline or all conversions?

7

Conversion was always in same loop so it should be able to serialize them but its not doing this.

8

I meant same inline asm, not same loop. And then reuse the registers for each. Just an attempt.

Or try to forbid unrolling the loop.

9

You mean 1 asm block containing both tmem-rmem movement and conversion? I don’t know how to move correctly using asm. cute::copy knows how to. When I copy that code from ptx, it doesn’t work.

10

yes, that was one idea. And putting several iterations into the inline asm block, but reusing the register names.

The other idea is to deactivate unrolling for the respective loop, either with CUTE_NO_UNROLL or #pragma unroll 1

(both ideas are kind of the opposite of each other, but could both work)

11

I will try the asm block when I learn more about cutlass / ptx. Thank you for help. The other I already tried (no unroll + comparison to compile-time constants to select compile-time register array index to avoid spill) it didnt stop ptxas from batching all.

My work involves keeping track of multiple accumulators, and there’s little space in registers already, then copying makes it worse until I reduce tile size considerably (that causes 5-10% lower tensor core utilization).

12

It is a pity that there is not more control.

You can introduce an artifical dependency.

Like a dynamically created false boolean, which conditionally executes instructions that would change registers (like memory indices or data).

E.g. a dependency between LDTM and F2FP instructions (the result of which can be reinterpreted as integer) to enforce an order.

13

Maybe role-switching warps can do this. Run some chunks on 1 warp, other chunks on other warp. Needs more warpgroups, limits usage of other roles, requires splitting branches with less registers on one branch, but better than nothing.