[mlir][sparse] support dynamic sparse tensor slices.

[Peiming]

add check test

[Peiming]

add matmul dynamic slice integrate test

[Peiming]

rebase.

[Peiming]

rebase.

[Peiming]

rebase

[Peiming]

fix crash.

[Peiming]

reabse

[Peiming]

rebase.

[Peiming]

rebase

[Peiming]

fix typos

[aartbik]

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

500

while you are here:

that do not alter (not does)

510–511

can we make this / an and
(it reads like a division right now)

510–511

I would continue this comment on the same line L585, since it just continues the reasoning

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

342

Generates code

347

Generates code

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

93

put this before the if

(so that First, Second, Third are at same leve)

106

put this Third before the if

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

193

Generates

274

related

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

103

an attribute

110

directly return

in fact, write this as

if (Value v = ...)

return v

...

670

xxx -> op?

1067–1071

we cannot generate

[Peiming]

address comments.

[Peiming]

revert unintended changes.

[wrengr]

A bunch of repetitive comments about using the proper types. If you rebase this after D145756, then that handles a bunch of them.

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

561

ditto

571

ditto

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

345

If this is actually a dimension, then use the Dimension type. (Otherwise, should rename the variable)

350

If this is actually a dimension, then use the Dimension type. (Otherwise, should rename the variable)

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

46

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

52

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

59

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

66–70

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

79

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

86

I've noticed you use this variable name a lot in this file. What exactly do you mean by it? Seems like there should be a more informative name

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

197

Is this actually (1) a memref of all the coordinates associated with a single element; in which case it should be "coords"; or (2) a single coordinate value; in which case it should be "crd"?

198

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

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

105

I'm pretty sure there's already something that does this in StaticValueUtils.h

[wrengr]

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

565

This should just check the offset local variable, rather than calling the function again. And for brevity it ought to use the syntax if (auto offset = ...) return ...

575

ditto

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

48

You should add a FIXME comment about how toOrigDim is deprecated.

54

ditto

115

Should just use a foreach loop instead; i.e., for (auto cond : ArrayRef(conds).drop_front())

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

193

"are"

194

"in"

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

1049

Instead of doing this manually, it would be a lot easier to follow if you use for (... : llvm::enumerate(llvm::zip(...))) instead

Also, should use the Dimension type

1062–1063

yes please :)

[Peiming]

address comments.

[Peiming]

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

86

it means transformed coordinates.

115

conds is an SmallVector, and it does not provide a drop_front, I prefer to start with 1 here instead of casting it to ArrayRef before iteration.

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

198

The header for Level is not included yet, will do it later to avoid more conflicts.

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

105

Thanks for letting me know.

[Peiming]

add bazel dependence.

[wrengr]

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

115

I still think the foreach loop is cleaner, since it captures the intention directly and avoids introducing extraneous details like the index variable and getting the size.

233–239

Elsewhere the code always uses the name "pHi", so that name should be used here too

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

1049

Elsewhere throughout the MLIR code folks name this "it". Also, you should make it auto& to avoid an unnecessary copy

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

344

Did you mean to add this back in?

Back when I renamed stuff in this file I cleaned all those up (since at that point, at least, the StorageSpecifierKind argument always took a Level), so the "[CLARIFY_DIM_LVL]" note no longer exists. If this CL reintroduces ambiguity/confusion issues, then you should give a longer comment explaining what the problem is

[Peiming]

address comments.

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

115

Okay. Changed

233–239

This is slightly different here actually, maybe I will just call it lvlSz to be more precise.

[wrengr]

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

233–239

hmm, but it comes from createOrFoldDimOp(d) so shouldn't it be dimSz? Then again, since it's being stored into lvlSizes, that does suggest it ought to be lvlSz... So can you add a comment explaining why createOrFoldDimOp(d) is safe/correct to use here? (And presumably it'll be a TODO/FIXME, since I'm guessing this is just a stand-in for once I finish up the proper dim<->lvl conversion class)

[Peiming]

