TMA: CuTe DSL V.S. CuTe C++

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

Python DSL Code
      ↓
   cute.jit  (遍历装饰器修饰的func)
      ↓
 cute.compile (JIT 入口)
      ↓
   MLIR IR  (多层 Dialect)
      ↓
   LLVM IR  (含内联 PTX)
      ↓
   PTX ASM
      ↓
   .CUBIN
      ↓
 GPU Execution

1

from cutlass.cute.nvgpu import cpasync

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

@cute.jit
    def __call__(
        self,
        a: cute.Tensor,
        b: cute.Tensor,
        c: cute.Tensor,
        stream: cuda.CUstream,
        epilogue_op: cutlass.Constexpr = lambda x: x,
    ):
    # Setup store for C
    tma_atom_c = None
    tma_tensor_c = None
    if cutlass.const_expr(self.use_tma_store):
        c_cta_v_layout = cute.composition(
            cute.make_identity_layout(c.shape), self.epi_tile
        )
        epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            c,
            epi_smem_layout,
            c_cta_v_layout,
        )
 
     # Compute grid size
     grid = self._compute_grid(c, self.cta_tile_shape_mnk, self.cluster_shape_mn)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

@dsl_user_op
def make_tiled_tma_atom(
    op: Union[
        CopyBulkTensorTileG2SOp,
        CopyBulkTensorTileG2SMulticastOp,
        CopyBulkTensorTileS2GOp,
        CopyReduceBulkTensorTileS2GOp,
    ],
    gmem_tensor: Tensor,
    smem_layout: Union[Layout, core.ComposedLayout],
    cta_tiler: Tiler,
    num_multicast: int = 1,
    *,
    internal_type: Optional[Type[Numeric]] = None,
    loc=None,
    ip=None,
) -> Tuple[core.CopyAtom, Tensor]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

  # ...
    # Configure gemm kernel
    gemm = DenseGemmKernel(
        acc_dtype,
        use_2cta_instrs,
        mma_tiler_mn,
        cluster_shape_mn,
        use_tma_store,
    )

    torch_stream = torch.cuda.Stream()
    stream = cuda.CUstream(torch_stream.cuda_stream)
    # Compile gemm kernel
    compiled_gemm = cute.compile(gemm, a_tensor, b_tensor, c_tensor, stream)

    # Launch GPU kernel
    # Warm up
    for i in range(warmup_iterations):
        compiled_gemm(a_tensor, b_tensor, c_tensor, stream)
    # Execution
    for i in range(iterations):
        compiled_gemm(a_tensor, b_tensor, c_tensor, stream)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

# Python DSL
cute.copy(
    tma_atom_a,
    tAgA[(None, ab_producer_state.count)],
    tAsA[(None, ab_producer_state.index)],
    tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
    mcast_mask=a_full_mcast_mask,
)

# 转换为 MLIR Op
%result = cute_nvgpu.arch.copy.SM100.tma_load(
    mode = #cute_nvgpu.tma_load_mode<tiled>,
    num_cta = 1,
    src_desc = %tma_descriptor_ptr,
    dsmem_data_addr = %smem_ptr,
    dsmem_bar_addr = %barrier_ptr,
    coord = [%coord_0, %coord_1],
    offsets = [],
    multicast_mask = %mask
) : (...) -> ()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

 def compile(
        self,
        module,
        pipeline: str,
        cuda_toolkit: str = "",
        arch: str = "",
        enable_verifier=False,
    ):
        """Compiles the module by invoking the pipeline."""
        try:
            pm = self.passmanager.PassManager.parse(pipeline)
            pm.enable_verifier(enable_verifier)

            # Enable IR dump for dubug
            pm.enable_ir_printing(print_after_change=True, enable_debug_info=True)
            
            pm.run(module.operation)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

