[mlir][sparse] introduce a higher-order tensor mapping

[aartbik]

Summary

This extension to the sparse tensor type system in MLIR
opens up a whole new set of sparse storage schemes, such as
block sparse storage (e.g. BCSR) and ELL (aka jagged diagonals).

This revision merely introduces the type extension and
initial documentation. The actual interpretation of the type
(reading in tensors, lowering to code, etc.) will follow.

[aartbik]

fixed typo

[Peiming]

LGTM, but perhaps I should wait for some one more experienced to accept the patch?

[Peiming]

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

71

It took me a while to understand the sentence.

Is For block-sparse (or block-sparsity), blocks are typically xxx better?

[aartbik]

In D135206#3835507, @Peiming wrote:

LGTM, but perhaps I should wait for some one more experienced to accept the patch?

I think we are all sufficiently (in)experienced with this ;-)

The higher order trick for block sparsity/ELL has been discussed in various publications, but never worked out in the nitty gritty details.
So we all learn as we go on this one ;-)

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

71

reordered

[aartbik]

addressed comment: reordered words

[aartbik]

clang

[Peiming]

In D135206#3835668, @aartbik wrote:

In D135206#3835507, @Peiming wrote:

LGTM, but perhaps I should wait for some one more experienced to accept the patch?

I think we are all sufficiently (in)experienced with this ;-)

The higher order trick for block sparsity/ELL has been discussed in various publications, but never worked out in the nitty gritty details.
So we all learn as we go on this one ;-)

Okay, then I will faithfully accept the patch ;-)

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

 /// Checks whether the given attribute is a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED bool
 mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr);
 
 /// Creates a sparse_tensor.encoding attribute with the given parameters.
 MLIR_CAPI_EXPORTED MlirAttribute mlirSparseTensorEncodingAttrGet(
     MlirContext ctx, intptr_t numDimLevelTypes,
     enum MlirSparseTensorDimLevelType const *dimLevelTypes,
-    MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth);
+    MlirAffineMap dimOrdering, MlirAffineMap higherOrdering,
+    int pointerBitWidth, int indexBitWidth);
 
 /// Returns the number of dim level types in a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED intptr_t
 mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr);
 
 /// Returns a specified dim level type in a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED enum MlirSparseTensorDimLevelType
 mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos);
 
 /// Returns the dimension ordering in a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED MlirAffineMap
 mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr);
 
+/// Returns the higher ordering in a sparse_tensor.encoding attribute.
+MLIR_CAPI_EXPORTED MlirAffineMap
+mlirSparseTensorEncodingAttrGetHigherOrdering(MlirAttribute attr);
+
 /// Returns the pointer bit width in a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED int
 mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr);
 
 /// Returns the index bit width in a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED int
 mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr);
 
 #ifdef __cplusplus
 }
 #endif
 
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.capi.h.inc"
 
 #endif // MLIR_C_DIALECT_SPARSETENSOR_H

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

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
         - **compressed** : dimension is sparse, only nonzeros along this dimensions
           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.
+
     - An optional 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.
-      By default, an identify mapping is used, which implies that the original
-      indices directly correspond to stored indices (viz. "(i,j,...) -> (i,j,...)").
+      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.  By default, an identify mapping is used,
+      which implies that the original indices directly correspond to stored
+      indices.
+
+    - An optional higher-ordering mapping from the original index space of
+      the tensor to a higher-order index space, used to define block-sparse
+      storage or ELL (jagged diagonal) storage. For example, for a 2-d tensor,
+      the mapping `(i, j) -> (i floordiv 2, j floordiv 3, i mod 2, j mod 3)`
+      imposes an higher-order partitioning into 2x3 blocks along the matrix
+      layout. A dimension ordering can be used to define a desired ordering
+      on this higher-order index space. Likewise, the dimension level types
+      define dense or compressed storage along this higher-order index space.
+      For block-sparse, blocks are typically stored with compression while
+      dense storage is used within each block (although hybrid schemes are
+      possible as well). The higher-order mapping also provides a notion of
+      "counting a dimension", where every stored element with the same index
+      is mapped to a new slice. For instance, ELL storage of a 2-d tensor can
+      be defined with the mapping `(i, j) -> (#i, i, j)` using the notation
+      of [Chou20]. Lacking the `#` symbol in MLIR's affine mapping, we use
+      a free symbol `c` to define such counting, together with a constant
+      that denotes the number of resulting slices. For example, the mapping
+      `(i, j)[c] -> (c * 3 * i, i, j)` with the first two higher-order indices
+      stored dense and the innermost compressed denotes ELL storage with
+      three jagged diagonals that count the dimension `i`.
+
+      TODO: introduce a real counting symbol to MLIR's mapping, since an
+            expression like 3*c*i has no direct interpretation?
+
     - 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
       required range (viz. the maximum number of stored entries over all indirection
       dimensions). The choices are `8`, `16`, `32`, `64`, or, the default, `0` to
       indicate the native bit width.
