This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2 --- a/lib/Analysis/Utility.cpp +++ b/lib/Analysis/Utility.cpp @@ -1197,7 +1197,7 @@ if (k < 256 / aElemTy.getIntOrFloatBitWidth()) return false; if (!(retShapePerCTA[rank - 2] % 64 == 0 && - retShapePerCTA[rank - 1] % 8 == 0)) + retShapePerCTA[rank - 1] % 16 == 0)) return false; if (aElemTy.isF64() || bElemTy.isF64() || retType.getElementType().isF64()) { @@ -1222,7 +1222,7 @@ if (rank == 3) return false; if (!(numWarps % 4 == 0 && retShapePerCTA[rank - 2] % 64 == 0 && - retShapePerCTA[rank - 1] % 8 == 0 && + retShapePerCTA[rank - 1] % 16 == 0 && (llvm::isa(aElemTy) || aElemTy.isInteger(8) || aElemTy.isF16() || aElemTy.isBF16() || aElemTy.isF32()))) { --- a/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp +++ b/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp @@ -184,13 +184,21 @@ auto kReg = StringAttr::get(ctx, "register"); LinearLayout quot; int bestContig = 1; + bool isDim0 = false; for (int contig = 1; bitwidth * contig <= 32; contig *= 2) { auto maybeQuot = divideLeft( ll, LinearLayout::identity1D(contig, rowColDims[1], dims[1])); + bool currDim0 = false; + if (!maybeQuot) { + maybeQuot = divideLeft( + ll, LinearLayout::identity1D(contig, rowColDims[1], dims[0])); + currDim0 = true; + } if (!maybeQuot) break; quot = *maybeQuot; bestContig = contig; + isDim0 = currDim0; } // Pack contiguous elements @@ -200,7 +208,7 @@ bitwidth * bestContig); if (!ret) return ret; - auto castbbitwidth = LinearLayout::identity1D(bestContig, kReg, dims[1]); + auto castbbitwidth = LinearLayout::identity1D(bestContig, kReg, isDim0 ? dims[0] : dims[1]); return castbbitwidth * ret.value(); } if (auto maybeQuot = divideLeft( @@ -241,10 +249,7 @@ auto nColsTile = tile.getOutDimSize(rowColDims[1]); auto nColsLL = ll.getInDimSize(rowColDims[1]); - auto nColsMissing = nColsLL / nColsTile; - if (nColsMissing == 0) { - return std::nullopt; - } + auto nColsMissing = std::max(nColsLL / nColsTile, 1); auto kReg = StringAttr::get(ctx, "register"); auto kLane = StringAttr::get(ctx, "lane"); auto kWarp = StringAttr::get(ctx, "warp"); @@ -261,7 +266,9 @@ // basis as that would disallow the resulting layout to be left-divisible by // the tile auto trivialBlock = LinearLayout::identity1D(1, kBlock, kBlock); - auto comp = (tile * trivialBlock) + auto tileMult = tile * trivialBlock; + auto tileMultShrunk = tileMult.resizeOutDim(rowColDims[0], std::min(ll.getInDimSize(rowColDims[0]), tileMult.getOutDimSize(rowColDims[0]))).resizeOutDim(rowColDims[1], std::min(ll.getInDimSize(rowColDims[1]), tileMult.getOutDimSize(rowColDims[1]))); + auto comp = tileMultShrunk .compose(ll) .sublayout({kReg, kLane}, to_vector(ll.getOutDimNames())); if (instr32Rows) { @@ -269,7 +276,7 @@ // basis comp = comp.resizeInDim(kLane, comp.getInDimSize(kLane) / 2); } - if (!comp.isInjective()) + if (!comp.isInjective() && nColsLL > 32) return std::nullopt; // Fit the warp bases either tiling on the RHS or in row=16 @@ -324,9 +331,8 @@ // Add CTAs as a trivial map auto nCTAs = ll.getInDimSize(kBlock); tile *= LinearLayout::identity1D(nCTAs, kBlock, kBlock); - assert(tile.getOutDimSize(rowColDims[1]) == ll.getInDimSize(rowColDims[1])); - - auto ret = tile.compose(ll); + auto tileShrunk = tile.resizeOutDim(rowColDims[0], std::min(ll.getInDimSize(rowColDims[0]), tile.getOutDimSize(rowColDims[0]))).resizeOutDim(rowColDims[1], std::min(ll.getInDimSize(rowColDims[1]), tile.getOutDimSize(rowColDims[1]))); + auto ret = tileShrunk.compose(ll); return ret; } @@ -349,6 +355,18 @@ auto *ctx = memType.getContext(); auto layout = getDistributedLayoutForTmemLdSt( memType, TMemAccessAtom::I32x32b, numWarps); + if (!layout) { + layout = getDistributedLayoutForTmemLdSt( + memType, TMemAccessAtom::I16x256b, numWarps); + } + if (!layout) { + layout = getDistributedLayoutForTmemLdSt( + memType, TMemAccessAtom::I16x128b, numWarps); + } + if (!layout) { + layout = getDistributedLayoutForTmemLdSt( + memType, TMemAccessAtom::I16x64b, numWarps); + } assert(layout); return LinearEncodingAttr::get(ctx, std::move(*layout)); } @@ -422,6 +440,10 @@ bool isDistributedLayoutTMemCompatible(Operation *op, RankedTensorType tensorType, gpu::MemDescType memType) { + for (auto l : getTmemCompatibleLayouts(op, tensorType, memType)) { + if (l == tensorType.getEncoding()) + return true; + } auto maxnreg = getContextualMaxNReg(op); return succeeded(computeTMemLdStEncodingInfo(tensorType, memType, maxnreg)); } --- a/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp +++ b/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp @@ -117,7 +117,7 @@ } FailureOr -lowerTMemLdSt(const LinearLayout &cvt, int maxnreg, int bitwidth, +lowerTMemLdSt(const LinearLayout &cvt, int maxnreg, int bitwidth, bool isScales, std::function emitError, bool unpacked = false) { // We will fill in the returned value recursively (if it exists) @@ -126,7 +126,7 @@ auto removeBroadcastSrc = actionRemoveBroadcastedRegs(cvt); if (!removeBroadcastSrc.isIdentity()) return lowerTMemLdSt(removeBroadcastSrc.apply(cvt), maxnreg, bitwidth, - emitError, unpacked); + isScales, emitError, unpacked); auto *ctx = cvt.getInDimNames().begin()->getContext(); auto S = [ctx](StringRef str) { return StringAttr::get(ctx, str); }; auto kReg = S("register"); @@ -139,6 +139,9 @@ for (int contig = 1; bitwidth * contig <= 32; contig *= 2) { auto maybeQuot = divideLeft(cvt, LinearLayout::identity1D(contig, kReg, kCol)); + if (!maybeQuot) { + maybeQuot = divideLeft(cvt, LinearLayout::identity1D(contig, kReg, kRow)); + } if (!maybeQuot) break; quot = *maybeQuot; @@ -182,7 +185,8 @@ if (unpacked) { quot = LinearLayout::zeros1D(1, kReg, kCol, 32 / bitwidth) * quot; } - auto info = lowerTMemLdSt(quot, maxnreg, newBitwidth, emitError, unpacked); + auto info = lowerTMemLdSt(quot, maxnreg, newBitwidth, isScales, emitError, + unpacked); if (failed(info)) return failure(); if (bestContig > 1) { @@ -231,6 +235,19 @@ auto row = reps.getBasis(kLane, 4, kRow); auto col = reps.getBasis(kLane, 4, kCol); secondHalfOffset = (row << 16) | col; + if (*secondHalfOffset == 0) { + // Workaround for ptxas bug, we cannot use secondHalfOffset = 0 to write + // only 16 elements. We use secondHalfOffset = 1 instead and we pad the + // allocation. + if (!isScales) { + if (emitError) { + emitError() + << "Only supported for scales as we pad the allocation."; + } + return failure(); + } + secondHalfOffset = 1; + } // We "quotient it out", meaning we remove the last basis from reps auto basis = reps.getBases(); basis[kLane][4] = {0, 0}; @@ -295,8 +312,9 @@ cvt = LinearLayout(std::move(bases), cvt.getOutDims(), /*isSurjective=*/cvt.isSurjective()); + bool isScales = isa(memTy.getEncoding()); int bitwidth = memTy.getElementTypeBitWidth(); - return lowerTMemLdSt(cvt, maxnreg, bitwidth, emitError); + return lowerTMemLdSt(cvt, maxnreg, bitwidth, isScales, emitError); } bool supportsTMemLoadReduce(RankedTensorType regTy, MemDescType memTy, --- a/python/test/gluon/test_fpsan.py +++ b/python/test/gluon/test_fpsan.py @@ -1870,6 +1870,8 @@ def _native_mma_k(type_a: str) -> int: + if is_cuda() and type_a == "f32": + return 16 return 256 // _float_dtype_info(type_a)[0] --- a/python/test/unit/language/test_compile_only.py +++ b/python/test/unit/language/test_compile_only.py @@ -213,7 +213,7 @@ assert re.search(pattern, str(ttgir)), "The TTGIR does not match the expected pattern." ptx = k.asm["ptx"] - pattern = (r"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32") + pattern = (r"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X") assert re.search(pattern, str(ptx)), "The PTX does not match the expected pattern." assert k.asm["cubin"] != b"" --- a/python/test/unit/language/test_core.py +++ b/python/test/unit/language/test_core.py @@ -3673,10 +3673,7 @@ if not is_cuda(): return [] return [(2, 4, 32, 1, False, False, 'None', 'ieee', 'float16', 'float32', 1, None), - (1, 2, 32, 1, False, False, 'None', 'ieee', 'float8e5', 'float32', 1, None), - # N=8: TF32 K=8 (wgmma.m64n8k8, sm90+) and FP16 K=16 (wgmma.m64n8k16) - (64, 8, 8, 4, False, False, 'None', 'tf32', 'float32', 'float32', 1, None), - (64, 8, 16, 4, False, False, 'None', 'ieee', 'float16', 'float32', 1, None)] + (1, 2, 32, 1, False, False, 'None', 'ieee', 'float8e5', 'float32', 1, None)] @pytest.mark.interpreter @@ -3911,21 +3908,19 @@ # make sure ld/st are vectorized ptx = pgm.asm['ptx'] - # XXX: skip small sizes because they are not vectorized; with runtime - # strides, v4 needs the contiguous dim >= 16 (K for loads, N for stores). - enough_work = (M * N // (num_warps * 32) >= 4) and (K > 16 or N > 16 or M > 16) - if enough_work and K >= 16: - if 'float64' in in_dtype: - assert 'ld.global.v2.b64' in ptx - else: - assert 'ld.global.v4' in ptx - if enough_work and N >= 16: - if 'float8' in in_dtype: - assert 'st.global.v2' in ptx - elif 'float64' in in_dtype: - assert 'st.global.v2.b64' in ptx - else: - assert 'st.global.v4' in ptx + if (K > 16 or N > 16 or M > 16) and (M * N // (num_warps * 32) >= 4): + if K >= 16: + if 'float64' in in_dtype: + assert 'ld.global.v2.b64' in ptx + else: + assert 'ld.global.v4' in ptx + if N >= 16: + if 'float8' in in_dtype: + assert 'st.global.v2' in ptx + elif 'float64' in in_dtype: + assert 'st.global.v2.b64' in ptx + else: + assert 'st.global.v4' in ptx is_tcgen5 = (capability[0] == 10) and (num_warps % 4) == 0 and (M % 64) == 0 and (N % 8) == 0 --- a/python/test/unit/language/test_matmul.py +++ b/python/test/unit/language/test_matmul.py @@ -402,6 +402,8 @@ pytest.skip(f"CDNA4 does not support {BLOCK_K=} for scaled mfma {nonKDim=} variants") if (BLOCK_M == 256 or BLOCK_N == 256) and BLOCK_K == 256: pytest.skip("Config requires too much shared memory") + if is_cuda() and torch.cuda.get_device_capability()[0] >= 10 and BLOCK_N < 32: + pytest.skip("Blackwell TMEM scaling instructions require aligned BLOCK_N >= 32") if not is_rubin() and BLOCK_N == 256 and BLOCK_K == 256: NUM_STAGES = min(NUM_STAGES, 2) --- a/python/tutorials/09-persistent-matmul.py +++ b/python/tutorials/09-persistent-matmul.py @@ -62,7 +62,7 @@ def is_hopper(): - return torch.cuda.get_device_capability()[0] == 9 + return torch.cuda.get_device_capability()[0] >= 9 def supports_ws(): --- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir +++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir @@ -302,7 +302,7 @@ // CHECK-LABEL: @tc_gen5_mma_scaled_a_tmem // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32 // CHECK: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,r,l,r,r,r,b,b" %[[TMEM_BASE]], %[[A_BASE]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,r,l,r,r,r,b,b" %[[TMEM_BASE]], %[[A_BASE]] tt.func @tc_gen5_mma_scaled_a_tmem( %a: !ttg.memdesc<128x256xf8E5M2, #tmem, #ttng.tensor_memory>, %b: !ttg.memdesc<256x64xf8E5M2, #shared, #ttg.shared_memory>, @@ -337,10 +337,10 @@ // CHECK: %[[P1:.+]] = llvm.and %{{.*}}, %[[P0]] : i1 // CHECK: llvm.cond_br %[[P1]] // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144708608 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5 + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5 // CHECK: %[[TRUE:.+]] = llvm.mlir.constant(true) : i1 // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681579536 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]], %{{.+}}, %{{.+}}, %[[TRUE]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]], %{{.+}}, %{{.+}}, %[[TRUE]] tt.func @tc_gen5_mma_block_scale(%a: !ttg.memdesc<128x64xf8E4M3FN, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<32x128xi8, #shared1, #ttg.shared_memory>, %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, @@ -371,13 +371,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { // CHECK-LABEL: @tc_gen5_mma_block_scale_fp4_a // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144769664 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]] // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681640592 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]] // CHECK: %[[DESC2:.+]] = llvm.mlir.constant(1218511520 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]] // CHECK: %[[DESC3:.+]] = llvm.mlir.constant(1755382448 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]] tt.func @tc_gen5_mma_block_scale_fp4_a(%a: !ttg.memdesc<128x64xi8, #shared1, #ttg.shared_memory>, %b: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>, %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, @@ -540,8 +540,8 @@ // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4 // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<3> to i32 // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(138413184 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]] - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]] tt.func @tc_gen5_mma_block_scale_nvfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>, %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>, @@ -573,9 +573,9 @@ // CHECK-LABEL: @tc_gen5_mma_block_scale_mxfp4 // CHECK-DAG: %[[TMEM_BASE:.+]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<3> to i32 // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(146801792 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]] // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(1220543648 : i32) : i32 - // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]] + // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]] tt.func @tc_gen5_mma_block_scale_mxfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>, %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>, @@ -1099,7 +1099,7 @@ // CHECK-LABEL: @load_store_16x32bx1_broadcast tt.func private @load_store_16x32bx1_broadcast(%arg0: !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x8xi8, #linear>) { %true = arith.constant true - // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 0, {$2} + // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 1, {$2} ttng.tmem_store %arg1, %arg0, %true : tensor<16x8xi8, #linear> -> !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable> tt.return } @@ -1114,7 +1114,7 @@ // CHECK-LABEL: @load_store_16x2_scales_uses_zero_second_half tt.func private @load_store_16x2_scales_uses_zero_second_half(%arg0: !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x2xi8, #linear>) { %true = arith.constant true - // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 0, {$2} + // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 1, {$2} // CHECK: llvm.return ttng.tmem_store %arg1, %arg0, %true : tensor<16x2xi8, #linear> -> !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable> tt.return @@ -1178,22 +1178,22 @@ %pred: i1) { // CHECK: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg4 : !llvm.ptr<3> to i32 // CHECK: %[[A_E4M3_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32 - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 8 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 16 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 24 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]] // CHECK: %[[A_E5M2_BASE:.+]] = llvm.ptrtoint %arg1 : !llvm.ptr<3> to i32 - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 8 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 16 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 24 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]] ttng.tc_gen5_mma_scaled %a_e4m3, %b_e5m2, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e4m3 rhs = e5m2 : !ttg.memdesc<128x128xf8E4M3FN, #tmem_fp8_lhs, #ttng.tensor_memory>, @@ -1221,13 +1221,13 @@ %pred: i1) { // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32 // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32 - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 8 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 16 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 24 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e5m2 : !ttg.memdesc<128x64xi8, #tmem_fp4_padded_lhs, #ttng.tensor_memory>, @@ -1257,13 +1257,13 @@ %pred: i1) { // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32 // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32 - // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 0 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 8 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 16 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] - // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2 + // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 24 ], $2 // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]] ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e2m1 : !ttg.memdesc<128x128xi8, #tmem_fp4_dense_lhs, #ttng.tensor_memory>, --- a/test/Conversion/tritongpu_to_llvm_rubin.mlir +++ b/test/Conversion/tritongpu_to_llvm_rubin.mlir @@ -8,7 +8,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} { // CHECK-LABEL: @tc_gen5_mma_block_scale_fp8_sm107 // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32 - // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 + // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X // CHECK-NOT: tcgen05.mma tt.func @tc_gen5_mma_block_scale_fp8_sm107(%a: !ttg.memdesc<128x128xf8E4M3FN, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<64x128xi8, #shared1, #ttg.shared_memory>, @@ -41,7 +41,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} { // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4_sm107_m256 // CHECK: %[[DESC:.+]] = llvm.mlir.constant(136316040 : i32) : i32 - // CHECK-COUNT-4: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block32 + // CHECK-COUNT-4: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::2X // CHECK-NOT: tcgen05.mma tt.func @tc_gen5_mma_block_scale_nvfp4_sm107_m256(%a: !ttg.memdesc<256x128xi8, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<128x128xi8, #shared1, #ttg.shared_memory>, @@ -103,7 +103,7 @@ #tmem_scales1 = #ttng.tensor_memory_scales_encoding module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107", "ttg.threads-per-warp" = 32 : i32, "ttng.two-ctas" = true} { // CHECK-LABEL: @tc_gen5_mma_scaled_fp4_padded_shared_linear_a - // CHECK-COUNT-4: tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale.block32 + // CHECK-COUNT-4: tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale.scale_vec::1X // CHECK-NOT: tcgen05.mma tt.func @tc_gen5_mma_scaled_fp4_padded_shared_linear_a( %a: !ttg.memdesc<256x64xi8, #shared_fp4_padded_sl, #smem, mutable>, @@ -158,8 +158,8 @@ #tmem_scales = #ttng.tensor_memory_scales_encoding<> module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} { // CHECK-LABEL: @tc_gen5_mma_block_scale_breuse - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32.collector::b::fill - // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32.collector::b::lastuse + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X.collector::b::fill + // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X.collector::b::lastuse tt.func @tc_gen5_mma_block_scale_breuse(%a: !ttg.memdesc<256x64xf8E4M3FN, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<32x128xi8, #shared1, #ttg.shared_memory>, %c: !ttg.memdesc<256x128xf32, #tmem, #ttng.tensor_memory, mutable>, @@ -195,7 +195,7 @@ // scale kind in the descriptor for both K=128 sub-instructions of a // logical BLOCK_K=256 tile. // CHECK: %[[DESC:.+]] = llvm.mlir.constant(153093256 : i32) : i32 - // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 + // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X // CHECK-NOT: tcgen05.mma tt.func @tc_gen5_mma_block_scale_nvfp4_ue5m3_sm107(%a: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>, %b: !ttg.memdesc<128x128xi8, #shared1, #ttg.shared_memory>, --- a/test/TritonGPU/accelerate-matmul.mlir +++ b/test/TritonGPU/accelerate-matmul.mlir @@ -1052,40 +1052,6 @@ // ----- -// Verify TF32 dot with N=8, K=8 (native WGMMA tile) selects MMAv3 on sm90. -// CHECK: #[[$MMA:.+]] = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 8, 8]}> -#blocked_tf32_n8 = #ttg.blocked<{sizePerThread = [2, 2], threadsPerWarp = [8, 4], warpsPerCTA = [4, 1], order = [1, 0]}> -module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { - // CHECK-LABEL: tf32_n8k8_sm90 - tt.func public @tf32_n8k8_sm90( - %a: tensor<64x8xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked_tf32_n8}>>, - %b: tensor<8x8xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked_tf32_n8}>>) -> tensor<64x8xf32, #blocked_tf32_n8> { - %cst = arith.constant dense<0.000000e+00> : tensor<64x8xf32, #blocked_tf32_n8> - // CHECK: ttng.warp_group_dot {{.*}} -> tensor<64x8xf32, #[[$MMA]]> - %d = tt.dot %a, %b, %cst, inputPrecision = tf32 : tensor<64x8xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked_tf32_n8}>> * tensor<8x8xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked_tf32_n8}>> -> tensor<64x8xf32, #blocked_tf32_n8> - tt.return %d : tensor<64x8xf32, #blocked_tf32_n8> - } -} - -// ----- - -// Verify FP16 dot with N=8, K=16 (native WGMMA tile) selects MMAv3 on sm90. -// CHECK: #[[$MMA:.+]] = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 8, 16]}> -#blocked_fp16_n8 = #ttg.blocked<{sizePerThread = [2, 2], threadsPerWarp = [8, 4], warpsPerCTA = [4, 1], order = [1, 0]}> -module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { - // CHECK-LABEL: fp16_n8k16_sm90 - tt.func public @fp16_n8k16_sm90( - %a: tensor<64x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked_fp16_n8}>>, - %b: tensor<16x8xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked_fp16_n8}>>) -> tensor<64x8xf32, #blocked_fp16_n8> { - %cst = arith.constant dense<0.000000e+00> : tensor<64x8xf32, #blocked_fp16_n8> - // CHECK: ttng.warp_group_dot {{.*}} -> tensor<64x8xf32, #[[$MMA]]> - %d = tt.dot %a, %b, %cst : tensor<64x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked_fp16_n8}>> * tensor<16x8xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked_fp16_n8}>> -> tensor<64x8xf32, #blocked_fp16_n8> - tt.return %d : tensor<64x8xf32, #blocked_fp16_n8> - } -} - -// ----- - #blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [0, 1]}> #blocked1 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [16, 2], warpsPerCTA = [4, 1], order = [1, 0]}> #blocked2 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [2, 16], warpsPerCTA = [4, 1], order = [1, 0]}> --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -29,8 +29,6 @@ return (1, 1, 32) elif lhs_bitwidth == 64: return (1, 1, 4) - elif lhs_bitwidth == 32: - return (1, 1, 8) else: return (1, 1, 16) --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp @@ -398,12 +398,12 @@ std::string opcode = "tcgen05.mma.cta_group::" + std::to_string(twoCTAs ? 2 : 1) + ".kind::"; if (mxfpInstKind == mxfpKind::mxf8f6f4) { - opcode += "mxf8f6f4.block_scale.block32"; + opcode += "mxf8f6f4.block_scale.scale_vec::1X"; } else if (mxfpInstKind == mxfpKind::mxf4) { - opcode += "mxf4.block_scale.block32"; + opcode += "mxf4.block_scale.scale_vec::2X"; } else if (mxfpInstKind == mxfpKind::mxf4nvf4) { - opcode += getScaleVecSize(op) == 32 ? "mxf4nvf4.block_scale.block32" - : "mxf4nvf4.block_scale.block16"; + opcode += getScaleVecSize(op) == 32 ? "mxf4nvf4.block_scale.scale_vec::2X" + : "mxf4nvf4.block_scale.scale_vec::4X"; } else { assert(0 && "Unsupported mxfp kind."); } @@ -881,15 +881,18 @@ int scaleFactorColsPerSet = getScaleFactorColsPerSet(mxfpInstKind, op, dot.mmaSizeK); int numRepKWords = ceil(numRepK, 4 / scaleFactorColsPerSet); + int aMultibuffers = cast(op.getAScale().getType()).getRank() > 2 ? cast(op.getAScale().getType()).getDimSize(0) : 1; + int bMultibuffers = cast(op.getBScale().getType()).getRank() > 2 ? cast(op.getBScale().getType()).getDimSize(0) : 1; int numColPerScaleBlockA = ceil( ttng::getTmemAllocSizes(cast(op.getAScale().getType())) - .numCols, + .numCols / aMultibuffers, numRepM * numRepKWords); int numColPerScaleBlockB = ceil( ttng::getTmemAllocSizes(cast(op.getBScale().getType())) - .numCols, + .numCols / bMultibuffers, numRepN * numRepKWords); - numColPerScaleBlockB = std::max(numColPerScaleBlockB, 2); + if (desc.mmaSizeN >= 32) + numColPerScaleBlockB = std::max(numColPerScaleBlockB, 2); int subWordIdx = k % (4 / scaleFactorColsPerSet); int wordIdx = k / (4 / scaleFactorColsPerSet); int scaleIdxA = linearizeScaleBlockIdx(op.getAScale(), m, wordIdx, numRepM,