address comment.

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 = [{
     %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

 
 //===----------------------------------------------------------------------===//
 // SparseTensorDialect Types.
 //===----------------------------------------------------------------------===//
 
 /// We normalized sparse tensor encoding attribute by always using
 /// ordered/unique DLT such that "compressed-nu-no" and "compressed-nu" (as well
 /// as other variants) lead to the same storage specifier type, and stripping
 /// irrelevant fields that does not alter the sparse tensor memory layout.
 static SparseTensorEncodingAttr
 getNormalizedEncodingForSpecifier(SparseTensorEncodingAttr enc) {
   SmallVector<DimLevelType> dlts;
   for (auto dlt : enc.getDimLevelType())
     dlts.push_back(*getDimLevelType(*getLevelFormat(dlt), true, true));
 
   return SparseTensorEncodingAttr::get(
       enc.getContext(), dlts,
       AffineMap(), // dimOrdering (irrelavant to storage speicifer)
       AffineMap(), // highLvlOrdering (irrelavant to storage specifer)
-      enc.getPointerBitWidth(), enc.getIndexBitWidth(),
-      // FIXME: we should keep the slice information, for now it is okay as only
-      // constant can be used for slice
-      ArrayRef<SparseTensorDimSliceAttr>{} /*enc.getDimSlices()*/);
+      enc.getPointerBitWidth(), enc.getIndexBitWidth(), enc.getDimSlices());
 }
 
 StorageSpecifierType
 StorageSpecifierType::get(MLIRContext *ctx, SparseTensorEncodingAttr encoding) {
   return Base::get(ctx, getNormalizedEncodingForSpecifier(encoding));
 }
 
 IntegerType StorageSpecifierType::getSizesType() const {

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

Value genToIndices(OpBuilder &builder, Location loc, Value tensor, uint64_t d,
                    uint64_t cooStart);
 
 /// Infers the result type and generates ToValuesOp.
 Value genToValues(OpBuilder &builder, Location loc, Value tensor);
 
 /// Generates code to retrieve the values size for the sparse tensor.
 Value genValMemSize(OpBuilder &builder, Location loc, Value tensor);
 
+/// Generate scode to retrieve the slice offset for the sparse tensor slice,
+/// return a constant if the offset is statically known.
+Value createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc, Value tensor,
+                                unsigned dim);
+
+/// Generate scode to retrieve the slice slice for the sparse tensor slice,
+/// return a constant if the offset is statically known.
+Value createOrFoldSliceStrideOp(OpBuilder &builder, Location loc, Value tensor,
+                                unsigned dim);
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

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

Value sparse_tensor::genToValues(OpBuilder &builder, Location loc,
   RankedTensorType srcTp = tensor.getType().cast<RankedTensorType>();
   Type valTp = get1DMemRefType(srcTp.getElementType(),
                                /*withLayout=*/false);
   return builder.create<ToValuesOp>(loc, valTp, tensor);
 }
 
 Value sparse_tensor::genValMemSize(OpBuilder &builder, Location loc,
                                    Value tensor) {
-  SmallVector<Value> fields;
-  auto desc = getMutDescriptorFromTensorTuple(tensor, fields);
+  auto desc = getDescriptorFromTensorTuple(tensor);
   return desc.getValMemSize(builder, loc);
-}
+}
+
+Value sparse_tensor::createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc,
+                                               Value tensor, unsigned dim) {
+  auto enc = getSparseTensorEncoding(tensor.getType());
+  assert(enc && enc.isSlice());
+  std::optional<unsigned> offset = enc.getStaticDimSliceOffset(dim);
+  if (enc.getStaticDimSliceOffset(dim))
+    return constantIndex(builder, loc, *offset);
+  return builder.create<ToSliceOffsetOp>(loc, tensor, APInt(64, dim));
+}
+
+Value sparse_tensor::createOrFoldSliceStrideOp(OpBuilder &builder, Location loc,
+                                               Value tensor, unsigned dim) {
+  auto enc = getSparseTensorEncoding(tensor.getType());
+  assert(enc && enc.isSlice());
+  std::optional<unsigned> offset = enc.getStaticDimSliceStride(dim);
+  if (enc.getStaticDimSliceStride(dim))
+    return constantIndex(builder, loc, *offset);
+  return builder.create<ToSliceStrideOp>(loc, tensor, APInt(64, dim));
+}

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

struct LoopLevelInfo {
     Block *const userCodeBlock; // the block holding user' generated code.
     const Value iv;             // the induction variable for the loop
   };
 
   /// Linearizes address for dense dimension (i.e., p = (i * d0) + j).
   Value genAddress(OpBuilder &builder, Location loc, size_t tid, size_t dim,
                    Value iv);
 
+  /// Generate a predicate to determine whether the tranformed coordinates is
+  /// on the given slice.
+  /// Returns std::pair<Transformed coordinates, Predicate>
+  std::pair<Value, Value> genSliceLegitPredicate(OpBuilder &builder,
+                                                 Location loc, Value coord,
+                                                 unsigned tid, unsigned lvl);
   bool isOutputTensor(size_t tid) {
     return hasOutput && tid == tensors.size() - 1;
   }
 
   bool isSparseOutput(size_t tid) { return isOutputTensor(tid) && isSparseOut; }
 
   /// Setups [lo, hi] for iterating tensor[dim], it assumes that tensor[0
   /// ...dims-1] has already been setup.
private:
   /// loops that operates on sparse tensors more easily.
   StringAttr loopTag;
   /// Whether the loop emitter needs to treat the last tensor as the output
   /// tensor.
   bool hasOutput;
   bool isSparseOut;
   /// Input and (optional) output tensors.
   std::vector<Value> tensors;
+  /// Values realted to slices.
   std::vector<bool> isSparseSlices;