+
     - The required bit width for "index" storage (elements of the coordinates of
       stored entries). A narrow width reduces the memory footprint of overhead
       storage, as long as the width suffices to define the total required range
       (viz. the maximum value of each tensor index over all dimensions). The
-      choices are `8`, `16`, `32`, `64`, or, the default, `0` to indicate the
+      choices are `8`, `16`, `32`, `64`, or, the default, `0` to indicate a
       native bit width.
 
     Examples:
 
     ```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)>,
+      dimOrdering = affine_map<(i, j) -> (j, i)>,
       pointerBitWidth = 32,
       indexBitWidth = 8
     }>
     ... tensor<8x8xf64, #DCSC> ...
+
+    // Block sparse row storage (2x3 blocks).
+    #BCSR = #sparse_tensor.encoding<{
+      dimLevelType = [ "compressed", "compressed", "dense", "dense" ],
+      dimOrdering  = affine_map<(ii, jj, i, j) -> (ii, jj, i, j)>,
+      higherOrdering = affine_map<(i, j) -> (i floordiv 2, j floordiv 3, i mod 2, j mod 3)>
+    }>
+    ... tensor<20x30xf32, #BCSR> ...
+
+    // ELL storage (4 jagged diagonals, i.e., at most 4 nonzeros per row).
+    #ELL = #sparse_tensor.encoding<{
+      dimLevelType = [ "dense", "dense", "compressed" ],
+      dimOrdering  = affine_map<(ii, i, j) -> (ii, i, j)>,
+      higherOrdering = affine_map<(i, j)[c] -> (c * 4 * i, i, j)>
+    }>
+    ... tensor<?x?xf64, #ELL> ...
     ```
   }];
 
   // Data in sparse tensor encoding.
   let parameters = (
     ins
     // A dimension level type for each dimension of the tensor type.
     ArrayRefParameter<
       "SparseTensorEncodingAttr::DimLevelType",
       "per dimension level type"
       >: $dimLevelType,
     // A dimension order on the indices of this tensor type.
-    // TODO: block structure with higher-dim inputs
     "AffineMap":$dimOrdering,
+    // A mapping between the original and higher-ordering index space.
+    "AffineMap":$higherOrdering,
     // The required bit width for pointer storage.
     "unsigned":$pointerBitWidth,
     // The required bit width for index storage.
     "unsigned":$indexBitWidth
   );
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;

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

let description = [{
     PhD thesis, Leiden University, May 1996.
     * [Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman,
     Nicolas Vasilache, Bixia Zheng, and Fredrik Kjolstad. Compiler Support
     for Sparse Tensor Computations in MLIR. ACM Transactions on Architecture
     and Code Optimization, June, 2022. See: https://dl.acm.org/doi/10.1145/3544559
     * [Chou18] Stephen Chou, Fredrik Berg Kjolstad, and Saman Amarasinghe.
     Format Abstraction for Sparse Tensor Algebra Compilers. Proceedings of
     the ACM on Programming Languages, October 2018.
+    * [Chou20] Stephen Chou, Fredrik Berg Kjolstad, and Saman Amarasinghe.
+    Automatic Generation of Efficient Sparse Tensor Format Conversion Routines.
+    Proceedings of the 41st ACM SIGPLAN Conference on Programming Language
+    Design and Implementation, June, 2020.
     * [Gustavson72] Fred G. Gustavson. Some basic techniques for solving
     sparse systems of linear equations. In Sparse Matrices and Their
     Applications, pages 41–52. Plenum Press, New York, 1972.
     * [Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David
     Lugato, and Saman Amarasinghe. The Tensor Algebra Compiler. Proceedings of
     the ACM on Programming Languages, October 2017.
     * [Kjolstad19] Fredrik Berg Kjolstad, Peter Ahrens, Shoaib Ashraf Kamil,
     and Saman Amarasinghe. Tensor Algebra Compilation with Workspaces,

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local())
       .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,
+             llvm::Optional<MlirAffineMap> dimOrdering,
+             llvm::Optional<MlirAffineMap> higherOrdering, int pointerBitWidth,
              int indexBitWidth, MlirContext context) {
             return cls(mlirSparseTensorEncodingAttrGet(
                 context, dimLevelTypes.size(), dimLevelTypes.data(),
                 dimOrdering ? *dimOrdering : MlirAffineMap{nullptr},
+                higherOrdering ? *higherOrdering : MlirAffineMap{nullptr},
                 pointerBitWidth, indexBitWidth));
           },
           py::arg("cls"), py::arg("dim_level_types"), py::arg("dim_ordering"),
-          py::arg("pointer_bit_width"), py::arg("index_bit_width"),
-          py::arg("context") = py::none(),
+          py::arg("higher_ordering"), py::arg("pointer_bit_width"),
+          py::arg("index_bit_width"), py::arg("context") = py::none(),
           "Gets a sparse_tensor.encoding from parameters.")
       .def_property_readonly(
           "dim_level_types",
           [](MlirAttribute self) {
             std::vector<MlirSparseTensorDimLevelType> ret;
             for (int i = 0,
                      e = mlirSparseTensorEncodingGetNumDimLevelTypes(self);
                  i < e; ++i)
               ret.push_back(
                   mlirSparseTensorEncodingAttrGetDimLevelType(self, i));
             return ret;
           })
       .def_property_readonly(
           "dim_ordering",
           [](MlirAttribute self) -> llvm::Optional<MlirAffineMap> {
             MlirAffineMap ret =
                 mlirSparseTensorEncodingAttrGetDimOrdering(self);
             if (mlirAffineMapIsNull(ret))
               return {};
             return ret;
           })
       .def_property_readonly(
+          "higher_ordering",
+          [](MlirAttribute self) -> llvm::Optional<MlirAffineMap> {
+            MlirAffineMap ret =
+                mlirSparseTensorEncodingAttrGetHigherOrdering(self);
+            if (mlirAffineMapIsNull(ret))
+              return {};
+            return ret;
+          })
+      .def_property_readonly(
           "pointer_bit_width",
           [](MlirAttribute self) {
             return mlirSparseTensorEncodingAttrGetPointerBitWidth(self);
           })
       .def_property_readonly("index_bit_width", [](MlirAttribute self) {
         return mlirSparseTensorEncodingAttrGetIndexBitWidth(self);
       });
 }
 
 PYBIND11_MODULE(_mlirDialectsSparseTensor, m) {
   m.doc() = "MLIR SparseTensor dialect.";
   populateDialectSparseTensorSubmodule(m);
 }

