[Mlir-commits] [mlir] [MLIR][NVVM] Add tcgen05.mma MLIR Ops (PR #164356)
Durgadoss R
llvmlistbot at llvm.org
Fri Oct 24 03:17:26 PDT 2025
================
@@ -4586,6 +4586,567 @@ def NVVM_ClusterLaunchControlQueryCancelOp
}];
}
+//===----------------------------------------------------------------------===//
+// NVVM tcgen05.mma attributes
+//===----------------------------------------------------------------------===//
+
+def Tcgen05MMAKindF16 : I32EnumAttrCase<"F16", 0, "f16">;
+def Tcgen05MMAKindTF32 : I32EnumAttrCase<"TF32", 1, "tf32">;
+def Tcgen05MMAKindF8F6F4 : I32EnumAttrCase<"F8F6F4", 2, "f8f6f4">;
+def Tcgen05MMAKindINT8 : I32EnumAttrCase<"I8", 3, "i8">;
+
+def Tcgen05MMAKind : I32EnumAttr<
+ "Tcgen05MMAKind",
+ "tcgen05 MMA Supported Types",
+ [Tcgen05MMAKindF8F6F4, Tcgen05MMAKindINT8, Tcgen05MMAKindF16,
+ Tcgen05MMAKindTF32]> {
+ let cppNamespace = "::mlir::NVVM";
+ let genSpecializedAttr = 0;
+}
+
+def Tcgen05MMAKindAttr : EnumAttr<NVVM_Dialect, Tcgen05MMAKind, "tcgen05_mma_kind"> {
+ let description = [{
+ The Tcgen05MMAKind attribute describes the allowed set of types for matrix A and B in the tcgen05.mma.{sp} Op. The following are supported types for each kind:
+
+ ```
+ +-------------+--------------------------------------------+
+ | Matrix Kind | supported types for A / B |
+ +-------------+--------------------------------------------+
+ | f16 | f16, bf16 |
+ | tf32 | tf32 |
+ | f8f6f4 | e4m3, e5m2, e2m3, e3m2, e2m1 |
+ | i8 | unsigned 8b, signed 8b |
+ +-------------+--------------------------------------------+
+ ```
+ }];
+ let assemblyFormat = "`<` $value `>`";
+}
+
+def Tcgen05MMACollectorOpDiscard : I32EnumAttrCase<"DISCARD", 0, "discard">;
+def Tcgen05MMACollectorOpLastUse : I32EnumAttrCase<"LASTUSE", 1, "lastuse">;
+def Tcgen05MMACollectorOpFill : I32EnumAttrCase<"FILL", 2, "fill">;
+def Tcgen05MMACollectorOpUse : I32EnumAttrCase<"USE", 3, "use">;
+
+def Tcgen05MMACollectorOp : I32EnumAttr<
+ "Tcgen05MMACollectorOp",
+ "tcgen05.mma Collector Buffer Operation",
+ [Tcgen05MMACollectorOpDiscard,
+ Tcgen05MMACollectorOpLastUse,
+ Tcgen05MMACollectorOpFill,
+ Tcgen05MMACollectorOpUse]> {
+ let cppNamespace = "::mlir::NVVM";
+ let genSpecializedAttr = 0;
+}
+
+def Tcgen05MMACollectorOpAttr : EnumAttr<NVVM_Dialect, Tcgen05MMACollectorOp, "tcgen05_mma_collectorop"> {
+ let description = [{
+ Tcgen05MMACollectorOp attribute specifies the collector buffer operations.
+ The following are the supported operations:
+ * discard : Release buffer after use (default)
+ * lastuse : Mark buffer for last use
+ * fill : Fill buffer
+ * use : Use buffer without modification
+ }];
+ let assemblyFormat = "`<` $value `>`";
+}
+
+//===----------------------------------------------------------------------===//
+// NVVM tcgen05.mma Ops.
+//===----------------------------------------------------------------------===//
+def NVVM_Tcgen05MMAOp : NVVM_Op<"tcgen05.mma",
+ [AttrSizedOperandSegments,
+ NVVMRequiresSMa<[100, 110]>]> {
+ let summary = "Performs MMA operation on 5th-gen tensor cores";
+
+ let description = [{
+ The `tcgen05.mma` operation is an asynchronous tensor core instruction that
+ 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 + (D * 2^ -scaleInputD) // if `scaleInputD` is provided
+ 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)
+
+ Optional Operands:
+ - `scaleInputD` is an Immediate value operand used for scaling D matrix by 2 ^ (-scaleInputD). The valid range is [0, 15]
+
+ - `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:
+ - `kind` is a Tcgen05MMAKind attribute
+
+ - `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:
+ - collectorOp is a Tcgen05MMACollectorOp attribute with matrix A as the collector buffer
+
+ - `aShift` shifts the rows of the A matrix down by one row and can only be
+ applied if A is in tensor memory
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma)
+ }];
+
+ let arguments = (ins
+ Tcgen05MMAKindAttr:$kind,
+ CTAGroupKindAttr:$ctaGroup,
+ DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
+ "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
+ UnitAttr:$aShift,
+ LLVM_PointerTensor:$matrixD,
+ AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
+ I64:$matrixB,
+ I32:$idesc,
+ I1:$enableInputD,
+ Optional<I64>:$scaleInputD,
+ Optional<FixedVectorOfLengthAndType<[4, 8], [I32]>>:$disableOutputLane
+ );
+
+ let assemblyFormat = [{
+ $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD (`scale` `=` $scaleInputD^)?
+ (`mask` `=` $disableOutputLane^)? attr-dict `:` `(` type(operands) `)`
+ }];
+
+ let hasVerifier = true;
+
+ let extraClassDeclaration = [{
+ static mlir::NVVM::IDArgPair getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder);
+ }];
+
+ let llvmBuilder = [{
+ auto [ID, args] = NVVM::Tcgen05MMAOp::getIntrinsicIDAndArgs(
+ *op, moduleTranslation, builder);
+ createIntrinsicCall(builder, ID, args);
+ }];
+}
+
+def NVVM_Tcgen05MMASparseOp : NVVM_Op<"tcgen05.mma.sp",
+ [AttrSizedOperandSegments,
+ NVVMRequiresSMa<[100, 110]>]> {
+ let summary = "Performs MMA operation with sparse A matrix on 5th-gen tensor cores";
+
+ let description = [{
+ The `tcgen05.mma.sp` operation is an asynchronous tensor core instruction
+ that performs matrix multiplication, accumulation with sparse `A` matrix in
+ a single fused operation. It targets 5th-generation tensor cores, providing
+ developers with fine-grained control over execution and scheduling.
+
+ ```
+ D = A * B + (D * 2^ -scaleInputD) // if `scaleInputD` is provided
+ D = A * B // if `enableInputD` is false
+ D = A * B + D // otherwise
+ ```
+
+ 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
+ - sparseMetadata located in tensor memory specifies the mapping of the `K / 2`
+ non-zero elements to the K elements before performing the MMA operation
+
+ Other attributes and operands are similar to that of tcgen05.mma Op
+
+ [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-sp)
+ }];
+
+ let arguments = (ins
+ Tcgen05MMAKindAttr:$kind,
+ CTAGroupKindAttr:$ctaGroup,
+ DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
+ "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
+ UnitAttr:$aShift,
+ LLVM_PointerTensor:$matrixD,
+ AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
+ I64:$matrixB,
+ I32:$idesc,
+ I1:$enableInputD,
+ LLVM_PointerTensor:$sparseMetadata,
+ Optional<I64>:$scaleInputD,
+ Optional<FixedVectorOfLengthAndType<[4, 8], [I32]>>:$disableOutputLane
+ );
+
+ let assemblyFormat = [{
+ $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD `,` $sparseMetadata (`scale` `=` $scaleInputD^)? (`mask` `=` $disableOutputLane^)? attr-dict `:` `(` type(operands) `)`
+ }];
----------------
durga4github wrote:
Thanks for moving the asm-format right next to the args. It is much easier to read and relate quickly
https://github.com/llvm/llvm-project/pull/164356
More information about the Mlir-commits
mailing list