cant make a vectorize load register to global memory

[N1GHTR0]

I’m trying to figure out why I can’t get a proper 128-bit vectorized store from registers to global memory using CuTe. When I load from global memory into registers, the compiler emits a nice LDG.E.128, like:

LDG.E.128 R8, [R8.64]
So far so good.

But when I try to store back to global memory from registers, even though I’m using a float4 (and the data should be 128-bit aligned), the SASS shows two separate instructions:

STG.E.64 [R2.64], R8
STG.E.64 [R2.64+0x8], R10
Not only is this not vectori

zed, but it also breaks memory coalescing across threads. I thought STG.E.128 was a valid instruction — am I misremembering? Or there is not support for LDG.E.128 in CuTe (very unlikely), or am i mising something
I’m trying to figure out why I can’t get a proper 128-bit vectorized store from registers to global memory using CuTe. When I load from global memory into registers, the compiler emits a nice LDG.E.128, like:

LDG.E.128 R8, [R8.64]
So far so good.

But when I try to store back to global memory from registers, even though I’m using a float4 (and the data should be 128-bit aligned), the SASS shows two separate instructions:

STG.E.64 [R2.64], R8
STG.E.64 [R2.64+0x8], R10
Not only is this not vectorized, but it also breaks memory coalescing across threads. I thought STG.E.128 was a valid instruction — am I misremembering? Or there is not support for LDG.E.128 in CuTe (very unlikely), or am i mising something

[thakkarV]

what's your input layouts? are you using the vectorizing copy policy?

[N1GHTR0]

My tensor layout has a 4:1 shape-to-stride ratio, which means it’s essentially a 4x1 row-major layout, as shown in the attached picture. I’m copying data from registers to global memory. I’m not explicitly using a vectorizing copy policy here; I assume that AutoVectorizingCopyWithAssumedAlignment handles that automatically. I also tried using copy(Copy_Atom<UniversalCopy<uint128_t>, float>{}, rthr_trms, thr_trmsn) and got the same result. (Is this the correct usage? I'm not sure.)

[thakkarV]

can you please include your full gmem tensor layout and your TV partitioner that created the RMEM tensor? If the copy is not vectorizing, that means it cannot provably find those contiguous elements. something is wrong with your layouts.

[N1GHTR0]

My full global memory layout is (2, 24, 512, 128) with strides (1572864, 65536, 128, 1).
I create one local_tile which gives me a (1, 1, 8, 128) layout, and I coalesce it to make it shape (8, 128).
I take one of the 8 rows based on the warp index and then create another local_tile so that each thread gets 4 values.
So the per-thread shape is 4.
I use make_tensor_like to create the same shape, copy from global memory into that register, and this copy is fully vectorized.
Everything works fine up to this point.

However, after performing calculations and copying back, the store is not vectorized.

there is my full code if needed

template <class Valtensor>
__device__ float thrRmsshfl(Valtensor thrRv)
{
    float thr_local_sum = thrRv(0) * thrRv(0) + 
                          thrRv(1) * thrRv(1) + 
                          thrRv(2) * thrRv(2) + 
                          thrRv(3) * thrRv(3);

    for (int offset = 16; offset > 0; offset /= 2)
        thr_local_sum += __shfl_down_sync(0xffffffff, thr_local_sum, offset);

    float variance = __shfl_sync(0xffffffff, thr_local_sum, 0) / 128.0f;

    float inv_std = rsqrtf(variance + 1e-6f);

    return (inv_std);
}


template <class Layoutrmsg, class Layoutscales, class bTiler, class wTiler>
__global__  void rmsn_kernel(float *rmsVals, float *rmsnVals, float *scales ,Layoutrmsg lrm, Layoutscales lscales, bTiler blockTile, wTiler warpTile) 
{
    using namespace cute;

    Tensor trms = make_tensor(make_gmem_ptr(rmsVals), lrm);
    Tensor trmsn = make_tensor(make_gmem_ptr(rmsnVals), lrm);

    auto bcoord = make_coord(blockIdx.x, blockIdx.y, blockIdx.z, 0);

    Tensor blk_trms = coalesce( local_tile(trms, blockTile, bcoord) );      // 8, 128 : 128, 1
    Tensor blk_trmsn = coalesce( local_tile(trmsn, blockTile, bcoord) );    // 8, 128 : 128, 1

    // block divisions
    int warpidx_in_block = threadIdx.x / 32;

    Tensor warp_trms = blk_trms(warpidx_in_block, _);                       // 128 : 1
    Tensor warp_trmsn = blk_trmsn(warpidx_in_block, _);                     // 128 : 1

    // warp divisions
    auto wcoord = make_coord(threadIdx.x % 32);

    Tensor tscales = make_tensor(make_gmem_ptr(scales), lscales);           // 128 : 1

    Tensor thr_trms = local_tile(warp_trms, warpTile, wcoord);              // 4 : 1  
    Tensor thr_trmsn = local_tile(warp_trmsn, warpTile, wcoord);            // 4 : 1
    Tensor thr_scales = local_tile(tscales, warpTile, wcoord);              // 4 : 1

    // thread claculations
    Tensor rthr_trms = make_tensor_like(thr_trms);                          // 4 : 1
    clear(rthr_trms);

    copy(thr_trms, rthr_trms);

    float inv_std = thrRmsshfl(rthr_trms);

    rthr_trms(0) = (rthr_trms(0) * inv_std) * thr_scales(0);
    rthr_trms(1) = (rthr_trms(1) * inv_std) * thr_scales(1);
    rthr_trms(2) = (rthr_trms(2) * inv_std) * thr_scales(2);
    rthr_trms(3) = (rthr_trms(3) * inv_std) * thr_scales(3);

    copy(rthr_trms, thr_trmsn);
}

int main()
{
    using namespace cute;

    thrust::host_vector<float> h_A(2*24*512*128);
    thrust::host_vector<float> h_scales(128);

    for (int i = 0; i < 128; ++i) h_scales[i] = 2.0f * (static_cast<float>(rand()) / static_cast<float>(RAND_MAX)) - 1.0f;
    for (int j = 0; j < 2*24*512*128; ++j) h_A[j] = 2.0f * (static_cast<float>(rand()) / static_cast<float>(RAND_MAX)) - 1.0f;

    dim3 dimBlock(256);
    dim3 dimGrid(2, 24, 64);

    thrust::device_vector<float> d_A = h_A;
    thrust::device_vector<float> d_scales = h_scales;

    Layout lscales = make_layout(make_shape(Int<128>{}), make_stride(Int<1>{}));
    Layout lrm = make_layout(make_shape(Int<2>{}, Int<24>{}, Int<512>{}, Int<128>{}), make_stride(Int<1572864>{}, Int<65536>{}, Int<128>{}, Int<1>{}));
    auto blockTile = make_shape(Int<1>{}, Int<1>{}, Int<8>{}, Int<128>{});
    auto warpTile = make_shape(Int<4>{});

    // Kernel launch
    rmsn_kernel<<<dimGrid, dimBlock>>>(d_A.data().get() ,d_scales.data().get() ,lrm, lscales, blockTile, warpTile);
    cudaGetLastError();
    cudaDeviceSynchronize();

    thrust::host_vector<float> cute_result = d_A;
}

[thakkarV]

copy(rthr_trms, thr_trmsn);

you are not using a copy policy for autovec with assumed alignment. what is your gmem tensor type? does it have a dynamic shape. if so you will have to use the assumed alignment copy policy because the copy cannot vectorize if it does not know the shape of the gmem tensor.