mlir/lib/CAPI/Dialect/SparseTensor.cpp

 
 bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {
   return unwrap(attr).isa<SparseTensorEncodingAttr>();
 }
 
 MlirAttribute mlirSparseTensorEncodingAttrGet(
     MlirContext ctx, intptr_t numDimLevelTypes,
     MlirSparseTensorDimLevelType const *dimLevelTypes,
-    MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth) {
+    MlirAffineMap dimOrdering, MlirAffineMap higherOrdering,
+    int pointerBitWidth, int indexBitWidth) {
   SmallVector<SparseTensorEncodingAttr::DimLevelType> cppDimLevelTypes;
   cppDimLevelTypes.resize(numDimLevelTypes);
   for (intptr_t i = 0; i < numDimLevelTypes; ++i)
     cppDimLevelTypes[i] =
         static_cast<SparseTensorEncodingAttr::DimLevelType>(dimLevelTypes[i]);
-  return wrap(SparseTensorEncodingAttr::get(unwrap(ctx), cppDimLevelTypes,
-                                            unwrap(dimOrdering),
-                                            pointerBitWidth, indexBitWidth));
+  return wrap(SparseTensorEncodingAttr::get(
+      unwrap(ctx), cppDimLevelTypes, unwrap(dimOrdering),
+      unwrap(higherOrdering), pointerBitWidth, indexBitWidth));
 }
 
 MlirAffineMap mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr) {
   return wrap(unwrap(attr).cast<SparseTensorEncodingAttr>().getDimOrdering());
 }
 
