[mlir][linalg] Use OpBuilder in rewriteAsPaddedOp (NFC).

[gysit]

Summary

Adapt the rewriteAsPaddedOp method to use the OpBuilder instead of the PatterRewriter.

[gysit]

Rebase.

[gysit]

Rebase.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

164

this older exit condition seems fishy, can this be cleaned up ?

[gysit]

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

164

If paddingFunc returns failure we do not have a padding value. I agree the clean solution should be return failure in that case. But the exit condition will stay there. I will clean this up.

[gysit]

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

164

Turns out it is still needed. We use paddingFunc also to control if a specific operand shall be padded. So if there is no padding value this is actually not a failure but a success. We can refactor this once we use all the new flags / padding control.

[gysit]

Rebase.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// `linalg.generic`.
 struct PadTensorOpTransformationPattern : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const override;
 };
 
-/// Try to create a static bounding box around each operand of `opToPad`.
-/// If successful, `paddedOp` will be updated to the cloned static form.
-LogicalResult
-rewriteAsPaddedOp(PatternRewriter &rewriter, LinalgOp opToPad,
+/// Pad the operands of `opToPad` to a static bounding box. Use `paddingFunc`
+/// and `nofoldFunc` to set the padding value and the nofold attribute of the
+/// introduced PadTensorOps, respectively. Update `paddedOp` to the cloned
+/// statically shaped operation and return the extracted dynamically shaped
+/// results. If padding fails, return failure.
+FailureOr<SmallVector<Value>>
+rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
                   const PaddingValueComputationFunction &paddingFunc,
                   const PaddingNoFoldComputationFunction &nofoldFunc,
                   LinalgOp &paddedOp);
 
 using OptimizeCopyFn =
     std::function<LogicalResult(PatternRewriter &, PadTensorOp, Value)>;
 
 /// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp and

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

 }
 
 /// Helper function that tries to pad `opOperand`. Exit early and return success
 /// for scalar operands or if `paddingFunc` returns failure. Otherwise, try to
 /// pad the operand even if it already has a static shape. Set `result` to the
 /// result of the created PadTensorOp or return failure if the operand cannot be
 /// padded to a static shape.
 static LogicalResult padOperandToSmallestStaticBoundingBox(
-    PatternRewriter &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,
+    OpBuilder &b, linalg::LinalgOp opToPad, OpOperand *opOperand,
     const PaddingValueComputationFunction &paddingFunc,
     const PaddingNoFoldComputationFunction &nofoldFunc, Value &result) {
   // Can't pad scalars.
   if (opToPad.getShape(opOperand).empty())
     return success();
   // Can't pad if no padding value is known.
-  FailureOr<Value> paddingValue = paddingFunc(rewriter, *opOperand);
+  FailureOr<Value> paddingValue = paddingFunc(b, *opOperand);
   if (failed(paddingValue))
     return success();
   auto sliceOp = opOperand->get().getDefiningOp<tensor::ExtractSliceOp>();
   // Not a slice op, cannot construct a static bounding box.
   if (!sliceOp)
     return failure();
   SmallVector<int64_t> staticSizes;
   staticSizes.reserve(opToPad.getRank(opOperand));
   auto shapedOp = cast<OffsetSizeAndStrideOpInterface>(sliceOp.getOperation());
   for (auto size : shapedOp.getMixedSizes()) {
     auto indexAttr = size.is<Attribute>()
                          ? size.get<Attribute>().dyn_cast<IntegerAttr>()
                          : linalg::getSmallestBoundingIndex(size.get<Value>());
     // SmallestBoundingIndex must exist for all sizes.
     // For now return an error if we can't find it.
-    if (!indexAttr)
-      return rewriter.notifyMatchFailure(
-          opToPad, "No constant bounding box can be found for padding");
+    if (!indexAttr) {
+      LLVM_DEBUG(DBGS() << "No constant bounding box can be found for padding");
+      return failure();
+    }
     staticSizes.push_back(indexAttr.getInt());
   }
   auto staticTensorType = RankedTensorType::get(
       staticSizes, getElementTypeOrSelf(opOperand->get()));
   bool nofold = nofoldFunc ? nofoldFunc(*opOperand) : false;
   result = linalg::PadTensorOp::createPadHighOp(
       staticTensorType, opOperand->get(), paddingValue.getValue(),
-      /*nofold=*/nofold, opToPad->getLoc(), rewriter);
+      /*nofold=*/nofold, opToPad->getLoc(), b);
   return success();
 }
 
