[mlir][sparse] introduce new 2:4 block sparsity level type.

[aartbik]

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

138

can we use the NEW_SYNTAX for this

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

176

use property to encode ...
or
use property for encoding ...

238

no : in syntax, but you can pick

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

22

no :

[Peiming]

address comments.

[Peiming]

fix naming inconsistency.

mlir/include/mlir-c/Dialect/SparseTensor.h

 
 /// Dimension level types (and properties) that define sparse tensors.
 /// See the documentation in SparseTensorAttrDefs.td for their meaning.
 ///
 /// These correspond to SparseTensorEncodingAttr::DimLevelType in the C++ API.
 /// If updating, keep them in sync and update the static_assert in the impl
 /// file.
 enum MlirSparseTensorDimLevelType {
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE = 4,                     // 0b0001_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED = 8,                // 0b0010_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU = 9,             // 0b0010_01
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO = 10,            // 0b0010_10
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO = 11,         // 0b0010_11
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON = 16,                // 0b0100_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU = 17,             // 0b0100_01
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO = 18,             // 0b0100_10
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO = 19,          // 0b0100_11
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI = 32,       // 0b1000_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU = 33,    // 0b1000_01
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NO = 34,    // 0b1000_10
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU_NO = 35, // 0b1000_11
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE = 4,                     // 0b00001_00
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED = 8,                // 0b00010_00
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU = 9,             // 0b00010_01
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO = 10,            // 0b00010_10
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO = 11,         // 0b00010_11
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON = 16,                // 0b00100_00
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU = 17,             // 0b00100_01
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO = 18,             // 0b00100_10
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO = 19,          // 0b00100_11
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI = 32,       // 0b01000_00
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU = 33,    // 0b01000_01
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NO = 34,    // 0b01000_10
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU_NO = 35, // 0b01000_11
+  MLIR_SPARSE_TENSOR_TWO_OUT_OF_FOUR = 64,                    // 0b10000_00
 };
 
 //===----------------------------------------------------------------------===//
 // SparseTensorEncodingAttr
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given attribute is a `sparse_tensor.encoding` attribute.
 MLIR_CAPI_EXPORTED bool

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

 /// versions; consequently, client code should use the predicate functions
 /// defined below, rather than relying on knowledge about the particular
 /// binary encoding.
 ///
 /// The `Undef` "format" is a special value used internally for cases
 /// where we need to store an undefined or indeterminate `DimLevelType`.
 /// It should not be used externally, since it does not indicate an
 /// actual/representable format.
