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

[Peiming]

rebase

[Peiming]

fix bug

[Peiming]

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

616

This is enabled in the next patch

[bixia]

mlir/test/Dialect/SparseTensor/roundtrip.mlir

217

I am still trying to put everything into a consistent picture, but here are my questions so far:

(1) the "slice" in sparse_tensor.encoding here has three fields, offset/size/stride, but the specifier only has offset/stride (the size is represented in dim size). Maybe sparse_tensor.encoding shouldn't encode "size" info? I meant, you can use the same sparse_tensor.encoding to annotate different sizes of slices, similar to the fact that you can use the same sparse_tensor.encoding to annotate sparse tensors of the same layout but different sizes?

(2) when the "slice" in sparse_tensor.encoding is not provided, would that be equivalent to having offset/stride == 0/1 ? In this PR, we issue an error if "slice" is not in the encoding we we query offset/stride.

[Peiming]

mlir/test/Dialect/SparseTensor/roundtrip.mlir

217

For 1) You are right, and I thought about removing it too (it somehow overlapped with the tensor DimOp) .

But I decided to keep it for future extension, it could be used to stored the actual sizes of the underlying tensor. Then it can be used to remove the check if the slice size equals to the tensor size.

For 2), it is because I want to reduce the number of fields for non slices after lowering (so that we don’t need actual fields to store offset/stride array for non slices). We could of course skip the field for static slice but it would also require us to fill the hole when casting static slice to a dynamic one. See my comment in next patch in ExtractSliceOpConverter.

Both problems are really some design choices, and we can have further discussion on them

[Peiming]

rebase

[aartbik]

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

323

Can you please update the doc in SparseTensorStorageLayout.h as part of this revision.
Good to keep that up to date!

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

363

We probably need to show CSR/CSR_SLICE as part of the example now
(typically, we omit CSR def since it is so widespread, but here it seems essential)

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

619

hehe, okay ;-)

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

40

period at end

[wrengr]

Please be sure to rebase this patch onto D144773, because otherwise this is coig to cause a lot of patch conflicts.

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

323

Is this *really* for dimensions? Since everything else in StorageSpecifierKind is for levels, I'm pretty sure this should be for levels too. (N.B., the name of DimSize is corrected in D144773)

324

Ditto

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

355

ditto

[Peiming]

rebase

[Peiming]

address comments.

[Peiming]

fix indention.

[Peiming]

revert unintended changes.

[Peiming]

rebase.

[aartbik]

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

356

mark this with TODO or FIXME so we can easily find what we promised later ;-)

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

120

We copy

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

76 ↗

(On Diff #503955)

These part needs to move ABOVE examples, to L61

then, add a slide example also to doc

[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,

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

let description = [{
 
     ```
   }];
   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^)? `:` (`from` qualified(type($source))^ `to`)?"
+                                                        " qualified(type($result))";
 }
 
 def SparseTensor_GetStorageSpecifierOp : SparseTensor_Op<"storage_specifier.get", [Pure]>,
     Arguments<(ins SparseTensorStorageSpecifier:$specifier,
                    SparseTensorStorageSpecifierKindAttr:$specifierKind,
                    OptionalAttr<IndexAttr>:$dim)>,
     Results<(outs Index:$result)> {
   let summary = "";

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

 static LogicalResult verifySparsifierGetterSetter(
     StorageSpecifierKind mdKind, std::optional<APInt> lvl,
     TypedValue<StorageSpecifierType> md, Operation *op) {
   if (mdKind == StorageSpecifierKind::ValMemSize && lvl) {
     return op->emitError(
         "redundant level argument for querying value memory size");
   }
 
   const auto enc = md.getType().getEncoding();
   const Level lvlRank = enc.getLvlRank();
 
+  // TODO:
+  //  if (mdKind == StorageSpecifierKind::DimOffset ||
+  //      mdKind == StorageSpecifierKind::DimStride)
+  //    if (!enc.isSlice())
+  //      return op->emitError("requested slice data on non-slice tensor");
+
   if (mdKind != StorageSpecifierKind::ValMemSize) {
     if (!lvl)
       return op->emitError("missing level argument");
 
     const Level l = lvl.value().getZExtValue();
     if (l >= lvlRank)
       return op->emitError("requested level out of bound");
 

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();
   const Level lvlRank = enc.getLvlRank();
 
-  SmallVector<Type, 2> result;
+  SmallVector<Type, 4> result;
   // TODO: how can we get the lowering type for index type in the later pipeline
   // to be consistent? LLVM::StructureType does not allow index fields.
   auto indexType = IntegerType::get(tp.getContext(), 64);
   auto dimSizes = LLVM::LLVMArrayType::get(ctx, indexType, lvlRank);
   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, lvlRank);
+    auto dimStride = LLVM::LLVMArrayType::get(ctx, indexType, lvlRank);
+
+    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 {
 private:
   Value extractField(OpBuilder &builder, Location loc,
-                     ArrayRef<int64_t> indices) {
+                     ArrayRef<int64_t> indices) const {
     return genCast(builder, loc,
                    builder.create<LLVM::ExtractValueOp>(loc, value, indices),
                    builder.getIndexType());
   }
 
   void insertField(OpBuilder &builder, Location loc, ArrayRef<int64_t> indices,
                    Value v) {
     value = builder.create<LLVM::InsertValueOp>(
         loc, value, genCast(builder, loc, v, builder.getIntegerType(64)),
         indices);
   }
 
 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 extractField(builder, loc,
                       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) {
 
   insertField(builder, loc, ArrayRef<int64_t>{kDimSizePosInSpecifier, dim},
               size);
 }
 
+/// Builds IR extracting the pos-th stride from the descriptor.
+Value SpecifierStructBuilder::dimStride(OpBuilder &builder, Location loc,
+                                        unsigned dim) const {
+  return extractField(builder, loc,
+                      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) {
+  insertField(builder, loc, ArrayRef<int64_t>{kDimStridePosInSpecifier, dim},
+              size);
+}
+
 /// Builds IR extracting the pos-th memory size into the descriptor.
 Value SpecifierStructBuilder::memSize(OpBuilder &builder, Location loc,
-                                      unsigned pos) {
+                                      unsigned pos) const {
   return extractField(builder, loc,
                       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) {
   insertField(builder, loc, ArrayRef<int64_t>{kMemSizePosInSpecifier, pos},
               size);
 }
 
+/// 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);
       std::optional<unsigned> dim;
       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 @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> index {
   // expected-error@+1 {{redundant level argument for querying value memory size}}
   %0 = sparse_tensor.storage_specifier.get %arg0 val_mem_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : index
 }
 
+//// -----
+//
+//#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>) -> index {
   // expected-error@+1 {{missing level argument}}
   %0 = sparse_tensor.storage_specifier.get %arg0 idx_mem_sz
        : !sparse_tensor.storage_specifier<#SparseVector>

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 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 : from !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]] : index
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> index {
   %0 = sparse_tensor.storage_specifier.get %arg0 dim_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : index
 }
 
 // -----
 
+#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]] : index
+func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector_Slice>) -> index {
+  %0 = sparse_tensor.storage_specifier.get %arg0 dim_offset at 0
+       : !sparse_tensor.storage_specifier<#SparseVector_Slice>
+  return %0 : index
+}
+
+// -----
+
+#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]] : index
+func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> index {
+  %0 = sparse_tensor.storage_specifier.get %arg0 dim_stride at 0
+       : !sparse_tensor.storage_specifier<#SparseVector>
+  return %0 : index
+}
+
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_set_md(
 //  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>,
 //  CHECK-SAME: %[[I:.*]]: index)
 //       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: index)