[mlir][sparse] introduce operations to query sparse tensor slice offset/strides at the given dimenion

[Peiming]

rebase.

[Peiming]

[mlir][sparse] extends storage specifier operations for slices.

[Peiming]

revert.

[aartbik]

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

212

can you add a bit more in the description other than an example
Having a simple illustrative example in terms of a tensor value helps too

[aartbik]

will these be implemented both along lib path as well as direct IR path?

[wrengr]

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

209

Do you really mean "dimension" here? or is this actually a "level"?

224

ditto

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

675

use getRankedTensorType

682

ditto

[Peiming]

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

209

I did mean dimension. Because slice should be applied on the user-faced shape (if you have high order mapping, the slice is not applied on storage level).

[wrengr]

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

209

Cool cool. I just wanted to make sure we got the right thing, so we don't need to change it later :)

[Peiming]

address comments.

[aartbik]

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

215

are still a work in progress

242

same

[Peiming]

address comments.

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

 
 // The C++ enum for Storage Specifier kind.
 def SparseTensorStorageSpecifierKindEnum
     : I32EnumAttr<"StorageSpecifierKind", "sparse tensor storage specifier kind", [
         I32EnumAttrCase<"DimSize",    0, "dim_sz">,
         I32EnumAttrCase<"PtrMemSize", 1, "ptr_mem_sz">,
         I32EnumAttrCase<"IdxMemSize", 2, "idx_mem_sz">,
         I32EnumAttrCase<"ValMemSize", 3, "val_mem_sz">,
+        I32EnumAttrCase<"DimOffset",  4, "dim_offset">,
+        I32EnumAttrCase<"DimStride",  5, "dim_stride">,
       ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = SparseTensor_Dialect.cppNamespace;
 }
 
 // Define the enum StorageSpecifier kind attribute.
 def SparseTensorStorageSpecifierKindAttr
     : EnumAttr<SparseTensor_Dialect, SparseTensorStorageSpecifierKindEnum,
                "SparseTensorStorageSpecifierKind"> {
    let mnemonic = "kind";
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Traits.
 //===----------------------------------------------------------------------===//
 
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
+def IsSparseTensorSlicePred
+  : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self) && "
+          "  ::mlir::sparse_tensor::getSparseTensorEncoding($_self).isSlice()">;
+
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,
 // `RankedTensorOf`, `AnyRankedTensor`.
 
 class SparseTensorOf<list<Type> allowedTypes>
   : TensorOf<allowedTypes, [IsSparseTensorPred], "sparse tensor">;
 
+class SparseTensorSliceOf<list<Type> allowedTypes>
+  : TensorOf<allowedTypes, [IsSparseTensorSlicePred], "sparse tensor slice">;
+
 def AnySparseTensor : SparseTensorOf<[AnyType]>;
+def AnySparseTensorSlice : SparseTensorSliceOf<[AnyType]>;
 
 class RankedSparseTensorOf<list<Type> allowedTypes>
   : RankedTensorOf<allowedTypes, [IsSparseTensorPred], "ranked sparse tensor">;
 
 def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
 
 #endif // SPARSETENSOR_ATTRDEFS

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

