[mlir][sparse] add more dimension level types and properties

[aartbik]

Summary

We recently removed the singleton dimension level type (see the revision
https://reviews.llvm.org/D131002) since it was unimplemented but also
incomplete (properties were missing). This revision add singleton back as
extra dimension level type, together with properties ordered/not-ordered
and unique/not-unique. Even though still not lowered to actual code, this
provides a complete way of defining many more sparse storage schemes (in
the long run, we want to support even dimension level types and properties
using the additional extensions proposed in [Chou]).

Note that the current solution of using suffixes for the properties is not
ideal, but keeps the extension relatively simple with respect to parsing and
printing. Furthermore, it is rather consistent with the TACO implementation
which uses things like Compressed-Unique as well. Nevertheless, we probably
want to separate dimension level types from properties when we add more types
and properties.

[wrengr]

While using string suffixes is ugly for now, I don't think it's too bad. That is, in the long term I think we'll want to replace the string-based specification of dimleveltypes with an AST-based specification. Whenever (if ever) that happens, then we can treat the variable properties as arguments/attributes of the base dlt; for example compressed { nu, no } (or whatever custom syntax we prefer). Whereas if some dlt fixes a particular property (e.g., dense fixes both the sortedness and uniqueness properties), then it wouldn't take an argument for the property since it would have it as an intrinsic attribute instead.

[aartbik]

In D132897#3756933, @wrengr wrote:

While using string suffixes is ugly for now, I don't think it's too bad. That is, in the long term I think we'll want to replace the string-based specification of dimleveltypes with an AST-based specification. Whenever (if ever) that happens, then we can treat the variable properties as arguments/attributes of the base dlt; for example compressed { nu, no } (or whatever custom syntax we prefer). Whereas if some dlt fixes a particular property (e.g., dense fixes both the sortedness and uniqueness properties), then it wouldn't take an argument for the property since it would have it as an intrinsic attribute instead.

Agreed. And yes, we want to move towards something like "compressed { nu, no }" in the longer run.
This will be very similar to the C++ implementation in TACO, where properties are passed into the format construction:

ModeFormat::Singleton(ModeFormat::NOT_UNIQUE)

[aartbik]

minor edits

[aartbik]

fixed typo

[Peiming]

I am wondering how long will current string-encoding stand? Agree with Wren that current implementation is ugly (and not extensible)

[aartbik]

In D132897#3758919, @Peiming wrote:

I am wondering how long will current string-encoding stand? Agree with Wren that current implementation is ugly (and not extensible)

We can work on beautifying the parsing/printing/encoding next. This revision, however, will get the ball rolling on supporting singleton properly within sparse compiler itself.

[aartbik]

rebased with main

[stella.stamenova]

Looks like this caused failures on the windows mlir bot: https://lab.llvm.org/buildbot/#/builders/13/builds/25183/steps/6/logs/stdio

[Peiming]

In D132897#3759118, @stella.stamenova wrote:

Looks like this caused failures on the windows mlir bot: https://lab.llvm.org/buildbot/#/builders/13/builds/25183/steps/6/logs/stdio

It should already be fixed by https://reviews.llvm.org/D132961

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

 #include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
 
-/// Dimension level types that define sparse tensors:
-///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE - dimension is dense, every
-///   entry is stored
-///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED - dimension is sparse,
-///   only nonzeros are stored (no duplicates).
+/// 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,
   MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO,
+  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO,
 };
 
 //===----------------------------------------------------------------------===//
 // SparseTensorEncodingAttr
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given attribute is a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED bool

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

 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
 include "mlir/IR/TensorEncoding.td"
 
 // All of the Tensor attributes will extend this class.
 class SparseTensor_Attr<string name,
                         list<Trait> traits = []>