-  /// The dim type array for each tensor.
+  std::vector<std::vector<Value>> sliceOffsets;
+  std::vector<std::vector<Value>> sliceStrides;
+  /// The dim and dim type array for each tensor.
+  std::vector<std::vector<Value>> dims;
   std::vector<std::vector<DimLevelType>> dimTypes;
   /// Sparse iteration information (by tensor and dim). These arrays
   /// are updated to remain current within the current loop.
   std::vector<std::vector<Value>> pidxs;
   std::vector<std::vector<Value>> coord;
   std::vector<std::vector<Value>> highs;
   std::vector<std::vector<Value>> ptrBuffer; // to_pointers
   std::vector<std::vector<Value>> idxBuffer; // to_indices
   std::vector<Value> valBuffer;              // to_value
 
   // Loop Stack, stores the information of all the nested loops that are
   // alive.
   std::vector<LoopLevelInfo> loopStack;
 
   // Loop Sequence Stack, stores the unversial index for the current loop
   // sequence.
   std::vector<Value> loopSeqStack;
 
   // Maps AffineDimExpr to the index of the loop in loopStack.

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

if (!load.getType().isa<IndexType>()) {
     if (load.getType().getIntOrFloatBitWidth() < 64)
       load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
     load =
         builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
   }
   return load;
 }
 
-// TODO: Support dynamic sized slice.
-static Value getSliceOffset(OpBuilder &builder, Location loc,
-                            SparseTensorEncodingAttr enc, unsigned lvl) {
-  return constantIndex(builder, loc, *enc.getStaticLvlSliceOffset(lvl));
-}
-
-static Value getSliceSize(OpBuilder &builder, Location loc,
-                          SparseTensorEncodingAttr enc, unsigned lvl) {
-  return constantIndex(builder, loc, *enc.getStaticLvlSliceSize(lvl));
+static Value genSliceOffset(OpBuilder &builder, Location loc, Value tensor,
+                            unsigned lvl) {
+  auto enc = getSparseTensorEncoding(tensor.getType());
+  return createOrFoldSliceOffsetOp(builder, loc, tensor, toOrigDim(enc, lvl));
 }
 
-static Value getSliceStride(OpBuilder &builder, Location loc,
-                            SparseTensorEncodingAttr enc, unsigned lvl) {
-  return constantIndex(builder, loc, *enc.getStaticLvlSliceStride(lvl));
+static Value genSliceStride(OpBuilder &builder, Location loc, Value tensor,
+                            unsigned lvl) {
+  auto enc = getSparseTensorEncoding(tensor.getType());
+  return createOrFoldSliceStrideOp(builder, loc, tensor, toOrigDim(enc, lvl));
 }
 
 static Value toSliceCoord(OpBuilder &builder, Location loc, Value v,
-                          SparseTensorEncodingAttr enc, unsigned lvl) {
-
-  Value stride = getSliceStride(builder, loc, enc, lvl);
-  Value offset = getSliceOffset(builder, loc, enc, lvl);
+                          Value offset, Value stride, Value tensor,
+                          unsigned lvl) {
   // iv = iv * stride + offset
   v = builder.create<arith::MulIOp>(loc, v, stride);
   v = builder.create<arith::AddIOp>(loc, v, offset);
   return v;
 }
 
 static std::pair<Value, Value> fromSliceCoord(OpBuilder &builder, Location loc,
-                                              Value v,
-                                              SparseTensorEncodingAttr enc,
+                                              Value iv, Value offset,
+                                              Value stride, Value tensor,
                                               unsigned lvl) {
-  Value stride = getSliceStride(builder, loc, enc, lvl);
-  Value offset = getSliceOffset(builder, loc, enc, lvl);
   // iv = (iv - offset) / stride
-  v = builder.create<arith::SubIOp>(loc, v, offset);
-  Value rem = builder.create<arith::RemUIOp>(loc, v, stride);
-  v = builder.create<arith::DivUIOp>(loc, v, stride);
-  return std::make_pair(v, rem);
+  iv = builder.create<arith::SubIOp>(loc, iv, offset);
+  Value rem = builder.create<arith::RemUIOp>(loc, iv, stride);
+  iv = builder.create<arith::DivUIOp>(loc, iv, stride);
+  return std::make_pair(iv, rem);
+}
+
+std::pair<Value, Value>
+LoopEmitter::genSliceLegitPredicate(OpBuilder &builder, Location loc,
+                                    Value coord, unsigned tid, unsigned lvl) {
+  assert(isSparseSlices[tid]);
+  Value slice = tensors[tid];
+  Value offset = sliceOffsets[tid][lvl];
+  Value stride = sliceStrides[tid][lvl];
+  auto enc = getSparseTensorEncoding(slice.getType());
+
+  std::pair<Value, Value> trans =
+      fromSliceCoord(builder, loc, coord, offset, stride, slice, lvl);
+
+  SmallVector<Value, 3> conds; // at most 3 conditions
+
+  if (auto staticOffset = enc.getStaticLvlSliceOffset(lvl);
+      !(staticOffset.has_value() && *staticOffset == 0)) {
+    // 1st, coord >= offset (skip the check if offset is known to be 0).
+    auto geOffset = builder.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::uge, coord, offset);
+    conds.push_back(geOffset);
   }
 
-static std::pair<Value, Value>
-genSliceLegitPredicate(OpBuilder &builder, Location loc, Value coord,
-                       SparseTensorEncodingAttr enc, unsigned lvl) {
-  std::pair<Value, Value> trans = fromSliceCoord(builder, loc, coord, enc, lvl);
-  // 1st, coord >= offset (TODO: seems unsigned >= 0 won't be folded, skip
-  // the check if the offset is zero).
-  auto geOffset =
-      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge, coord,
-                                    getSliceOffset(builder, loc, enc, lvl));
   // 2nd, coord_in_slice < length
-  auto ltLength =
-      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, trans.first,
-                                    getSliceSize(builder, loc, enc, lvl));
-
-  // 3rd, rem == 0; confirmed that (a % 1) will be folded to 0
-  auto fitStride =
-      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, trans.second,
-                                    constantIndex(builder, loc, 0));
+  auto ltLength = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
+                                                trans.first, dims[tid][lvl]);
+  conds.push_back(ltLength);
+
+  if (auto staticStride = enc.getStaticLvlSliceStride(lvl);
+      !(staticStride.has_value() && *staticStride == 1)) {
+    // 3rd, rem == 0 (skip the check if stride is known to be 1).
+    auto fitStride = builder.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::eq, trans.second,
+    conds.push_back(fitStride);
+  }
+
+  // Must meet all condition to be a valid coordinate in slice.
+  auto pred = conds.front();
+  for (unsigned i = 1, e = conds.size(); i < e; i++)
+    pred = builder.create<arith::AndIOp>(loc, pred, conds[i]);
 
