[QST] CuteDSL in memory pass

Here I write a simple cuteDSL program in order to perform cast from fp32 tensor to bf16 tensor:

import argparse
import math
import torch
import triton
from typing import Tuple, Type, Callable, Optional, Union, Literal

import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.cute.testing as testing
import cutlass.pipeline as pipeline

from cutlass.pipeline.helpers import MbarrierArray

import cutlass.torch as cutlass_torch
import cutlass.utils.hopper_helpers as sm90_utils
from cutlass.cute.runtime import from_dlpack
from cutlass import Int32, Boolean, const_expr
from cutlass.cute.nvgpu import cpasync, warp, warpgroup

from typing import Tuple, Type

class S2R:
    def __init__(self):
        pass
    
     
    @cute.jit
    def make_smem_layouts(
        self,
        mA_dtype: cutlass.Numeric,
        mA_layout: utils.LayoutEnum,
        mB_dtype: cutlass.Numeric,
        mB_layout: utils.LayoutEnum
    ):
        a_is_k_major = mA_layout.sm90_mma_major_mode() == warpgroup.OperandMajorMode.K
        b_is_k_major = mB_layout.sm90_mma_major_mode() == warpgroup.OperandMajorMode.K
        
        # m = const_expr()
        # n = const_expr()
        a_smem_shape = (m, n)
        
        a_major_mode_size = const_expr(n if a_is_k_major else m)
        print(f"w_major_mode_size: {a_major_mode_size}")
        
        smem_layout_atom_a = sm90_utils.get_smem_layout_atom(mA_layout, mA_dtype, a_major_mode_size)
        print(f"smem_layout_atom_a {smem_layout_atom_a}")
        
        a_smem_layout_atom = warpgroup.make_smem_layout_atom(
            smem_layout_atom_a,
            mA_dtype,
        )
        print(f"a_smem_layout_atom: {a_smem_layout_atom}")
        a_smem_layout = cute.tile_to_shape(
            a_smem_layout_atom,
            a_smem_shape,
            order=(0, 1) if a_is_k_major else (1, 0),
        )
        
        b_smem_shape = (m, n)
        b_major_mode_size = const_expr(n if a_is_k_major else m)
        smem_layout_atom_b = sm90_utils.get_smem_layout_atom(mB_layout, mB_dtype, b_major_mode_size)
        
        print(f"smem_layout_atom_b: {smem_layout_atom_b}")
        b_smem_layout_atom = warpgroup.make_smem_layout_atom(
            smem_layout_atom_b,
            mB_dtype,
        )
        print(f"b_smem_layout_atom: {b_smem_layout_atom}")
        
        b_smem_layout = cute.tile_to_shape(
            b_smem_layout_atom,
            b_smem_shape,
            order=(0, 1) if b_is_k_major else (1, 0),
        )

        print(f"bsmem_layout_staged: {b_smem_layout}")
        
        
        return (
            a_smem_layout,
            b_smem_layout,
        )
        
    def make_tiled_copy(self, dtype: cute.Numeric, major_mode_size: int, num_copy_bits: int=128):
        assert dtype in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]
        threads_per_cta = 128

        copy_op = cute.nvgpu.CopyUniversalOp()
        copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
        
        # value layout:
        elems = num_copy_bits // dtype.width
        val_layout = cute.make_layout((elems, 1), stride=(1, 0))
        
        # thread layout: 
        assert major_mode_size == 64
        threads_per_col = const_expr(major_mode_size // elems)
        rows_per_block = threads_per_cta // threads_per_col
        thr_layout = cute.make_layout((threads_per_col, rows_per_block), stride=(1, threads_per_col))
        
        tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
        
        print(tiled_copy)
        return tiled_copy
        
    
    @cute.jit
    def __call__(self, mA: cute.Tensor, mB: cute.Tensor, m: int, n: int):
        threads_per_cta = 128
        a_dtype = mA.element_type
        b_dtype = mB.element_type
        assert a_dtype == cutlass.Float32
        assert b_dtype == cutlass.BFloat16
        tiled_copy_a = self.make_tiled_copy(a_dtype, 64, 128)
        tiled_copy_b = self.make_tiled_copy(b_dtype, 64, 64)
        a_smem_layout, b_smem_layout = self.make_smem_layouts(mA.element_type, utils.LayoutEnum.from_tensor(mA), mB.element_type, utils.LayoutEnum.from_tensor(mB))
        a_smem_size = cute.cosize(a_smem_layout)
        b_smem_size = cute.cosize(b_smem_layout)
        @cute.struct
        class SharedStorage:
            sA: cute.struct.Align[
                cute.struct.MemRange[a_dtype, a_smem_size],
                1024,
            ]
            sB: cute.struct.Align[
                cute.struct.MemRange[b_dtype, b_smem_size],
                1024,
            ]
            
        self.shared_storage = SharedStorage
        
        self.kernel(
            tiled_copy_a,
            tiled_copy_b,    
            mA, mB, 
            a_smem_layout, b_smem_layout
        ).launch(
            block=(threads_per_cta, 1, 1),
            grid=(1,1,1)
        )
    
    @cute.kernel
    def kernel(self, tiled_copy_a: cute.TiledCopy, tiled_copy_b: cute.TiledCopy, mA: cute.Tensor, mB: cute.Tensor, a_smem_layout: cute.ComposedLayout, b_smem_layout: cute.ComposedLayout):
        tidx, _, _ = cute.arch.thread_idx()
        thr_copy_a = tiled_copy_a.get_slice(tidx)
        thr_copy_b = tiled_copy_b.get_slice(tidx)
        
        smem = cutlass.utils.SmemAllocator()
        
        storage = smem.allocate(self.shared_storage)
        
        sA = storage.sA.get_tensor(
            a_smem_layout.outer, swizzle=a_smem_layout.inner
        )
        
        sB = storage.sB.get_tensor(
            b_smem_layout.outer, swizzle=b_smem_layout.inner
        )
        
        tXgA = thr_copy_a.partition_S(mA)
        tXgB = thr_copy_b.partition_D(mB)
        
        print(f"tXgA: {tXgA}")
        print(f"tXgB: {tXgB}")
        
        tXsA = thr_copy_a.partition_S(sA)
        tXsB = thr_copy_b.partition_D(sB)
    
        print(f"sA layout: {sA.layout}")
        print(f"sB layout: {sB.layout}")
        print(f"mA: {mA}")
        print(f"mB: {mB}")
        print(f"tXsA: {tXsA}")
        print(f"tXsB: {tXsB}")
        
        # cute.autovec_copy(tXgA, tXsA)
        cute.copy(thr_copy_a, tXgA, tXsA)
        
        tXrA = cute.make_fragment_like(tXsA)
        tXrB = cute.make_fragment_like(tXsB)
        
        print(f"tXrA: {tXrA}")
        tXrA = thr_copy_a.retile(tXrA)
        print(f"tXrA: {tXrA}")
        cute.autovec_copy(tXsA, tXrA)
    
        tXrB.store(tXrA.load().to(cutlass.BFloat16))
        
        cute.autovec_copy(tXrB, tXsB)
        cute.autovec_copy(tXsB, tXgB)
        
m = 64
n = 16
dtype = torch.float32
x = torch.randn((n, m), dtype=dtype, device="cuda").transpose_(0, 1)

print(x.shape)
print(x.is_contiguous())

y = torch.zeros((n, m), dtype=torch.bfloat16, device="cuda").transpose_(0,1)
print(y.shape)
print(y.is_contiguous())

keep = x.clone()
print(keep)
# print(2 * keep)
kernel = S2R()
kernel(from_dlpack(x, assumed_align=16), from_dlpack(y, assumed_align=16), m, n)

from cutlass.cute import KeepPTX, KeepCUBIN
hello_world_compiled_ptx_on = cute.compile[KeepPTX, KeepCUBIN](kernel,from_dlpack(x, assumed_align=16), from_dlpack(y, assumed_align=16), m, n)
# print(x)
print(y)

But the ptx seems confusing:

//
// Generated by NVIDIA NVVM Compiler
//
// Compiler Build ID: CL-36006120
// Cuda compilation tools, release 12.9, V12.9.83
// Based on NVVM 20.0.0
//

.version 8.8
.target sm_90a
.address_size 64

	// .globl	kernel_cutlass_kernel___main__S2R_object_at__TiledCopy_TilerMN64181_TVLayouttiled128441_CopyAtom_ThrID10_TVLayoutSrc1401_TVLayoutDst1401_Valuetypef32_TiledCopy_TilerMN64181_TVLayouttiled1_0
.extern .shared .align 1024 .b8 __dynamic_shmem__0[];

.visible .entry kernel_cutlass_kernel___main__S2R_object_at__TiledCopy_TilerMN64181_TVLayouttiled128441_CopyAtom_ThrID10_TVLayoutSrc1401_TVLayoutDst1401_Valuetypef32_TiledCopy_TilerMN64181_TVLayouttiled1_0(
	.param .align 8 .b8 kernel_cutlass_kernel___main__S2R_object_at__TiledCopy_TilerMN64181_TVLayouttiled128441_CopyAtom_ThrID10_TVLayoutSrc1401_TVLayoutDst1401_Valuetypef32_TiledCopy_TilerMN64181_TVLayouttiled1_0_param_0[8],
	.param .align 8 .b8 kernel_cutlass_kernel___main__S2R_object_at__TiledCopy_TilerMN64181_TVLayouttiled128441_CopyAtom_ThrID10_TVLayoutSrc1401_TVLayoutDst1401_Valuetypef32_TiledCopy_TilerMN64181_TVLayouttiled1_0_param_1[8]
)
.reqntid 128, 1, 1
{
	.reg .b16 	%rs<13>;
	.reg .b32 	%r<29>;
	.reg .f32 	%f<13>;
	.reg .b64 	%rd<7>;

	ld.param.u64 	%rd1, [kernel_cutlass_kernel___main__S2R_object_at__TiledCopy_TilerMN64181_TVLayouttiled128441_CopyAtom_ThrID10_TVLayoutSrc1401_TVLayoutDst1401_Valuetypef32_TiledCopy_TilerMN64181_TVLayouttiled1_0_param_1];
	ld.param.u64 	%rd2, [kernel_cutlass_kernel___main__S2R_object_at__TiledCopy_TilerMN64181_TVLayouttiled128441_CopyAtom_ThrID10_TVLayoutSrc1401_TVLayoutDst1401_Valuetypef32_TiledCopy_TilerMN64181_TVLayouttiled1_0_param_0];
	mov.u32 	%r1, %tid.x;
	shl.b32 	%r2, %r1, 2;
	shl.b32 	%r3, %r1, 3;
	and.b32  	%r4, %r3, 896;
	shl.b32 	%r5, %r1, 4;
	and.b32  	%r6, %r5, 112;
	or.b32  	%r7, %r6, %r4;
	shl.b32 	%r8, %r1, 8;
	and.b32  	%r9, %r8, 2048;
	or.b32  	%r10, %r9, %r7;
	mov.u32 	%r11, __dynamic_shmem__0;
	add.s32 	%r12, %r11, %r10;
	add.s32 	%r13, %r11, %r3;
	add.s32 	%r14, %r13, 4096;
	mul.wide.u32 	%rd3, %r2, 4;
	add.s64 	%rd4, %rd2, %rd3;
	shr.u32 	%r15, %r12, 3;
	and.b32  	%r16, %r15, 112;
	xor.b32  	%r17, %r16, %r12;
	ld.global.v4.f32 	{%f1, %f2, %f3, %f4}, [%rd4];
	st.shared.v4.f32 	[%r17], {%f1, %f2, %f3, %f4};
	add.s32 	%r18, %r12, 1024;
	shr.u32 	%r19, %r18, 3;
	and.b32  	%r20, %r19, 112;
	xor.b32  	%r21, %r20, %r18;
	ld.global.v4.f32 	{%f5, %f6, %f7, %f8}, [%rd4+2048];
	st.shared.v4.f32 	[%r21], {%f5, %f6, %f7, %f8};
	ld.shared.v4.f32 	{%f9, %f10, %f11, %f12}, [%r17];
	cvt.rn.bf16.f32 	%rs1, %f12;
	cvt.rn.bf16.f32 	%rs2, %f11;
	cvt.rn.bf16.f32 	%rs3, %f10;
	cvt.rn.bf16.f32 	%rs4, %f9;
	cvt.rn.bf16.f32 	%rs5, %f8;
	cvt.rn.bf16.f32 	%rs6, %f7;
	cvt.rn.bf16.f32 	%rs7, %f6;
	cvt.rn.bf16.f32 	%rs8, %f5;
	shr.u32 	%r22, %r14, 3;
	and.b32  	%r23, %r22, 112;
	xor.b32  	%r24, %r23, %r14;
	st.shared.v4.b16 	[%r24], {%rs4, %rs3, %rs2, %rs1};
	add.s32 	%r25, %r13, 5120;
	shr.u32 	%r26, %r25, 3;
	and.b32  	%r27, %r26, 112;
	xor.b32  	%r28, %r27, %r25;
	st.shared.v4.b16 	[%r28], {%rs8, %rs7, %rs6, %rs5};
	mul.wide.u32 	%rd5, %r2, 2;
	add.s64 	%rd6, %rd1, %rd5;
	ld.shared.v4.b16 	{%rs9, %rs10, %rs11, %rs12}, [%r24];
	st.global.v4.b16 	[%rd6], {%rs9, %rs10, %rs11, %rs12};
	st.global.v4.b16 	[%rd6+1024], {%rs8, %rs7, %rs6, %rs5};
	ret;

}

There is only one ld.shared.v4, it seems to do a memory optimize pass to reuse the register. But My question is why only one ld.shared.v4 is optimized, it is clearly that we can eliminate both ld.shared.v4 in convert since it must go to register first!

Start by reproducing the provided CuteDSL S2R program and inspecting the generated PTX around the shared-memory loads and FP32-to-BF16 conversions. Compare the two conversion paths and determine why only one shared load is eliminated, then document whether both loads can be removed or what constraint prevents that optimization.