-  : AttrDef<SparseTensor_Dialect, name, traits>;
+    : AttrDef<SparseTensor_Dialect, name, traits>;
 
 // Sparse tensor encoding attribute.
 def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
          [ DeclareAttrInterfaceMethods<VerifiableTensorEncoding> ] > {
   let mnemonic = "encoding";
 
   let description = [{
     An attribute to encode TACO-style information on sparsity properties
     of tensors. The encoding is eventually used by a **sparse compiler**
     pass to generate sparse code fully automatically for all tensor
     expressions that involve tensors with a sparse encoding. Compiler
     passes that run before this sparse compiler pass need to be
     aware of the semantics of tensor types with such an encoding.
 
     The attribute consists of the following fields.
     - Dimension level type for each dimension of a tensor type:
         - **dense** : dimension is dense, all entries along this dimension
-    are stored.
+          are stored
         - **compressed** : dimension is sparse, only nonzeros along this dimensions
-    are stored, without duplicates, i.e., compressed (unique).
+          are stored
+        - **singleton** : dimension stores individual indices with no siblings
+      By default, each dimension level types has the property of being unique
+      (no duplicates at that level) and ordered (indices appear sorted at that
+      level). The following two suffixes can be used to make the last two
+      dimension level types not-unique (duplicates may appear) and not-ordered
+      (indices may appear unsorted).
+        - **-nu** : not unique
+        - **-no** : not ordered
+      Currently, these suffixes, is present, should appear in this order.
+      In the future, we may introduce many more dimension level types and
+      properties, and separate specifying the two completely rather than
+      using this suffix mechanism.
     - Dimension ordering on the indices of this tensor type. Unlike dense
       storage, most sparse storage schemes do not provide fast random access.
       This affine map specifies the order of dimensions that should be supported
       by the sparse storage scheme. For example, for a 2-d tensor, "(i,j) -> (i,j)"
       requests row-wise storage and "(i,j) -> (j,i)" requests column-wise storage.
     - The required bit width for "pointer" storage (integral offsets into
       the sparse storage scheme). A narrow width reduces the memory footprint
       of overhead storage, as long as the width suffices to define the total
let description = [{
 
     ```mlir
     // Sparse vector.
     #SparseVector = #sparse_tensor.encoding<{
       dimLevelType = [ "compressed" ]
     }>
     ... tensor<?xf32, #SparseVector> ...
 
+    // Sorted Coordinate Scheme.
+    #SortedCOO = #sparse_tensor.encoding<{
+      dimLevelType = [ "compressed-nu", "singleton" ]
+    }>
+    ... tensor<?x?xf64, #SortedCOO> ...
+
     // Doubly compressed sparse column storage with specific bitwidths.
     #DCSC = #sparse_tensor.encoding<{
       dimLevelType = [ "compressed", "compressed" ],
       dimOrdering = affine_map<(i,j) -> (j,i)>,
       pointerBitWidth = 32,
       indexBitWidth = 8
     }>
     ... tensor<8x8xf64, #DCSC> ...
     ```
   }];
 
   // Data in sparse tensor encoding.
   let parameters = (
     ins
-    // A dimension level type for each dimension of a tensor type.
+    // A dimension level type for each dimension of the tensor type.
     ArrayRefParameter<
       "SparseTensorEncodingAttr::DimLevelType",
-      "Per-dimension level type (dense or compressed)"
+      "per dimension level type"
       >: $dimLevelType,
     // A dimension order on the indices of this tensor type.
     // TODO: block structure with higher-dim inputs
     "AffineMap":$dimOrdering,
     // The required bit width for pointer storage.
     "unsigned":$pointerBitWidth,
     // The required bit width for index storage.
     "unsigned":$indexBitWidth
   );
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 
   let extraClassDeclaration = [{
-    // Dimension level types.
+    // Dimension level types. By default, each type has the unique and
+    // ordered properties. Alternatives properties are indicated by
+    // Nu (not-unique) and No (not-ordered).
+    //
+    // TODO: separate type and property in encoding
+    //
     enum class DimLevelType {
-      Dense, Compressed
+      Dense,
+      Compressed, CompressedNu, CompressedNo, CompressedNuNo,
+      Singleton, SingletonNu, SingletonNo, SingletonNuNo,
     };
   }];
 }
 
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 /// The actions performed by @newSparseTensor.
 enum class Action : uint32_t {
   kEmpty = 0,
   kFromFile = 1,
   kFromCOO = 2,
   kSparseToSparse = 3,
   kEmptyCOO = 4,
   kToCOO = 5,
-  kToIterator = 6
+  kToIterator = 6,
 };
 
 /// This enum mimics `SparseTensorEncodingAttr::DimLevelType` for
 /// breaking dependency cycles.  `SparseTensorEncodingAttr::DimLevelType`
 /// is the source of truth and this enum should be kept consistent with it.
 enum class DimLevelType : uint8_t {
   kDense = 0,
   kCompressed = 1,
-  kSingleton = 2
+  kCompressedNu = 2,
+  kCompressedNo = 3,
+  kCompressedNuNo = 4,
+  kSingleton = 5,
+  kSingletonNu = 6,
+  kSingletonNo = 7,
+  kSingletonNuNo = 8,
 };
 
 //===----------------------------------------------------------------------===//
 //
 // Public functions which operate on MLIR buffers (memrefs) to interact
 // with sparse tensors (which are only visible as opaque pointers externally).
 // Because these functions deal with memrefs, they should only be used
 // by MLIR compiler-generated code (or code similarly guaranteed to remain

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("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED);
+      .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_attribute_subclass(m, "EncodingAttr",
                           mlirAttributeIsASparseTensorEncodingAttr)
       .def_classmethod(
           "get",
           [](py::object cls,
              std::vector<MlirSparseTensorDimLevelType> dimLevelTypes,
              llvm::Optional<MlirAffineMap> dimOrdering, int pointerBitWidth,

mlir/lib/CAPI/Dialect/SparseTensor.cpp

MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor,
                                       mlir::sparse_tensor::SparseTensorDialect)
 
 // Ensure the C-API enums are int-castable to C++ equivalents.
 static_assert(
     static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) ==
             static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Dense) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) ==
             static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::Compressed),
+                SparseTensorEncodingAttr::DimLevelType::Compressed) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::CompressedNu) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::CompressedNo) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::CompressedNuNo) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::Singleton) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::SingletonNu) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::SingletonNo) &&
+        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO) ==
+            static_cast<int>(
+                SparseTensorEncodingAttr::DimLevelType::SingletonNuNo),
     "MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch");
 
 bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {
   return unwrap(attr).isa<SparseTensorEncodingAttr>();
 }
 
 MlirAttribute mlirSparseTensorEncodingAttrGet(
     MlirContext ctx, intptr_t numDimLevelTypes,

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

if (attr.getName() == "dimLevelType") {
                            "expected a string value in dimension level types");
           return {};
         }
         auto strVal = strAttr.getValue();
         if (strVal == "dense") {
           dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Dense);
         } else if (strVal == "compressed") {
           dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Compressed);