+MlirAffineMap
+mlirSparseTensorEncodingAttrGetHigherOrdering(MlirAttribute attr) {
+  return wrap(
+      unwrap(attr).cast<SparseTensorEncodingAttr>().getHigherOrdering());
+}
+
 intptr_t mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr) {
   return unwrap(attr).cast<SparseTensorEncodingAttr>().getDimLevelType().size();
 }
 
 MlirSparseTensorDimLevelType
 mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos) {
   return static_cast<MlirSparseTensorDimLevelType>(
       unwrap(attr).cast<SparseTensorEncodingAttr>().getDimLevelType()[pos]);

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

Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
   DictionaryAttr dict;
   if (failed(parser.parseAttribute(dict)))
     return {};
   if (failed(parser.parseGreater()))
     return {};
   // Process the data from the parsed dictionary value into struct-like data.
   SmallVector<SparseTensorEncodingAttr::DimLevelType, 4> dlt;
   AffineMap dimOrd = {};
+  AffineMap higherOrd = {};
   unsigned ptr = 0;
   unsigned ind = 0;
   for (const NamedAttribute &attr : dict) {
     if (attr.getName() == "dimLevelType") {
       auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
       if (!arrayAttr) {
         parser.emitError(parser.getNameLoc(),
                          "expected an array for dimension level types");
for (const NamedAttribute &attr : dict) {
     } else if (attr.getName() == "dimOrdering") {
       auto affineAttr = attr.getValue().dyn_cast<AffineMapAttr>();
       if (!affineAttr) {
         parser.emitError(parser.getNameLoc(),
                          "expected an affine map for dimension ordering");
         return {};
       }
       dimOrd = affineAttr.getValue();
+    } else if (attr.getName() == "higherOrdering") {
+      auto affineAttr = attr.getValue().dyn_cast<AffineMapAttr>();
+      if (!affineAttr) {
+        parser.emitError(parser.getNameLoc(),
+                         "expected an affine map for higher ordering");
+        return {};
+      }
+      higherOrd = affineAttr.getValue();
     } else if (attr.getName() == "pointerBitWidth") {
       auto intAttr = attr.getValue().dyn_cast<IntegerAttr>();
       if (!intAttr) {
         parser.emitError(parser.getNameLoc(),
                          "expected an integral pointer bitwidth");
         return {};
       }
       ptr = intAttr.getInt();
     } else if (attr.getName() == "indexBitWidth") {
       auto intAttr = attr.getValue().dyn_cast<IntegerAttr>();
       if (!intAttr) {
         parser.emitError(parser.getNameLoc(),
                          "expected an integral index bitwidth");
         return {};
       }
       ind = intAttr.getInt();
     } else {
       parser.emitError(parser.getNameLoc(), "unexpected key: ")
           << attr.getName().strref();
       return {};
     }
   }
   // Construct struct-like storage for attribute.
-  return parser.getChecked<SparseTensorEncodingAttr>(parser.getContext(), dlt,
-                                                     dimOrd, ptr, ind);
+  return parser.getChecked<SparseTensorEncodingAttr>(
+      parser.getContext(), dlt, dimOrd, higherOrd, ptr, ind);
 }
 
 void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
   // Print the struct-like storage in dictionary fashion.
   printer << "<{ dimLevelType = [ ";
   for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
     switch (getDimLevelType()[i]) {
     case DimLevelType::Dense:
for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
     }
     if (i != e - 1)
       printer << ", ";
   }
   printer << " ]";
   // Print remaining members only for non-default values.
   if (getDimOrdering() && !getDimOrdering().isIdentity())
     printer << ", dimOrdering = affine_map<" << getDimOrdering() << ">";
+  if (getHigherOrdering())
+    printer << ", higherOrdering = affine_map<" << getHigherOrdering() << ">";
   if (getPointerBitWidth())
     printer << ", pointerBitWidth = " << getPointerBitWidth();
   if (getIndexBitWidth())
     printer << ", indexBitWidth = " << getIndexBitWidth();
   printer << " }>";
 }
 
 LogicalResult SparseTensorEncodingAttr::verify(
     function_ref<InFlightDiagnostic()> emitError,
     ArrayRef<DimLevelType> dimLevelType, AffineMap dimOrdering,
-    unsigned pointerBitWidth, unsigned indexBitWidth) {
+    AffineMap higherOrdering, unsigned pointerBitWidth,
+    unsigned indexBitWidth) {
   if (!acceptBitWidth(pointerBitWidth))
     return emitError() << "unexpected pointer bitwidth: " << pointerBitWidth;
   if (!acceptBitWidth(indexBitWidth))
     return emitError() << "unexpected index bitwidth: " << indexBitWidth;
   if (dimOrdering) {
     if (!dimOrdering.isPermutation())
       return emitError()
              << "expected a permutation affine map for dimension ordering";
     if (dimOrdering.getNumResults() != dimLevelType.size())
       return emitError() << "unexpected mismatch in ordering and dimension "
                             "level types size";
   }
+  if (higherOrdering) {
+    if (higherOrdering.getNumDims() >= higherOrdering.getNumResults())
+      return emitError() << "unexpected higher ordering mapping from "
+                         << higherOrdering.getNumDims() << " to "
+                         << higherOrdering.getNumResults();
+    if (higherOrdering.getNumResults() != dimLevelType.size())
+      return emitError() << "unexpected mismatch in higher ordering and "
+                            "dimension level types size";
+  }
   return success();
 }
 
 LogicalResult SparseTensorEncodingAttr::verifyEncoding(
     ArrayRef<int64_t> shape, Type elementType,
     function_ref<InFlightDiagnostic()> emitError) const {
   // Check structural integrity.
   if (failed(verify(emitError, getDimLevelType(), getDimOrdering(),
-                    getPointerBitWidth(), getIndexBitWidth())))
+                    getHigherOrdering(), getPointerBitWidth(),
+                    getIndexBitWidth())))
     return failure();
   // Check integrity with tensor type specifics. Dimension ordering is optional,
   // but we always should have dimension level types for the full rank.
   unsigned size = shape.size();
   if (size == 0)
     return emitError() << "expected non-scalar sparse tensor";
+  if (getHigherOrdering()) {
+    if (getHigherOrdering().getNumDims() != size)
+      return emitError() << "expected an affine map of size " << size
+                         << " for higher ordering";
+
+    // TODO: verification of higher ordering contents
+
+    size = getHigherOrdering().getNumResults(); // higher-order size!
+  }
   if (getDimOrdering() && getDimOrdering().getNumResults() != size)
     return emitError() << "expected an affine map of size " << size
                        << " for dimension ordering";
   if (getDimLevelType().size() != size)
     return emitError() << "expected an array of size " << size
                        << " for dimension level types";
   return success();
 }

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

genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
 
   unsigned srcRank = srcTp.getRank();
   unsigned dstRank = dstTp.getRank();
   Type elemTp = srcTp.getElementType();
   assert(elemTp == dstTp.getElementType() &&
          "reshape should not change element type");
   // Start an iterator over the source tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