-  auto pred = builder.create<arith::AndIOp>(loc, geOffset, ltLength);
-  pred = builder.create<arith::AndIOp>(loc, pred, fitStride);
   return {trans.first, pred};
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 Value LoopEmitter::genAddress(OpBuilder &builder, Location loc, size_t tid,
                               size_t dim, Value iv) {
   Value p = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
   Value mul = builder.create<arith::MulIOp>(loc, highs[tid][dim], p);
-  if (isSparseSlices[tid]) {
-    auto enc = getSparseTensorEncoding(tensors[tid].getType());
-    iv = toSliceCoord(builder, loc, iv, enc, dim);
-  }
+  if (isSparseSlices[tid])
+    iv = toSliceCoord(builder, loc, iv, sliceOffsets[tid][dim],
+                      sliceStrides[tid][dim], tensors[tid], dim);
   Value add = builder.create<arith::AddIOp>(loc, mul, iv);
   return add;
 }
 
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                          bool isSparseOut, ArrayRef<unsigned> topSort) {
   initialize(tensors, loopTag, hasOutput, isSparseOut, topSort);
 }
 
 void LoopEmitter::initialize(ValueRange tensors, StringAttr loopTag,
                              bool hasOutput, bool isSparseOut,
                              ArrayRef<unsigned> topSort) {
   // First initializes fields.
   this->loopTag = loopTag;
   this->hasOutput = hasOutput;
   this->isSparseOut = isSparseOut;
   this->tensors.assign(tensors.begin(), tensors.end());
   this->isSparseSlices.assign(tensors.size(), false);
+  this->sliceOffsets.assign(tensors.size(), std::vector<Value>());
+  this->sliceStrides.assign(tensors.size(), std::vector<Value>());
+  this->dims.assign(tensors.size(), std::vector<Value>());
   this->dimTypes.assign(tensors.size(), std::vector<DimLevelType>());
   this->pidxs.assign(tensors.size(), std::vector<Value>());
   this->coord.assign(tensors.size(), std::vector<Value>());
   this->highs.assign(tensors.size(), std::vector<Value>());
   this->ptrBuffer.assign(tensors.size(), std::vector<Value>());
   this->idxBuffer.assign(tensors.size(), std::vector<Value>());
   this->valBuffer.assign(tensors.size(), nullptr);
   this->loopStack.reserve(topSort.size());
for (size_t tid = 0, e = tensors.size(); tid < e; tid++) {
     if (enc && !(isOutputTensor(tid) && isSparseOut)) {
       isSparseSlices[tid] = enc.isSlice();
       for (auto dimTp : enc.getDimLevelType())
         dimTypes[tid].push_back(dimTp);
     } else
       dimTypes[tid].assign(rank, DimLevelType::Dense);
 
     // Initialize using empty value.
+    sliceOffsets[tid].assign(rank, Value());
+    sliceStrides[tid].assign(rank, Value());
+    dims[tid].assign(rank, Value());
     pidxs[tid].assign(rank, Value());
     coord[tid].assign(rank, Value());
     highs[tid].assign(rank, Value());
     ptrBuffer[tid].assign(rank, Value());
     idxBuffer[tid].assign(rank, Value());
   }
 
   // FIXME: This map should be maintained outside loop emitter.
   for (unsigned i = 0, e = topSort.size(); i < e; i++) {
     // This is an inverse map of the topologically sorted loop index from
-    // sparsifier. This is needed to map the AffineDimExpr back to the loopStack
-    // index used in loop emitter.
+    // sparsifier. This is needed to map the AffineDimExpr back to the
+    // loopStack index used in loop emitter.
     sparsiferLoopLvlMap[topSort[i]] = i;
   }
 }
 
 void LoopEmitter::initializeLoopEmit(OpBuilder &builder, Location loc,
                                      LoopEmitter::OutputUpdater updater) {
   // For every tensor, find lower and upper bound on dimensions, set the
   // same bounds on loop indices, and obtain dense or sparse buffer(s).
for (int64_t d = 0; d < rank; d++) {
         // Singleton dimension, fetch indices.
         idxBuffer[t][d] = genToIndices(builder, loc, tensor, d, cooStart);
       } else {
         // Dense dimension, nothing to fetch.
         assert(isDenseDLT(dimTypes[t][d]));
       }
 
       // Find upper bound in current dimension.
       unsigned p = toOrigDim(enc, d);
       Value up = mlir::linalg::createOrFoldDimOp(builder, loc, tensor, p);
-      highs[t][d] = up;
+      highs[t][d] = dims[t][d] = up;
+      if (isSparseSlices[t]) {
+        sliceOffsets[t][d] = genSliceOffset(builder, loc, tensors[t], d);
+        sliceStrides[t][d] = genSliceStride(builder, loc, tensors[t], d);
+      }
     }
 
     // Perform the required bufferization. Dense inputs materialize
