[mlir][sparse] introduce a new compressed(hi) dimension level type

[Peiming]

Summary

compressed(hi) is similar to compressed, but instead of reusing the previous position high as the current position low, it uses a pair of positions for each sparse index.

The patch only introduces the definition (syntax) but does not provide codegen implementation.

[Peiming]

add roundtrip test.

[aartbik]

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

30

introducing parenthesis now seems a bit taking a new path.
Would compressed-hi-xxx a bit safer with the current (granted ugly) solution?

[Peiming]

address comments.

[wrengr]

I disagree with the operational description of the new level-type as given in the CL description. Like I mentioned in the chat: I think that it'd make more sense to implement the new level-type by having two separate arrays for the position-lows and position-highs, rather than interleaving them into a single array. At the very least, the external description/definition of the level-type should allow for such an implementation, rather than specifying the interleaved implementation.

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

38–41

If we go with the two-arrays implementation, then these could be given the values 0b1010_XX: i.e., where the 0bX010_XX part indicates the presence of the position-lows array, and the 0b1XXX_XX part indicates the presence of the position-highs array.

The reason I mention this is because that sort of encoding will make sense in the future once we add ELL support; i.e., we'd have the bits (hasLowPositions, hasHighPositions, hasCoordinates) and then we'd have the level types CompressedWithHi = (1,1,1); Compressed = (1,0,1); Singleton = (0,0,1); Counted = (1,0,0); Dense = (0,0,0). Of course, since this changes how the Compressed/Singleton/Dense types are encoded relative to one another, we can always wait until we add the Counted level-type to make these changes.

[Peiming]

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

38–41

Make sense, let's wait then ;-)

Though, strictly speaking, compressed level has both lo and hi for the first index on the level (or you can say the last index on the level).

[Peiming]

rebase

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,             // 0b001_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED = 8,        // 0b010_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU = 9,     // 0b010_01
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO = 10,    // 0b010_10
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO = 11, // 0b010_11
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON = 16,        // 0b100_00
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU = 17,     // 0b100_01
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO = 18,     // 0b100_10
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO = 19,  // 0b100_11
+  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
 };
 
 //===----------------------------------------------------------------------===//
 // SparseTensorEncodingAttr
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given attribute is a `sparse_tensor.encoding` attribute.
 MLIR_CAPI_EXPORTED bool

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

 /// 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.
 enum class DimLevelType : uint8_t {
-  Undef = 0,           // 0b000_00
-  Dense = 4,           // 0b001_00
-  Compressed = 8,      // 0b010_00
-  CompressedNu = 9,    // 0b010_01
-  CompressedNo = 10,   // 0b010_10
-  CompressedNuNo = 11, // 0b010_11
-  Singleton = 16,      // 0b100_00
-  SingletonNu = 17,    // 0b100_01
-  SingletonNo = 18,    // 0b100_10
-  SingletonNuNo = 19,  // 0b100_11
+  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
 };
 
 /// This enum defines all the storage formats supported by the sparse compiler,
 /// without the level properties.
 enum class LevelFormat : uint8_t {
-  Dense = 4,      // 0b001_00
-  Compressed = 8, // 0b010_00
-  Singleton = 16, // 0b100_00
+  Dense = 4,             // 0b0001_00
+  Compressed = 8,        // 0b0010_00
+  Singleton = 16,        // 0b0100_00
+  CompressedWithHi = 32, // 0b1000_00
 };
 
 /// Returns string representation of the given dimension level type.
 inline std::string toMLIRString(DimLevelType dlt) {
   switch (dlt) {
   // TODO: should probably raise an error instead of printing it...
   case DimLevelType::Undef:
     return "undef";
inline std::string toMLIRString(DimLevelType dlt) {
   case DimLevelType::Singleton:
     return "singleton";
   case DimLevelType::SingletonNu:
     return "singleton-nu";
   case DimLevelType::SingletonNo:
     return "singleton-no";
   case DimLevelType::SingletonNuNo:
     return "singleton-nu-no";
+  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";
   }
   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) ? (propertyBits == 0)
-                           : (formatBits == 2 || formatBits == 4);
+  return (formatBits <= 1)
+             ? (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;
 }
 
 // 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);
 }
 
