[mlir][sparse] Adding SparseTensorEncodingAttr::getDimSlice

[wrengr]

Summary

This helps catch segfaults and OOB.

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

let extraClassDeclaration = [{
 
     bool isDenseLvl(::mlir::sparse_tensor::Level l) const { return isDenseDLT(getLvlType(l)); }
     bool isCompressedLvl(::mlir::sparse_tensor::Level l) const { return isCompressedDLT(getLvlType(l)); }
     bool isCompressedWithHiLvl(::mlir::sparse_tensor::Level l) const { return isCompressedWithHiDLT(getLvlType(l)); }
     bool isSingletonLvl(::mlir::sparse_tensor::Level l) const { return isSingletonDLT(getLvlType(l)); }
     bool isOrderedLvl(::mlir::sparse_tensor::Level l) const { return isOrderedDLT(getLvlType(l)); }
     bool isUniqueLvl(::mlir::sparse_tensor::Level l) const { return isUniqueDLT(getLvlType(l)); }
 
-    bool isSlice() const {
-      return !getDimSlices().empty();
-    }
+    bool isSlice() const;
+
+    ::mlir::sparse_tensor::SparseTensorDimSliceAttr getDimSlice(::mlir::sparse_tensor::Dimension dim) const;
 
     std::optional<uint64_t> getStaticDimSliceOffset(::mlir::sparse_tensor::Dimension dim) const;
     std::optional<uint64_t> getStaticDimSliceSize(::mlir::sparse_tensor::Dimension dim) const;
     std::optional<uint64_t> getStaticDimSliceStride(::mlir::sparse_tensor::Dimension dim) const;
     std::optional<uint64_t> getStaticLvlSliceOffset(::mlir::sparse_tensor::Level lvl) const;
     std::optional<uint64_t> getStaticLvlSliceSize(::mlir::sparse_tensor::Level lvl) const;
     std::optional<uint64_t> getStaticLvlSliceStride(::mlir::sparse_tensor::Level lvl) const;
   }];

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

 
 DimLevelType SparseTensorEncodingAttr::getLvlType(Level l) const {
   if (!getImpl())
     return DimLevelType::Dense;
   assert(l < getLvlRank() && "Level is out of bounds");
   return getDimLevelType()[l];
 }
 
+bool SparseTensorEncodingAttr::isSlice() const {
+  assert(getImpl() && "Uninitialized SparseTensorEncodingAttr");
+  return !getDimSlices().empty();
+}
+
+SparseTensorDimSliceAttr
+SparseTensorEncodingAttr::getDimSlice(Dimension dim) const {
+  assert(isSlice() && "Is not a slice");
+  const auto dimSlices = getDimSlices();
+  assert(dim < dimSlices.size() && "Dimension is out of bounds");
+  return dimSlices[dim];
+}
+
 std::optional<uint64_t>
 SparseTensorEncodingAttr::getStaticDimSliceOffset(Dimension dim) const {
-  return getDimSlices()[dim].getStaticOffset();
+  return getDimSlice(dim).getStaticOffset();
 }
 
 std::optional<uint64_t>
 SparseTensorEncodingAttr::getStaticDimSliceSize(Dimension dim) const {
-  return getDimSlices()[dim].getStaticSize();
+  return getDimSlice(dim).getStaticSize();
 }
 
 std::optional<uint64_t>
 SparseTensorEncodingAttr::getStaticDimSliceStride(Dimension dim) const {
-  return getDimSlices()[dim].getStaticStride();
+  return getDimSlice(dim).getStaticStride();
 }
 
 std::optional<uint64_t>
 SparseTensorEncodingAttr::getStaticLvlSliceOffset(Level lvl) const {
   // FIXME: `toOrigDim` is deprecated.
   return getStaticDimSliceOffset(toOrigDim(*this, lvl));
 }