[mlir][Linalg] DoNotCommit - Rewrite PadTensor before bufferizing

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

let extraClassDeclaration = [{
 
     RankedTensorType getSourceType() {
       return source().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
+    // Infer the dynamic shape of the result tensor along each dim
+    SmallVector<Value> getResultTypeShapes(OpBuilder &b);
+
     // Infer the shape of the result tensor given the static shapes
     // and element type of the result tensor.
     static RankedTensorType inferResultType(RankedTensorType sourceType,
                                 ArrayRef<int64_t> staticLow,
                                 ArrayRef<int64_t> staticHigh);
 
     // Return a PadTensorOp that pads `source` to `type` size where the static
     // sizes are assumed to be greater than the dynamic sizes. The op performs

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

for (int i = 0; i < rank; ++i) {
     auto highValue = builder.createOrFold<SubIOp>(loc, resultDimSize, dimOp);
     high.push_back(highValue);
     low.push_back(builder.createOrFold<ConstantIndexOp>(loc, 0));
   }
   return PadTensorOp::createPadScalarOp(type, source, pad, low, high, loc,
                                         builder);
 }
 
-LogicalResult PadTensorOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+SmallVector<Value> PadTensorOp::getResultTypeShapes(OpBuilder &b) {
   Location loc = getLoc();
   auto lowPad = getMixedLowPad();
   auto highPad = getMixedHighPad();
   SmallVector<Value> shapes;
   for (auto dim : llvm::seq<int64_t>(0, getSourceType().getRank())) {
     // Shape along each dimension is source dim + low pad + high pad.
     SmallVector<Value> mapOperands;
     mapOperands.push_back(b.createOrFold<memref::DimOp>(loc, source(), dim));
auto addOpFoldResult = [&](OpFoldResult valueOrAttr) {
           valueOrAttr.get<Attribute>().cast<IntegerAttr>().getInt();
       expr = expr + staticValue;
     };
     addOpFoldResult(lowPad[dim]);
     addOpFoldResult(highPad[dim]);
     shapes.push_back(applyMapToValues(
         b, loc, AffineMap::get(1, numSymbols, expr), mapOperands)[0]);
   }
-  reifiedReturnShapes.emplace_back(std::move(shapes));
+  return shapes;
+}
+
+LogicalResult PadTensorOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  reifiedReturnShapes.emplace_back(getResultTypeShapes(b));
   return success();
 }
 
 namespace {
 // Folds linalg.pad_tensor when padding is static zeros.
 struct FoldStaticZeroPadding : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 

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

 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferUtils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
FunctionType bufferizedFuncType = getOrCreateBufferizedFunctionType(
       bufferizedFunctionTypes);
   OpBuilder b(returnOp);
   b.create<ReturnOp>(returnOp.getLoc(), returnValues);
   returnOp->erase();
 
   // 3. Rewrite the bbArgs.
   // Iterate on the original `numArgs` and replace them in order.
   // This guarantees the argument order still matches after the rewrite.
   Block &frontBlock = funcOp.body().front();
   unsigned numArgs = frontBlock.getNumArguments();
   for (unsigned idx = 0; idx < numArgs; ++idx) {
     auto bbArg = frontBlock.getArgument(0);
     auto tensorType = bbArg.getType().dyn_cast<TensorType>();
     // Non-tensor types are just forwarded.
     if (!tensorType) {
       frontBlock.addArgument(bbArg.getType());
       bbArg.replaceAllUsesWith(frontBlock.getArguments().back());
struct LinalgComprehensiveModuleBufferize
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<linalg::LinalgDialect, memref::MemRefDialect>();
   }
 };
 } // end namespace
 
+/// Returns the static/dynamic mixed sizes of the memref.
+static SmallVector<OpFoldResult> getMixedSizes(OpBuilder &b, Location loc,
+                                               Value memref) {
+  auto inputType = memref.getType().cast<ShapedType>();
+  auto inputShape = inputType.getShape();
+  SmallVector<OpFoldResult> sizeMixedValues;
+  for (int64_t i = 0; i < inputType.getRank(); ++i) {
+    if (inputShape[i] == ShapedType::kDynamicSize) {
+      Value dim = b.create<tensor::DimOp>(loc, memref, i);
+      sizeMixedValues.push_back(dim);
+    } else {
+      sizeMixedValues.push_back(b.getI64IntegerAttr(inputShape[i]));
+    }
+  }
+  return sizeMixedValues;
+}
+
+/// Conversion pattern that bufferizes `linalg.pad_tensor` operation.
+class PadTensorOpPattern : public OpRewritePattern<PadTensorOp> {
+public:
+  using OpRewritePattern<PadTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(PadTensorOp op,
+                                PatternRewriter &rewriter) const {
+    Location loc = op->getLoc();
+    Value sourceTensor = op.source();
+    auto sourceTensorType = sourceTensor.getType().cast<RankedTensorType>();
+
+    // Allocate the destination buffer
+    SmallVector<Value> resultShape = op.getResultTypeShapes(rewriter);
+    Value resultTensor = rewriter.create<linalg::InitTensorOp>(
+        loc, resultShape, sourceTensorType.getElementType());
+
+    // Get padding value and fill the destination buffer.
+    auto yieldOps = op.region().getOps<linalg::YieldOp>();
+    if (std::distance(yieldOps.begin(), yieldOps.end()) != 1) {
+      return rewriter.notifyMatchFailure(op,
+                                         "linalg.pad_tensor with more than one "
+                                         "padding value is not supported");
+    }
+    Value paddingValue = (*yieldOps.begin()).values()[0];
+    auto constOp = paddingValue.getDefiningOp<ConstantOp>();
+    if (!constOp) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "linalg.pad_tensor with non-constant padding value is not supported");
+    }
+    if (constOp.getValue().isa<DenseElementsAttr>()) {
+      return rewriter.notifyMatchFailure(
+          op, "linalg.pad_tensor with non-scalar constant padding value is not "
+              "supported");
+    }
+    resultTensor =
+        rewriter.create<linalg::FillOp>(loc, paddingValue, resultTensor)
+            ->getResult(0);
+
+    // Get the interior region.
+    SmallVector<OpFoldResult> sizes =
+        getMixedSizes(rewriter, loc, sourceTensor);
+    SmallVector<OpFoldResult> strides(sourceTensorType.getRank(),
+                                      rewriter.getI64IntegerAttr(1));
+
+    // Copy input into the interior region
+    resultTensor = rewriter.create<tensor::InsertSliceOp>(
+        loc, sourceTensor, resultTensor, op.getMixedLowPad(), sizes, strides);
+
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getResultType(),
+                                                resultTensor);
+    return success();
+  }
+};
+
+static void applyEnablingTransformations(ModuleOp moduleOp) {
+  RewritePatternSet patterns(moduleOp.getContext());
+  patterns.add<PadTensorOpPattern>(moduleOp.getContext());
+  (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
+}
+
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
   ModuleOp moduleOp = getOperation();
+  applyEnablingTransformations(moduleOp);
 
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return signalPassFailure();
 
   GlobalCreator globalCreator(moduleOp);