[Mlir-commits] [mlir] [MLIR][NVVM] Support tcgen05.mma{.block_scale}.decompress_b Ops (PR #218354)

[Kirill Vedernikov]

================
@@ -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
----------------
kvederni wrote:

I'd like to use the names that are already used in other tcgen05.mma Ops. We can change all the names in a separate NFC PR.

https://github.com/llvm/llvm-project/pull/218354