[mlir][sparse] make foreach operation support sparse tensor slices.

[Peiming]

clean up.

[Peiming]

update comments.

[Peiming]

rebase

[Peiming]

rebase

[bixia]

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

61

This function looks more like fromSliceCoord to me, where "v" is the coordinate in the slice view. Am I right?

66

iv = slice_iv * stride + offset? I am not sure whether we should also add a period in this case?

72

Related to the above, this looks like toSliceCoord to me.

78

slice_iv = (iv - offset)/stride?

379

confirm?

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

962

Converter?

1022–1029

I don't see FileCheck test for this pattern, can we add such test?

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

6

Do we really need this lib?

50–53

We should avoid having the second part repeat the first part, maybe even try to avoid having the same NZ values in two different positions.

[Peiming]

address comments.

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

61

Hmmm, yeah, updated.

379

No, I confirmed that it is folded.

But it does not really matter, I manually skip the check anyway in the following revision.

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

1022–1029

This is a nop right now (the check test is added in the dynamic slice revision, which actually implements it).

[aartbik]

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

420

so what happens if you do "the right thing" here already?

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

61

Typo: Converts

62

tensor (singular)?

74

Typo: Converts

369

Using numbers and code is hard to read.
Just say

First, coord >= ...
Seconds, ...
etc.

401

insertion point

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

251

needs its own comment, or be moved some place else?

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

974

TODO, not "probably" ;-)

[Peiming]

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

420

I remember that it crashed, as there is no fields for dimSliceAttr in specifier yet.

[Peiming]

address comments.

[Peiming]

address comments.

[Peiming]

rebase

[Peiming]

disable test case with memory leak (for now).

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

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());
+      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<SparseTensorEncodingAttr::DimSlice>{} /*enc.getDimSlices()*/);
 }
 
 StorageSpecifierType
 StorageSpecifierType::get(MLIRContext *ctx, SparseTensorEncodingAttr encoding) {
   return Base::get(ctx, getNormalizedEncodingForSpecifier(encoding));
 }
 
 IntegerType StorageSpecifierType::getSizesType() const {

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

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;
+  std::vector<bool> isSparseSlices;
   /// The dim type array for each tensor.
   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

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.getLvlSliceOffset(lvl));
+}
+
+static Value getSliceLength(OpBuilder &builder, Location loc,
+                            SparseTensorEncodingAttr enc, unsigned lvl) {
+  return constantIndex(builder, loc, *enc.getLvlSliceLength(lvl));
+}
+
+static Value getSliceStride(OpBuilder &builder, Location loc,
+                            SparseTensorEncodingAttr enc, unsigned lvl) {
+  return constantIndex(builder, loc, *enc.getLvlSliceStride(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);
+  // 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,
+                                              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);
+}
+
+static Value maybeToSliceCoord(OpBuilder &builder, Location loc, Value v,
+                               SparseTensorEncodingAttr enc, unsigned lvl) {
+  if (!enc || !enc.isSlice())
+    return v;
+  return toSliceCoord(builder, loc, v, enc, lvl);
+}
+
 //===----------------------------------------------------------------------===//
 // 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);
+  }
   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->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());
   this->sparsiferLoopLvlMap.assign(topSort.size(), 0);
 
   for (size_t tid = 0, e = tensors.size(); tid < e; tid++) {
     auto t = tensors[tid];
     // a scalar or 0-dimension tensors
     if (isZeroRankedTensorOrScalar(t.getType()))
       continue;
     auto rtp = t.getType().cast<RankedTensorType>();
     auto rank = static_cast<size_t>(rtp.getRank());
     auto enc = getSparseTensorEncoding(rtp);
     // We always treat sparse output tensor as dense so that we always iterate
     // it based on dim size.
-    if (enc && !(isOutputTensor(tid) && isSparseOut))
+    if (enc && !(isOutputTensor(tid) && isSparseOut)) {
+      isSparseSlices[tid] = enc.isSlice();
       for (auto dimTp : enc.getDimLevelType())
         dimTypes[tid].push_back(dimTp);
-    else
+    } else
       dimTypes[tid].assign(rank, DimLevelType::Dense);
 
     // Initialize using empty 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());
Value LoopEmitter::genAffine(OpBuilder &builder, AffineExpr a, Location loc) {
   }
 }
 
 Operation *LoopEmitter::enterLoopOverTensorAtDim(
     OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
     ArrayRef<size_t> dims, MutableArrayRef<Value> reduc, bool isParallel) {
   // TODO: support multiple return on parallel for?
   assert(!isParallel || reduc.size() <= 1);
-
   bool isSparseInput = false;
   size_t tid = tids.front(), dim = dims.front();
   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.
     assert(!isSparseInput || !isSparse);
     if (isSparse) {
       tid = t;
       dim = d;
     }
     isSparseInput = isSparseInput || isSparse;
   }
 
-  Value step = constantIndex(builder, loc, 1);
-  Value lo = isSparseInput ? pidxs[tid][dim]      // current offset
-                           : loopSeqStack.back(); // univeral tid
-  Value hi = highs[tid][dim];
+  auto enc = getSparseTensorEncoding(tensors[tid].getType());
+  // Use the stride as the step for a dense level of a sparse tensor slices.
+  // TODO: support dynamic slices.
+  Value step = isSparseSlices[tid] && !isSparseInput
+                   ? constantIndex(builder, loc, *enc.getLvlSliceStride(dim))
+                   : constantIndex(builder, loc, 1);
+  Value lo = isSparseInput ? pidxs[tid][dim] // current offset
+                           : maybeToSliceCoord(builder, loc,
+                                               loopSeqStack.back(), enc, dim);
+  Value hi = isSparseInput
+                 ? highs[tid][dim]
+                 : maybeToSliceCoord(builder, loc, highs[tid][dim], enc, 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];
if (isParallel) {
     // 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);
     loop = forOp;
   }
   assert(loop && iv);
 