builtin.module(
  cute-to-nvvm{
    cubin-format=bin 
    opt-level=3 
    enable-device-assertions=false 
    toolkitPath=/usr/local/cuda 
    cubin-chip=sm_100a
  },
  external-kernel-for-gpu-launch
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

Python DSL 
  ↓
1. CuteRemoveStaticArgs
2. CuteExpandOps
3. ConvertCuteAlgoToArch
4. CuteFoldStatic
5. InsertRangeInformation


6. Canonicalizer
7. CSE
8. LoopInvariantCodeMotion


9. ConvertCuteTypesInScfOps


10. ArithExpandOpsPass
11. SCFToControlFlow


12. ConvertToLLVMPass


13. SCFToControlFlow
14. Canonicalizer
15. CSE


16. ConvertVectorToLLVMPass
17. ConvertNVVMToLLVMPass
18. GpuNVVMAttachTarget
19. ExpandStridedMetadata
20. StripDebugInfo
21. ConvertGpuOpsToNVVMOps
22. GpuToLLVMConversionPass


23. Canonicalizer


24. ReconcileUnrealizedCasts
25. GpuModuleToBinaryPass
26. ExternalKernelForGpuLaunchPass
  ↓
LLVM IR

1
2
3

// 生成的 TMA descriptor 初始化代码
%llvm.getelementptr %89[0, 0] : (!llvm.ptr) -> !llvm.ptr
llvm.store %140, %141 : i64, !llvm.ptr  // 存储各种 TMA 参数

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

--- a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py
+++ b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py
@@ -352,6 +352,12 @@ class JitExecutor:
             cubin_data = cubin_data.split(b'bin = "')[1].split(b'">')[0]
             cubin_data = self._get_escaped_cubin_bytes(cubin_data)
             callback(sym, func_sym, cubin_data)
+
+            # dump cubin to file
+            import os
+            cubin_file_path = os.path.join("/workspace/cutlass/examples/python
/CuTeDSL/blackwell", f"{func_sym}.cubin")
+            with open(cubin_file_path, "wb") as f:
+                f.write(cubin_data)
             return ir.WalkResult.ADVANCE
 
         module.operation.walk(walk_gpu_binary_op)

1

cuobjdump -sass kernel_xxx.cubin > kernel_xxx.sass

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

// cute_nvgpu dialect
cute_nvgpu.arch.copy.SM100.tma_load(
    mode = #cute_nvgpu.tma_load_mode<tiled>,
    num_cta = 1,
    src_desc = %desc_ptr,
    dsmem_data_addr = %smem_addr,
    dsmem_bar_addr = %bar_addr,
    coord = [%x, %y],
    multicast_mask = %mask
)

// Lower 到 NVVM Dialect
nvvm.cp.async.bulk.tensor.2d.shared.cluster.global.tile.mbarrier.arrive
    %desc_ptr, [%smem_addr], %bar_addr, [%x, %y], %mask

1
2
3
4
5
6
7

// cute.Tensor 的 Layout 信息
!cute.layout<"((128,128),1,1):((65536,1),0,0)">

// 展开为地址计算
%offset = llvm.mul %i, 65536
%offset2 = llvm.add %offset, %j
%addr = llvm.getelementptr %base[%offset2]

1
2
3
4
5
6

// tcgen05.mma
cute.gemm(%tiled_mma, %acc, %a_frag, %b_frag, %acc)

// Lower 到 NVVM
nvvm.tcgen05.mma.ss.ss.t.f32.f16.f16.f32
    %acc, %a_smem_desc, %b_smem_desc, %acc_tmem_desc

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

Python 代码 (dense_gemm.py)
    │
    ├─> cute.compile(gemm, ...)
    │       │
    │       └─> compiler.py::compile()
    │               │
    │               ├─> CompileOptions 解析
    │               ├─> AST 预处理（可选）
    │               └─> dsl.py::_func()
    │                       │
    │                       ├─> 参数类型转换
    │                       │   └─> cute.Tensor → memref<...>
    │                       │
    │                       ├─> generate_original_ir()
    │                       │   │
    │                       │   ├─> 创建 ir.Module
    │                       │   ├─> 构建 GPU Module
    │                       │   │   └─> @cute.kernel 装饰的函数
    │                       │   │       └─> gpu.func @kernel_...
    │                       │   │
    │                       │   ├─> 构建 Host Function
    │                       │   │   └─> func.func @dense_gemm
    │                       │   │       └─> gpu.launch_func
    │                       │   │
    │                       │   └─> 执行用户代码 funcBody(...)
    │                       │       └─> 生成 CUTE Dialect IR
    │                       │           ├─> cute.copy → TMA ops
    │                       │           ├─> cute.gemm → MMA ops
    │                       │           └─> cute.barrier → sync ops
    │                       │
    │                       ├─> 计算 module_hash (SHA256)
    │                       │
    │                       ├─> 检查缓存
    │                       │   ├─> Hit  → 返回 JitExecutor
    │                       │   └─> Miss → compile_and_cache()
    │                       │
    │                       └─> compile_and_cache()
    │                               │
    │                               ├─> preprocess_pipeline()
    │                               │   └─> 添加 toolkitPath, sm-arch
    │                               │
    │                               └─> compile_and_jit()
    │                                       │
    │                                       └─> Compiler.compile()
    │                                           │
    │                                           ├─> PassManager.parse(pipeline)
    │                                           │   │
    │                                           │   └─> Pass Chain:
    │                                           │       │
    │                                           │       ├─> cute-to-linalg
    │                                           │       │   └─> CUTE ops → Linalg ops
    │                                           │       │
    │                                           │       ├─> linalg-to-gpu
    │                                           │       │   └─> Linalg → GPU Dialect
    │                                           │       │
    │                                           │       ├─> gpu-to-nvvm
    │                                           │       │   └─> GPU → NVVM (含 tcgen05)
    │                                           │       │
    │                                           │       ├─> nvvm-to-llvm
    │                                           │       │   └─> NVVM → LLVM IR
    │                                           │       │
    │                                           │       └─> llvm-to-cubin
    │                                           │           └─> LLVM → CUBIN
    │                                           │
    │                                           └─> ExecutionEngine.create()
    │                                               │
    │                                               └─> 生成 CUBIN 二进制
    │
    └─> 返回 JitExecutor
        │
        └─> compiled_gemm(a, b, c, stream)
            │
            └─> ctypes 调用 C API
                └─> CUDA Driver 加载并执行 kernel

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

// Host 调用
Copy_Atom tma_atom_A = make_tma_atom(
  SM90_TMA_LOAD{},        // Operation Tag
  mA,                     // GMEM Tensor
  sA_layout,              // SMEM Layout  
  select<0,2>(mma_tiler)  // Tiler
);

// 模板展开
template <class Operation, class GTensor, class SLayout, class TiledShape>
auto make_tma_atom(Operation const& op, 
                   GTensor const& gtensor,
                   SLayout const& slayout, 
                   TiledShape const& tiled_shape)
{
  // 1. 编译期类型推导
  using GEngine = typename GTensor::engine_type;
  using SEngine = typename SLayout::engine_type;
  using TValue  = typename GTensor::value_type;
  
  // 2. 编译期布局分析
  auto tma_shape = transform_shape(tiled_shape, slayout);
  auto tma_stride = transform_stride(gtensor.stride(), slayout);
  
  // 3. 生成TMA描述符结构体（编译期计算偏移和配置）
  auto tma_desc = make_tma_descriptor<Operation, TValue>(
    gtensor.data(),           // 基地址
    tma_shape,                // 传输shape
    tma_stride,               // 步幅
    get_swizzle_bits(slayout) // Swizzle模式
  );
  
  // 4. 返回Copy_Atom类型（包含TMA描述符）
  return Copy_Atom<Operation, TValue, decltype(tma_desc)>{tma_desc};
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

auto [tAgA, tAsA] = tma_partition(
  tma_atom_A,                // Copy_Atom<SM90_TMA_LOAD, ...>
  Int<0>{},                  // 非多播模式
  Layout<_1>{},              // 非多播布局
  group_modes<0,3>(tCsA),   // SMEM目标
  group_modes<0,3>(tCgA)    // GMEM源
);

// 模板实例化
template <class TMAAtom, class MulticastMode, class MulticastLayout,
          class STensor, class GTensor>
auto tma_partition(TMAAtom const& atom,
                   MulticastMode const& mcast_mode,
                   MulticastLayout const& mcast_layout,
                   STensor const& stensor,
                   GTensor const& gtensor)
{
  // 编译期分支：是否多播
  if constexpr (is_constant<0, MulticastMode>::value) {
    // 非多播路径
    return tma_partition_no_multicast(atom, stensor, gtensor);
  } else {
    // 多播路径（Tutorial 03）
    return tma_partition_with_multicast(atom, mcast_mode, mcast_layout, 
                                        stensor, gtensor);
  }
}

// 非多播实现
template <class TMAAtom, class STensor, class GTensor>
auto tma_partition_no_multicast(TMAAtom const& atom,
                                STensor const& stensor, 
                                GTensor const& gtensor)
{
  // 1. 提取TMA描述符配置
  constexpr auto tma_shape   = atom.shape();   // (_128, _64)
  constexpr auto tma_stride  = atom.stride();  // (_1, _128)
  
  // 2. 重排GMEM张量为TMA期望的格式
  //    mode-0: 单次TMA传输的完整块 (flatten)
  //    mode-1+: 多次传输的迭代维度
  auto tAgA_reordered = recast<tma_shape>(gtensor);
  //    ((TmaShape_M, TmaShape_K), RestModes...)
  
  // 3. 重排SMEM张量为扁平化视图
  auto tAsA_flat = flatten_to_tma_tile(stensor, tma_shape);
  //    (TmaSize_total, 1)
  
  return cute::make_tuple(tAgA_reordered, tAsA_flat);
}

1
2
3
4
5
6
7
8
9

// 输入 tCgA: ((_128,_16),_1,_4,4)
group_modes<0,3>(tCgA)  // → ((_128,_16,_4),4)
                        //    ^^^^^^^^^^^^^ mode-0: MMA Tile MK
                        //                ^ mode-1: Tiles_K

// TMA分区后 tAgA: (((_64,_128),_1),4)
//                   ^^^^^^^^^^^^ mode-0: 单次TMA传输块（flatten后8192元素）
//                            ^  mode-1: 内部虚拟维度（用于对齐）
//                               ^ mode-2: K tile迭代

1
2
3
4
5

if (elect_one_warp && elect_one_thr) {
  copy(tma_atom_A.with(shared_storage.tma_barrier), 
       tAgA(_,k_tile), 
       tAsA);
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

template <class... CopyArgs,
          class SrcEngine, class SrcLayout,
          class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<CopyArgs...>       const& copy_atom,
     Tensor<SrcEngine, SrcLayout> const& src,       // (V,Rest...)
     Tensor<DstEngine, DstLayout>      & dst)       // (V,Rest...){
     
     copy_atom.call(src, dst);
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

  // Check and call instruction, or recurse
  template <class SEngine, class SLayout,
            class DEngine, class DLayout>
  CUTE_HOST_DEVICE
  void
  call(Tensor<SEngine,SLayout> const& src,
       Tensor<DEngine,DLayout>      & dst) const
  {
    static_assert(SLayout::rank == 1, "Expected rank-1 src tensor");
    static_assert(DLayout::rank == 1, "Expected rank-1 dst tensor");

    if constexpr (is_constant<NumValSrc, decltype(size(src))>::value ||
                  is_constant<NumValDst, decltype(size(dst))>::value) {
      // Dispatch to unpack to execute instruction
      return copy_unpack(static_cast<Traits const&>(*this), src, dst);
    } else if constexpr (is_tuple<decltype(shape(src))>::value &&
                         is_tuple<decltype(shape(dst))>::value) {
      // If the size of the src/dst doesn't match the instruction,
      //   recurse this rank-1 layout by peeling off the mode
      //   ((A,B,C,...)) -> (A,B,C,...)
      return copy(*this, tensor<0>(src), tensor<0>(dst));
    } else {
      static_assert(dependent_false<SEngine>,
                    "CopyAtom: Src/Dst partitioning does not match the instruction requirement.");
    }
  }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

template <class TValue, class TMADescriptor, class Barrier,
          class SrcTensor, class DstTensor>
CUTE_DEVICE
void Copy_Atom<SM90_TMA_LOAD, TValue, TMADescriptor>::call(
    Barrier const& barrier,
    SrcTensor const& src,  // tAgA(_,k_tile)
    DstTensor      & dst)  // tAsA
{
  // 编译期检查
  static_assert(is_smem<DstTensor>::value, "Dst must be SMEM");
  static_assert(is_gmem<SrcTensor>::value, "Src must be GMEM");
  
  // 提取地址和大小（运行时）
  uint64_t gmem_addr = reinterpret_cast<uint64_t>(src.data().get());
  uint32_t smem_addr = cast_smem_ptr_to_uint(dst.data().get());
  uint32_t tma_bytes = sizeof(TValue) * cosize_v<DstTensor>;
  
  // 调用底层TMA copy函数
  SM90_TMA_LOAD::copy(
    tma_desc_,      // TMA描述符（编译期常量内存）
    barrier,        // Barrier地址（运行时）
    smem_addr,      // SMEM目标地址（运行时）
    gmem_addr,      // GMEM源地址（运行时）
    tma_bytes       // 传输字节数（编译期常量）
  );
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

struct SM90_TMA_LOAD_1D
{
  CUTE_HOST_DEVICE static void
  copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
       void      * smem_ptr,
       int32_t const& crd0)
  {
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
    uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
    cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
#if defined(CUTE_ARCH_TMA_SM120_ENABLED)
    asm volatile (
      "cp.async.bulk.tensor.1d.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint"
      " [%0], [%1, {%3}], [%2], %4;"
      :
      : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
        "r"(crd0), "l"(cache_hint)
      : "memory");
#else
    asm volatile (
      "cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
      " [%0], [%1, {%3}], [%2], %4;"
      :
      : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
        "r"(crd0), "l"(cache_hint)
      : "memory");
#endif
#else
    CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
  }
  
 // ...
 }

1
2
3
4
5
6
7
8

➜  cutlass git:(main) ✗ nm -D /root/wkspace/cutlass/_mlir/_mlir_libs/_mlirExecutionEngine.cpython-313-x86_64-linux-gnu.so | grep -i "execution\|nvptx\|ptx" | head -20

0000000000008038 T PyInit__mlirExecutionEngine
                 U mlirExecutionEngineCreate
                 U mlirExecutionEngineDestroy
                 U mlirExecutionEngineDumpToObjectFile
                 U mlirExecutionEngineLookupPacked
                 U mlirExecutionEngineRegisterSymbol

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

[Debug] Using compilation pipeline: builtin.module(
                                            cute-to-nvvm{
                                                cubin-format=bin 
                                                opt-level=3 
                                                enable-device-assertions=false 
                                                link-libraries= 
                                                toolkitPath=/usr/local/cuda 
                                                cubin-chip=sm_100a 
                                            },
                                            external-kernel-for-gpu-launch
                                    )