more edits for n:m

Author: yinying-lisa-li

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

@@ -145,7 +145,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     - **compressed** : only nonzeros along this level are stored
     - **loose_compressed** : as compressed, but allows for free space between regions
     - **singleton** : a variant of the compressed format, where coordinates have no siblings
-    - **block[2, 4]** : the compression uses a 2:4 encoding per 1x4 block
+    - **block[n, m]** : the compression uses a n:m encoding per 1xm block
 
     For a compressed level, each position interval is represented in a compact
     way with a lowerbound `pos(i)` and an upperbound `pos(i+1) - 1`, which implies

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

@@ -123,7 +123,7 @@ class SparseTensorStorageBase {
   /// Safely checks if the level uses singleton storage.
   bool isSingletonLvl(uint64_t l) const { return isSingletonLT(getLvlType(l)); }
 
-  /// Safely checks if the level uses 2 out of 4 storage.
+  /// Safely checks if the level uses n out of m storage.
   bool isNOutOfMLvl(uint64_t l) const { return isNOutOfMLT(getLvlType(l)); }
 
   /// Safely checks if the level is ordered.
@@ -792,7 +792,7 @@ SparseTensorStorage<P, C, V>::SparseTensorStorage(
     } else if (isSingletonLvl(l)) {
       assert(0 && "general singleton not supported yet");
     } else if (isNOutOfMLvl(l)) {
-      assert(0 && "2Out4 not supported yet");
+      assert(0 && "n ouf of m not supported yet");
     } else {
       assert(isDenseLvl(l));
     }

mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp

@@ -40,7 +40,7 @@ FailureOr<uint64_t> LvlTypeParser::parseLvlType(AsmParser &parser) const {
   if (base.compare("block") == 0) {
     ParseResult res = parser.parseCommaSeparatedList(
         mlir::OpAsmParser::Delimiter::OptionalSquare,
-        [&]() -> ParseResult { return parseBlockSize(parser, &blockSizes); },
+        [&]() -> ParseResult { return parseBlockSizes(parser, &blockSizes); },
         " in block n out of m");
     FAILURE_IF_FAILED(res)
   }
@@ -95,8 +95,8 @@ ParseResult LvlTypeParser::parseProperty(AsmParser &parser,
 }
 
 ParseResult
-LvlTypeParser::parseBlockSize(AsmParser &parser,
-                              SmallVector<unsigned> *blockSizes) const {
+LvlTypeParser::parseBlockSizes(AsmParser &parser,
+                               SmallVector<unsigned> *blockSizes) const {
   int intVal;
   auto loc = parser.getCurrentLocation();
   OptionalParseResult intValParseResult = parser.parseOptionalInteger(intVal);

mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h

@@ -22,8 +22,8 @@ class LvlTypeParser {
 
 private:
   ParseResult parseProperty(AsmParser &parser, uint64_t *properties) const;
-  ParseResult parseBlockSize(AsmParser &parser,
-                             SmallVector<unsigned> *blockSizes) const;
+  ParseResult parseBlockSizes(AsmParser &parser,
+                              SmallVector<unsigned> *blockSizes) const;
 };
 
 } // namespace ir_detail

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

@@ -254,3 +254,21 @@ func.func private @NV_24(%arg0: tensor<?x?x?xf64, #NV_24>) {
 func.func private @NV_24(%arg0: tensor<?x?x?xf64, #NV_24>) {
   return
 }
+
+// -----
+
+#NOutOfM = #sparse_tensor.encoding<{
+  map = ( i, j, k ) ->
+  ( i            : dense,
+    k floordiv 8 : dense,
+    j            : dense,
+    k mod 8      : block[5, 8]
+  )
+}>
+
+// CHECK-DAG: #[[$NOutOfM:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d2 floordiv 8 : dense, d1 : dense, d2 mod 8 : block[5, 8]) }>
+// CHECK-LABEL: func private @NOutOfM(
+// CHECK-SAME: tensor<?x?x?xf64, #[[$NOutOfM]]>
+func.func private @NOutOfM(%arg0: tensor<?x?x?xf64, #NOutOfM>) {
+  return
+}