+        } else if (strVal == "compressed-nu") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNu);
+        } else if (strVal == "compressed-no") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNo);
+        } else if (strVal == "compressed-nu-no") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNuNo);
+        } else if (strVal == "singleton") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Singleton);
+        } else if (strVal == "singleton-nu") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNu);
+        } else if (strVal == "singleton-no") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNo);
+        } else if (strVal == "singleton-nu-no") {
+          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNuNo);
         } else {
           parser.emitError(parser.getNameLoc(),
                            "unexpected dimension level type: ")
               << strVal;
           return {};
         }
       }
     } else if (attr.getName() == "dimOrdering") {
void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
   for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
     switch (getDimLevelType()[i]) {
     case DimLevelType::Dense:
       printer << "\"dense\"";
       break;
     case DimLevelType::Compressed:
       printer << "\"compressed\"";
       break;
+    case DimLevelType::CompressedNu:
+      printer << "\"compressed-nu\"";
+      break;
+    case DimLevelType::CompressedNo:
+      printer << "\"compressed-no\"";
+      break;
+    case DimLevelType::CompressedNuNo:
+      printer << "\"compressed-nu-no\"";
+      break;
+    case DimLevelType::Singleton:
+      printer << "\"singleton\"";
+      break;
+    case DimLevelType::SingletonNu:
+      printer << "\"singleton-nu\"";
+      break;
+    case DimLevelType::SingletonNo:
+      printer << "\"singleton-no\"";
+      break;
+    case DimLevelType::SingletonNuNo:
+      printer << "\"singleton-nu-no\"";
+      break;
     }
     if (i != e - 1)
       printer << ", ";
   }
   printer << " ]";
   // Print remaining members only for non-default values.
   if (getDimOrdering() && !getDimOrdering().isIdentity())
     printer << ", dimOrdering = affine_map<" << getDimOrdering() << ">";

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

 
 DimLevelType mlir::sparse_tensor::dimLevelTypeEncoding(
     SparseTensorEncodingAttr::DimLevelType dlt) {
   switch (dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Dense:
     return DimLevelType::kDense;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return DimLevelType::kCompressed;
+  case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
+    return DimLevelType::kCompressedNu;
+  case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
+    return DimLevelType::kCompressedNo;
+  case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
+    return DimLevelType::kCompressedNuNo;
+  case SparseTensorEncodingAttr::DimLevelType::Singleton:
+    return DimLevelType::kSingleton;
+  case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
+    return DimLevelType::kSingletonNu;
+  case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
+    return DimLevelType::kSingletonNo;
+  case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
+    return DimLevelType::kSingletonNuNo;
   }
   llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
 }
 
 //===----------------------------------------------------------------------===//
 // Misc code generators.
 //===----------------------------------------------------------------------===//
 

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

case SparseTensorEncodingAttr::DimLevelType::Compressed:
       if (alreadyCompressed)
         return false; // Multiple compressed dimensions not yet supported.
       alreadyCompressed = true;
       break;
     case SparseTensorEncodingAttr::DimLevelType::Dense:
       if (alreadyCompressed)
         return false; // Dense after Compressed not yet supported.
       break;
+    default: // TODO: investigate
+      return false;
     }
   }
   return true;
 }
 
 /// Helper method to translate indices during a reshaping operation.
 /// TODO: provide as general utility to MLIR at large?
 static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

#DCSC = #sparse_tensor.encoding<{
   dimOrdering = affine_map<(i,j) -> (j,i)>,
   pointerBitWidth = 0,
   indexBitWidth = 64
 }>
 
 // CHECK-LABEL: func private @sparse_dcsc(
 // CHECK-SAME: tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, indexBitWidth = 64 }>>)
 func.func private @sparse_dcsc(tensor<?x?xf32, #DCSC>)
+
+// -----
+
+#COO = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu-no", "singleton-no" ]
+}>
+
+// 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>)
+
+// -----
+
+#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>)