+///
+// TODO: We should generalize TwoOutOfFour to N out of M and use property to
+// encode the value of N and M.
 enum class DimLevelType : uint8_t {
-  Undef = 0,                 // 0b0000_00
-  Dense = 4,                 // 0b0001_00
-  Compressed = 8,            // 0b0010_00
-  CompressedNu = 9,          // 0b0010_01
-  CompressedNo = 10,         // 0b0010_10
-  CompressedNuNo = 11,       // 0b0010_11
-  Singleton = 16,            // 0b0100_00
-  SingletonNu = 17,          // 0b0100_01
-  SingletonNo = 18,          // 0b0100_10
-  SingletonNuNo = 19,        // 0b0100_11
-  CompressedWithHi = 32,     // 0b1000_00
-  CompressedWithHiNu = 33,   // 0b1000_01
-  CompressedWithHiNo = 34,   // 0b1000_10
-  CompressedWithHiNuNo = 35, // 0b1000_11
+  Undef = 0,                 // 0b00000_00
+  Dense = 4,                 // 0b00001_00
+  Compressed = 8,            // 0b00010_00
+  CompressedNu = 9,          // 0b00010_01
+  CompressedNo = 10,         // 0b00010_10
+  CompressedNuNo = 11,       // 0b00010_11
+  Singleton = 16,            // 0b00100_00
+  SingletonNu = 17,          // 0b00100_01
+  SingletonNo = 18,          // 0b00100_10
+  SingletonNuNo = 19,        // 0b00100_11
+  CompressedWithHi = 32,     // 0b01000_00
+  CompressedWithHiNu = 33,   // 0b01000_01
+  CompressedWithHiNo = 34,   // 0b01000_10
+  CompressedWithHiNuNo = 35, // 0b01000_11
+  TwoOutOfFour = 64,         // 0b10000_00
 };
 
 /// This enum defines all the storage formats supported by the sparse compiler,
 /// without the level properties.
 enum class LevelFormat : uint8_t {
-  Dense = 4,             // 0b0001_00
-  Compressed = 8,        // 0b0010_00
-  Singleton = 16,        // 0b0100_00
-  CompressedWithHi = 32, // 0b1000_00
+  Dense = 4,             // 0b00001_00
+  Compressed = 8,        // 0b00010_00
+  Singleton = 16,        // 0b00100_00
+  CompressedWithHi = 32, // 0b01000_00
+  TwoOutOfFour = 64,     // 0b10000_00
 };
 
 /// Returns string representation of the given dimension level type.
 constexpr const char *toMLIRString(DimLevelType dlt) {
   switch (dlt) {
   // TODO: should probably raise an error instead of printing it...
   case DimLevelType::Undef:
     return "undef";
constexpr const char *toMLIRString(DimLevelType dlt) {
   case DimLevelType::CompressedWithHi:
     return "compressed-hi";
   case DimLevelType::CompressedWithHiNu:
     return "compressed-hi-nu";
   case DimLevelType::CompressedWithHiNo:
     return "compressed-hi-no";
   case DimLevelType::CompressedWithHiNuNo:
     return "compressed-hi-nu-no";
+  case DimLevelType::TwoOutOfFour:
+    return "compressed24";
   }
   return "";
 }
 
 /// Check that the `DimLevelType` contains a valid (possibly undefined) value.
 constexpr bool isValidDLT(DimLevelType dlt) {
   const uint8_t formatBits = static_cast<uint8_t>(dlt) >> 2;
   const uint8_t propertyBits = static_cast<uint8_t>(dlt) & 3;
   // If undefined or dense, then must be unique and ordered.
   // Otherwise, the format must be one of the known ones.
-  return (formatBits <= 1)
+  return (formatBits <= 1 || formatBits == 16)
              ? (propertyBits == 0)
              : (formatBits == 2 || formatBits == 4 || formatBits == 8);
 }
 
 /// Check if the `DimLevelType` is the special undefined value.
 constexpr bool isUndefDLT(DimLevelType dlt) {
   return dlt == DimLevelType::Undef;
 }
 
 /// Check if the `DimLevelType` is dense.
 constexpr bool isDenseDLT(DimLevelType dlt) {
   return dlt == DimLevelType::Dense;
 }
 
+/// Check if the `DimLevelType` is 2:4
+constexpr bool isTwoOutOfFourDLT(DimLevelType dlt) {
+  return dlt == DimLevelType::TwoOutOfFour;
+}
+
 // We use the idiom `(dlt & ~3) == format` in order to only return true
 // for valid DLTs.  Whereas the `dlt & format` idiom is a bit faster but
 // can return false-positives on invalid DLTs.
 
 /// Check if the `DimLevelType` is compressed (regardless of properties).
 constexpr bool isCompressedDLT(DimLevelType dlt) {
   return (static_cast<uint8_t>(dlt) & ~3) ==
          static_cast<uint8_t>(DimLevelType::Compressed);
static_assert(
      *getLevelFormat(DimLevelType::SingletonNuNo) == LevelFormat::Singleton),
     "getLevelFormat conversion is broken");
 
 static_assert(
     (buildLevelType(LevelFormat::Dense, false, true) == std::nullopt &&
      buildLevelType(LevelFormat::Dense, true, false) == std::nullopt &&
      buildLevelType(LevelFormat::Dense, false, false) == std::nullopt &&
      *buildLevelType(LevelFormat::Dense, true, true) == DimLevelType::Dense &&
+     buildLevelType(LevelFormat::TwoOutOfFour, false, true) == std::nullopt &&
+     buildLevelType(LevelFormat::TwoOutOfFour, true, false) == std::nullopt &&
+     buildLevelType(LevelFormat::TwoOutOfFour, false, false) == std::nullopt &&
+     *buildLevelType(LevelFormat::TwoOutOfFour, true, true) ==
+         DimLevelType::TwoOutOfFour &&
      *buildLevelType(LevelFormat::Compressed, true, true) ==
          DimLevelType::Compressed &&
      *buildLevelType(LevelFormat::Compressed, true, false) ==
          DimLevelType::CompressedNu &&
      *buildLevelType(LevelFormat::Compressed, false, true) ==
          DimLevelType::CompressedNo &&
      *buildLevelType(LevelFormat::Compressed, false, false) ==
          DimLevelType::CompressedNuNo &&
static_assert((isValidDLT(DimLevelType::Undef) &&
                isValidDLT(DimLevelType::CompressedNuNo) &&
                isValidDLT(DimLevelType::Singleton) &&
                isValidDLT(DimLevelType::SingletonNu) &&
                isValidDLT(DimLevelType::SingletonNo) &&
                isValidDLT(DimLevelType::SingletonNuNo) &&
                isValidDLT(DimLevelType::CompressedWithHi) &&
                isValidDLT(DimLevelType::CompressedWithHiNu) &&
                isValidDLT(DimLevelType::CompressedWithHiNo) &&
-               isValidDLT(DimLevelType::CompressedWithHiNuNo)),
+               isValidDLT(DimLevelType::CompressedWithHiNuNo) &&
+               isValidDLT(DimLevelType::TwoOutOfFour)),
               "isValidDLT definition is broken");
 
 static_assert((!isCompressedDLT(DimLevelType::Dense) &&
                isCompressedDLT(DimLevelType::Compressed) &&
                isCompressedDLT(DimLevelType::CompressedNu) &&
                isCompressedDLT(DimLevelType::CompressedNo) &&
                isCompressedDLT(DimLevelType::CompressedNuNo) &&
                !isCompressedDLT(DimLevelType::Singleton) &&
static_assert((!isSingletonDLT(DimLevelType::Dense) &&
                !isSingletonDLT(DimLevelType::CompressedNuNo) &&
                isSingletonDLT(DimLevelType::Singleton) &&
                isSingletonDLT(DimLevelType::SingletonNu) &&
                isSingletonDLT(DimLevelType::SingletonNo) &&
                isSingletonDLT(DimLevelType::SingletonNuNo)),
               "isSingletonDLT definition is broken");
 
 static_assert((isOrderedDLT(DimLevelType::Dense) &&
+               isOrderedDLT(DimLevelType::TwoOutOfFour) &&
                isOrderedDLT(DimLevelType::Compressed) &&
                isOrderedDLT(DimLevelType::CompressedNu) &&
                !isOrderedDLT(DimLevelType::CompressedNo) &&
                !isOrderedDLT(DimLevelType::CompressedNuNo) &&
                isOrderedDLT(DimLevelType::Singleton) &&
                isOrderedDLT(DimLevelType::SingletonNu) &&
                !isOrderedDLT(DimLevelType::SingletonNo) &&
                !isOrderedDLT(DimLevelType::SingletonNuNo) &&
                isOrderedDLT(DimLevelType::CompressedWithHi) &&
                isOrderedDLT(DimLevelType::CompressedWithHiNu) &&
                !isOrderedDLT(DimLevelType::CompressedWithHiNo) &&
                !isOrderedDLT(DimLevelType::CompressedWithHiNuNo)),
               "isOrderedDLT definition is broken");
 
 static_assert((isUniqueDLT(DimLevelType::Dense) &&
+               isUniqueDLT(DimLevelType::TwoOutOfFour) &&
                isUniqueDLT(DimLevelType::Compressed) &&
                !isUniqueDLT(DimLevelType::CompressedNu) &&
                isUniqueDLT(DimLevelType::CompressedNo) &&
                !isUniqueDLT(DimLevelType::CompressedNuNo) &&
                isUniqueDLT(DimLevelType::Singleton) &&
                !isUniqueDLT(DimLevelType::SingletonNu) &&
                isUniqueDLT(DimLevelType::SingletonNo) &&
                !isUniqueDLT(DimLevelType::SingletonNuNo) &&

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

let extraClassDeclaration = [{
     //
 
     /// Safely looks up the level-type for the requested level.  (Returns
     /// `DimLevelType::Dense` for the null encoding, since dense-tensors
     /// are always all-dense.)
     ::mlir::sparse_tensor::DimLevelType getLvlType(::mlir::sparse_tensor::Level l) const;
 
     bool isDenseLvl(::mlir::sparse_tensor::Level l) const { return isDenseDLT(getLvlType(l)); }
+    bool isTwoOutOfFourLvl(::mlir::sparse_tensor::Level l) const { return isTwoOutOfFourDLT(getLvlType(l)); }
     bool isCompressedLvl(::mlir::sparse_tensor::Level l) const { return isCompressedDLT(getLvlType(l)); }
     bool isCompressedWithHiLvl(::mlir::sparse_tensor::Level l) const { return isCompressedWithHiDLT(getLvlType(l)); }
     bool isSingletonLvl(::mlir::sparse_tensor::Level l) const { return isSingletonDLT(getLvlType(l)); }
     bool isOrderedLvl(::mlir::sparse_tensor::Level l) const { return isOrderedDLT(getLvlType(l)); }
     bool isUniqueLvl(::mlir::sparse_tensor::Level l) const { return isUniqueDLT(getLvlType(l)); }
 
     /// Returns true if every level is dense.  Also returns true for
     /// the null encoding (since dense-tensors are always all-dense).

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

 namespace py = pybind11;
 using namespace llvm;
 using namespace mlir;
 using namespace mlir::python::adaptors;
 
 static void populateDialectSparseTensorSubmodule(const py::module &m) {
   py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local())
       .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE)
+      .value("compressed24", MLIR_SPARSE_TENSOR_TWO_OUT_OF_FOUR)
       .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED)
       .value("compressed-nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU)
       .value("compressed-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO)
       .value("compressed-nu-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO)
       .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON)
       .value("singleton-nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU)
       .value("singleton-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO)
       .value("singleton-nu-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO)

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

#define FOREVERY_LEVELTYPE(DO)                                                 \
   DO(DimLevelType::CompressedNuNo)                                             \
   DO(DimLevelType::Singleton)                                                  \
   DO(DimLevelType::SingletonNu)                                                \
   DO(DimLevelType::SingletonNo)                                                \
   DO(DimLevelType::SingletonNuNo)                                              \
   DO(DimLevelType::CompressedWithHi)                                           \
   DO(DimLevelType::CompressedWithHiNu)                                         \
   DO(DimLevelType::CompressedWithHiNo)                                         \
-  DO(DimLevelType::CompressedWithHiNuNo)
+  DO(DimLevelType::CompressedWithHiNuNo)                                       \
+  DO(DimLevelType::TwoOutOfFour)
 #define LEVELTYPE_INITLIST_ELEMENT(lvlType)                                    \
   std::make_pair(StringRef(toMLIRString(lvlType)), lvlType),
 #define LEVELTYPE_INITLIST                                                     \
   { FOREVERY_LEVELTYPE(LEVELTYPE_INITLIST_ELEMENT) }
 
 // TODO(wrengr): Since this parser is non-trivial to construct, is there
 // any way to hook into the parsing process so that we construct it only once
 // at the begining of parsing and then destroy it once parsing has finished?

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 
 std::optional<uint64_t>
 SparseTensorEncodingAttr::getStaticLvlSliceStride(Level lvl) const {
   // FIXME: `toOrigDim` is deprecated.
   return getStaticDimSliceStride(toOrigDim(*this, lvl));
 }
 
 const static DimLevelType validDLTs[] = {DimLevelType::Dense,
+                                         DimLevelType::TwoOutOfFour,
                                          DimLevelType::Compressed,
                                          DimLevelType::CompressedNu,
                                          DimLevelType::CompressedNo,
                                          DimLevelType::CompressedNuNo,
                                          DimLevelType::Singleton,
                                          DimLevelType::SingletonNu,
                                          DimLevelType::SingletonNo,
                                          DimLevelType::SingletonNuNo,

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

#CSR_SLICE = #sparse_tensor.encoding<{
   dimSlices = [ (1, ?, 1), (?, 4, 2) ]
 }>
 
 // CHECK-LABEL: func private @sparse_slice(
 // CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ], dimSlices = [ (1, ?, 1), (?, 4, 2) ] }>>
 func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
 
 
+// -----
+
+// TODO: It is probably better to use [dense, dense, 2:4] (see NV_24 defined using new syntax
+// below) to encode a 2D matrix, but it would require dim2lvl mapping which is not ready yet.
+// So we take the simple path for now.
+#NV_24= #sparse_tensor.encoding<{
+  lvlTypes = [ "dense", "compressed24" ],
+}>
+
+// CHECK-LABEL: func private @sparse_2_out_of_4(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed24" ] }>>
+func.func private @sparse_2_out_of_4(tensor<?x?xf64, #NV_24>)
+
 ///////////////////////////////////////////////////////////////////////////////
 // Migration plan for new STEA surface syntax,
 // use the NEW_SYNTAX on selected examples
 // and then TODO: remove when fully migrated
 ///////////////////////////////////////////////////////////////////////////////
 
 // -----
 

 
 // FIXME: should not have to use 4 dims ;-)
 //
 // CHECK-LABEL: func private @foo(
 // CHECK-SAME: tensor<?x?x?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "dense", "dense" ] }>>
 func.func private @foo(%arg0: tensor<?x?x?x?xf64, #BCSR_explicit>) {
   return
 }
+
+// -----
+
+#NV_24 = #sparse_tensor.encoding<{
+  NEW_SYNTAX =
+  ( i, j ) ->
+  ( i            : dense,
+    j floordiv 4 : dense,
+    j mod 4      : compressed24
+  )
+}>
+
+//
+// CHECK-LABEL: func private @foo_2_out_of_4(
+// CHECK-SAME: tensor<?x?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "dense", "compressed24" ] }>>
+func.func private @foo_2_out_of_4(%arg0: tensor<?x?x?xf64, #NV_24>) {
+  return
+}