| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # this module is named _tcgen05 to avoid name collision with cute.nvgpu.tcgen05 |
| 4 | |
| 5 | import cutlass |
| 6 | from cutlass import Boolean, Float32, Int32, Uint32, Uint64, cute |
| 7 | from cutlass._mlir import ir |
| 8 | from cutlass._mlir.dialects import llvm, nvvm, vector |
| 9 | from cutlass.cutlass_dsl import dsl_user_op |
| 10 | |
| 11 | NVVM_CTA_GROUP_MAP = [ |
| 12 | None, |
| 13 | nvvm.CTAGroupKind.CTA_1, |
| 14 | nvvm.CTAGroupKind.CTA_2, |
| 15 | ] |
| 16 | LDST_MAP = { |
| 17 | "32x32b": (nvvm.Tcgen05LdStShape.SHAPE_32X32B, 1), |
| 18 | "16x128b": (nvvm.Tcgen05LdStShape.SHAPE_16X128B, 2), |
| 19 | "16x256b": (nvvm.Tcgen05LdStShape.SHAPE_16X256B, 4), |
| 20 | } |
| 21 | |
| 22 | |
| 23 | def _make_tmem_llvm_ptr(addr, *, loc=None, ip=None): |
| 24 | ptr_ty = llvm.PointerType.get(cute.AddressSpace.tmem.value) |
| 25 | val = Int32(addr).ir_value(loc=loc, ip=ip) |
| 26 | return llvm.inttoptr(ptr_ty, val, loc=loc, ip=ip) |
| 27 | |
| 28 | |
| 29 | @dsl_user_op |
| 30 | def alloc( |
| 31 | taddr: cute.Pointer, |
| 32 | cta_group: int = 1, |
| 33 | *, |
| 34 | loc=None, |
| 35 | ip=None, |
| 36 | ) -> None: |
| 37 | nvvm.tcgen05_alloc( |
| 38 | taddr.to_llvm_ptr(loc=loc, ip=ip), |
| 39 | Uint32(512).ir_value(loc=loc, ip=ip), |
| 40 | group=NVVM_CTA_GROUP_MAP[cta_group], |
| 41 | loc=loc, |
| 42 | ip=ip, |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | @dsl_user_op |
| 47 | def dealloc(cta_group: int = 1, *, loc=None, ip=None) -> None: |
| 48 | nvvm.tcgen05_dealloc( |
| 49 | _make_tmem_llvm_ptr(0, loc=loc, ip=ip), |
| 50 | Int32(512).ir_value(loc=loc, ip=ip), |
| 51 | group=NVVM_CTA_GROUP_MAP[cta_group], |
| 52 | loc=loc, |
| 53 | ip=ip, |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def make_bf16_idesc( |
| 58 | MMA_M: int, |
| 59 | MMA_N: int, |
| 60 | *, |
| 61 | negate_A: bool = False, |
| 62 | negate_B: bool = False, |
| 63 | transpose_A: bool = False, |
| 64 | transpose_B: bool = False, |
| 65 | ): |
| 66 | idesc = Uint32( |
| 67 | (1 << 4) | (1 << 7) | (1 << 10) | ((MMA_N >> 3) << 17) | ((MMA_M >> 4) << 24) |
| 68 | ) |
| 69 | idesc |= Uint32(negate_A) << 13 |
| 70 | idesc |= Uint32(negate_B) << 14 |
| 71 | idesc |= Uint32(transpose_A) << 15 |
| 72 | idesc |= Uint32(transpose_B) << 16 |
| 73 | return idesc |
| 74 | |
| 75 | |
| 76 | def make_sdesc_128B_swizzle(LBO: int): |
| 77 | SBO = 8 * 128 |
| 78 | return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61)) |
| 79 | |
| 80 | |
| 81 | def make_mxfp8_idesc(MMA_M: int, MMA_N: int): |
| 82 | return Uint32((MMA_N >> 3 << 17) | (1 << 23) | (MMA_M >> 7 << 27)) |
| 83 | |
| 84 | |
| 85 | @dsl_user_op |
| 86 | def mma_f16( |
| 87 | d_tmem, |
| 88 | a_desc, |
| 89 | b_desc, |
| 90 | idesc, |
| 91 | enable_input_d, |
| 92 | cta_group: int = 1, |
| 93 | *, |
| 94 | loc=None, |
| 95 | ip=None, |
| 96 | ) -> None: |
| 97 | with cute.arch.elect_one(): |
| 98 | nvvm.tcgen05_mma( |
| 99 | nvvm.Tcgen05MMAKind.F16, |
| 100 | NVVM_CTA_GROUP_MAP[cta_group], |
| 101 | _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), |
| 102 | Uint64(a_desc).ir_value(loc=loc, ip=ip), |
| 103 | Uint64(b_desc).ir_value(loc=loc, ip=ip), |
| 104 | Int32(idesc).ir_value(loc=loc, ip=ip), |
| 105 | Boolean(enable_input_d).ir_value(loc=loc, ip=ip), |
| 106 | loc=loc, |
| 107 | ip=ip, |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | @dsl_user_op |
| 112 | def mma_ts_f16( |
| 113 | d_tmem, |
| 114 | a_tmem, |
| 115 | b_desc, |
| 116 | idesc, |
| 117 | enable_input_d, |
| 118 | cta_group: int = 1, |
| 119 | *, |
| 120 | loc=None, |
| 121 | ip=None, |
| 122 | ) -> None: |
| 123 | with cute.arch.elect_one(): |
| 124 | nvvm.tcgen05_mma( |
| 125 | nvvm.Tcgen05MMAKind.F16, |
| 126 | NVVM_CTA_GROUP_MAP[cta_group], |
| 127 | _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), |
| 128 | _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip), |
| 129 | Uint64(b_desc).ir_value(loc=loc, ip=ip), |
| 130 | Int32(idesc).ir_value(loc=loc, ip=ip), |
| 131 | Boolean(enable_input_d).ir_value(loc=loc, ip=ip), |
| 132 | loc=loc, |
| 133 | ip=ip, |
| 134 | ) |
| 135 | |
| 136 | |
| 137 | @dsl_user_op |
| 138 | def mma_mxfp8( |
| 139 | d_tmem, |
| 140 | a_desc, |
| 141 | b_desc, |
| 142 | idesc, |
| 143 | sfa_tmem, |
| 144 | sfb_tmem, |
| 145 | enable_input_d, |
| 146 | cta_group: int = 1, |
| 147 | *, |
| 148 | loc=None, |
| 149 | ip=None, |
| 150 | ) -> None: |
| 151 | # CuTe DSL 4.8 renamed the block-scale enum and keyword. |
| 152 | if hasattr(nvvm, "Tcgen05MMABlockScale"): |
| 153 | scale_kwargs = {"block_scale": nvvm.Tcgen05MMABlockScale.DEFAULT} |
| 154 | else: |
| 155 | scale_kwargs = {"scale_vec_size": nvvm.Tcgen05MMAScaleVecSize.X1} |
| 156 | with cute.arch.elect_one(): |
| 157 | nvvm.tcgen05_mma_block_scale( |
| 158 | nvvm.Tcgen05MMAKind.MXF8F6F4, |
| 159 | NVVM_CTA_GROUP_MAP[cta_group], |
| 160 | _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), |
| 161 | Uint64(a_desc).ir_value(loc=loc, ip=ip), |
| 162 | Uint64(b_desc).ir_value(loc=loc, ip=ip), |
| 163 | Int32(idesc).ir_value(loc=loc, ip=ip), |
| 164 | Boolean(enable_input_d).ir_value(loc=loc, ip=ip), |
| 165 | _make_tmem_llvm_ptr(sfa_tmem, loc=loc, ip=ip), |
| 166 | _make_tmem_llvm_ptr(sfb_tmem, loc=loc, ip=ip), |
| 167 | loc=loc, |
| 168 | ip=ip, |
| 169 | **scale_kwargs, |
| 170 | ) |
| 171 | |
| 172 | |
| 173 | @dsl_user_op |
| 174 | def cp( |
| 175 | tmem, |
| 176 | sdesc, |
| 177 | shape: str, |
| 178 | mcast: str, |
| 179 | cta_group: int = 1, |
| 180 | *, |
| 181 | loc=None, |
| 182 | ip=None, |
| 183 | ) -> None: |
| 184 | SHAPE_MAP = { |
| 185 | "128x256b": nvvm.Tcgen05CpShape.SHAPE_128x256b, |
| 186 | "4x256b": nvvm.Tcgen05CpShape.SHAPE_4x256b, |
| 187 | "128x128b": nvvm.Tcgen05CpShape.SHAPE_128x128b, |
| 188 | "64x128b": nvvm.Tcgen05CpShape.SHAPE_64x128b, |
| 189 | "32x128b": nvvm.Tcgen05CpShape.SHAPE_32x128b, |
| 190 | } |
| 191 | MCAST_MAP = { |
| 192 | "none": nvvm.Tcgen05CpMulticast.NONE, |
| 193 | "warpx2::02_13": nvvm.Tcgen05CpMulticast.WARPX2_02_13, |
| 194 | "warpx2::01_23": nvvm.Tcgen05CpMulticast.WARPX2_01_23, |
| 195 | "warpx4": nvvm.Tcgen05CpMulticast.WARPX4, |
| 196 | } |
| 197 | with cute.arch.elect_one(): |
| 198 | nvvm.tcgen05_cp( |
| 199 | SHAPE_MAP[shape], |
| 200 | _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip), |
| 201 | Uint64(sdesc).ir_value(loc=loc, ip=ip), |
| 202 | group=NVVM_CTA_GROUP_MAP[cta_group], |
| 203 | multicast=MCAST_MAP[mcast], |
| 204 | loc=loc, |
| 205 | ip=ip, |
| 206 | ) |
| 207 | |
| 208 | |
| 209 | @dsl_user_op |
| 210 | def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None): |
| 211 | mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip) |
| 212 | group = NVVM_CTA_GROUP_MAP[cta_group] |
| 213 | if cutlass.const_expr(cta_mask is not None): |
| 214 | with cute.arch.elect_one(): |
| 215 | nvvm.tcgen05_commit( |
| 216 | mbar_llvm, |
| 217 | multicast_mask=cta_mask.ir_value(loc=loc, ip=ip), |
| 218 | group=group, |
| 219 | loc=loc, |
| 220 | ip=ip, |
| 221 | ) |
| 222 | else: |
| 223 | with cute.arch.elect_one(): |
| 224 | nvvm.tcgen05_commit(mbar_llvm, group=group, loc=loc, ip=ip) |
| 225 | |
| 226 | |
| 227 | @dsl_user_op |
| 228 | def ld(row, col, shape: str, num: int, *, loc=None, ip=None): |
| 229 | nvvm_shape, regs_per_num = LDST_MAP[shape] |
| 230 | num_regs = regs_per_num * num |
| 231 | tmem = (Int32(row) << Int32(16)) | Int32(col) |
| 232 | tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip) |
| 233 | |
| 234 | if num_regs == 1: |
| 235 | reg = nvvm.tcgen05_ld(Int32.mlir_type, nvvm_shape, tmem_ptr, loc=loc, ip=ip) |
| 236 | reg_f32 = llvm.bitcast(Float32.mlir_type, reg, loc=loc, ip=ip) |
| 237 | return Float32(reg_f32) |
| 238 | |
| 239 | else: |
| 240 | vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc) |
| 241 | vec_f32_ty = ir.VectorType.get([num_regs], Float32.mlir_type, loc=loc) |
| 242 | regs = nvvm.tcgen05_ld(vec_i32_ty, nvvm_shape, tmem_ptr, loc=loc, ip=ip) |
| 243 | regs_f32 = llvm.bitcast(vec_f32_ty, regs, loc=loc, ip=ip) |
| 244 | return cute.TensorSSA(regs_f32, (num_regs,), Float32) |
| 245 | |
| 246 | |
| 247 | @dsl_user_op |
| 248 | def st(row, col, shape: str, num: int, vals, *, loc=None, ip=None) -> None: |
| 249 | # if input is TensorSSA, convert to Tensor so we can bitcast |
| 250 | if isinstance(vals, cute.TensorSSA): |
| 251 | vals_ = cute.make_rmem_tensor_like(vals) |
| 252 | vals_.store(vals) |
| 253 | vals = vals_ |
| 254 | |
| 255 | # bitcast to Int32 |
| 256 | vals = cute.recast_tensor(vals, Int32) |
| 257 | |
| 258 | nvvm_shape, regs_per_num = LDST_MAP[shape] |
| 259 | num_regs = regs_per_num * num |
| 260 | tmem = (Int32(row) << Int32(16)) | Int32(col) |
| 261 | tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip) |
| 262 | |
| 263 | if num_regs == 1: |
| 264 | nvvm.tcgen05_st( |
| 265 | nvvm_shape, |
| 266 | tmem_ptr, |
| 267 | vals[0].ir_value(loc=loc, ip=ip), |
| 268 | loc=loc, |
| 269 | ip=ip, |
| 270 | ) |
| 271 | else: |
| 272 | vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc) |
| 273 | val_vec = vector.from_elements( |
| 274 | vec_i32_ty, |
| 275 | [vals[i].ir_value(loc=loc, ip=ip) for i in range(num_regs)], |
| 276 | loc=loc, |
| 277 | ip=ip, |
| 278 | ) |
| 279 | nvvm.tcgen05_st(nvvm_shape, tmem_ptr, val_vec, loc=loc, ip=ip) |
| 280 | |
| 281 | |
| 282 | @dsl_user_op |
| 283 | def fence_after_thread_sync(*, loc=None, ip=None): |
| 284 | nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip) |
| 285 | |
| 286 | |
| 287 | @dsl_user_op |
| 288 | def fence_before_thread_sync(*, loc=None, ip=None): |
| 289 | nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip) |
| 290 | |
| 291 | |
| 292 | @dsl_user_op |
| 293 | def wait_ld(*, loc=None, ip=None): |
| 294 | nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.LOAD, loc=loc, ip=ip) |
| 295 | |
| 296 | |
| 297 | @dsl_user_op |
| 298 | def wait_st(*, loc=None, ip=None): |
| 299 | nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.STORE, loc=loc, ip=ip) |