-    // from the input tensors. Sparse inputs use sparse primitives to obtain the
-    // values.
-    // Delegates extra output initialization to clients.
+    // from the input tensors. Sparse inputs use sparse primitives to obtain
+    // the values. Delegates extra output initialization to clients.
     bool isOutput = isOutputTensor(t);
     Type elementType = rtp.getElementType();
     if (!enc) {
       // Non-annotated dense tensors.
       auto denseTp = MemRefType::get(shape, elementType);
       Value denseVal =
           builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
       // Dense outputs need special handling.
       if (isOutput && updater)
         denseVal = updater(builder, loc, denseVal, tensor);
 
       valBuffer[t] = denseVal;
     } else {
       // Annotated sparse tensors.
-      // We also need the value buffer for annotated all dense `sparse` tensor.
+      // We also need the value buffer for annotated all dense `sparse`
+      // tensor.
       auto dynShape = {ShapedType::kDynamic};
       auto sparseTp = MemRefType::get(dynShape, elementType);
       valBuffer[t] = builder.create<ToValuesOp>(loc, sparseTp, tensor);
     }
     // NOTE: we can also prepare for 0 dim here in advance, this will hosit
-    // some loop preparation from tensor iteration, but will also (undesirably)
-    // hosit the code ouside if conditions.
+    // some loop preparation from tensor iteration, but will also
+    // (undesirably) hosit the code ouside if conditions.
   }
 }
 
 void LoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
                                   ArrayRef<size_t> tids,
                                   ArrayRef<size_t> dims) {
   // Universal Index start from 0
   assert(loopSeqStack.size() == loopStack.size());
Operation *LoopEmitter::enterLoopOverTensorAtDim(
   for (auto [t, d] : llvm::zip(tids, dims)) {
     assert(dimTypes[t].size() > d); // Must be a valid tid, dim pair
     assert(!coord[t][d]);           // We cannot re-enter the same level
     auto dimType = dimTypes[t][d];
     // Must be a recognizable DLT.
     assert(isDenseDLT(dimType) || isCompressedDLT(dimType) ||
            isSingletonDLT(dimType));
     bool isSparse = isCompressedDLT(dimType) || isSingletonDLT(dimType);
-    // We can at most have one sparse input, otherwise, a while loop is required
-    // to co-iterate multiple sparse tensors.
+    // We can at most have one sparse input, otherwise, a while loop is
+    // required to co-iterate multiple sparse tensors.
     assert(!isSparseInput || !isSparse);
     if (isSparse) {
       tid = t;
       dim = d;
     }
     isSparseInput = isSparseInput || isSparse;
   }
 
-  auto enc = getSparseTensorEncoding(tensors[tid].getType());
   // TODO: support dynamic slices.
   Value step = constantIndex(builder, loc, 1);
   Value lo = isSparseInput ? pidxs[tid][dim]      // current offset
                            : loopSeqStack.back(); // universal index
   Value hi = highs[tid][dim];
 
   Operation *loop = nullptr;
   Value iv;
   if (isParallel) {
     scf::ParallelOp parOp =
         builder.create<scf::ParallelOp>(loc, lo, hi, step, reduc);
     builder.setInsertionPointToStart(parOp.getBody());
     assert(parOp.getNumReductions() == reduc.size());
     iv = parOp.getInductionVars()[0];
 
     // In-place update on the reduction variable vector.
-    // Note that the init vals is not the actual reduction variables but instead
-    // used as a `special handle` to (temporarily) represent them. The
+    // Note that the init vals is not the actual reduction variables but
+    // instead used as a `special handle` to (temporarily) represent them. The
     // expression on init vals will be moved into scf.reduce and replaced with
     // the block arguments when exiting the loop (see exitForLoop). This is
-    // needed as we can not build the actual reduction block and get the actual
-    // reduction varaible before users fill parallel loop body.
+    // needed as we can not build the actual reduction block and get the
+    // actual reduction varaible before users fill parallel loop body.
     for (int i = 0, e = reduc.size(); i < e; i++)
       reduc[i] = parOp.getInitVals()[i];
     loop = parOp;
   } else {
     scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
     builder.setInsertionPointToStart(forOp.getBody());
     iv = forOp.getInductionVar();
 
Operation *LoopEmitter::enterLoopOverTensorAtDim(
 
   if (isSparseSlices[tid] && isSparseInput) {
     // For sparse level slices, we need to filter out invalid coordinates that
     // are not included in the slice.
     SmallVector<Type> types;
     for (Value red : reduc)
       types.push_back(red.getType());
 
-    auto [trans, pred] = genSliceLegitPredicate(builder, loc, c, enc, dim);
+    auto [trans, pred] = genSliceLegitPredicate(builder, loc, c, tid, dim);
     bool hasReduc = !types.empty();
     scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, pred,
                                                /*else*/ hasReduc);
     if (hasReduc) {
       // scf.for (a) -> v
       //  %s = scf.if (a) -> v
       //    user-generated code.
       //  else
Operation *LoopEmitter::enterFilterLoopOverTensorAtDim(
 
   Value step = constantIndex(builder, loc, 1);
 
   Value lo = pidxs[tid][dim];
   Value hi = highs[tid][dim];
 
   // TODO: We should instead use a whileOp for filter loop to allow early
   // break when exceeding (for ordered dimensions).
-  // TODO: There are many other potiential opportunities that we might apply in
-  // the future. E.g., we could use binary search to located the pointer index.
+  // TODO: There are many other potiential opportunities that we might apply
+  // in the future. E.g., we could use binary search to located the pointer
+  // index.
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
 
   // In-place update on the reduction variable vector.
   assert(forOp.getNumRegionIterArgs() == reduc.size());
   for (int i = 0, e = reduc.size(); i < e; i++)
     reduc[i] = forOp.getRegionIterArg(i);
 
   builder.setInsertionPointToStart(forOp.getBody());
for (auto [tid, dim] : llvm::zip(tids, dims)) {
     // Prepares for next level.
     if (isCompressedDLT(dimTypes[tid][dim]) ||
         isSingletonDLT(dimTypes[tid][dim])) {
       Value ptr = idxBuffer[tid][dim];
       Value s = pidxs[tid][dim];
       Value load = genIndexLoad(builder, loc, ptr, s);
       coord[tid][dim] = load;
       if (isSparseSlices[tid]) {
-        auto enc = getSparseTensorEncoding(tensors[tid].getType());
         auto [trans, pred] =
-            genSliceLegitPredicate(builder, loc, load, enc, dim);
+            genSliceLegitPredicate(builder, loc, load, tid, dim);
         slicesPreds.emplace_back(pred, i);
         // Updates to the relative coordinate to the slice.
         coord[tid][dim] = trans;
       }
       i++;
     }
   }
 
   if (!slicesPreds.empty()) {
     // Skip invalid loop iteration when slice coordinate is inapplicable.
     SmallVector<Value> yields(after->getArguments());
     // Generates a list of if statments
     //  pidx = in_slice ? pidx : pidx + 1
     // TODO: instead of always picking pidx + 1, we should set pidx = high to
-    // break to loop the coordinates is larger than the slice size.
+    // break to loop if the coordinates is larger than the slice size.
     for (auto [pred, idx] : slicesPreds) {
       Value nextPidx = builder.create<arith::AddIOp>(
           loc, yields[idx], constantIndex(builder, loc, 1));
       yields[idx] =
           builder.create<arith::SelectOp>(loc, pred, yields[idx], nextPidx);
     }
 
     Value pred = slicesPreds.front().first;

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/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

static scf::ForOp createFor(OpBuilder &builder, Location loc, Value upper,
   Value one = constantOne(builder, loc, indexType);
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lower, upper, one, fields);
   for (unsigned i = 0, e = fields.size(); i < e; i++)
     fields[i] = forOp.getRegionIterArg(i);
   builder.setInsertionPointToStart(forOp.getBody());
   return forOp;
 }
 
+/// Returns a value from the mixed slice info (either a attribute for statically
+/// known value or a dynamic value). Creates a constantOp if needed.
+static Value valueFromMixedSliceInfo(OpBuilder &builder, Location loc,
+                                     OpFoldResult ofr) {
+  Value v = ofr.dyn_cast<Value>();
+  if (!v) {
+    auto attr = ofr.get<Attribute>().cast<IntegerAttr>();
+    v = builder.create<arith::ConstantIndexOp>(loc, attr.getInt());
+  }
+  return v;
+}
+
 /// Gets the dimension size for the given sparse tensor at the given
 /// original dimension 'dim'. Returns std::nullopt if no sparse encoding is
 /// attached to the given tensor type.
 static std::optional<Value>
 sizeFromTensorAtDim(OpBuilder &builder, Location loc,
                     const SparseTensorDescriptor &desc, unsigned dim) {
   RankedTensorType rtp = desc.getTensorType();
   // Access into static dimension can query original type directly.
matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
     if (!sz)
       return failure();
 
     rewriter.replaceOp(op, *sz);
     return success();
   }
 };
 
+template <typename Op, StorageSpecifierKind kind>
+class SparseSliceGetterOpConver : public OpConversionPattern<Op> {
+public:
+  using OpConversionPattern<Op>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(Op op, typename Op::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Simply lowers to specifer.get xxx.
+    auto desc = getDescriptorFromTensorTuple(adaptor.getSlice());
+    auto v = desc.getSpecifierField(rewriter, op.getLoc(), kind,
+                                    op.getDim().getZExtValue());
+
+    rewriter.replaceOp(op, v);
+    return success();
+  }
+};
+
 /// Sparse codegen rule for trivial tensor casts.
 class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::CastOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Only rewrite identically annotated source/dest.
matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
     // Replace the requested pointer access with corresponding field.
     // The cast_op is inserted by type converter to intermix 1:N type
     // conversion.
     Location loc = op.getLoc();
     auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     uint64_t dim = op.getDimension().getZExtValue();
     Value field = desc.getIdxMemRefOrView(rewriter, loc, dim);
 
-    // Insert a cast to bridge the actual type to the user expected type. If the
-    // actual type and the user expected type aren't compatible, the compiler or
-    // the runtime will issue an error.
+    // Insert a cast to bridge the actual type to the user expected type. If
+    // the actual type and the user expected type aren't compatible, the
+    // compiler or the runtime will issue an error.
     Type resType = op.getResult().getType();
     if (resType != field.getType())
       field = rewriter.create<memref::CastOp>(loc, resType, field);
     rewriter.replaceOp(op, field);
 
     return success();
   }
 };

 
 class SparseExtractSliceCoverter
     : public OpConversionPattern<tensor::ExtractSliceOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::ExtractSliceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    MLIRContext *ctx = op.getContext();
     auto srcEnc = getSparseTensorEncoding(op.getSourceType());
     auto dstEnc = getSparseTensorEncoding(op.getResult().getType());
     if (!srcEnc && !dstEnc)
       return failure();
 
-    // We should probably check these in ExtractSliceOp::verify.
+    // FIXME: We should probably check these in ExtractSliceOp::verify.
     assert(srcEnc && dstEnc && dstEnc.isSlice());
     assert(srcEnc.getDimLevelType() == dstEnc.getDimLevelType());
     assert(srcEnc.getDimOrdering() == dstEnc.getDimOrdering());
     assert(srcEnc.getHigherOrdering() == dstEnc.getHigherOrdering());
     assert(srcEnc.getPointerBitWidth() == dstEnc.getPointerBitWidth());
     assert(srcEnc.getIndexBitWidth() == dstEnc.getIndexBitWidth());
 
-    // TODO: support dynamic slices.
-    for (int i = 0, e = op.getSourceType().getRank(); i < e; i++) {
-      assert(op.getStaticStrides()[i] == dstEnc.getStaticDimSliceStride(i));
-      assert(op.getStaticOffsets()[i] == dstEnc.getStaticDimSliceOffset(i));
-      assert(op.getStaticSizes()[i] == dstEnc.getStaticDimSliceSize(i));
-    }
+    SmallVector<Value> fields;
+    auto desc = getMutDescriptorFromTensorTuple(adaptor.getSource(), fields);
 
-    // TODO: create a new specifer for slices (need to encode slice metadata).
-    // It does not matter now because only constant offset/stride are allowed.
-    rewriter.replaceOp(op, adaptor.getSource());
+    auto newSpec = rewriter.create<StorageSpecifierInitOp>(
+        loc, StorageSpecifierType::get(ctx, dstEnc), desc.getSpecifier());
+    desc.setSpecifier(newSpec);
+
+    // Fills in slice information.
+    unsigned dim = 0;
+    for (auto [offset, size, stride] : llvm::zip(
+             op.getMixedOffsets(), op.getMixedSizes(), op.getMixedStrides())) {
+      Value offsetV = valueFromMixedSliceInfo(rewriter, loc, offset);
+      Value sizeV = valueFromMixedSliceInfo(rewriter, loc, size);
+      Value strideV = valueFromMixedSliceInfo(rewriter, loc, stride);
+      // TODO: We could probably only set dynamic value here. But it would
+      // requires us to fill the hole when casting a static slice to dynamic
+      // slice.
+      desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::DimOffset,
+                             dim, offsetV);
+      desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::DimSize, dim,
+                             sizeV);
+      desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::DimStride,
+                             dim, strideV);
+      dim++;
+    }
+
+    // NOTE: we can not generates tuples directly from descriptor here, as the
+    // descriptor is holding the original type, yet we want the slice type
+    // here (they shared every memref but with an updated specifier).
+    rewriter.replaceOp(op, genTuple(rewriter, loc, op.getResult().getType(),
+                                    desc.getFields()));
     return success();
   }
 };
 
 /// Sparse codegen rule for number of entries operator.
 class SparseNumberOfEntriesConverter
     : public OpConversionPattern<NumberOfEntriesOp> {
 public:

 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     bool enableBufferInitialization) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseTensorDeallocConverter,
                SparseExtractSliceCoverter, SparseTensorLoadConverter,
                SparseExpandConverter, SparseCompressConverter,
