[mlir][sparse] support parsing slices in sparse tensor encoding attribute

[aartbik]

roundtrip and validation tests? also makes it easier to see what it will look like

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

187

This needs a lot more documentation what the triplets mean per dimension ;-)

[Peiming]

rebase

[Peiming]

support parsing in dynamic slices.

[Peiming]

rebase.

[bixia]

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

27

We normally end a comment with a period, see also SparseTensorOps.td.

33–36

Do we intend to add the description?

46

offset/size/stride to be consistent with ordering of the parameters in the code?

47

Would it be better if we use mlir::ShapedType::kDynamic instead?
Do you know why mlir::ShapedType::kDynamic was change from -1 to std::numeric_limits<int64_t>::min()?

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

65

Would it be better if we say "positive value or ? for" here?

[Peiming]

address comments

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

33–36

No, Sorry that I forget to add the description here...

47

I do not know why it is changed, any negative value should be okay though?

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

65

SG

[Peiming]

code simplification.

[Peiming]

revert unintended changes.

[Peiming]

fix bug

[Peiming]

fix typo.

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

 //===----------------------------------------------------------------------===//
 
 // All of the Tensor attributes will extend this class.
 class SparseTensor_Attr<string name,
                         list<Trait> 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
         - **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
ArrayRefParameter<
       >: $dimLevelType,
     // A dimension order on the indices of this tensor type.
     "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
+    "unsigned":$indexBitWidth,
+    // A dimension level type for each dimension of the tensor type.
+    ArrayRefParameter<
+      "std::tuple<int, int, int>",
+      "per dimension slice metadata"
+      >: $dimSlices
   );
 
+  let builders = [
+    AttrBuilder<(ins "ArrayRef<::mlir::sparse_tensor::DimLevelType>":$dimLevelType,
+                     "AffineMap":$dimOrdering,
+                     "AffineMap":$higherOrdering,
+                     "unsigned":$pointerBitWidth,
+                     "unsigned":$indexBitWidth), [{
+      return $_get($_ctxt, dimLevelType,
+                         dimOrdering,
+                         higherOrdering,
+                         pointerBitWidth,
+                         indexBitWidth,
+                         ArrayRef<std::tuple<int, int, int>>{});
+    }]>
+  ];
+
   let extraClassDeclaration = [{
+    ///
+    using DimSlice = std::tuple<int, int, int>;
+
     /// Returns the type for pointer storage based on pointerBitWidth
     Type getPointerType() const;
 
     /// Returns the type for index storage based on indexBitWidth
     Type getIndexType() const;
 
     /// Constructs a new encoding with the dimOrdering and higherOrdering
     /// reset to the default/identity.
     SparseTensorEncodingAttr withoutOrdering() const;
 
     /// Return true if every level is dense in the encoding.
     bool isAllDense() const;
 
     /// Return true if the encoding has an identity dimension ordering.
     bool hasIdDimOrdering() const;
+
+    bool isSlice() const {
+      return !getDimSlices().empty();
+    }
+
+    std::optional<int> getDimSliceOffset(unsigned dim) const;
+    std::optional<int> getDimSliceLength(unsigned dim) const;
+    std::optional<int> getDimSliceStride(unsigned dim) const;
+    std::optional<int> getLvlSliceOffset(unsigned lvl) const;
+    std::optional<int> getLvlSliceLength(unsigned lvl) const;
+    std::optional<int> getLvlSliceStride(unsigned lvl) const;
   }];
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Storage Specifier Enum Attribute

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

Type SparseTensorEncodingAttr::getIndexType() const {
   return idxWidth ? IntegerType::get(getContext(), idxWidth) : indexType;
 }
 
 SparseTensorEncodingAttr SparseTensorEncodingAttr::withoutOrdering() const {
   return SparseTensorEncodingAttr::get(
       getContext(), getDimLevelType(), AffineMap(), AffineMap(),
       getPointerBitWidth(), getIndexBitWidth());
 }
 
 bool SparseTensorEncodingAttr::isAllDense() const {
   return llvm::all_of(getDimLevelType(), isDenseDLT);
 }
 
 bool SparseTensorEncodingAttr::hasIdDimOrdering() const {
   return !getDimOrdering() || getDimOrdering().isIdentity();
 }
 
