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

[Peiming]

rebase

[Peiming]

fix bug

[Peiming]

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

620

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

371

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

386

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

623

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

371

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)

372

Ditto

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

367

ditto

[Peiming]

rebase

[Peiming]

address comments.

[Peiming]

fix indention.

[Peiming]

revert unintended changes.

[Peiming]

rebase.

[aartbik]

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

371

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

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

123

We copy

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

81

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<"LvlSize",    0, "lvl_sz">,
         I32EnumAttrCase<"PosMemSize", 1, "pos_mem_sz">,
         I32EnumAttrCase<"CrdMemSize", 2, "crd_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 level-sizes, position arrays, coordinate arrays,
     and the value array.
+    If this is a specifier for slices, it also holds the extra strides/offsets
+    for each tensor dimension.
+
+    TODO: The sparse tensor slice support is currently in a unstable state, and
+    is subject to change in the future.
 
     Example:
 
     ```mlir
+    #CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ]}>
+    #CSR_SLICE = #sparse_tensor.encoding<{
+      dimLevelType = [ "dense", "compressed" ],
+      slice = [ (1, 4, 1), (1, 4, 2) ]
+    }>
+
     %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<LevelAttr>:$level)>,
     Results<(outs Index:$result)> {
   let summary = "";

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

 static LogicalResult verifySparsifierGetterSetter(
     StorageSpecifierKind mdKind, std::optional<Level> 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();
     if (l >= lvlRank)
       return op->emitError("requested level is out of bounds");
 

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 sizeType = IntegerType::get(tp.getContext(), 64);
   auto lvlSizes = LLVM::LLVMArrayType::get(ctx, sizeType, lvlRank);
   auto memSizes = LLVM::LLVMArrayType::get(ctx, sizeType,
                                            getNumDataFieldsFromEncoding(enc));
   result.push_back(lvlSizes);
   result.push_back(memSizes);
+
+  if (enc.isSlice()) {
+    // Extra fields are required for the slice information.
+    auto dimOffset = LLVM::LLVMArrayType::get(ctx, sizeType, lvlRank);
+    auto dimStride = LLVM::LLVMArrayType::get(ctx, sizeType, 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 kLvlSizePosInSpecifier = 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 level-sizes, all zero values for memory-sizes.
-  static Value getInitValue(OpBuilder &builder, Location loc, Type structType);
+  // Undef value for dimension sizes, all zero value for memory sizes.
+  static Value getInitValue(OpBuilder &builder, Location loc, Type structType,
+                            Value source);
 
-  Value lvlSize(OpBuilder &builder, Location loc, Level lvl);
+  Value lvlSize(OpBuilder &builder, Location loc, Level lvl) const;
   void setLvlSize(OpBuilder &builder, Location loc, Level lvl, Value size);
 
-  Value memSize(OpBuilder &builder, Location loc, FieldIndex fidx);
+  Value dimOffset(OpBuilder &builder, Location loc, Dimension dim) const;
+  void setDimOffset(OpBuilder &builder, Location loc, Dimension dim,
+                    Value size);
+
+  Value dimStride(OpBuilder &builder, Location loc, Dimension dim) const;
+  void setDimStride(OpBuilder &builder, Location loc, Dimension dim,
+                    Value size);
+
+  Value memSize(OpBuilder &builder, Location loc, FieldIndex fidx) const;
   void setMemSize(OpBuilder &builder, Location loc, FieldIndex fidx,
                   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 copy 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,
+                                        Dimension dim) const {
+  return builder.create<LLVM::ExtractValueOp>(
+      loc, value,
+      ArrayRef<int64_t>({kDimOffsetPosInSpecifier, static_cast<int64_t>(dim)}));
+}
+
+/// Builds IR inserting the pos-th offset into the descriptor.
+void SpecifierStructBuilder::setDimOffset(OpBuilder &builder, Location loc,
+                                          Dimension dim, Value size) {
+  value = builder.create<LLVM::InsertValueOp>(
+      loc, value, size,
+      ArrayRef<int64_t>({kDimOffsetPosInSpecifier, static_cast<int64_t>(dim)}));
+}
+
 /// Builds IR extracting the `lvl`-th level-size from the descriptor.
 Value SpecifierStructBuilder::lvlSize(OpBuilder &builder, Location loc,
-                                      Level lvl) {
+                                      Level lvl) const {
   // This static_cast makes the narrowing of `lvl` explicit, as required
   // by the braces notation for the ctor.
   return extractField(
       builder, loc,
       ArrayRef<int64_t>{kLvlSizePosInSpecifier, static_cast<int64_t>(lvl)});
 }
 
 /// Builds IR inserting the `lvl`-th level-size into the descriptor.
 void SpecifierStructBuilder::setLvlSize(OpBuilder &builder, Location loc,
                                         Level lvl, Value size) {
   // This static_cast makes the narrowing of `lvl` explicit, as required
   // by the braces notation for the ctor.
   insertField(
       builder, loc,
       ArrayRef<int64_t>{kLvlSizePosInSpecifier, static_cast<int64_t>(lvl)},
       size);
 }
 
-/// Builds IR extracting the `fidx`-th memory-size from the descriptor.
+/// Builds IR extracting the pos-th stride from the descriptor.
+Value SpecifierStructBuilder::dimStride(OpBuilder &builder, Location loc,
+                                        Dimension dim) const {
+  return extractField(
+      builder, loc,
+      ArrayRef<int64_t>{kDimStridePosInSpecifier, static_cast<int64_t>(dim)});
+}
+
+/// Builds IR inserting the pos-th stride into the descriptor.
+void SpecifierStructBuilder::setDimStride(OpBuilder &builder, Location loc,
+                                          Dimension dim, Value size) {
+  insertField(
+      builder, loc,
+      ArrayRef<int64_t>{kDimStridePosInSpecifier, static_cast<int64_t>(dim)},
+      size);
+}
+
+/// Builds IR extracting the pos-th memory size into the descriptor.
 Value SpecifierStructBuilder::memSize(OpBuilder &builder, Location loc,
-                                      FieldIndex fidx) {
-  return extractField(builder, loc,
-                      ArrayRef<int64_t>{kMemSizePosInSpecifier, fidx});
+                                      FieldIndex fidx) const {
+  return extractField(
+      builder, loc,
+      ArrayRef<int64_t>{kMemSizePosInSpecifier, static_cast<int64_t>(fidx)});
 }
 
 /// Builds IR inserting the `fidx`-th memory-size into the descriptor.
 void SpecifierStructBuilder::setMemSize(OpBuilder &builder, Location loc,
                                         FieldIndex fidx, Value size) {
-  insertField(builder, loc, ArrayRef<int64_t>{kMemSizePosInSpecifier, fidx},
-              size);
+  insertField(
+      builder, loc,
+      ArrayRef<int64_t>{kMemSizePosInSpecifier, static_cast<int64_t>(fidx)},
 }
 
+/// 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::LvlSize) {
-      assert(op.getLevel().has_value());
-      v = Base::onLvlSize(rewriter, op, spec, op.getLevel().value());
-    } else {
+    switch (op.getSpecifierKind()) {
+    case StorageSpecifierKind::LvlSize: {
+      Value v = Base::onLvlSize(rewriter, op, spec, (*op.getLevel()));
+      rewriter.replaceOp(op, v);
+      return success();
+    }
+    case StorageSpecifierKind::DimOffset: {
+      Value v = Base::onDimOffset(rewriter, op, spec, (*op.getLevel()));
+      rewriter.replaceOp(op, v);
+      return success();
+    }
+    case StorageSpecifierKind::DimStride: {
+      Value v = Base::onDimStride(rewriter, op, spec, (*op.getLevel()));
+      rewriter.replaceOp(op, v);
+      return success();
+    }
+    case StorageSpecifierKind::CrdMemSize:
+    case StorageSpecifierKind::PosMemSize:
+    case StorageSpecifierKind::ValMemSize: {
       auto enc = op.getSpecifier().getType().getEncoding();
       StorageLayout layout(enc);
-      FieldIndex fidx =
-          layout.getMemRefFieldIndex(op.getSpecifierKind(), op.getLevel());
-      v = Base::onMemSize(rewriter, op, spec, fidx);
-    }
-
+      std::optional<unsigned> lvl;
+      if (op.getLevel())
+        lvl = (*op.getLevel());
+      unsigned idx = layout.getMemRefFieldIndex(op.getSpecifierKind(), lvl);
+      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 onLvlSize(OpBuilder &builder, SetStorageSpecifierOp op,
                          SpecifierStructBuilder &spec, Level lvl) {
     spec.setLvlSize(builder, op.getLoc(), lvl, op.getValue());
     return spec;
   }
 
+  static Value onDimOffset(OpBuilder &builder, SetStorageSpecifierOp op,
+                           SpecifierStructBuilder &spec, Dimension d) {
+    spec.setDimOffset(builder, op.getLoc(), d, op.getValue());
+    return spec;
+  }
+
+  static Value onDimStride(OpBuilder &builder, SetStorageSpecifierOp op,
+                           SpecifierStructBuilder &spec, Dimension d) {
+    spec.setDimStride(builder, op.getLoc(), d, op.getValue());
+    return spec;
+  }
+
   static Value onMemSize(OpBuilder &builder, SetStorageSpecifierOp op,
                          SpecifierStructBuilder &spec, FieldIndex fidx) {
     spec.setMemSize(builder, op.getLoc(), fidx, op.getValue());
     return spec;
   }
 };
 
 struct StorageSpecifierGetOpConverter
     : public SpecifierGetterSetterOpConverter<StorageSpecifierGetOpConverter,
                                               GetStorageSpecifierOp> {
   using SpecifierGetterSetterOpConverter::SpecifierGetterSetterOpConverter;
+
   static Value onLvlSize(OpBuilder &builder, GetStorageSpecifierOp op,
                          SpecifierStructBuilder &spec, Level lvl) {
     return spec.lvlSize(builder, op.getLoc(), lvl);
   }
+
+  static Value onDimOffset(OpBuilder &builder, GetStorageSpecifierOp op,
+                           const SpecifierStructBuilder &spec, Dimension d) {
+    return spec.dimOffset(builder, op.getLoc(), d);
+  }
+
+  static Value onDimStride(OpBuilder &builder, GetStorageSpecifierOp op,
+                           const SpecifierStructBuilder &spec, Dimension d) {
+    return spec.dimStride(builder, op.getLoc(), d);
+  }
+
   static Value onMemSize(OpBuilder &builder, GetStorageSpecifierOp op,
                          SpecifierStructBuilder &spec, FieldIndex fidx) {
     return spec.memSize(builder, op.getLoc(), fidx);
   }
 };
 
 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/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

 // In addition, for a "trailing COO region", defined as a compressed level
 // followed by one or more singleton levels, the default SOA storage that
 // is inherent to the TACO format is optimized into an AOS storage where
 // all coordinates of a stored element appear consecutively.  In such cases,
 // a special operation (sparse_tensor.coordinates_buffer) must be used to
 // access the AOS coordinates array. In the code below, the method `getCOOStart`
 // is used to find the start of the "trailing COO region".
 //
+// If the sparse tensor is a slice (produced by `tensor.extract_slice`
+// operation), instead of allocating a new sparse tensor for it, it reuses the
+// same sets of MemRefs but attaching a additional set of slicing-metadata for
+// per-dimension slice offset and stride.
+//
 // Examples.
 //
 // #CSR storage of 2-dim matrix yields
 //   memref<?xindex>                           ; positions-1
 //   memref<?xindex>                           ; coordinates-1
 //   memref<?xf64>                             ; values
 //   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
 //
 // #COO storage of 2-dim matrix yields
 //   memref<?xindex>,                          ; positions-0, essentially [0,sz]
 //   memref<?xindex>                           ; AOS coordinates storage
 //   memref<?xf64>                             ; values
 //   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
 //
+// Slice on #COO storage of 2-dim matrix yields
+//   ;; Inherited from the original sparse tensors
+//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
+//   memref<?xindex>                           ; AOS coordinates storage
+//   memref<?xf64>                             ; values
+//   struct<(array<2 x i64>, array<3 x i64>,   ; lvl0, lvl1, 3xsizes
+//   ;; Extra slicing-metadata
+//           array<2 x i64>, array<2 x i64>)>) ; dim offset, dim stride.
+//
 //===----------------------------------------------------------------------===//
 
 enum class SparseTensorFieldKind : uint32_t {
   StorageSpec = 0,
   PosMemRef = 1,
   CrdMemRef = 2,
   ValMemRef = 3
 };

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 crd_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]] lvl_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 lvl_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]] lvl_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)