-      op->getContext(), encSrc.getDimLevelType(), AffineMap(),
+      op->getContext(), encSrc.getDimLevelType(), AffineMap(), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
   SmallVector<Value, 4> srcSizes;
   SmallVector<Value, 8> params;
   sizesFromPtr(rewriter, srcSizes, loc, encSrc, srcTp, adaptor.getSrc());
   newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, srcSizes,
             adaptor.getSrc());
   Value iter = genNewCall(rewriter, loc, params);
   // Start a new COO for the destination tensor.
static void genSparseCOOIterationLoop(
   auto enc = getSparseTensorEncoding(tensorTp);
   assert(enc && "Generating Sparse Tensor COO Loop on a Dense Tensor!");
 
   unsigned rank = tensorTp.getRank();
   Type elemTp = tensorTp.getElementType();
 
   // Start an iterator over the tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
-      rewriter.getContext(), enc.getDimLevelType(), AffineMap(),
+      rewriter.getContext(), enc.getDimLevelType(), AffineMap(), AffineMap(),
       enc.getPointerBitWidth(), enc.getIndexBitWidth());
   SmallVector<Value, 4> sizes;
   SmallVector<Value, 8> params;
   sizesFromPtr(rewriter, sizes, loc, noPerm, tensorTp, t);
   newParams(rewriter, params, loc, tensorTp, noPerm, Action::kToIterator, sizes,
             t);
   Value iter = genNewCall(rewriter, loc, params);
 
if (encDst && encSrc) {
                   sizes, src);
         rewriter.replaceOp(op, genNewCall(rewriter, loc, params));
       } else { // use via-COO conversion.
         // Set up encoding with right mix of src and dst so that the two
         // method calls can share most parameters, while still providing
         // the correct sparsity information to either of them.
         auto enc = SparseTensorEncodingAttr::get(
             op->getContext(), encDst.getDimLevelType(), encDst.getDimOrdering(),
-            encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
+            encDst.getHigherOrdering(), encSrc.getPointerBitWidth(),
+            encSrc.getIndexBitWidth());
         newParams(rewriter, params, loc, stp, enc, Action::kToCOO, sizes, src);
         Value coo = genNewCall(rewriter, loc, params);
         params[3] = constantPointerTypeEncoding(rewriter, loc, encDst);
         params[4] = constantIndexTypeEncoding(rewriter, loc, encDst);
         params[6] = constantAction(rewriter, loc, Action::kFromCOO);
         params[7] = coo;
         Value dst = genNewCall(rewriter, loc, params);
         genDelCOOCall(rewriter, loc, stp.getElementType(), coo);
if (!encDst && encSrc) {
       Type elemTp = dstTensorTp.getElementType();
       // Fabricate a no-permutation encoding for newParams().
       // The pointer/index types must be those of `src`.
       // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
           op->getContext(),
           SmallVector<SparseTensorEncodingAttr::DimLevelType>(
               rank, SparseTensorEncodingAttr::DimLevelType::Dense),
-          AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
+          AffineMap(), AffineMap(), encSrc.getPointerBitWidth(),
+          encSrc.getIndexBitWidth());
       SmallVector<Value, 4> sizes;
       SmallVector<Value, 8> params;
       sizesFromPtr(rewriter, sizes, loc, encSrc, srcTensorTp, src);
       newParams(rewriter, params, loc, dstTensorTp, encDst, Action::kToIterator,
                 sizes, src);
       Value iter = genNewCall(rewriter, loc, params);
       Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
matchAndRewrite(OutOp op, OpAdaptor adaptor,
     ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
     // Convert to default permuted COO.
     Value src = adaptor.getOperands()[0];
     auto encSrc = getSparseTensorEncoding(srcType);
     SmallVector<Value, 4> sizes;
     SmallVector<Value, 8> params;
     sizesFromPtr(rewriter, sizes, loc, encSrc, srcType, src);
     auto enc = SparseTensorEncodingAttr::get(
-        op->getContext(), encSrc.getDimLevelType(), AffineMap(),
+        op->getContext(), encSrc.getDimLevelType(), AffineMap(), AffineMap(),
         encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
     newParams(rewriter, params, loc, srcType, enc, Action::kToCOO, sizes, src);
     Value coo = genNewCall(rewriter, loc, params);
     // Then output the tensor to external file with indices in the externally
     // visible lexicographic index order. A sort is required if the source was
     // not in that order yet (note that the sort can be dropped altogether if
     // external format does not care about the order at all, but here we assume
     // it does).

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

static RankedTensorType getUnorderedCOOFromType(RankedTensorType src) {
   // TODO: only if all the inputs (for concatentate) are unique at the last
   // level should the COO has a unique level at the end. Ends by a unordered
   // unique singleton level.
   dims.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNo);
   // TODO: Maybe pick the bitwidth based on input/output tensors (probably the
   // largest one among them) in the original operation instead of using the
   // default value.
   auto enc = SparseTensorEncodingAttr::get(
-      ctx, dims, AffineMap::getMultiDimIdentityMap(rank, ctx), 0, 0);
+      ctx, dims, AffineMap::getMultiDimIdentityMap(rank, ctx), AffineMap(), 0,
+      0);
   return RankedTensorType::get(src.getShape(), src.getElementType(), enc);
 }
 
 //===---------------------------------------------------------------------===//
 // The actual sparse tensor rewriting rules.
 //===---------------------------------------------------------------------===//
 
 namespace {

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

for (OpOperand *t : op.getInputOperands()) {
       // TODO: investigate fusing the conversion with computation,
       //       especially if it is a direct yield!
       //
       auto srcTp = tval.getType().cast<RankedTensorType>();
       auto dstEnc = SparseTensorEncodingAttr::get(
           op->getContext(), srcEnc.getDimLevelType(),
           permute(getContext(), op.getMatchingIndexingMap(t),
                   topSort), // new order
-          srcEnc.getPointerBitWidth(), srcEnc.getIndexBitWidth());
+          srcEnc.getHigherOrdering(), srcEnc.getPointerBitWidth(),
+          srcEnc.getIndexBitWidth());
       auto dstTp = RankedTensorType::get(srcTp.getShape(),
                                          srcTp.getElementType(), dstEnc);
       auto convert = rewriter.create<ConvertOp>(tval.getLoc(), dstTp, tval);
       op->setOperand(tensor, convert);
       rewriter.setInsertionPointAfter(op);
       rewriter.create<bufferization::DeallocTensorOp>(tval.getLoc(), convert);
       return success();
     }

mlir/test/CAPI/sparse_tensor.c

 // CHECK-LABEL: testRoundtripEncoding()
 static int testRoundtripEncoding(MlirContext ctx) {
   fprintf(stderr, "testRoundtripEncoding()\n");
   // clang-format off
   const char *originalAsm =
     "#sparse_tensor.encoding<{ "
     "dimLevelType = [ \"dense\", \"compressed\", \"compressed\"], "
     "dimOrdering = affine_map<(d0, d1, d2) -> (d0, d1, d2)>, "
+    "higherOrdering = affine_map<(d0, d1)[s0] -> (s0, d0, d1)>, "
     "pointerBitWidth = 32, indexBitWidth = 64 }>";
   // clang-format on
   MlirAttribute originalAttr =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm));
   // CHECK: isa: 1
   fprintf(stderr, "isa: %d\n",
           mlirAttributeIsASparseTensorEncodingAttr(originalAttr));
   MlirAffineMap dimOrdering =
       mlirSparseTensorEncodingAttrGetDimOrdering(originalAttr);
   // CHECK: (d0, d1, d2) -> (d0, d1, d2)
   mlirAffineMapDump(dimOrdering);
