[Mlir-commits] [mlir] [MLIR][NVVM] Support sp{de}compress MLIR Ops (PR #232731)

[Durgadoss R]

================
@@ -7200,6 +7200,282 @@ def NVVM_Tcgen05MMABlockScaleDecompressBOp :
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM spcompress and spdecompress Ops
+//===----------------------------------------------------------------------===//
+
+// Structured sparsity factor: spXtoY keeps X out of every Y elements.
+def SPFactor1To2  : I32EnumAttrCase<"SP1TO2",  0, "sp1to2">;
+def SPFactor1To4  : I32EnumAttrCase<"SP1TO4",  1, "sp1to4">;
+def SPFactor1To8  : I32EnumAttrCase<"SP1TO8",  2, "sp1to8">;
+def SPFactor1To16 : I32EnumAttrCase<"SP1TO16", 3, "sp1to16">;
+def SPFactor2To4  : I32EnumAttrCase<"SP2TO4",  4, "sp2to4">;
+def SPFactor2To8  : I32EnumAttrCase<"SP2TO8",  5, "sp2to8">;
+def SPFactor2To16 : I32EnumAttrCase<"SP2TO16", 6, "sp2to16">;
+def SPFactor4To8  : I32EnumAttrCase<"SP4TO8",  7, "sp4to8">;
+def SPFactor4To16 : I32EnumAttrCase<"SP4TO16", 8, "sp4to16">;
+def SPFactor : I32EnumAttr<"SPFactor", "Structured sparsity factor", [
+    SPFactor1To2,
+    SPFactor1To4,
+    SPFactor1To8,
+    SPFactor1To16,
+    SPFactor2To4,
+    SPFactor2To8,
+    SPFactor2To16,
+    SPFactor4To8,
+    SPFactor4To16,
+  ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::NVVM";
+}
+def SPFactorAttr : EnumAttr<NVVM_Dialect, SPFactor, "sp_factor">;
+
+// Structured sparsity metadata index size
+def SPIndexSizeB2 : I32EnumAttrCase<"B2", 2, "b2">;
+def SPIndexSizeB4 : I32EnumAttrCase<"B4", 4, "b4">;
+def SPIndexSize : I32EnumAttr<"SPIndexSize",
+    "Structured sparsity metadata index size", [
+    SPIndexSizeB2,
+    SPIndexSizeB4,
+  ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::NVVM";
+}
+def SPIndexSizeAttr : EnumAttr<NVVM_Dialect, SPIndexSize, "sp_index_size">;
+
+// Structured sparsity element size
+def SPElemSizeB8  : I32EnumAttrCase<"B8",  8,  "b8">;
+def SPElemSizeB16 : I32EnumAttrCase<"B16", 16, "b16">;
+def SPElemSize : I32EnumAttr<"SPElemSize", "Structured sparsity element size", [
+    SPElemSizeB8,
+    SPElemSizeB16,
+  ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::NVVM";
+}
+def SPElemSizeAttr : EnumAttr<NVVM_Dialect, SPElemSize, "sp_elem_size">;
+
+// Structured sparsity repetition factor
+def SPRepFactorX1  : I32EnumAttrCase<"X1",  1,  "x1">;
+def SPRepFactorX2  : I32EnumAttrCase<"X2",  2,  "x2">;
+def SPRepFactorX4  : I32EnumAttrCase<"X4",  4,  "x4">;
+def SPRepFactorX8  : I32EnumAttrCase<"X8",  8,  "x8">;
+def SPRepFactorX16 : I32EnumAttrCase<"X16", 16, "x16">;
+def SPRepFactorX32 : I32EnumAttrCase<"X32", 32, "x32">;
+def SPRepFactorX64 : I32EnumAttrCase<"X64", 64, "x64">;
+def SPRepFactor : I32EnumAttr<"SPRepFactor",
+    "Structured sparsity repetition factor",
+    [SPRepFactorX1, SPRepFactorX2, SPRepFactorX4, SPRepFactorX8,
+     SPRepFactorX16, SPRepFactorX32, SPRepFactorX64]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::NVVM";
+}
+def SPRepFactorAttr : EnumAttr<NVVM_Dialect, SPRepFactor, "sp_rep_factor">;
+
+// spcompress only supports the 2:4 sparsity factor.
+def SPCompressFactorAttr : ConfinedAttr<SPFactorAttr, [
+    EnumAttrIsOneOf<SPFactorAttr, [
+      SPFactor2To4,
+    ]>
+  ]>;
+
+// Register bundles are i32 or fixed-length, rank-1 vectors of i32.
+defvar SPDataType = FixedVectorOfRankAndLengthAndType<[1],
+    [2, 4, 8, 16, 32, 64, 128], [I32]>;
+defvar SPRegType = AnyTypeOf<[I32,
+    FixedVectorOfRankAndLengthAndType<[1], [2, 4, 8, 16, 32, 64], [I32]>]>;
+
+def NVVM_SPCompressOp : NVVM_Op<"spcompress",
----------------
durga4github wrote:

ok, looking at the existing Ops, (tcgen05-family, dot-accumulate family etc), we have used "." as a separator. So, we can do the same here for consistency.

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