[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

241

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

23

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

29–32

Do we intend to add the description?

42

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

43

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

29–32

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

43

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

 //===-- SparseTensorAttrDefs.td - attributes definitions ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPARSETENSOR_ATTRDEFS
 #define SPARSETENSOR_ATTRDEFS
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
 include "mlir/IR/TensorEncoding.td"
 
-//===----------------------------------------------------------------------===//
-// Sparse Tensor Type Encoding Attribute
-//===----------------------------------------------------------------------===//
-
 // All of the Tensor attributes will extend this class.
 class SparseTensor_Attr<string name,
                         list<Trait> traits = []>
     : AttrDef<SparseTensor_Dialect, name, traits>;
 
+//===----------------------------------------------------------------------===//
+// Sparse Tensor Dimension Slice Attribute
+//===----------------------------------------------------------------------===//
+
+def SparseTensorDimSliceAttr : SparseTensor_Attr<"SparseTensorDimSlice", []> {
+  let mnemonic = "slice";
+
+  // TODO
+  let description = [{
+  abc
+  }];
+
+  let parameters = (
+    ins
+    "int64_t" : $offset,
+    "int64_t" : $size,
+    "int64_t" : $stride
+  );
+
+  let extraClassDeclaration = [{
+    /// Special value for dynamic offset/stride/size.
+    static constexpr int64_t kDynamic = -1;
+
+    static bool isDynamic(int64_t v) {
+      return v == kDynamic;
+    }
+
+    std::optional<uint64_t> getStaticOffset() const {
+      if (isDynamic(getOffset()))
+        return std::nullopt;
+      return static_cast<uint64_t>(getOffset());
+    };
+
+    std::optional<uint64_t> getStaticStride() const {
+      if (isDynamic(getStride()))
+        return std::nullopt;
+      return static_cast<uint64_t>(getStride());
+    }
+
+    std::optional<uint64_t> getStaticSize() const {
+      if (isDynamic(getSize()))
+        return std::nullopt;
+      return static_cast<uint64_t>(getSize());
+    }
+
+    bool isCompletelyDynamic() const {
+      return isDynamic(getOffset()) && isDynamic(getStride()) && isDynamic(getSize());
+    };
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Sparse Tensor Type Encoding Attribute
+//===----------------------------------------------------------------------===//
+
 // 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**
#ELL = #sparse_tensor.encoding<{
       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<
       "::mlir::sparse_tensor::DimLevelType",
       "per dimension level type"
       >: $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<
+      "::mlir::sparse_tensor::SparseTensorDimSliceAttr",
+      "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<::mlir::sparse_tensor::SparseTensorDimSliceAttr>{});
+    }]>
+  ];
+
   let extraClassDeclaration = [{
     /// 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<uint64_t> getStaticDimSliceOffset(unsigned dim) const;
+    std::optional<uint64_t> getStaticDimSliceSize(unsigned dim) const;
+    std::optional<uint64_t> getStaticDimSliceStride(unsigned dim) const;
+    std::optional<uint64_t> getStaticLvlSliceOffset(unsigned lvl) const;
+    std::optional<uint64_t> getStaticLvlSliceSize(unsigned lvl) const;
+    std::optional<uint64_t> getStaticLvlSliceStride(unsigned lvl) const;
   }];
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Storage Specifier Enum Attribute

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

static bool acceptBitWidth(unsigned bitWidth) {
   case 32:
   case 64:
     return true;
   default:
     return false;
   }
 }
 