+  MlirAffineMap higherOrdering =
+      mlirSparseTensorEncodingAttrGetHigherOrdering(originalAttr);
+  // CHECK: (d0, d1)[s0] -> (s0, d0, d1)
+  mlirAffineMapDump(higherOrdering);
   // CHECK: level_type: 0
   // CHECK: level_type: 1
   // CHECK: level_type: 1
   int numLevelTypes = mlirSparseTensorEncodingGetNumDimLevelTypes(originalAttr);
   enum MlirSparseTensorDimLevelType *levelTypes =
       malloc(sizeof(enum MlirSparseTensorDimLevelType) * numLevelTypes);
   for (int i = 0; i < numLevelTypes; ++i) {
     levelTypes[i] =
         mlirSparseTensorEncodingAttrGetDimLevelType(originalAttr, i);
     fprintf(stderr, "level_type: %d\n", levelTypes[i]);
   }
   // CHECK: pointer: 32
   int pointerBitWidth =
       mlirSparseTensorEncodingAttrGetPointerBitWidth(originalAttr);
   fprintf(stderr, "pointer: %d\n", pointerBitWidth);
   // CHECK: index: 64
   int indexBitWidth =
       mlirSparseTensorEncodingAttrGetIndexBitWidth(originalAttr);
   fprintf(stderr, "index: %d\n", indexBitWidth);
 
   MlirAttribute newAttr = mlirSparseTensorEncodingAttrGet(
-      ctx, numLevelTypes, levelTypes, dimOrdering, pointerBitWidth,
-      indexBitWidth);
+      ctx, numLevelTypes, levelTypes, dimOrdering, higherOrdering,
+      pointerBitWidth, indexBitWidth);
   mlirAttributeDump(newAttr); // For debugging filecheck output.
   // CHECK: equal: 1
   fprintf(stderr, "equal: %d\n", mlirAttributeEqual(originalAttr, newAttr));
 
   free(levelTypes);
   return 0;
 }
 

