[Mlir-commits] [mlir] [MLIR][NVVM] Support tcgen05.mma{.block_scale}.decompress_b Ops (PR #218354)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 24 01:59:49 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm
Author: Kirill Vedernikov (kvederni)
<details>
<summary>Changes</summary>
This change adds support for `tcgen05.mma.decompress_b` and `tcgen05.mma.block_scale.decompress_b` MLIR Ops.
---
Patch is 127.92 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/218354.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+152)
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+122)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-block-scale-shared-decompress-b.mlir (+157)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-block-scale-tensor-decompress-b.mlir (+157)
- (modified) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir (+11)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-shared-decompress-b.mlir (+307)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-tensor-decompress-b.mlir (+307)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 6ac489f9ad2bc..f0161173421eb 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6816,6 +6816,158 @@ def NVVM_Tcgen05MMAWsSparseOp : NVVM_Op<"tcgen05.mma.ws.sp",
}];
}
+def NVVM_Tcgen05MMADecompressBOp :
+ NVVM_VoidIntrinsicOp<"tcgen05.mma.decompress_b",
+ [NVVMRequiresSMf<[107]>]> {
+ let summary = "Performs MMA operation with pre-compressed B matrix on 5th-gen tensor cores";
+
+ let description = [{
+ The `tcgen05.mma.decompress_b` operation is an asynchronous tensor core
+ instruction that decompresses the B matrix and performs matrix
+ multiplication, accumulation in a single fused operation. It targets
+ 5th-generation tensor cores, providing developers with fine-grained
+ control over execution and scheduling.
+
+ ```
+ D = A * B // if `enableInputD` is false
+ D = A * B + D // otherwise
+ ```
+
+ where:
+ - A is an `M x K` matrix in tensor memory or described using shared memory descriptor
+ - B is a `K x N` matrix described using shared memory descriptor
+ - D is an `M x N` accumulator matrix in tensor memory
+
+ The `shared memory descriptor` can be generated using `tcgen05.mma_smem_desc` Op
+
+ - `idesc` is a 32-bit value representing the [Instruction Descriptor](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instruction-descriptor)
+
+ - `decompressBMetadata` is a decompress metadata of B matrix
+
+ Optional Operands:
+ - `disableOutputLane` is a vector mask for selective output
+ * vector<4 x i32> when ctaGroup is CTA_1
+ * vector<8 x i32> when ctaGroup is CTA_2
+
+ Required Attributes:
+ - `ctaGroup` specifies CTA group configuration
+ * cta_1: MMA will be performed on the current thread's CTA
+ * cta_2: MMA will be performed on the current thread and it's peer CTA
+
+ Default Attributes:
+ - collectorOpA is a Tcgen05MMACollectorOp attribute with matrix A as the collector buffer
+ - collectorOpB is a Tcgen05MMACollectorOp attribute with matrix B as the collector buffer
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma)
+ }];
+
+ let arguments = (ins
+ LLVM_PointerTensor:$matrixD,
+ AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
+ I64:$matrixB,
+ I32:$idesc,
+ I1:$enableInputD,
+ LLVM_PointerTensor:$decompressBMetadata,
+ Optional<FixedVectorOfLengthAndType<[4, 8], [I32]>>:$disableOutputLane,
+ CTAGroupKindAttr:$ctaGroup,
+ DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
+ "Tcgen05MMACollectorOp::DISCARD">:$collectorOpA,
+ DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
+ "Tcgen05MMACollectorOp::DISCARD">:$collectorOpB
+ );
+
+ let assemblyFormat = [{
+ $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD
+ `,` $decompressBMetadata (`,` `mask` `=` $disableOutputLane^)?
+ `cta_group` `=` $ctaGroup
+ oilist(`collector_a` `=` $collectorOpA | `collector_b` `=` $collectorOpB)
+ attr-dict `:` `(` type(operands) `)`
+ }];
+
+ let hasVerifier = true;
+}
+
+defvar Tcgen05MMABlockScaleDecompressBScaleList = [
+ Tcgen05MMABlockScaleDefault,
+ Tcgen05MMABlockScaleBlock32
+ ];
+
+defvar Tcgen05MMABlockScaleDecompressBScaleAttr =
+ ConfinedAttr<Tcgen05MMABlockScaleAttr,
+ [EnumAttrIsOneOf<Tcgen05MMABlockScaleAttr,
+ Tcgen05MMABlockScaleDecompressBScaleList>]>;
+
+def NVVM_Tcgen05MMABlockScaleDecompressBOp :
+ NVVM_VoidIntrinsicOp<"tcgen05.mma.block_scale.decompress_b",
+ [NVVMRequiresSMf<[107]>]> {
+ let summary = "Performs block scaled MMA operation with pre-compressed B matrix on 5th-gen tensor cores";
+
+ let description = [{
+ The `tcgen05.mma.block_scale.decompress_b` operation is an asynchronous tensor core
+ instruction that decompresses the B matrix and performs matrix
+ multiplication, accumulation with block scaling in a single fused operation.
+ It targets 5th-generation tensor cores, providing developers with
+ fine-grained control over execution and scheduling.
+
+ ```
+ D = (A * scale_a) * (B * scale_b)` // if `enableInputD` is false
+ D = (A * scale_a) * (B * scale_b) + D`
+ ```
+
+ where:
+ - A is an M x (K / 2) matrix in tensor memory or described using shared memory descriptor
+ - B is a `K x N` matrix described using shared memory descriptor
+ - D is an `M x N` accumulator matrix in tensor memory
+ - `scale_a` and `scale_b` are matrices in tensor memory used to scale `A` and `B` respectively
+
+ The `shared memory descriptor` can be generated using `tcgen05.mma_smem_desc` Op
+
+ - `idesc` is a 32-bit value representing the [Instruction Descriptor](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instruction-descriptor)
+
+ - `decompressBMetadata` is a decompress metadata of B matrix
+
+ Required Attributes:
+ - `ctaGroup` specifies CTA group configuration
+ * cta_1: MMA will be performed on the current thread's CTA
+ * cta_2: MMA will be performed on the current thread and it's peer CTA
+
+ Default Attributes:
+ - collectorOpA is a Tcgen05MMACollectorOp attribute with matrix A as the collector buffer
+ - collectorOpB is a Tcgen05MMACollectorOp attribute with matrix B as the collector buffer
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma)
+ }];
+
+ let arguments = (ins
+ LLVM_PointerTensor:$matrixD,
+ AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
+ I64:$matrixB,
+ I32:$idesc,
+ I1:$enableInputD,
+ LLVM_PointerTensor:$scaleA,
+ LLVM_PointerTensor:$scaleB,
+ LLVM_PointerTensor:$decompressBMetadata,
+ CTAGroupKindAttr:$ctaGroup,
+ DefaultValuedAttr<Tcgen05MMABlockScaleDecompressBScaleAttr,
+ "Tcgen05MMABlockScale::DEFAULT">:$blockScale,
+ DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
+ "Tcgen05MMACollectorOp::DISCARD">:$collectorOpA,
+ DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
+ "Tcgen05MMACollectorOp::DISCARD">:$collectorOpB
+ );
+
+ let assemblyFormat = [{
+ $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD
+ `,` $scaleA `,` $scaleB `,` $decompressBMetadata
+ `cta_group` `=` $ctaGroup
+ oilist(
+ `block_scale` `=` $blockScale
+ | `collector_a` `=` $collectorOpA
+ | `collector_b` `=` $collectorOpB
+ ) attr-dict `:` `(` type(operands) `)`
+ }];
+}
+
def SIMTFloatType : AnyTypeOf<[F16, BF16, F32, F64,
VectorOfLengthAndType<[2], [F16, BF16, F32, F64]>]>;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 4dddb33fd050e..ef807e8b8f77d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -6513,6 +6513,128 @@ mlir::NVVM::IDArgPair Tcgen05MMAWsSparseOp::getIntrinsicIDAndArgs(
return {ID, args};
}
+//===----------------------------------------------------------------------===//
+// NVVM tcgen05.mma.decompress_b functions
+//===----------------------------------------------------------------------===//
+
+mlir::NVVM::IDArgPair Tcgen05MMADecompressBOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<Tcgen05MMADecompressBOp>(op);
+ llvm::SmallVector<llvm::Value *> args;
+
+ args.push_back(mt.lookupValue(thisOp.getMatrixD()));
+
+ llvm::Value *A = mt.lookupValue(thisOp.getMatrixA());
+ const bool isATensor = isa<llvm::PointerType>(A->getType());
+ args.push_back(A);
+
+ args.push_back(mt.lookupValue(thisOp.getMatrixB()));
+ args.push_back(mt.lookupValue(thisOp.getIdesc()));
+ args.push_back(mt.lookupValue(thisOp.getEnableInputD()));
+ args.push_back(mt.lookupValue(thisOp.getDecompressBMetadata()));
+
+ llvm::Value *DisableOutputLane =
+ mt.lookupValue(thisOp.getDisableOutputLane());
+ bool hasDisableOutputLane = DisableOutputLane != nullptr;
+
+ NVVM::CTAGroupKind ctaGroup = thisOp.getCtaGroup();
+
+ using namespace llvm::Intrinsic;
+ ID intrinsicID = not_intrinsic;
+
+ if (hasDisableOutputLane) {
+ if (ctaGroup == NVVM::CTAGroupKind::CTA_1) {
+ intrinsicID =
+ isATensor
+ ? nvvm_tcgen05_mma_tensor_f8f6f4_disable_output_lane_cg1_decompress_b
+ : nvvm_tcgen05_mma_shared_f8f6f4_disable_output_lane_cg1_decompress_b;
+ } else if (ctaGroup == NVVM::CTAGroupKind::CTA_2) {
+ intrinsicID =
+ isATensor
+ ? nvvm_tcgen05_mma_tensor_f8f6f4_disable_output_lane_cg2_decompress_b
+ : nvvm_tcgen05_mma_shared_f8f6f4_disable_output_lane_cg2_decompress_b;
+ } else {
+ llvm_unreachable("Unknown ctaGroup for tcgen05.mma.decompress_b");
+ }
+ } else {
+ intrinsicID = isATensor ? nvvm_tcgen05_mma_tensor_f8f6f4_decompress_b
+ : nvvm_tcgen05_mma_shared_f8f6f4_decompress_b;
+ }
+
+ assert(intrinsicID != not_intrinsic &&
+ "Invalid intrinsic for Tcgen05MMADecompressBOp.");
+
+ if (hasDisableOutputLane)
+ args.push_back(DisableOutputLane);
+ else
+ args.push_back(
+ builder.getInt32(static_cast<unsigned>(getNVVMCtaGroupKind(ctaGroup))));
+
+ args.push_back(
+ builder.getInt32(static_cast<unsigned>(thisOp.getCollectorOpA())));
+ args.push_back(
+ builder.getInt32(static_cast<unsigned>(thisOp.getCollectorOpB())));
+
+ return {intrinsicID, args};
+}
+
+LogicalResult Tcgen05MMADecompressBOp::verify() {
+ mlir::Value disableOutputLane = getDisableOutputLane();
+
+ if (disableOutputLane) {
+ NVVM::CTAGroupKind ctaGroup = getCtaGroup();
+
+ mlir::VectorType disableOutputLaneType =
+ cast<mlir::VectorType>(disableOutputLane.getType());
+ if ((ctaGroup == NVVM::CTAGroupKind::CTA_1 &&
+ disableOutputLaneType.getNumElements() != 4) ||
+ (ctaGroup == NVVM::CTAGroupKind::CTA_2 &&
+ disableOutputLaneType.getNumElements() != 8))
+ return emitOpError() << "Disable Output Lane of length "
+ << disableOutputLaneType.getNumElements()
+ << " is incompatible with CtaGroupAttr";
+ }
+
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// NVVM tcgen05.mma.block_scale.decompress_b functions
+//===----------------------------------------------------------------------===//
+
+mlir::NVVM::IDArgPair Tcgen05MMABlockScaleDecompressBOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<Tcgen05MMABlockScaleDecompressBOp>(op);
+ llvm::SmallVector<llvm::Value *> args;
+
+ args.push_back(mt.lookupValue(thisOp.getMatrixD()));
+
+ llvm::Value *A = mt.lookupValue(thisOp.getMatrixA());
+ const bool isATensor = isa<llvm::PointerType>(A->getType());
+ args.push_back(A);
+
+ args.push_back(mt.lookupValue(thisOp.getMatrixB()));
+ args.push_back(mt.lookupValue(thisOp.getIdesc()));
+ args.push_back(mt.lookupValue(thisOp.getEnableInputD()));
+ args.push_back(mt.lookupValue(thisOp.getScaleA()));
+ args.push_back(mt.lookupValue(thisOp.getScaleB()));
+ args.push_back(mt.lookupValue(thisOp.getDecompressBMetadata()));
+ args.push_back(builder.getInt32(
+ static_cast<unsigned>(getNVVMCtaGroupKind(thisOp.getCtaGroup()))));
+ args.push_back(
+ builder.getInt32(static_cast<unsigned>(thisOp.getCollectorOpA())));
+ args.push_back(
+ builder.getInt32(static_cast<unsigned>(thisOp.getCollectorOpB())));
+
+ using namespace llvm::Intrinsic;
+ ID intrinsicID =
+ isATensor
+ ? nvvm_tcgen05_mma_tensor_mxf8f6f4_block_scale_block32_decompress_b
+ : nvvm_tcgen05_mma_shared_mxf8f6f4_block_scale_block32_decompress_b;
+
+ return {intrinsicID, args};
+}
+
//===----------------------------------------------------------------------===//
// NVVM tcgen05.ld.red functions
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-block-scale-shared-decompress-b.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-block-scale-shared-decompress-b.mlir
new file mode 100644
index 0000000000000..5c394a465e290
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-block-scale-shared-decompress-b.mlir
@@ -0,0 +1,157 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_block_scale_shared_decompress_b_cta_1
+llvm.func @nvvm_tcgen05_mma_block_scale_shared_decompress_b_cta_1(
+ %d_tmem : !llvm.ptr<6>,
+ %a_desc : i64,
+ %b_desc : i64,
+ %idesc : i32,
+ %enable_input_d : i1,
+ %decompress_metadata: !llvm.ptr<6>,
+ %scale_a : !llvm.ptr<6>,
+ %scale_b : !llvm.ptr<6>) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=lastuse */ i32 1)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> block_scale = block32 collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=fill */ i32 2)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> collector_b = fill : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=use */ i32 3)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> block_scale = block32 collector_b = use : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=lastuse */ i32 1, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> collector_a = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=lastuse */ i32 1, /* collector_b=lastuse */ i32 1)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> block_scale = block32 collector_a = lastuse collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=lastuse */ i32 1, /* collector_b=fill */ i32 2)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> collector_a = lastuse collector_b = fill : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=lastuse */ i32 1, /* collector_b=use */ i32 3)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> block_scale = block32 collector_a = lastuse collector_b = use : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> collector_a = fill : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=lastuse */ i32 1)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> block_scale = block32 collector_a = fill collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale.block32.decompress_b(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=fill */ i32 2)
+ nvvm.tcgen05.mma.block_scale.decompress_b %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b, %decompress_metadata
+ cta_group = <cta_1> collector_a = fill collector_b = fill : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared.mxf8f6f4.block_scale...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/218354
More information about the Mlir-commits
mailing list