-               SparseInsertConverter, SparseToPointersConverter,
-               SparseToIndicesConverter, SparseToIndicesBufferConverter,
-               SparseToValuesConverter, SparseConvertConverter,
-               SparseNumberOfEntriesConverter>(typeConverter,
-                                               patterns.getContext());
+               SparseInsertConverter,
+               SparseSliceGetterOpConver<ToSliceOffsetOp,
+                                         StorageSpecifierKind::DimOffset>,
+               SparseSliceGetterOpConver<ToSliceStrideOp,
+                                         StorageSpecifierKind::DimStride>,
+               SparseToPointersConverter, SparseToIndicesConverter,
+               SparseToIndicesBufferConverter, SparseToValuesConverter,
+               SparseConvertConverter, SparseNumberOfEntriesConverter>(
+      typeConverter, patterns.getContext());
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

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

public:
   }
 
   unsigned getNumFields() const { return fields.size(); }
 
   ///
   /// Getters: get the value for required field.
   ///
 
+  Value getSpecifier() const { return fields.back(); }
+
   Value getSpecifierField(OpBuilder &builder, Location loc,
                           StorageSpecifierKind kind,
                           std::optional<unsigned> dim) const {
     SparseTensorSpecifier md(fields.back());
     return md.getSpecifierField(builder, loc, kind, dim);
   }
 
   Value getDimSize(OpBuilder &builder, Location loc, unsigned dim) const {
void setMemRefField(unsigned fidx, Value v) {
     fields[fidx] = v;
   }
 
   void setField(unsigned fidx, Value v) {
     assert(fidx < fields.size());
     fields[fidx] = v;
   }
 
+  void setSpecifier(Value newSpec) { fields.back() = newSpec; }
+
   void setSpecifierField(OpBuilder &builder, Location loc,
                          StorageSpecifierKind kind, std::optional<unsigned> dim,
                          Value v) {
     SparseTensorSpecifier md(fields.back());
     md.setSpecifierField(builder, loc, v, kind, dim);
     fields.back() = md;
   }
 
   void setDimSize(OpBuilder &builder, Location loc, unsigned dim, Value v) {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_foreach_slices.mlir

#CSR = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ]
 }>
 
 #CSR_SLICE = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   slice = [ (1, 4, 1), (1, 4, 2) ]
 }>
 