+/// Check if the `DimLevelType` is compressed (regardless of properties).
+constexpr bool isCompressedWithHiDLT(DimLevelType dlt) {
+  return (static_cast<uint8_t>(dlt) & ~3) ==
+         static_cast<uint8_t>(DimLevelType::CompressedWithHi);
+}
+
 /// Check if the `DimLevelType` is singleton (regardless of properties).
 constexpr bool isSingletonDLT(DimLevelType dlt) {
   return (static_cast<uint8_t>(dlt) & ~3) ==
          static_cast<uint8_t>(DimLevelType::Singleton);
 }
 
 /// Check if the `DimLevelType` is ordered (regardless of storage format).
 constexpr bool isOrderedDLT(DimLevelType dlt) {
static_assert((isValidDLT(DimLevelType::Undef) &&
                isValidDLT(DimLevelType::Dense) &&
                isValidDLT(DimLevelType::Compressed) &&
                isValidDLT(DimLevelType::CompressedNu) &&
                isValidDLT(DimLevelType::CompressedNo) &&
                isValidDLT(DimLevelType::CompressedNuNo) &&
                isValidDLT(DimLevelType::Singleton) &&
                isValidDLT(DimLevelType::SingletonNu) &&
                isValidDLT(DimLevelType::SingletonNo) &&
-               isValidDLT(DimLevelType::SingletonNuNo)),
+               isValidDLT(DimLevelType::SingletonNuNo) &&
+               isValidDLT(DimLevelType::CompressedWithHi) &&
+               isValidDLT(DimLevelType::CompressedWithHiNu) &&
+               isValidDLT(DimLevelType::CompressedWithHiNo) &&
+               isValidDLT(DimLevelType::CompressedWithHiNuNo)),
               "isValidDLT definition is broken");
 
 static_assert((!isCompressedDLT(DimLevelType::Dense) &&
                isCompressedDLT(DimLevelType::Compressed) &&
                isCompressedDLT(DimLevelType::CompressedNu) &&
                isCompressedDLT(DimLevelType::CompressedNo) &&
                isCompressedDLT(DimLevelType::CompressedNuNo) &&
                !isCompressedDLT(DimLevelType::Singleton) &&
                !isCompressedDLT(DimLevelType::SingletonNu) &&
                !isCompressedDLT(DimLevelType::SingletonNo) &&
                !isCompressedDLT(DimLevelType::SingletonNuNo)),
               "isCompressedDLT definition is broken");
 
+static_assert((!isCompressedWithHiDLT(DimLevelType::Dense) &&
+               isCompressedWithHiDLT(DimLevelType::CompressedWithHi) &&
+               isCompressedWithHiDLT(DimLevelType::CompressedWithHiNu) &&
+               isCompressedWithHiDLT(DimLevelType::CompressedWithHiNo) &&
+               isCompressedWithHiDLT(DimLevelType::CompressedWithHiNuNo) &&
+               !isCompressedWithHiDLT(DimLevelType::Singleton) &&
+               !isCompressedWithHiDLT(DimLevelType::SingletonNu) &&
+               !isCompressedWithHiDLT(DimLevelType::SingletonNo) &&
+               !isCompressedWithHiDLT(DimLevelType::SingletonNuNo)),
+              "isCompressedWithHiDLT definition is broken");
+
 static_assert((!isSingletonDLT(DimLevelType::Dense) &&
                !isSingletonDLT(DimLevelType::Compressed) &&
                !isSingletonDLT(DimLevelType::CompressedNu) &&
                !isSingletonDLT(DimLevelType::CompressedNo) &&
                !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::Compressed) &&
                isOrderedDLT(DimLevelType::CompressedNu) &&
                !isOrderedDLT(DimLevelType::CompressedNo) &&
                !isOrderedDLT(DimLevelType::CompressedNuNo) &&
                isOrderedDLT(DimLevelType::Singleton) &&
                isOrderedDLT(DimLevelType::SingletonNu) &&
                !isOrderedDLT(DimLevelType::SingletonNo) &&
-               !isOrderedDLT(DimLevelType::SingletonNuNo)),
+               !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::Compressed) &&
                !isUniqueDLT(DimLevelType::CompressedNu) &&
                isUniqueDLT(DimLevelType::CompressedNo) &&
                !isUniqueDLT(DimLevelType::CompressedNuNo) &&
                isUniqueDLT(DimLevelType::Singleton) &&
                !isUniqueDLT(DimLevelType::SingletonNu) &&
                isUniqueDLT(DimLevelType::SingletonNo) &&