+  Value c;
   if (isSparseInput) {
     pidxs[tid][dim] = iv;
     // Generating a load on the indices array yields the coordinate.
     Value ptr = idxBuffer[tid][dim];
-    coord[tid][dim] = genIndexLoad(builder, loc, ptr, iv);
+    c = genIndexLoad(builder, loc, ptr, iv);
   } else {
     // Dense tensor, the coordinates is the inducation variable.
-    coord[tid][dim] = iv;
+    c = iv;
+  }
+
+  if (isSparseSlices[tid]) {
+    // For sparse level slices, we need to filter out invalid coordinates that
+    // are not included in the slice.
+    std::pair<Value, Value> trans = fromSliceCoord(builder, loc, c, enc, dim);
+    if (isSparseInput) {
+      SmallVector<Type> types;
+      for (Value red : reduc)
+        types.push_back(red.getType());
+
+      // 1st, coord >= offset
+      auto geOff =
+          builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge, c,
+                                        getSliceOffset(builder, loc, enc, dim));
+      // 2nd, coords < length
+      auto ltLen = builder.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::ult, trans.first,
+          getSliceLength(builder, loc, enc, dim));
+
+      // 3rd, rem == 0;
+      auto fitStride = builder.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq, trans.second,
+          constantIndex(builder, loc, 0));
+
+      auto pred = builder.create<arith::AndIOp>(loc, geOff, ltLen);
+      pred = builder.create<arith::AndIOp>(loc, pred, fitStride);
+      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
+        //    yield a
+        //  yield %s
+        builder.create<scf::YieldOp>(loc, ifOp.getResults());
+        builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+        // On mismatch.
+        builder.create<scf::YieldOp>(loc, reduc);
       }
+      // Set the insert point to matched branch.
+      builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+    }
+    c = trans.first;
+  }
+
+  assert(c);
+  coord[tid][dim] = c;
   // NOTE: we can also prepare for next dim here in advance
   // Push the loop into stack
   loopStack.emplace_back(ArrayRef<size_t>(tid), ArrayRef<size_t>(dim), loop,
                          coord[tid][dim], loopTag);
   // Emit extra locals.
   emitExtraLocalsForTensorsAtDenseDims(builder, loc, tids, dims);
 
   return loop;

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

if (encDst != encSrc) {
       // This should be handled by rewriting before codegen.
       return failure();
     }
     rewriter.replaceOp(op, adaptor.getSource());
     return success();
   }
 };
 
+class SparseExtractSliceCoverter
+    : public OpConversionPattern<tensor::ExtractSliceOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(tensor::ExtractSliceOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcEnc = getSparseTensorEncoding(op.getSourceType());
+    auto dstEnc = getSparseTensorEncoding(op.getResult().getType());
+    if (!srcEnc && !dstEnc)
+      return failure();
+
+    // 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.getDimSliceStride(i));
+      assert(op.getStaticOffsets()[i] == dstEnc.getDimSliceOffset(i));
+      assert(op.getStaticSizes()[i] == dstEnc.getDimSliceLength(i));
+    }
+
+    // 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());
+    return success();
+  }
+};
+
 /// Sparse codegen rule for number of entries operator.
 class SparseNumberOfEntriesConverter
     : public OpConversionPattern<NumberOfEntriesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {

 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     bool enableBufferInitialization) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
-               SparseCastConverter, SparseTensorDeallocConverter,
-               SparseTensorLoadConverter, SparseExpandConverter,
-               SparseCompressConverter, SparseInsertConverter,
-               SparseToPointersConverter, SparseToIndicesConverter,
-               SparseToValuesConverter, SparseConvertConverter,
-               SparseNumberOfEntriesConverter>(typeConverter,
-                                               patterns.getContext());
+               SparseExtractSliceCoverter, SparseCastConverter,
+               SparseTensorDeallocConverter, SparseTensorLoadConverter,
+               SparseExpandConverter, SparseCompressConverter,
+               SparseInsertConverter, SparseToPointersConverter,
+               SparseToIndicesConverter, SparseToValuesConverter,
+               SparseConvertConverter, SparseNumberOfEntriesConverter>(
+      typeConverter, patterns.getContext());
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

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

+// DEFINE: %{option} = enable-runtime-library=false
+// DEFINE: %{command} = mlir-opt %s --sparse-compiler=%{option} | \
+// DEFINE: mlir-cpu-runner \
+// DEFINE:  -e entry -entry-point-result=void  \
+// DEFINE:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_lib_dir/libmlir_runner_utils%shlibext | \
+// DEFINE: FileCheck %s
+//
+// RUN: %{command}
+//
+
+// TODO: support slices on lib path
+#CSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ]
+}>
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ [1, 4, 1], [1, 4, 2] ]
+}>
+
+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 @entry() {
+    %c0 = arith.constant 0 : 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 ],
+        [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 ]
+    ]> : tensor<8x8xf64>
+
+
+    %tmp = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #CSR>
+    %a = tensor.extract_slice %tmp[1, 1][4, 4][1, 2] : tensor<8x8xf64, #CSR> to
+                                                       tensor<4x4xf64, #CSR_SLICE>
+    // Foreach on sparse tensor slices directly
+    //
+    // 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(%a) : (tensor<4x4xf64, #CSR_SLICE>) -> ()
+
+    %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
+    // 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_non_slice(%b) : (tensor<4x4xf64, #CSR>) -> ()
+
+    bufferization.dealloc_tensor %tmp : tensor<8x8xf64, #CSR>
+    bufferization.dealloc_tensor %b : tensor<4x4xf64, #CSR>
+    return
+  }
+}