let description = [{
     ```mlir
     %1 = sparse_tensor.values %0 : tensor<64x64xf64, #CSR> to memref<?xf64>
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
   let hasVerifier = 1;
 }
 
+def SparseTensor_ToSliceOffsetOp : SparseTensor_Op<"slice.offset", [Pure]>,
+    Arguments<(ins AnySparseTensorSlice:$slice, IndexAttr:$dim)>,
+    Results<(outs Index:$offset)> {
+  let summary = "Extracts the offset of the sparse tensor slice at the given dimension";
+  let description = [{
+    Example:
+
+    ```mlir
+    %1 = sparse_tensor.slice.offset %0 at 1 : tensor<64x64xf64, #Slice>
+    ```
+  }];
+  let assemblyFormat = "$slice `at` $dim attr-dict `:` type($slice)";
+  let hasVerifier = 1;
+}
+
+def SparseTensor_ToSliceStrideOp : SparseTensor_Op<"slice.stride", [Pure]>,
+    Arguments<(ins AnySparseTensorSlice:$slice, IndexAttr:$dim)>,
+    Results<(outs Index:$stride)> {
+  let summary = "Extracts the stride of the sparse tensor slice at the given dimension";
+  let description = [{
+    Example:
+
+    ```mlir
+    %1 = sparse_tensor.slice.stride %0 at 1 : tensor<64x64xf64, #Slice>
+    ```
+  }];
+  let assemblyFormat = "$slice `at` $dim attr-dict `:` type($slice)";
+  let hasVerifier = 1;
+}
+
 def SparseTensor_StorageSpecifierInitOp : SparseTensor_Op<"storage_specifier.init", [Pure]>,
+    Arguments<(ins Optional<SparseTensorStorageSpecifier>:$source)>,
     Results<(outs SparseTensorStorageSpecifier:$result)> {
   let summary = "";
   let description = [{
     Returns an initial storage specifier value. A storage specifier value holds
     the sizes for tensor dimensions, pointer arrays, index arrays, and the value array.
+    If this is a specifier for slices, it also holds the extra strides/offsets for tensor
+    dimensions.
 
     Example:
 
     ```mlir
     %0 = sparse_tensor.storage_specifier.init :  !sparse_tensor.storage_specifier<#CSR>
+    %1 = sparse_tensor.storage_specifier.init with %src
+         : !sparse_tensor.storage_specifier<#CSR> to
+           !sparse_tensor.storage_specifier<#CSR_SLICE>
     ```
   }];
 
-  let assemblyFormat = "attr-dict `:` qualified(type($result))";
+  let builders = [
+    OpBuilder<(ins "Type":$result),
+    [{
+      build($_builder, $_state, result, Value());
+    }]>
+  ];
+
+  let assemblyFormat = "attr-dict (`with` $source^)? `:` qualified(type($result))"
+                                                  "(`to` qualified(type($source))^)?";
 }
 
 def SparseTensor_GetStorageSpecifierOp : SparseTensor_Op<"storage_specifier.get", [Pure]>,
     Arguments<(ins SparseTensorStorageSpecifier:$specifier,
                    SparseTensorStorageSpecifierKindAttr:$specifierKind,
                    OptionalAttr<IndexAttr>:$dim)>,
     Results<(outs AnyType:$result)> {
   let summary = "";

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

static LogicalResult verifySparsifierGetterSetter(
     StorageSpecifierKind mdKind, std::optional<APInt> dim,
     TypedValue<StorageSpecifierType> md, Operation *op) {
   if (mdKind == StorageSpecifierKind::ValMemSize && dim) {
     return op->emitError(
         "redundant dimension argument for querying value memory size");
   }
 
   auto enc = md.getType().getEncoding();
+  // TODO:
+  //  if (mdKind == StorageSpecifierKind::DimOffset ||
+  //      mdKind == StorageSpecifierKind::DimStride)
+  //    if (!enc.isSlice())
+  //      return op->emitError("requested slice data on non-slice tensor");
+
   ArrayRef<DimLevelType> dlts = enc.getDimLevelType();
   unsigned rank = dlts.size();
 
   if (mdKind != StorageSpecifierKind::ValMemSize) {
     if (!dim)
       return op->emitError("missing dimension argument");
 
     unsigned d = dim.value().getZExtValue();

 LogicalResult ToValuesOp::verify() {
   RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
   MemRefType mtp = getResult().getType().cast<MemRefType>();
   if (ttp.getElementType() != mtp.getElementType())
     return emitError("unexpected mismatch in element types");
   return success();
 }
 
+LogicalResult ToSliceOffsetOp::verify() {
+  auto rank = getSlice().getType().cast<RankedTensorType>().getRank();
+  if (rank <= getDim().getSExtValue() || getDim().getSExtValue() < 0)
+    return emitError("requested dimension out of bound");
+  return success();
+}
+
+LogicalResult ToSliceStrideOp::verify() {
+  auto rank = getSlice().getType().cast<RankedTensorType>().getRank();
+  if (rank <= getDim().getSExtValue() || getDim().getSExtValue() < 0)
+    return emitError("requested dimension out of bound");
+  return success();
+}
+
 LogicalResult GetStorageSpecifierOp::verify() {
   if (failed(verifySparsifierGetterSetter(getSpecifierKind(), getDim(),
                                           getSpecifier(), getOperation()))) {
     return failure();
   }
 
   // Checks the result type
   if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=

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

 using namespace sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
-static SmallVector<Type, 2> getSpecifierFields(StorageSpecifierType tp) {
+static SmallVector<Type, 4> getSpecifierFields(StorageSpecifierType tp) {
   MLIRContext *ctx = tp.getContext();
   auto enc = tp.getEncoding();
   unsigned rank = enc.getDimLevelType().size();
 
-  SmallVector<Type, 2> result;
+  SmallVector<Type, 4> result;
   auto indexType = tp.getSizesType();
   auto dimSizes = LLVM::LLVMArrayType::get(ctx, indexType, rank);
   auto memSizes = LLVM::LLVMArrayType::get(ctx, indexType,
                                            getNumDataFieldsFromEncoding(enc));
   result.push_back(dimSizes);
   result.push_back(memSizes);
+
+  if (enc.isSlice()) {
+    // Extra fields are required for the slice information
+    auto dimOffset = LLVM::LLVMArrayType::get(ctx, indexType, rank);
+    auto dimStride = LLVM::LLVMArrayType::get(ctx, indexType, rank);
+
+    result.push_back(dimOffset);
+    result.push_back(dimStride);
+  }
+
   return result;
 }
 
 static Type convertSpecifier(StorageSpecifierType tp) {
   return LLVM::LLVMStructType::getLiteral(tp.getContext(),
                                           getSpecifierFields(tp));
 }
 
 //===----------------------------------------------------------------------===//
 // Specifier struct builder.
 //===----------------------------------------------------------------------===//
 
 constexpr uint64_t kDimSizePosInSpecifier = 0;
 constexpr uint64_t kMemSizePosInSpecifier = 1;
+constexpr uint64_t kDimOffsetPosInSpecifier = 2;
+constexpr uint64_t kDimStridePosInSpecifier = 3;
 
 class SpecifierStructBuilder : public StructBuilder {
 public:
   explicit SpecifierStructBuilder(Value specifier) : StructBuilder(specifier) {
     assert(value);
   }
 
   // Undef value for dimension sizes, all zero value for memory sizes.
-  static Value getInitValue(OpBuilder &builder, Location loc, Type structType);
+  static Value getInitValue(OpBuilder &builder, Location loc, Type structType,
+                            Value source);
 
-  Value dimSize(OpBuilder &builder, Location loc, unsigned dim);
+  Value dimOffset(OpBuilder &builder, Location loc, unsigned dim) const;
+  void setDimOffset(OpBuilder &builder, Location loc, unsigned dim, Value size);
+
+  Value dimSize(OpBuilder &builder, Location loc, unsigned dim) const;
   void setDimSize(OpBuilder &builder, Location loc, unsigned dim, Value size);
 
-  Value memSize(OpBuilder &builder, Location loc, unsigned pos);
+  Value dimStride(OpBuilder &builder, Location loc, unsigned dim) const;
+  void setDimStride(OpBuilder &builder, Location loc, unsigned dim, Value size);
+
+  Value memSize(OpBuilder &builder, Location loc, unsigned pos) const;
   void setMemSize(OpBuilder &builder, Location loc, unsigned pos, Value size);
+
+  Value memSizeArray(OpBuilder &builder, Location loc) const;
+  void setMemSizeArray(OpBuilder &builder, Location loc, Value array);
 };
 
 Value SpecifierStructBuilder::getInitValue(OpBuilder &builder, Location loc,
-                                           Type structType) {
+                                           Type structType, Value source) {
   Value metaData = builder.create<LLVM::UndefOp>(loc, structType);
   SpecifierStructBuilder md(metaData);
+  if (!source) {
     auto memSizeArrayType = structType.cast<LLVM::LLVMStructType>()
                                 .getBody()[kMemSizePosInSpecifier]
                                 .cast<LLVM::LLVMArrayType>();
 
     Value zero = constantZero(builder, loc, memSizeArrayType.getElementType());
     // Fill memSizes array with zero.
     for (int i = 0, e = memSizeArrayType.getNumElements(); i < e; i++)
       md.setMemSize(builder, loc, i, zero);
-
+  } else {
+    // We copies non-slice information (memory sizes array) from source
+    SpecifierStructBuilder sourceMd(source);
+    md.setMemSizeArray(builder, loc, sourceMd.memSizeArray(builder, loc));
+  }
   return md;
 }
 
+/// Builds IR extracting the pos-th offset from the descriptor.
+Value SpecifierStructBuilder::dimOffset(OpBuilder &builder, Location loc,
+                                        unsigned dim) const {
+  return builder.create<LLVM::ExtractValueOp>(
+      loc, value, ArrayRef<int64_t>({kDimOffsetPosInSpecifier, dim}));
+}
+
+/// Builds IR inserting the pos-th offset into the descriptor.
+void SpecifierStructBuilder::setDimOffset(OpBuilder &builder, Location loc,
+                                          unsigned dim, Value size) {
+  value = builder.create<LLVM::InsertValueOp>(
+      loc, value, size, ArrayRef<int64_t>({kDimOffsetPosInSpecifier, dim}));
+}
+
 /// Builds IR inserting the pos-th size into the descriptor.
 Value SpecifierStructBuilder::dimSize(OpBuilder &builder, Location loc,
-                                      unsigned dim) {
+                                      unsigned dim) const {
   return builder.create<LLVM::ExtractValueOp>(
       loc, value, ArrayRef<int64_t>({kDimSizePosInSpecifier, dim}));
 }
 
 /// Builds IR inserting the pos-th size into the descriptor.
 void SpecifierStructBuilder::setDimSize(OpBuilder &builder, Location loc,
                                         unsigned dim, Value size) {
   value = builder.create<LLVM::InsertValueOp>(
       loc, value, size, ArrayRef<int64_t>({kDimSizePosInSpecifier, dim}));
 }
 
+/// Builds IR extracting the pos-th stride from the descriptor.
+Value SpecifierStructBuilder::dimStride(OpBuilder &builder, Location loc,
+                                        unsigned dim) const {
+  return builder.create<LLVM::ExtractValueOp>(
+      loc, value, ArrayRef<int64_t>({kDimStridePosInSpecifier, dim}));
+}
+
+/// Builds IR inserting the pos-th stride into the descriptor.
+void SpecifierStructBuilder::setDimStride(OpBuilder &builder, Location loc,
+                                          unsigned dim, Value size) {
+  value = builder.create<LLVM::InsertValueOp>(
+      loc, value, size, ArrayRef<int64_t>({kDimStridePosInSpecifier, dim}));
+}
+
 /// Builds IR extracting the pos-th memory size into the descriptor.
 Value SpecifierStructBuilder::memSize(OpBuilder &builder, Location loc,
-                                      unsigned pos) {
+                                      unsigned pos) const {
   return builder.create<LLVM::ExtractValueOp>(
       loc, value, ArrayRef<int64_t>({kMemSizePosInSpecifier, pos}));
 }
 
 /// Builds IR inserting the pos-th memory size into the descriptor.
 void SpecifierStructBuilder::setMemSize(OpBuilder &builder, Location loc,
                                         unsigned pos, Value size) {
   value = builder.create<LLVM::InsertValueOp>(
       loc, value, size, ArrayRef<int64_t>({kMemSizePosInSpecifier, pos}));
 }
 
+/// Builds IR extracting the memory size array from the descriptor.
+Value SpecifierStructBuilder::memSizeArray(OpBuilder &builder,
+                                           Location loc) const {
+  return builder.create<LLVM::ExtractValueOp>(loc, value,
+                                              kMemSizePosInSpecifier);
+}
+
+/// Builds IR inserting the memory size array into the descriptor.
+void SpecifierStructBuilder::setMemSizeArray(OpBuilder &builder, Location loc,
+                                             Value array) {
+  value = builder.create<LLVM::InsertValueOp>(loc, value, array,
+                                              kMemSizePosInSpecifier);
+}
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // The sparse storage specifier type converter (defined in Passes.h).
 //===----------------------------------------------------------------------===//
 
 StorageSpecifierToLLVMTypeConverter::StorageSpecifierToLLVMTypeConverter() {
   addConversion([](Type type) { return type; });

 public:
   using OpAdaptor = typename SourceOp::Adaptor;
   using OpConversionPattern<SourceOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     SpecifierStructBuilder spec(adaptor.getSpecifier());
-    Value v;
-    if (op.getSpecifierKind() == StorageSpecifierKind::DimSize) {
-      v = Base::onDimSize(rewriter, op, spec,
-                          op.getDim().value().getZExtValue());
-    } else {
+    switch (op.getSpecifierKind()) {
+    case StorageSpecifierKind::DimOffset: {
+      Value v =
+          Base::onDimOffset(rewriter, op, spec, (*op.getDim()).getZExtValue());
+      rewriter.replaceOp(op, v);
+      return success();
+    }
+    case StorageSpecifierKind::DimSize: {
+      Value v =
+          Base::onDimSize(rewriter, op, spec, (*op.getDim()).getZExtValue());
+      rewriter.replaceOp(op, v);
+      return success();
+    }
+    case StorageSpecifierKind::DimStride: {
+      Value v =
+          Base::onDimStride(rewriter, op, spec, (*op.getDim()).getZExtValue());
+      rewriter.replaceOp(op, v);
+      return success();
+    }
+    case StorageSpecifierKind::IdxMemSize:
+    case StorageSpecifierKind::PtrMemSize:
+    case StorageSpecifierKind::ValMemSize: {
       auto enc = op.getSpecifier().getType().getEncoding();
       StorageLayout layout(enc);
       Optional<unsigned> dim = std::nullopt;
       if (op.getDim())
-        dim = op.getDim().value().getZExtValue();
+        dim = (*op.getDim()).getZExtValue();
       unsigned idx = layout.getMemRefFieldIndex(op.getSpecifierKind(), dim);
-      v = Base::onMemSize(rewriter, op, spec, idx);
-    }
-
+      Value v = Base::onMemSize(rewriter, op, spec, idx);
       rewriter.replaceOp(op, v);
       return success();
     }
+    }
+    llvm_unreachable("unrecognized specifer kind");
+  }
 };
 
 struct StorageSpecifierSetOpConverter
     : public SpecifierGetterSetterOpConverter<StorageSpecifierSetOpConverter,
                                               SetStorageSpecifierOp> {
   using SpecifierGetterSetterOpConverter::SpecifierGetterSetterOpConverter;
+
+  static Value onDimOffset(OpBuilder &builder, SetStorageSpecifierOp op,
+                           SpecifierStructBuilder &spec, unsigned d) {
+    spec.setDimOffset(builder, op.getLoc(), d, op.getValue());
+    return spec;
+  }
+
   static Value onDimSize(OpBuilder &builder, SetStorageSpecifierOp op,
                          SpecifierStructBuilder &spec, unsigned d) {
     spec.setDimSize(builder, op.getLoc(), d, op.getValue());
     return spec;
   }
 
+  static Value onDimStride(OpBuilder &builder, SetStorageSpecifierOp op,
+                           SpecifierStructBuilder &spec, unsigned d) {
+    spec.setDimStride(builder, op.getLoc(), d, op.getValue());
+    return spec;
+  }
+
   static Value onMemSize(OpBuilder &builder, SetStorageSpecifierOp op,
                          SpecifierStructBuilder &spec, unsigned i) {
     spec.setMemSize(builder, op.getLoc(), i, op.getValue());
     return spec;
   }
 };
 
 struct StorageSpecifierGetOpConverter
     : public SpecifierGetterSetterOpConverter<StorageSpecifierGetOpConverter,
                                               GetStorageSpecifierOp> {
   using SpecifierGetterSetterOpConverter::SpecifierGetterSetterOpConverter;
+
+  static Value onDimOffset(OpBuilder &builder, GetStorageSpecifierOp op,
+                           const SpecifierStructBuilder &spec, unsigned d) {
+    return spec.dimOffset(builder, op.getLoc(), d);
+  }
+
   static Value onDimSize(OpBuilder &builder, GetStorageSpecifierOp op,
-                         SpecifierStructBuilder &spec, unsigned d) {
+                         const SpecifierStructBuilder &spec, unsigned d) {
     return spec.dimSize(builder, op.getLoc(), d);
   }
+
+  static Value onDimStride(OpBuilder &builder, GetStorageSpecifierOp op,
+                           const SpecifierStructBuilder &spec, unsigned d) {
+    return spec.dimStride(builder, op.getLoc(), d);
+  }
+
   static Value onMemSize(OpBuilder &builder, GetStorageSpecifierOp op,
-                         SpecifierStructBuilder &spec, unsigned i) {
+                         const SpecifierStructBuilder &spec, unsigned i) {
     return spec.memSize(builder, op.getLoc(), i);
   }
 };
 
 struct StorageSpecifierInitOpConverter
     : public OpConversionPattern<StorageSpecifierInitOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(StorageSpecifierInitOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type llvmType = getTypeConverter()->convertType(op.getResult().getType());
-    rewriter.replaceOp(op, SpecifierStructBuilder::getInitValue(
-                               rewriter, op.getLoc(), llvmType));
+    rewriter.replaceOp(
+        op, SpecifierStructBuilder::getInitValue(
+                rewriter, op.getLoc(), llvmType, adaptor.getSource()));
     return success();
   }
 };
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 void mlir::populateStorageSpecifierToLLVMPatterns(TypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<StorageSpecifierGetOpConverter, StorageSpecifierSetOpConverter,
                StorageSpecifierInitOpConverter>(converter,
                                                 patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/invalid.mlir

 func.func @mismatch_values_types(%arg0: tensor<?xf64, #SparseVector>) -> memref<?xf32> {
   // expected-error@+1 {{unexpected mismatch in element types}}
   %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf32>
   return %0 : memref<?xf32>
 }
 
 // -----
 
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+func.func @sparse_slice_offset(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  // expected-error@+1 {{requested dimension out of bound}}
+  %0 = sparse_tensor.slice.offset %arg0 at 2 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+func.func @sparse_slice_stride(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  // expected-error@+1 {{requested dimension out of bound}}
+  %0 = sparse_tensor.slice.stride %arg0 at 2 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   // expected-error@+1 {{redundant dimension argument for querying value memory size}}
   %0 = sparse_tensor.storage_specifier.get %arg0 val_mem_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }
 
+//// -----
+//
+//#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+//
+//func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
+//  // _e_xpected-error@+1 {{requested slice data on non-slice tensor}}
+//  %0 = sparse_tensor.storage_specifier.get %arg0 dim_offset at 0
+//       : !sparse_tensor.storage_specifier<#SparseVector> to i64
+//  return %0 : i64
+//}
+
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   // expected-error@+1 {{missing dimension argument}}
   %0 = sparse_tensor.storage_specifier.get %arg0 idx_mem_sz
        : !sparse_tensor.storage_specifier<#SparseVector> to i64

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: return %[[T]] : memref<?xf64>
 func.func @sparse_values(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
 // -----
 
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+// CHECK-LABEL: func @sparse_slice_offset(
+//  CHECK-SAME: %[[A:.*]]: tensor<2x8xf64, #{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.slice.offset %[[A]] at 1 : tensor<2x8xf64, #{{.*}}>
+//       CHECK: return %[[T]] : index
+func.func @sparse_slice_offset(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  %0 = sparse_tensor.slice.offset %arg0 at 1 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+// CHECK-LABEL: func @sparse_slice_stride(
+//  CHECK-SAME: %[[A:.*]]: tensor<2x8xf64, #{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.slice.stride %[[A]] at 1 : tensor<2x8xf64, #{{.*}}>
+//       CHECK: return %[[T]] : index
+func.func @sparse_slice_stride(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  %0 = sparse_tensor.slice.stride %arg0 at 1 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_metadata_init(
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#{{.*}}>
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_metadata_init() -> !sparse_tensor.storage_specifier<#SparseVector> {
   %0 = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : !sparse_tensor.storage_specifier<#SparseVector>
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector_Slice = #sparse_tensor.encoding<{
+  dimLevelType = ["compressed"],
+  slice = [ (?, ?, ?) ]
+}>
+
+// CHECK-LABEL: func @sparse_metadata_init(
+//  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>
+//       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.init with %[[A]] :
+//       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
+func.func @sparse_metadata_init(%src : !sparse_tensor.storage_specifier<#SparseVector>)
+                                    -> !sparse_tensor.storage_specifier<#SparseVector_Slice> {
+  %0 = sparse_tensor.storage_specifier.init with %src : !sparse_tensor.storage_specifier<#SparseVector>
+                                                     to !sparse_tensor.storage_specifier<#SparseVector_Slice>
+  return %0 : !sparse_tensor.storage_specifier<#SparseVector_Slice>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_get_md(
 //  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.get %[[A]] dim_sz at 0
 //       CHECK: return %[[T]] : i64
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   %0 = sparse_tensor.storage_specifier.get %arg0 dim_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }
 
 // -----
 
+#SparseVector_Slice = #sparse_tensor.encoding<{
+  dimLevelType = ["compressed"],
+  slice = [ (?, ?, ?) ]
+}>
+
+// CHECK-LABEL: func @sparse_get_md(
+//  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>
+//       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.get %[[A]] dim_offset at 0
+//       CHECK: return %[[T]] : i64
+func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector_Slice>) -> i64 {
+  %0 = sparse_tensor.storage_specifier.get %arg0 dim_offset at 0
+       : !sparse_tensor.storage_specifier<#SparseVector_Slice> to i64
+  return %0 : i64
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{
+  dimLevelType = ["compressed"],
+  slice = [ (?, ?, ?) ]
+}>
+
+// CHECK-LABEL: func @sparse_get_md(
+//  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>
+//       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.get %[[A]] dim_stride at 0
+//       CHECK: return %[[T]] : i64
+func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
+  %0 = sparse_tensor.storage_specifier.get %arg0 dim_stride at 0
+       : !sparse_tensor.storage_specifier<#SparseVector> to i64
+  return %0 : i64
+}
+
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_set_md(
 //  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>,
 //  CHECK-SAME: %[[I:.*]]: i64)
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.set %[[A]] dim_sz at 0 with %[[I]]
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_set_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>, %arg1: i64)