-               !isUniqueDLT(DimLevelType::SingletonNuNo)),
+               !isUniqueDLT(DimLevelType::SingletonNuNo) &&
+               isUniqueDLT(DimLevelType::CompressedWithHi) &&
+               !isUniqueDLT(DimLevelType::CompressedWithHiNu) &&
+               isUniqueDLT(DimLevelType::CompressedWithHiNo) &&
+               !isUniqueDLT(DimLevelType::CompressedWithHiNuNo)),
               "isUniqueDLT definition is broken");
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_IR_ENUMS_H

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 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)); }
 
     bool isSlice() const {
       return !getDimSlices().empty();
     }
 

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local())
       .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE)
       .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);
+      .value("singleton-nu-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO)
+      .value("compressed(hi)", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI)
+      .value("compressed(hi)-nu",
+             MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU)
+      .value("compressed(hi)-no",
+             MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NO)
+      .value("compressed(hi)-nu-no",
+             MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU_NO);
 
   mlir_attribute_subclass(m, "EncodingAttr",
                           mlirAttributeIsASparseTensorEncodingAttr)
       .def_classmethod(
           "get",
           [](py::object cls,
              std::vector<MlirSparseTensorDimLevelType> dimLevelTypes,
              std::optional<MlirAffineMap> dimOrdering,

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::Compressed,
-    DimLevelType::CompressedNu,   DimLevelType::CompressedNo,
-    DimLevelType::CompressedNuNo, DimLevelType::Singleton,
-    DimLevelType::SingletonNu,    DimLevelType::SingletonNo,
-    DimLevelType::SingletonNuNo};
+const static DimLevelType validDLTs[] = {DimLevelType::Dense,
+                                         DimLevelType::Compressed,
+                                         DimLevelType::CompressedNu,
+                                         DimLevelType::CompressedNo,
+                                         DimLevelType::CompressedNuNo,
+                                         DimLevelType::Singleton,
+                                         DimLevelType::SingletonNu,
+                                         DimLevelType::SingletonNo,
+                                         DimLevelType::SingletonNuNo,
+                                         DimLevelType::CompressedWithHi,
+                                         DimLevelType::CompressedWithHiNu,
+                                         DimLevelType::CompressedWithHiNo,
+                                         DimLevelType::CompressedWithHiNuNo};
 
 static std::optional<DimLevelType> parseDLT(StringRef str) {
   for (DimLevelType dlt : validDLTs)
     if (str == toMLIRString(dlt))
       return dlt;
   return std::nullopt;
 }
 

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

 }>
 
 // CHECK-LABEL: func private @sparse_coo(
 // CHECK-SAME: tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu-no", "singleton-no" ] }>>)
 func.func private @sparse_coo(tensor<?x?xf32, #COO>)
 
 // -----
 
+#BCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed(hi)-nu", "singleton" ]
+}>
+
+// CHECK-LABEL: func private @sparse_bcoo(
+// CHECK-SAME: tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed(hi)-nu", "singleton" ] }>>)
+func.func private @sparse_bcoo(tensor<?x?xf32, #BCOO>)
+
+// -----
+
 #SortedCOO = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ]
 }>
 
 // CHECK-LABEL: func private @sparse_sorted_coo(
 // CHECK-SAME: tensor<10x10xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>)
 func.func private @sparse_sorted_coo(tensor<10x10xf64, #SortedCOO>)