+std::optional<int>
+SparseTensorEncodingAttr::getDimSliceOffset(unsigned dim) const {
+  auto v = std::get<0>(getDimSlices()[dim]);
+  if (v == -1)
+    return std::nullopt;
+  return v;
+}
+
+std::optional<int>
+SparseTensorEncodingAttr::getDimSliceLength(unsigned dim) const {
+  auto v = std::get<1>(getDimSlices()[dim]);
+  if (v == -1)
+    return std::nullopt;
+  return v;
+}
+
+std::optional<int>
+SparseTensorEncodingAttr::getDimSliceStride(unsigned dim) const {
+  auto v = std::get<2>(getDimSlices()[dim]);
+  if (v == -1)
+    return std::nullopt;
+  return v;
+}
+
+std::optional<int>
+SparseTensorEncodingAttr::getLvlSliceOffset(unsigned lvl) const {
+  return getDimSliceOffset(toOrigDim(*this, lvl));
+}
+
+std::optional<int>
+SparseTensorEncodingAttr::getLvlSliceLength(unsigned lvl) const {
+  return getDimSliceLength(toOrigDim(*this, lvl));
+}
+
+std::optional<int>
+SparseTensorEncodingAttr::getLvlSliceStride(unsigned lvl) const {
+  return getDimSliceStride(toOrigDim(*this, lvl));
+}
+
 Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
   if (failed(parser.parseLess()))
     return {};
   // Parse the data as a dictionary.
   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<DimLevelType> dlt;
+  SmallVector<DimSlice> slices;
   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) {
for (const NamedAttribute &attr : dict) {
     } 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 if (attr.getName() == "slice") {
+      auto slicesAttr = attr.getValue().dyn_cast<ArrayAttr>();
+      if (!slicesAttr) {
+        parser.emitError(parser.getNameLoc(), "expected an array for slices");
+        return {};
+      }
+      for (auto a : slicesAttr) {
+        auto sliceAttr = a.dyn_cast<ArrayAttr>();
+        if (!sliceAttr || sliceAttr.size() != 3) {
+          parser.emitError(parser.getNameLoc(),
+                           "expected an array of size 3 for a slice");
+          return {};
+        }
+        auto offset = sliceAttr.getValue()[0].dyn_cast<IntegerAttr>();
+        auto length = sliceAttr.getValue()[1].dyn_cast<IntegerAttr>();
+        auto stride = sliceAttr.getValue()[2].dyn_cast<IntegerAttr>();
+
+        if (!offset || !length || !stride) {
+          parser.emitError(parser.getNameLoc(),
+                           "expected an integral slice offset/length/stride");
+          return {};
+        }
+        slices.push_back(
+            std::make_tuple(offset.getInt(), length.getInt(), stride.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, higherOrd, ptr, ind);
+      parser.getContext(), dlt, dimOrd, higherOrd, ptr, ind, slices);
 }
 
 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++) {
     printer << toMLIRString(getDimLevelType()[i]);
     if (i != e - 1)
       printer << ", ";
   }
   printer << " ]";
   // Print remaining members only for non-default values.
   if (!hasIdDimOrdering())
     printer << ", dimOrdering = affine_map<" << getDimOrdering() << ">";
   if (getHigherOrdering())
     printer << ", higherOrdering = affine_map<" << getHigherOrdering() << ">";
   if (getPointerBitWidth())
     printer << ", pointerBitWidth = " << getPointerBitWidth();
   if (getIndexBitWidth())
     printer << ", indexBitWidth = " << getIndexBitWidth();
+  if (!getDimSlices().empty()) {
+    printer << ", slice = [ ";
+    llvm::interleaveComma(getDimSlices(), printer, [&](DimSlice dimSlice) {
+      int offset = std::get<0>(dimSlice);
+      int length = std::get<1>(dimSlice);
+      int stride = std::get<2>(dimSlice);
+      printer << '[' << offset << ", " << length << ", " << stride << ']';
+    });
+    printer << " ]";
+  }
+
   printer << " }>";
 }
 
 LogicalResult SparseTensorEncodingAttr::verify(
     function_ref<InFlightDiagnostic()> emitError,
     ArrayRef<DimLevelType> dimLevelType, AffineMap dimOrdering,
-    AffineMap higherOrdering, unsigned pointerBitWidth,
-    unsigned indexBitWidth) {
+    AffineMap higherOrdering, unsigned pointerBitWidth, unsigned indexBitWidth,
+    ArrayRef<DimSlice> dimSlices) {
+  // TODO: verify slices
   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";

 }
 
 LogicalResult SparseTensorEncodingAttr::verifyEncoding(
     ArrayRef<int64_t> shape, Type elementType,
     function_ref<InFlightDiagnostic()> emitError) const {
   // Check structural integrity.
   if (failed(verify(emitError, getDimLevelType(), getDimOrdering(),
                     getHigherOrdering(), getPointerBitWidth(),
-                    getIndexBitWidth())))
+                    getIndexBitWidth(), getDimSlices())))
     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)