[mlir][sparse] handle padding on sparse levels. (#90527)

Author: Peiming Liu

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

@@ -75,6 +75,40 @@ static Value genSliceStride(OpBuilder &builder, Location loc, Value tensor,
   return createOrFoldSliceStrideOp(builder, loc, tensor, toDim(enc, lvl));
 }
 
+static bool isIntOrFPZero(Attribute attr) {
+  if (auto f = llvm::dyn_cast<FloatAttr>(attr); f && f.getValue().isZero())
+    return true;
+  if (auto i = llvm::dyn_cast<IntegerAttr>(attr); i && i.getValue().isZero())
+    return true;
+  return false;
+}
+
+static Value unFoldOpIntResult(OpBuilder &builder, Location loc,
+                               OpFoldResult ofr) {
+  if (std::optional<int64_t> i = getConstantIntValue(ofr); i.has_value())
+    return constantIndex(builder, loc, *i);
+  return ofr.get<Value>();
+}
+
+static Value tryFoldTensors(Value t) {
+  // TODO: this should be done through a folding pass after switching to
+  // `sparse_tensor.iterate`-based sparsification.
+  auto stt = tryGetSparseTensorType(t);
+  auto padOp = t.getDefiningOp<tensor::PadOp>();
+  if (padOp && stt.has_value() && stt->hasEncoding() &&
+      padOp.getSourceType().getEncoding() == stt->getEncoding() &&
+      stt->getEncoding().isIdentity()) {
+    // Try fusing padOp with zeros.
+    Attribute padCst;
+    if (matchPattern(padOp.getBody()->getTerminator(),
+                     m_Op<tensor::YieldOp>(m_Constant(&padCst))) &&
+        isIntOrFPZero(padCst)) {
+      return padOp.getSource();
+    }
+  }
+  return t;
+}
+
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
@@ -166,15 +200,30 @@ void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
 std::unique_ptr<SparseIterator>
 LoopEmitter::makeLevelIterator(OpBuilder &builder, Location loc, TensorId t,
                                Level l) {
+  Value tensor = tensors[t];
+  auto stt = getSparseTensorType(tensor);
   auto it = makeSimpleIterator(*lvls[t][l], emitStrategy);
-  auto stt = getSparseTensorType(tensors[t]);
+
+  Value folded = tryFoldTensors(tensor);
+  if (folded != tensor) {
+    auto padOp = tensor.getDefiningOp<tensor::PadOp>();
+    assert(padOp);
+    if (padOp.getPaddedDims().test(l)) {
+      Value low = unFoldOpIntResult(builder, loc, padOp.getMixedLowPad()[l]);
+      Value high = unFoldOpIntResult(builder, loc, padOp.getMixedHighPad()[l]);
+      auto padIt = makePaddedIterator(std::move(it), low, high, emitStrategy);
+      return padIt;
+    }
+  }
+
   if (stt.hasEncoding() && stt.getEncoding().isSlice()) {
-    Value offset = genSliceOffset(builder, loc, tensors[t], l);
-    Value stride = genSliceStride(builder, loc, tensors[t], l);
+    Value offset = genSliceOffset(builder, loc, tensor, l);
+    Value stride = genSliceStride(builder, loc, tensor, l);
     auto slicedIt = makeSlicedLevelIterator(
         std::move(it), offset, stride, lvls[t][l]->getSize(), emitStrategy);
     return slicedIt;
   }
+
   return it;
 }
 
@@ -200,7 +249,9 @@ void LoopEmitter::initializeLoopEmit(
   //     on positions.
   for (TensorId t = 0, numTensors = getNumManifestTensors(); t < numTensors;
        t++) {
-    const Value tensor = tensors[t];
+    // TODO: this should be done through a folding pass after switching to
+    // `sparse_tensor.iterate`-based sparsification.
+    const Value tensor = tryFoldTensors(tensors[t]);
     const auto rtp = dyn_cast<RankedTensorType>(tensor.getType());
     if (!rtp)
       // Skips only scalar, zero ranked tensor still need to be bufferized and
@@ -213,14 +264,6 @@ void LoopEmitter::initializeLoopEmit(
     const Level lvlRank = stt.getLvlRank();
     const auto shape = rtp.getShape();
 
-    SmallVector<Value> lvlSzs;
-    for (Level l = 0; l < stt.getLvlRank(); l++) {
-      if (stt.hasEncoding())
-        lvlSzs.push_back(builder.create<LvlOp>(loc, tensor, l));
-      else
-        lvlSzs.push_back(builder.create<tensor::DimOp>(loc, tensor, l));
-    }
-
     // Scan all levels of current tensor.
     for (Level l = 0; l < lvlRank; l++) {
       // Find upper bound in current dimension.

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp

@@ -462,11 +462,54 @@ class DedupIterator : public ConcreteIterator {
   Value posHi;
 };
 
+// A util base-iterator that delegates all methods to the wrapped iterator.
+class SimpleWrapIterator : public SparseIterator {
+public:
+  SimpleWrapIterator(std::unique_ptr<SparseIterator> &&wrap, IterKind kind)
+      : SparseIterator(kind, *wrap), wrap(std::move(wrap)) {}
+
+  SmallVector<Type> getCursorValTypes(OpBuilder &b) const override {
+    return wrap->getCursorValTypes(b);
+  }
+  bool isBatchIterator() const override { return wrap->isBatchIterator(); }
+  bool randomAccessible() const override { return wrap->randomAccessible(); };
+  bool iteratableByFor() const override { return wrap->iteratableByFor(); };
+  SmallVector<Value> serialize() const override { return wrap->serialize(); };
+  void deserialize(ValueRange vs) override { wrap->deserialize(vs); };
+  ValueRange getCurPosition() const override { return wrap->getCurPosition(); }
+  void genInitImpl(OpBuilder &b, Location l,
+                   const SparseIterator *parent) override {
+    wrap->genInit(b, l, parent);
+  }
+  Value genNotEndImpl(OpBuilder &b, Location l) override {
+    return wrap->genNotEndImpl(b, l);
+  }
+  ValueRange forwardImpl(OpBuilder &b, Location l) override {
+    return wrap->forward(b, l);
+  };
+  Value upperBound(OpBuilder &b, Location l) const override {
+    return wrap->upperBound(b, l);
+  };
+
+  Value derefImpl(OpBuilder &b, Location l) override {
+    return wrap->derefImpl(b, l);
+  }
+
+  void locateImpl(OpBuilder &b, Location l, Value crd) override {
+    return wrap->locate(b, l, crd);
+  }
+
+  SparseIterator &getWrappedIterator() const { return *wrap; }
+
+protected:
+  std::unique_ptr<SparseIterator> wrap;
+};
+
 //
 // A filter iterator wrapped from another iterator. The filter iterator update
 // the wrapped iterator *in-place*.
 //
-class FilterIterator : public SparseIterator {
+class FilterIterator : public SimpleWrapIterator {
   // Coorindate translation between crd loaded from the wrap iterator and the
   // filter iterator.
   Value fromWrapCrd(OpBuilder &b, Location l, Value wrapCrd) const {
@@ -487,8 +530,8 @@ class FilterIterator : public SparseIterator {
   // when crd always < size.
   FilterIterator(std::unique_ptr<SparseIterator> &&wrap, Value offset,
                  Value stride, Value size)
-      : SparseIterator(IterKind::kFilter, *wrap), offset(offset),
-        stride(stride), size(size), wrap(std::move(wrap)) {}
+      : SimpleWrapIterator(std::move(wrap), IterKind::kFilter), offset(offset),
+        stride(stride), size(size) {}
 
   // For LLVM-style RTTI.
   static bool classof(const SparseIterator *from) {
@@ -498,19 +541,10 @@ class FilterIterator : public SparseIterator {
   std::string getDebugInterfacePrefix() const override {
     return std::string("filter<") + wrap->getDebugInterfacePrefix() + ">";
   }
-  SmallVector<Type> getCursorValTypes(OpBuilder &b) const override {
-    return wrap->getCursorValTypes(b);
-  }
 
-  bool isBatchIterator() const override { return wrap->isBatchIterator(); }
-  bool randomAccessible() const override { return wrap->randomAccessible(); };
   bool iteratableByFor() const override { return randomAccessible(); };
   Value upperBound(OpBuilder &b, Location l) const override { return size; };
 
-  SmallVector<Value> serialize() const override { return wrap->serialize(); };
-  void deserialize(ValueRange vs) override { wrap->deserialize(vs); };
-  ValueRange getCurPosition() const override { return wrap->getCurPosition(); }
-
   void genInitImpl(OpBuilder &b, Location l,
                    const SparseIterator *parent) override {
     wrap->genInit(b, l, parent);
@@ -541,7 +575,47 @@ class FilterIterator : public SparseIterator {
   ValueRange forwardImpl(OpBuilder &b, Location l) override;
 
   Value offset, stride, size;
-  std::unique_ptr<SparseIterator> wrap;
+};
+
+//
+// A pad iterator wrapped from another iterator. The pad iterator updates
+// the wrapped iterator *in-place*.
+//
+class PadIterator : public SimpleWrapIterator {
+
+public:
+  PadIterator(std::unique_ptr<SparseIterator> &&wrap, Value padLow,
+              Value padHigh)
+      : SimpleWrapIterator(std::move(wrap), IterKind::kPad), padLow(padLow),
+        padHigh(padHigh) {
+    assert(!randomAccessible() && "Not implemented.");
+  }
+
+  // For LLVM-style RTTI.
+  static bool classof(const SparseIterator *from) {
+    return from->kind == IterKind::kPad;
+  }
+
+  std::string getDebugInterfacePrefix() const override {
+    return std::string("pad<") + wrap->getDebugInterfacePrefix() + ">";
+  }
+
+  // The upper bound after padding becomes `size + padLow + padHigh`.
+  Value upperBound(OpBuilder &b, Location l) const override {
+    return ADDI(ADDI(wrap->upperBound(b, l), padLow), padHigh);
+  };
+
+  // The pad_coord = coord + pad_lo
+  Value derefImpl(OpBuilder &b, Location l) override {
+    updateCrd(ADDI(wrap->deref(b, l), padLow));
+    return getCrd();
+  }
+
+  void locateImpl(OpBuilder &b, Location l, Value crd) override {
+    assert(randomAccessible());
+  }
+
+  Value padLow, padHigh;
 };
 
 class NonEmptySubSectIterator : public SparseIterator {
@@ -1408,10 +1482,19 @@ sparse_tensor::makeSlicedLevelIterator(std::unique_ptr<SparseIterator> &&sit,
   return ret;
 }
 
+std::unique_ptr<SparseIterator>
+sparse_tensor::makePaddedIterator(std::unique_ptr<SparseIterator> &&sit,
+                                  Value padLow, Value padHigh,
+                                  SparseEmitStrategy strategy) {
+  auto ret = std::make_unique<PadIterator>(std::move(sit), padLow, padHigh);
+  ret->setSparseEmitStrategy(strategy);
+  return ret;
+}
+
 static const SparseIterator *tryUnwrapFilter(const SparseIterator *it) {
   auto *filter = llvm::dyn_cast_or_null<FilterIterator>(it);
   if (filter)
-    return filter->wrap.get();
+    return &filter->getWrappedIterator();
   return it;
 }
 

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h

@@ -76,6 +76,7 @@ enum class IterKind : uint8_t {
   kSubSect,
   kNonEmptySubSect,
   kFilter,
+  kPad,
 };
 
 /// Helper class that generates loop conditions, etc, to traverse a
@@ -291,26 +292,32 @@ std::unique_ptr<SparseTensorLevel> makeSparseTensorLevel(OpBuilder &b,
 std::unique_ptr<SparseIterator> makeSimpleIterator(const SparseTensorLevel &stl,
                                                    SparseEmitStrategy strategy);
 
-/// Helper function to create a synthetic SparseIterator object that iterate
+/// Helper function to create a synthetic SparseIterator object that iterates
 /// over a dense space specified by [0,`sz`).
 std::pair<std::unique_ptr<SparseTensorLevel>, std::unique_ptr<SparseIterator>>
 makeSynLevelAndIterator(Value sz, unsigned tid, unsigned lvl,
                         SparseEmitStrategy strategy);
 
-/// Helper function to create a SparseIterator object that iterate over a
+/// Helper function to create a SparseIterator object that iterates over a
 /// sliced space, the orignal space (before slicing) is traversed by `sit`.
 std::unique_ptr<SparseIterator>
 makeSlicedLevelIterator(std::unique_ptr<SparseIterator> &&sit, Value offset,
                         Value stride, Value size, SparseEmitStrategy strategy);
 
+/// Helper function to create a SparseIterator object that iterates over a
+/// padded sparse level (the padded value must be zero).
+std::unique_ptr<SparseIterator>
+makePaddedIterator(std::unique_ptr<SparseIterator> &&sit, Value padLow,
+                   Value padHigh, SparseEmitStrategy strategy);
+
 /// Helper function to create a SparseIterator object that iterate over the
 /// non-empty subsections set.
 std::unique_ptr<SparseIterator> makeNonEmptySubSectIterator(
     OpBuilder &b, Location l, const SparseIterator *parent, Value loopBound,
     std::unique_ptr<SparseIterator> &&delegate, Value size, unsigned stride,
     SparseEmitStrategy strategy);
 
-/// Helper function to create a SparseIterator object that iterate over a
+/// Helper function to create a SparseIterator object that iterates over a
 /// non-empty subsection created by NonEmptySubSectIterator.
 std::unique_ptr<SparseIterator> makeTraverseSubSectIterator(
     OpBuilder &b, Location l, const SparseIterator &subsectIter,