-LogicalResult
-linalg::rewriteAsPaddedOp(PatternRewriter &rewriter, LinalgOp opToPad,
+FailureOr<SmallVector<Value>>
+linalg::rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
                           const PaddingValueComputationFunction &paddingFunc,
                           const PaddingNoFoldComputationFunction &nofoldFunc,
                           LinalgOp &paddedOp) {
   Location loc = opToPad->getLoc();
 
   // TODO: there are cases where we may still want to pad to larger sizes.
   assert(opToPad.hasTensorSemantics() &&
          "expected operation to have tensor semantics");
 
-  OpBuilder::InsertionGuard g(rewriter);
+  OpBuilder::InsertionGuard g(b);
   // Set IP after op because we also take the dims of the original output.
-  rewriter.setInsertionPointAfter(opToPad);
+  b.setInsertionPointAfter(opToPad);
   // Make a copy of the shaped operands and update it.
   SmallVector<Value> newOperands;
   newOperands.reserve(opToPad.getNumInputsAndOutputs());
   for (OpOperand *opOperand : opToPad.getInputAndOutputOperands()) {
     Value paddedOperand;
     // If padding was requested but the shape cannot be bounded statically then
     // the pattern fails to apply.
     if (failed(padOperandToSmallestStaticBoundingBox(
-            rewriter, opToPad, opOperand, paddingFunc, nofoldFunc,
-            paddedOperand)))
+            b, opToPad, opOperand, paddingFunc, nofoldFunc, paddedOperand)))
       return failure();
     newOperands.push_back(paddedOperand ? paddedOperand : opOperand->get());
   }
 
   SmallVector<SmallVector<Value>> reifiedResultShapes;
   if (failed(cast<ReifyRankedShapedTypeOpInterface>(opToPad.getOperation())
-                 .reifyResultShapes(rewriter, reifiedResultShapes)))
+                 .reifyResultShapes(b, reifiedResultShapes)))
     return failure();
   assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
          "expected same number of results");
 
   // Clone `opToPad` to operate on the statically padded shapes.
   auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumOutputs()).getTypes();
-  paddedOp = opToPad.clone(rewriter, loc, resultTensorTypes, newOperands);
+  paddedOp = opToPad.clone(b, loc, resultTensorTypes, newOperands);
 
   // Recover the slice out of the new static results. This keeps the original
   // linalg op around because it uses the dims of the original results.
   SmallVector<Value> paddedSubviewResults;
   paddedSubviewResults.reserve(opToPad->getNumResults());
   for (auto en : llvm::enumerate(paddedOp->getResults())) {
     Value paddedResult = en.value();
     int64_t resultNumber = en.index();
     int64_t rank = paddedResult.getType().cast<RankedTensorType>().getRank();
-    SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
     SmallVector<OpFoldResult> sizes;
     for (Value v : reifiedResultShapes[resultNumber])
       sizes.push_back(v);
-    SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
-    paddedSubviewResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
+    SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
+    paddedSubviewResults.push_back(b.create<tensor::ExtractSliceOp>(
         loc, paddedResult, offsets, sizes, strides));
   }
-  rewriter.replaceOp(opToPad, paddedSubviewResults);
-  return success();
+  return paddedSubviewResults;
 }
 
 /// Linalg base tiling pattern.
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     StringRef opName, MLIRContext *context, LinalgTilingOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : RewritePattern(opName, benefit, context), filter(filter),
       options(options) {}
if (!options.paddingValueComputationFunction ||
       !linalgOp.hasTensorSemantics()) {
     result = *res;
     return success();
   }
 
   // Try to pad on the fly by rewriting res->op as a padded op. If successful,
   // `res.op` is rewritten in static form with padded operands.
   LinalgOp paddedOp;
-  if (succeeded(rewriteAsPaddedOp(
-          rewriter, res->op, options.paddingValueComputationFunction,
-          options.paddingNoFoldComputationFunction, paddedOp))) {
+  FailureOr<SmallVector<Value>> newResults = rewriteAsPaddedOp(
+      options.paddingNoFoldComputationFunction, paddedOp);
+  if (succeeded(newResults)) {
+    rewriter.replaceOp(res->op, newResults.getValue());
     filter.replaceLinalgTransformationFilter(rewriter, paddedOp);
     res->op = paddedOp;
     result = *res;
     // Do not perform replacement of `linalgOp`, let the derived patterns
     // do this as they see fit, from the resulting TiledLinalgOp.
     return success();
   }
   // Set so RAII guard does not propagate TiledLinalgOp to `result`.