mlir/test/Dialect/SparseTensor/invalid_encoding.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 #a = #sparse_tensor.encoding<{dimLevelType = []}>
 func.func private @scalar(%arg0: tensor<f64, #a>) -> () // expected-error {{expected non-scalar sparse tensor}}
 
 // -----
 
 #a = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
-func.func private @tensor_size_mismatch(%arg0: tensor<8xi32, #a>) -> () // expected-error {{expected an array of size 1 for dimension level types}}
+func.func private @tensor_dimlevel_size_mismatch(%arg0: tensor<8xi32, #a>) -> () // expected-error {{expected an array of size 1 for dimension level types}}
 
 // -----
 
 #a = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"], dimOrdering = affine_map<(i) -> (i)>}> // expected-error {{unexpected mismatch in ordering and dimension level types size}}
 func.func private @tensor_sizes_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
 // -----
 
 #a = #sparse_tensor.encoding<{dimLevelType = [1]}> // expected-error {{expected a string value in dimension level types}}
 func.func private @tensor_type_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
 // -----
 
 #a = #sparse_tensor.encoding<{dimLevelType = ["strange"]}> // expected-error {{unexpected dimension level type: strange}}
 func.func private @tensor_value_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
 // -----
 
 #a = #sparse_tensor.encoding<{dimOrdering = "wrong"}> // expected-error {{expected an affine map for dimension ordering}}
-func.func private @tensor_order_mismatch(%arg0: tensor<8xi32, #a>) -> ()
+func.func private @tensor_dimorder_mismatch(%arg0: tensor<8xi32, #a>) -> ()
+
+// -----
+
+#a = #sparse_tensor.encoding<{higherOrdering = "wrong"}> // expected-error {{expected an affine map for higher ordering}}
+func.func private @tensor_highorder_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
 // -----
 
 #a = #sparse_tensor.encoding<{dimOrdering = affine_map<(i,j) -> (i,i)>}> // expected-error {{expected a permutation affine map for dimension ordering}}
 func.func private @tensor_no_permutation(%arg0: tensor<16x32xf32, #a>) -> ()
 
 // -----
 

 
 #a = #sparse_tensor.encoding<{indexBitWidth = 128}> // expected-error {{unexpected index bitwidth: 128}}
 func.func private @tensor_invalid_int_index(%arg0: tensor<16x32xf32, #a>) -> ()
 
 // -----
 
 #a = #sparse_tensor.encoding<{key = 1}> // expected-error {{unexpected key: key}}
 func.func private @tensor_invalid_key(%arg0: tensor<16x32xf32, #a>) -> ()
+
+// -----
+
+#a = #sparse_tensor.encoding<{dimLevelType = [ "compressed", "compressed", "dense", "dense" ], dimOrdering  = affine_map<(ii, jj, i, j) -> (ii, jj, i, j)>, higherOrdering = affine_map<(i, j) -> (j, i)>}> // expected-error {{unexpected higher ordering mapping from 2 to 2}}
+func.func private @tensor_invalid_key(%arg0: tensor<10x60xf32, #a>) -> ()
+

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

 
 #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>)
+
+// -----
+
+#BCSR = #sparse_tensor.encoding<{
+   dimLevelType = [ "compressed", "compressed", "dense", "dense" ],
+   dimOrdering  = affine_map<(ii, jj, i, j) -> (ii, jj, i, j)>,
+   higherOrdering = affine_map<(i, j) -> (i floordiv 2, j floordiv 3, i mod 2, j mod 3)>
+}>
+
+// CHECK-LABEL: func private @sparse_bcsr(
+// CHECK-SAME: tensor<10x60xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense", "dense" ], higherOrdering = affine_map<(d0, d1) -> (d0 floordiv 2, d1 floordiv 3, d0 mod 2, d1 mod 3)> }>>
+func.func private @sparse_bcsr(tensor<10x60xf64, #BCSR>)
+
+
+// -----
+
+#ELL = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "dense", "compressed" ],
+  dimOrdering  = affine_map<(ii, i, j) -> (ii, i, j)>,
+  higherOrdering = affine_map<(i,j)[c] -> (c*4*i, i, j)>
+}>
+
+// CHECK-LABEL: func private @sparse_ell(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ], higherOrdering = affine_map<(d0, d1)[s0] -> (d0 * (s0 * 4), d0, d1)> }>>
+func.func private @sparse_ell(tensor<?x?xf64, #ELL>)
+

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

with ir.Context() as ctx, ir.Location.unknown():
     vec_strategy = ['none', 'dense-inner-loop']
     for level in levels:
       for ordering in orderings:
         for pwidth in [32]:
           for iwidth in [32]:
             for vec in vec_strategy:
               for e in [True]:
                 vl = 1 if vec == 0 else 16
-                attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
+                attr = st.EncodingAttr.get(level, ordering, None, pwidth,
+                                           iwidth)
                 opt = (f'parallelization-strategy=none '
                        f'vectorization-strategy={vec} '
                        f'vl={vl} enable-simd-index32={e}')
                 compiler = sparse_compiler.SparseCompiler(
                     options=opt, opt_level=0, shared_libs=[support_lib])
                 build_compile_and_run_SDDMMM(attr, compiler)
                 count = count + 1
   # CHECK: Passed 16 tests
   print('Passed ', count, 'tests')
 
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

with ir.Context() as ctx, ir.Location.unknown():
     ]
     bitwidths = [0]
     compiler = sparse_compiler.SparseCompiler(
         options=opt, opt_level=0, shared_libs=[support_lib])
     for level in levels:
       for ordering in orderings:
         for pwidth in bitwidths:
           for iwidth in bitwidths:
-            attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
+            attr = st.EncodingAttr.get(level, ordering, None, pwidth, iwidth)
             build_compile_and_run_SpMM(attr, compiler)
             count = count + 1
     # CHECK: Passed 8 tests
     print('Passed ', count, 'tests')
 
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_output.py

orderings = [
         ir.AffineMap.get_permutation([1, 0])
     ]
     bitwidths = [8, 16, 32, 64]
     compiler = sparse_compiler.SparseCompiler(
         options='', opt_level=2, shared_libs=[support_lib])
     for level in levels:
       for ordering in orderings:
         for bwidth in bitwidths:
-          attr = st.EncodingAttr.get(level, ordering, bwidth, bwidth)
+          attr = st.EncodingAttr.get(level, ordering, None, bwidth, bwidth)
           build_compile_and_run_output(attr, compiler)
           count = count + 1
 
   # CHECK: Passed 16 tests
   print('Passed', count, 'tests')
 
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

orderings = list(map(ir.AffineMap.get_permutation,
       itertools.permutations(range(rank))))
     bitwidths = [0]
     # The first type must be a dense tensor for numpy conversion to work.
     types = [ir.RankedTensorType.get(shape, f64)]
     for level in levels:
       for ordering in orderings:
         for pwidth in bitwidths:
           for iwidth in bitwidths:
-            attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
+            attr = st.EncodingAttr.get(level, ordering, None, pwidth, iwidth)
             types.append(ir.RankedTensorType.get(shape, f64, attr))
     #
     # For exhaustiveness we should have one or more StressTest, such
     # that their paths cover all 2*n*(n-1) directed pairwise combinations
     # of the `types` set.  However, since n is already superexponential,
     # such exhaustiveness would be prohibitive for a test that runs on
     # every commit.  So for now we'll just pick one particular path that
     # at least hits all n elements of the `types` set.

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

class Format:
   def mlir_tensor_attr(self) -> Optional[sparse_tensor.EncodingAttr]:
     """Constructs the MLIR attributes for the tensor format."""
     order = (
         range(self.rank()) if
         (self.ordering is None) else self.ordering.ordering)
     mlir_storage_format = [f.value for f in self.format_pack.formats]
     return sparse_tensor.EncodingAttr.get(mlir_storage_format,
                                           ir.AffineMap.get_permutation(order),
-                                          _POINTER_BIT_WIDTH, _INDEX_BIT_WIDTH)
+                                          None, _POINTER_BIT_WIDTH,
+                                          _INDEX_BIT_WIDTH)
 
 
 def _make_format(formats: List[ModeFormat],
                  ordering: Optional[List[int]] = None) -> Format:
   """Constructs a format from a list of ModeFormat and an optional ordering.
 
   Args:
     formats: A list of ModeFormat, one for each dimension of a tensor.

mlir/test/python/dialects/sparse_tensor/dialect.py

with Context() as ctx:
     print(f"dim_level_types: {casted.dim_level_types}")
     # CHECK: dim_ordering: None
     print(f"dim_ordering: {casted.dim_ordering}")
     # CHECK: pointer_bit_width: 16
     print(f"pointer_bit_width: {casted.pointer_bit_width}")
     # CHECK: index_bit_width: 32
     print(f"index_bit_width: {casted.index_bit_width}")
 
-    created = st.EncodingAttr.get(casted.dim_level_types, None, 0, 0)
+    created = st.EncodingAttr.get(casted.dim_level_types, None, None, 0, 0)
     # CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
     print(created)
     # CHECK: created_equal: False
     print(f"created_equal: {created == casted}")
 
     # Verify that the factory creates an instance of the proper type.
     # CHECK: is_proper_instance: True
     print(f"is_proper_instance: {isinstance(created, st.EncodingAttr)}")
with Context() as ctx:
     # CHECK: dim_ordering: (d0, d1) -> (d1, d0)
     print(f"dim_ordering: {casted.dim_ordering}")
     # CHECK: pointer_bit_width: 8
     print(f"pointer_bit_width: {casted.pointer_bit_width}")
     # CHECK: index_bit_width: 32
     print(f"index_bit_width: {casted.index_bit_width}")
 
     created = st.EncodingAttr.get(casted.dim_level_types, casted.dim_ordering,
-                                  8, 32)
+                                  casted.higher_ordering, 8, 32)
     # CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 8, indexBitWidth = 32 }>
     print(created)
     # CHECK: created_equal: True
     print(f"created_equal: {created == casted}")
 
 
 # CHECK-LABEL: TEST: testEncodingAttrOnTensorType
 @run