+#CSR_SLICE_DYN = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (?, ?, ?), (?, ?, ?) ]
+}>
+
 module {
   func.func @foreach_print_non_slice(%A: tensor<4x4xf64, #CSR>) {
     sparse_tensor.foreach in %A : tensor<4x4xf64, #CSR> do {
     ^bb0(%1: index, %2: index, %v: f64) :
       vector.print %1: index
       vector.print %2: index
       vector.print %v: f64
     }
     return
   }
 
   func.func @foreach_print_slice(%A: tensor<4x4xf64, #CSR_SLICE>) {
     sparse_tensor.foreach in %A : tensor<4x4xf64, #CSR_SLICE> do {
     ^bb0(%1: index, %2: index, %v: f64) :
       vector.print %1: index
       vector.print %2: index
       vector.print %v: f64
     }
     return
   }
 
+  func.func @foreach_print_slice_dyn(%A: tensor<?x?xf64, #CSR_SLICE_DYN>) {
+    sparse_tensor.foreach in %A : tensor<?x?xf64, #CSR_SLICE_DYN> do {
+    ^bb0(%1: index, %2: index, %v: f64) :
+      vector.print %1: index
+      vector.print %2: index
+      vector.print %v: f64
+    }
+    return
+  }
+
   func.func @entry() {
     %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
+    %c4 = arith.constant 4 : index
+
     %sa = arith.constant dense<[
         [ 0.0, 2.1, 0.0, 0.0, 0.0, 6.1, 0.0, 0.0 ],
         [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
         [ 0.0, 2.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
         [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 ],
         [ 0.0, 2.1, 0.0, 0.0, 0.0, 6.1, 0.0, 0.0 ],
         [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
         [ 0.0, 2.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
func.func @entry() {
     // CHECK-NEXT: 0
     // CHECK-NEXT: 2.1
     // CHECK-NEXT: 3
     // CHECK-NEXT: 2
     // CHECK-NEXT: 6.1
     //
     call @foreach_print_slice(%a) : (tensor<4x4xf64, #CSR_SLICE>) -> ()
 
+    // The same slice, but with dynamic encoding.
+    %a_dyn = tensor.extract_slice %tmp[%c1, %c1][%c4, %c4][%c1, %c2] :
+          tensor<8x8xf64, #CSR> to tensor<?x?xf64, #CSR_SLICE_DYN>
+    //
+    // CHECK:      1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 2.3
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 3
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 3
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 2.1
+    // CHECK-NEXT: 3
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 6.1
+    //
+    call @foreach_print_slice_dyn(%a_dyn) : (tensor<?x?xf64, #CSR_SLICE_DYN>) -> ()
+
     %dense = tensor.extract_slice %sa[1, 1][4, 4][1, 2] : tensor<8x8xf64> to
                                                           tensor<4x4xf64>
     %b = sparse_tensor.convert %dense : tensor<4x4xf64> to tensor<4x4xf64, #CSR>
     // Foreach on sparse tensor instead of slice they should yield the same result.
     //
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 2.3