[NVVM][NVPTX] Add support for tcgen05.mma

This commit adds support for tcgen05.mma instructions in NVPTX which tests under CodeGen/NVPTX/tcgen05-mma*. This tcgen05.mma instructions are modeled as intrinsics with multiple arguments to model cta_group, mma kind, collector usage etc. The rationale for the design is present documented in NVPTXUsage.rst file

Author: schwarzschild-radius

llvm/docs/NVPTXUsage.rst

@@ -47,6 +47,15 @@ enum class CTAGroupKind : uint8_t {
   CG_2 = 2,    // cta_group::2 modifier
 };
 
+enum class Tcgen05MMAKind : uint8_t { F16 = 0, TF32 = 1, F8F6F4 = 2, I8 = 3 };
+
+enum class Tcgen05CollectorUsageOp : uint8_t {
+  DISCARD = 0,
+  LASTUSE = 1,
+  FILL = 2,
+  USE = 3,
+};
+
 inline bool FPToIntegerIntrinsicShouldFTZ(Intrinsic::ID IntrinsicID) {
   switch (IntrinsicID) {
   case Intrinsic::nvvm_f2i_rm_ftz:

llvm/include/llvm/IR/IntrinsicsNVVM.td

@@ -84,7 +84,44 @@ enum NodeType : unsigned {
   StoreV2,
   StoreV4,
   StoreV8,
-  LAST_MEMORY_OPCODE = StoreV8,
+  TCGEN05_MMA_SHARED_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SHARED_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SHARED_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SHARED_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SHARED_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SHARED_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_TENSOR_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_TENSOR_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_TENSOR_DISABLE_OUTPUT_LANE_CG1_ASHIFT,
+  TCGEN05_MMA_TENSOR_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
+  TCGEN05_MMA_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG1_ASHIFT,
+  TCGEN05_MMA_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
+  TCGEN05_MMA_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG1_ASHIFT,
+  TCGEN05_MMA_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
+  TCGEN05_MMA_SP_SHARED_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SP_SHARED_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SP_SHARED_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SP_SHARED_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SP_SHARED_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SP_SHARED_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SP_TENSOR_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SP_TENSOR_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SP_TENSOR_DISABLE_OUTPUT_LANE_CG1_ASHIFT,
+  TCGEN05_MMA_SP_TENSOR_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
+  TCGEN05_MMA_SP_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SP_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SP_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG1,
+  TCGEN05_MMA_SP_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2,
+  TCGEN05_MMA_SP_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG1_ASHIFT,
+  TCGEN05_MMA_SP_TENSOR_F16_SCALE_D_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
+  TCGEN05_MMA_SP_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG1_ASHIFT,
+  TCGEN05_MMA_SP_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
+  LAST_MEMORY_OPCODE =
+      TCGEN05_MMA_SP_TENSOR_TF32_SCALE_D_DISABLE_OUTPUT_LANE_CG2_ASHIFT,
 };
 }