+void SparseTensorDimSliceAttr::print(AsmPrinter &printer) const {
+  printer << "(";
+  printer << (getStaticOffset() ? std::to_string(*getStaticOffset()) : "?");
+  printer << ", ";
+  printer << (getStaticSize() ? std::to_string(*getStaticSize()) : "?");
+  printer << ", ";
+  printer << (getStaticStride() ? std::to_string(*getStaticStride()) : "?");
+  printer << ")";
+}
+
+static ParseResult parseOptionalStaticSlice(int64_t &result,
+                                            AsmParser &parser) {
+  auto parseResult = parser.parseOptionalInteger(result);
+  if (parseResult.has_value()) {
+    if (parseResult.value().succeeded() && result < 0) {
+      parser.emitError(
+          parser.getCurrentLocation(),
+          "expect positive value for static slice offset/size/stride");
+      return failure();
+    }
+    return parseResult.value();
+  }
+
+  // Else, and '?' which represented dynamic slice
+  result = SparseTensorDimSliceAttr::kDynamic;
+  return parser.parseQuestion();
+}
+
+Attribute SparseTensorDimSliceAttr::parse(AsmParser &parser, Type type) {
+  if (failed(parser.parseLParen()))
+    return {};
+
+  int64_t offset = -1;
+  if (failed(parseOptionalStaticSlice(offset, parser)))
+    return {};
+
+  if (parser.parseComma().failed())
+    return {};
+
+  int64_t size = -1;
+  if (failed(parseOptionalStaticSlice(size, parser)))
+    return {};
+
+  if (parser.parseComma().failed())
+    return {};
+
+  int64_t stride = -1;
+  if (failed(parseOptionalStaticSlice(stride, parser)))
+    return {};
+
+  if (failed(parser.parseRParen()))
+    return {};
+
+  return parser.getChecked<SparseTensorDimSliceAttr>(parser.getContext(),
+                                                     offset, size, stride);
+}
+
 Type SparseTensorEncodingAttr::getPointerType() const {
   unsigned ptrWidth = getPointerBitWidth();
   Type indexType = IndexType::get(getContext());
   return ptrWidth ? IntegerType::get(getContext(), ptrWidth) : indexType;
 }
 
 Type SparseTensorEncodingAttr::getIndexType() const {
   unsigned idxWidth = getIndexBitWidth();

 bool SparseTensorEncodingAttr::isAllDense() const {
   return llvm::all_of(getDimLevelType(), isDenseDLT);
 }
 
 bool SparseTensorEncodingAttr::hasIdDimOrdering() const {
   return !getDimOrdering() || getDimOrdering().isIdentity();
 }
 
+std::optional<uint64_t>
+SparseTensorEncodingAttr::getStaticDimSliceOffset(unsigned dim) const {
+  return getDimSlices()[dim].getStaticOffset();
+}
+
+std::optional<uint64_t>
+SparseTensorEncodingAttr::getStaticDimSliceSize(unsigned dim) const {
+  return getDimSlices()[dim].getStaticSize();
+}
+
+std::optional<uint64_t>
+SparseTensorEncodingAttr::getStaticDimSliceStride(unsigned dim) const {
+  return getDimSlices()[dim].getStaticStride();
+}
+
+std::optional<uint64_t>
+SparseTensorEncodingAttr::getStaticLvlSliceOffset(unsigned lvl) const {
+  return getStaticDimSliceOffset(toOrigDim(*this, lvl));
+}
+
+std::optional<uint64_t>
+SparseTensorEncodingAttr::getStaticLvlSliceSize(unsigned lvl) const {
+  return getStaticDimSliceSize(toOrigDim(*this, lvl));
+}
+
+std::optional<uint64_t>
+SparseTensorEncodingAttr::getStaticLvlSliceStride(unsigned lvl) const {
+  return getStaticDimSliceStride(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 {};
+#define RETURN_ON_FAIL(stmt)                                                   \
+  if (failed(stmt)) {                                                          \
+    return {};                                                                 \
+  }
+
+  RETURN_ON_FAIL(parser.parseLess())
+  RETURN_ON_FAIL(parser.parseLBrace())
+
   // Process the data from the parsed dictionary value into struct-like data.
   SmallVector<DimLevelType> dlt;
+  SmallVector<SparseTensorDimSliceAttr> 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>();
+
+  StringRef attrName;
+  // Exactly 6 keys.
+  SmallVector<StringRef, 6> keys = {"dimLevelType",   "dimOrdering",
+                                    "higherOrdering", "pointerBitWidth",
+                                    "indexBitWidth",  "slice"};
+  while (succeeded(parser.parseOptionalKeyword(&attrName))) {
+    if (!llvm::is_contained(keys, attrName)) {
+      parser.emitError(parser.getNameLoc(), "unexpected key: ") << attrName;
+      return {};
+    }
+
+    // Consume the `=` after keys
+    RETURN_ON_FAIL(parser.parseEqual())
+    if (attrName == "dimLevelType") {
+      Attribute attr;
+      RETURN_ON_FAIL(parser.parseAttribute(attr));
+      auto arrayAttr = attr.dyn_cast<ArrayAttr>();
       if (!arrayAttr) {
         parser.emitError(parser.getNameLoc(),
                          "expected an array for dimension level types");
         return {};
       }
       for (auto i : arrayAttr) {
         auto strAttr = i.dyn_cast<StringAttr>();
         if (!strAttr) {
if (attrName == "dimLevelType") {
           dlt.push_back(DimLevelType::SingletonNuNo);
         } else {
           parser.emitError(parser.getNameLoc(),
                            "unexpected dimension level type: ")
               << strVal;
           return {};
         }
       }
-    } else if (attr.getName() == "dimOrdering") {
-      auto affineAttr = attr.getValue().dyn_cast<AffineMapAttr>();
+    } else if (attrName == "dimOrdering") {
+      Attribute attr;
+      RETURN_ON_FAIL(parser.parseAttribute(attr))
+
+      auto affineAttr = attr.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>();
+    } else if (attrName == "higherOrdering") {
+      Attribute attr;
+      RETURN_ON_FAIL(parser.parseAttribute(attr))
+
+      auto affineAttr = attr.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>();
+    } else if (attrName == "pointerBitWidth") {
+      Attribute attr;
+      RETURN_ON_FAIL(parser.parseAttribute(attr))
+
+      auto intAttr = attr.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>();
+    } else if (attrName == "indexBitWidth") {
+      Attribute attr;
+      RETURN_ON_FAIL(parser.parseAttribute(attr))
+
+      auto intAttr = attr.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();
+    } else if (attrName == "slice") {
+      RETURN_ON_FAIL(parser.parseLSquare())
+      // Dispatches to DimSliceAttr to skip mnemonic
+      bool finished = false;
+      while (auto attr = SparseTensorDimSliceAttr::parse(parser, nullptr)) {
+        auto sliceAttr = attr.cast<SparseTensorDimSliceAttr>();
+        slices.push_back(sliceAttr);
+        if (parser.parseOptionalComma().failed()) {
+          finished = true;
+          break;
+        }
+      }
+      // Wrong when parsing slices
+      if (!finished)
         return {};
+      RETURN_ON_FAIL(parser.parseRSquare())
     }
+
+    // Only the last item can omit the comma
+    if (parser.parseOptionalComma().failed())
+      break;
   }
+
+  RETURN_ON_FAIL(parser.parseRBrace())
+  RETURN_ON_FAIL(parser.parseGreater())
+#undef RETURN_ON_FAIL
+
   // 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,
+                          [&](SparseTensorDimSliceAttr attr) {
+                            // Calls SparseTensorDimSliceAttr::print directly to
+                            // skip mnemonic.
+                            attr.print(printer);
+                          });
+    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<SparseTensorDimSliceAttr> 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)

mlir/test/Dialect/SparseTensor/invalid_encoding.mlir

 #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>) -> ()
 
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (-1, ?, 1), (?, 4, 2) ] // expected-error{{expect positive value for static slice offset/size/stride}}
+}>
+func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

#ELL = #sparse_tensor.encoding<{
   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>)
 
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+// CHECK-LABEL: func private @sparse_slice(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], slice = [ (1, 4, 1), (1, 4, 2) ] }>>
+func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
+
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+// CHECK-LABEL: func private @sparse_slice(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], slice = [ (1, 4, 1), (1, 4, 2) ] }>>
+func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
+
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, ?, 1), (?, 4, 2) ]
+}>
+
+// CHECK-LABEL: func private @sparse_slice(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], slice = [ (1, ?, 1), (?, 4, 2) ] }>>
+func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)