[mlir][Linalg] Replace SimplePad with PadTensor in hoist-padding

[hanchung]

Summary

This is the last revision to migrate using SimplePadOp to PadTensorOp, and the
SimplePadOp is removed in the patch. Update a bit in SliceAnalysis because the
PadTensorOp takes a region different from SimplePadOp. This is not covered by
LinalgOp because it is not a structured op.

Also, remove a duplicated comment from cpp file, which is already described in a
header file. And update the pseudo-mlir in the comment.

[hanchung]

This is as same as D95615 but fixing one dep in CMakeLists.txt. I couldn't repro the issue on local side, so would wait until the builds pass in remote lab.

[hanchung]

rebase

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

let extraClassDeclaration = [{
     static StringRef getStaticLowAttrName() {
       return "static_low";
     }
 
     static StringRef getStaticHighAttrName() {
       return "static_high";
     }
 
+    RankedTensorType getSourceType() {
+      return source().getType().cast<RankedTensorType>();
+    }
+    RankedTensorType getResultType() {
+      return getResult().getType().cast<RankedTensorType>();
+    }
+
     // 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
SmallVector<Value, 8> getRanges() {
           res.push_back(operand);
       return res;
     }
   }];
 
   let hasFolder = 1;
 }
 
-def Linalg_SimplePadOp : Linalg_Op<"simple_pad", [NoSideEffect]> {
-  let summary = "TODO: replace with pad_tensors when ready.";
-
-  let description = [{
-    `linalg.simple_pad` is a tmp placeholder for padding and packing on tensors.
-    Its semantics are to pad a partially dynamic tensor to a fully static tensor
-    where the static sizes are assumed to be greater than the dynamic sizes. The
-    op perforrms "high" padding (i.e. it adds trailing padding values until the 
-    desired size is met).
-  }];
-
-  let arguments = (ins AnyRankedTensor:$tensor, AnyType:$padding);
-  let results = (outs AnyRankedTensor:$result);
-
-  // TODO: verify all static result, some dynamic input, static shapes match,
-  // element types match, ranks match etc. Use pad_tensors when ready but for
-  // now just let it ne fully specified by traits.
-  let verifier = ?;
-
-  let extraClassDeclaration = [{
-    RankedTensorType getSourceType() {
-      return tensor().getType().cast<RankedTensorType>(); }
-    RankedTensorType getResultType() {
-      return getResult().getType().cast<RankedTensorType>(); }
-   }];
-
-  let assemblyFormat = [{
-    $tensor `pad` $padding attr-dict `:`
-      type($tensor) `to` type($result) `pad` type($padding)
-  }];
-}
-
 def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "Linalg yield operation";
   let description = [{
     `linalg.yield` is a special terminator operation for blocks inside regions
     in `linalg` generic ops. It returns values to the immediately enclosing
     `linalg` generic op.
 

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

 //===- Hoisting.h - Linalg hoisting transformations -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_TRANSFORMS_HOISTING_H_
 #define MLIR_DIALECT_LINALG_TRANSFORMS_HOISTING_H_
 
 namespace mlir {
 class FuncOp;
 struct LogicalResult;
 
 namespace linalg {
-class SimplePadOp;
+class PadTensorOp;
 
 /// Hoist alloc/dealloc pairs and alloca op out of immediately enclosing
 /// scf::ForOp if both conditions are true:
 ///   1. All operands are defined outside the loop.
 ///   2. All uses are ViewLikeOp or DeallocOp.
 // TODO: generalize on a per-need basis.
 void hoistViewAllocOps(FuncOp func);
 

 void hoistRedundantVectorTransfers(FuncOp func);
 
 /// Same behavior as `hoistRedundantVectorTransfers` but works on tensors
 /// instead of buffers.
 void hoistRedundantVectorTransfersOnTensor(FuncOp func);
 
 /// Mechanically hoist padding operations on tensors by `nLoops` into a new,
 /// generally larger tensor. This achieves packing of multiple padding ops into
-/// a larger tensor. On success, `simplePadOp` is replaced by the cloned version
+/// a larger tensor. On success, `padTensorOp` is replaced by the cloned version
 /// in the packing loop so the caller can continue reasoning about the padding
 /// operation.
 ///
 /// Example in pseudo-mlir:
 /// =======================
 ///
 /// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR.
 /// ```
 ///    scf.for (%i, %j, %k)
 ///      %st0 = subtensor f(%i, %k) : ... to tensor<?x?xf32>
-///      %0 = linalg.simple_pad %st0 pad %pad :
-///             tensor<?x?xf32> to tensor<4x8xf32>
+///      %0 = linalg.pad_tensor %st0 low[0, 0] high[...] {
+///      ^bb0( ... ):
+///        linalg.yield %pad
+///      } : tensor<?x?xf32> to tensor<4x8xf32>
 ///      compute(%0)
 /// ```
 ///
 /// IR resembling the following is produced:
 ///
 /// ```
 ///    scf.for (%i) {
 ///      %packed_init = linalg.init_tensor range(%j) : tensor<?x4x8xf32>
-///      %packed = scf.for (%k) iter_args(%p : %packed_init)
+///      %packed = scf.for (%k) iter_args(%p : %packed_init) {
 ///        %st0 = subtensor f(%i, %k) : ... to tensor<?x?xf32>
-///        %0 = linalg.simple_pad %st0 pad %pad :
-///               tensor<?x?xf32> to tensor<4x8xf32>
+///        %0 = linalg.pad_tensor %st0 low[0, 0] high[...] {
+///        ^bb0( ... ):
+///          linalg.yield %pad
+///        } : tensor<?x?xf32> to tensor<4x8xf32>
+///        %1 = subtensor_insert %0 ... : tensor<4x8xf32> to tensor<?x4x8xf32>
 ///        scf.yield %1: tensor<?x4x8xf32>
 ///      } -> tensor<?x4x8xf32>
 ///      scf.for (%j, %k) {
 ///        %st0 = subtensor %packed [%k, 0, 0][1, 4, 8][1, 1, 1] :
 ///                 tensor<?x4x8xf32> to tensor<4x8xf32>
 ///        compute(%st0)
 ///      }
 ///    }
 /// ```
-LogicalResult hoistPaddingOnTensors(SimplePadOp &simplePadOp, unsigned nLoops);
+LogicalResult hoistPaddingOnTensors(PadTensorOp &padTensorOp, unsigned nLoops);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_HOISTING_H_

mlir/lib/Analysis/CMakeLists.txt

add_mlir_library(MLIRLoopAnalysis
   DEPENDS
   mlir-headers
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
+  MLIRLinalg
   MLIRPresburger
   MLIRSCF
   )
 
 add_subdirectory(Presburger)

mlir/lib/Analysis/SliceAnalysis.cpp

 
 static void getBackwardSliceImpl(Operation *op,
                                  SetVector<Operation *> *backwardSlice,
                                  TransitiveFilter filter) {
   if (!op)
     return;
 
   assert((op->getNumRegions() == 0 ||
-          isa<AffineForOp, scf::ForOp, linalg::LinalgOp>(op)) &&
+          isa<AffineForOp, scf::ForOp, linalg::LinalgOp, linalg::PadTensorOp>(
+              op)) &&
          "unexpected generic op with regions");
 
   // Evaluate whether we should keep this def.
   // This is useful in particular to implement scoping; i.e. return the
   // transitive forwardSlice in the current scope.
   if (!filter(op)) {
     return;
   }

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

func.walk([&](vector::TransferReadOp transferRead) {
       // walk.
       return WalkResult::interrupt();
     });
   }
 }
 
 /// Ensure prerequisites that guarantee pad op hoisting can occur.
 /// Return failure in the cases when we cannot perform hoisting; i.e. if either:
-///   1. There exists a use of `simplePadOp` that is not a linalg input operand.
+///   1. There exists a use of `padTensorOp` that is not a linalg input operand.
 ///   2. There isn't an enclosing `outermostEnclosingForOp` loop.
 ///   3. There exists an op with a region that is dominated by
 ///   `outermostEnclosingForOp` and that isn't a LoopLikeInterface or a
 ///    LinalgOp.
 ///   3. There exists an op with side effects that is dominated by
 ///    `outermostEnclosingForOp` and that isn't a LoopLikeInterface.
 ///
 /// While ensuring prerequisites:
 ///   1. Fill the `backwardSlice` to contain the topologically sorted ops
 ///   dominated by `outermostEnclosingForOp`.
 ///   2. Fill the `packingLoops` to contain only the enclosing loops of
 ///   `backwardSlice` whose IV is actually used in computing padding. Loops that
 ///   remain in `backwardSlice` but that are not in `packingLoops` are
 ///   dimensions of reuse.
 static LogicalResult
-hoistPaddingOnTensorsPrerequisites(linalg::SimplePadOp simplePadOp, int nLevels,
+hoistPaddingOnTensorsPrerequisites(linalg::PadTensorOp padTensorOp, int nLevels,
                                    llvm::SetVector<Operation *> &backwardSlice,
                                    llvm::SetVector<Operation *> &packingLoops) {
   // Bail on any use that isn't an input of a Linalg op.
   // Hoisting of inplace updates happens after vectorization.
-  for (OpOperand &use : simplePadOp.result().getUses()) {
+  for (OpOperand &use : padTensorOp.result().getUses()) {
     auto linalgUser = dyn_cast<linalg::LinalgOp>(use.getOwner());
     if (!linalgUser || !linalgUser.isInputTensor(&use))
       return failure();
   }
 
   // Get at most nLevels of enclosing loops.
   SmallVector<LoopLikeOpInterface> reverseEnclosingLoops;
   Operation *outermostEnclosingForOp = nullptr,
             *nextEnclosingForOp =
-                simplePadOp->getParentOfType<LoopLikeOpInterface>();
+                padTensorOp->getParentOfType<LoopLikeOpInterface>();
   while (nLevels-- > 0 && nextEnclosingForOp) {
     outermostEnclosingForOp = nextEnclosingForOp;
     reverseEnclosingLoops.push_back(outermostEnclosingForOp);
     nextEnclosingForOp =
         nextEnclosingForOp->getParentOfType<LoopLikeOpInterface>();
   }
   if (!outermostEnclosingForOp)
     return failure();
 
-  // Get the backwards slice from `simplePadOp` that is dominated by the
+  // Get the backwards slice from `padTensorOp` that is dominated by the
   // outermost enclosing loop.
   DominanceInfo domInfo(outermostEnclosingForOp);
-  getBackwardSlice(simplePadOp, &backwardSlice, [&](Operation *op) {
+  getBackwardSlice(padTensorOp, &backwardSlice, [&](Operation *op) {
     return domInfo.dominates(outermostEnclosingForOp, op);
   });
 
-#if 0
-
-  // Bail on any op with a region that is not a LoopLikeInterface or a LinalgOp.
-  // Bail on any op with side effects that is not a LoopLikeInterface.
-  if (llvm::any_of(backwardSlice, [](Operation *op) {
-        if (isa<LoopLikeOpInterface>(op))
-          return false;
-        if (!MemoryEffectOpInterface::hasNoEffect(op))
-          return true;
-        return op->getNumRegions() > 0 && !isa<LinalgOp>(op);
-      }))
-    return failure();
-
-#else
-
   // Bail on any op with a region that is not a LoopLikeInterface or a LinalgOp.
   if (llvm::any_of(backwardSlice, [](Operation *op) {
         return op->getNumRegions() > 0 && !isa<LoopLikeOpInterface>(op) &&
                !isa<LinalgOp>(op);
       }))
     return failure();
 
-#endif
-
   // Filter out the loops whose induction variable is not used to compute the
   // padded result. As a first approximation, just look for IVs that have no use
   // in the backwardSlice.
   // These are the dimensions of reuse that we can exploit to reduce the amount
   // of work / memory.
   // TODO: would this optimization compose better as a canonicalization?
   for (LoopLikeOpInterface loop : reverseEnclosingLoops) {
     auto forOp = dyn_cast<scf::ForOp>(loop.getOperation());
static Value buildLoopTripCount(OpBuilder &b, Operation *op) {
   AffineExpr lb, ub, step = getAffineSymbolExpr(0, ctx);
   bindDims(ctx, lb, ub);
   scf::ForOp forOp = cast<scf::ForOp>(op);
   return b.create<AffineApplyOp>(
       op->getLoc(), AffineMap::get(2, 1, {(ub - lb).ceilDiv(step)}, ctx),
       ValueRange{forOp.lowerBound(), forOp.upperBound(), forOp.step()});
 }
 
-/// Mechanically hoist padding operations on tensors by at most `nLoops` into a
-/// new, generally larger tensor. This achieves packing of multiple padding ops
-/// into a larger tensor. On success, `simplePadOp` is replaced by the cloned
-/// version in the packing loop so the caller can continue reasoning about the
-/// padding operation.
-///
-/// Example in pseudo-mlir:
-/// =======================
-///
-/// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR.
-/// ```
-///    scf.for (%i, %j, %k)
-///      %st0 = subtensor f(%i, %k) : ... to tensor<?x?xf32>
-///      %0 = linalg.simple_pad %st0 pad %pad :
-///             tensor<?x?xf32> to tensor<4x8xf32>
-///      compute(%0)
-/// ```
-///
-/// IR resembling the following is produced:
-///
-/// ```
-///    scf.for (%i) {
-///      %packed_init = linalg.init_tensor range(%j) : tensor<?x4x8xf32>
-///      %packed = scf.for (%k) iter_args(%p : %packed_init)
-///        %st0 = subtensor f(%i, %k) : ... to tensor<?x?xf32>
-///        %0 = linalg.simple_pad %st0 pad %pad :
-///               tensor<?x?xf32> to tensor<4x8xf32>
-///        scf.yield %1: tensor<?x4x8xf32>
-///      } -> tensor<?x4x8xf32>
-///      scf.for (%j, %k) {
-///        %st0 = subtensor %packed [%k, 0, 0][1, 4, 8][1, 1, 1] :
-///                 tensor<?x4x8xf32> to tensor<4x8xf32>
-///        compute(%st0)
-///      }
-///    }
-/// ```
-LogicalResult mlir::linalg::hoistPaddingOnTensors(SimplePadOp &simplePadOp,
+LogicalResult mlir::linalg::hoistPaddingOnTensors(PadTensorOp &padTensorOp,
                                                   unsigned nLoops) {
   llvm::SetVector<Operation *> backwardSlice, packingLoops;
-  if (failed(hoistPaddingOnTensorsPrerequisites(simplePadOp, nLoops,
+  if (failed(hoistPaddingOnTensorsPrerequisites(padTensorOp, nLoops,
                                                 backwardSlice, packingLoops)))
     return failure();
 
   // Update actual number of loops, which may be smaller.
   nLoops = packingLoops.size();
 
-  Location loc = simplePadOp->getLoc();
-  RankedTensorType paddedTensorType = simplePadOp.getResultType();
+  Location loc = padTensorOp->getLoc();
+  RankedTensorType paddedTensorType = padTensorOp.getResultType();
   unsigned paddedRank = paddedTensorType.getRank();
 
   // Backward slice is a topologically sorted list of ops starting at
   // `outermostEnclosingForOp`.
   Operation *outermostEnclosingForOp = backwardSlice.front();
   // IP just before the outermost loop considered that we hoist above.
   OpBuilder b(outermostEnclosingForOp);
 
   // Create the packed tensor<?x?x..?xpadded_shape> into which we amortize
   // padding.
   SmallVector<int64_t> packedShape(nLoops, ShapedType::kDynamicSize);
-  // TODO: go grab dims when necessary, for now SimplePadOp returns a static
+  // TODO: go grab dims when necessary, for now PadTensorOp returns a static
   // tensor.
   llvm::append_range(packedShape, paddedTensorType.getShape());
   auto packedTensorType =
       RankedTensorType::get(packedShape, paddedTensorType.getElementType());
   auto dynamicSizes = llvm::to_vector<4>(llvm::map_range(
       packingLoops, [&](Operation *op) { return buildLoopTripCount(b, op); }));
   Value packedTensor = b.create<linalg::InitTensorOp>(
       loc, dynamicSizes, packedTensorType.getShape(),
       packedTensorType.getElementType());
 
   // Clone the operations involved in the backward slice, iteratively stepping
   // into the loops that we encounter.
   // The implementation proceeds in a stack-like fashion:
   //   1. Iteratively clone and step into the loops, pushing the `packedTensor`
   //      deeper in the stack.
   //   2. Create a SubTensorInsert at the top of the stack.
   //   3. Iteratively pop and yield the result of the SubTensorInsertOp across
   //     the cloned loops.
   SmallVector<Value> clonedLoopIvs;
   clonedLoopIvs.reserve(nLoops);
   BlockAndValueMapping bvm;
   // Stack step 1. iteratively clone loops and push `packedTensor`.
-  // Insert `simplePadOp` into the backwardSlice so we clone it too.
-  backwardSlice.insert(simplePadOp);
+  // Insert `padTensorOp` into the backwardSlice so we clone it too.
+  backwardSlice.insert(padTensorOp);
   for (Operation *op : backwardSlice) {
-    if (op->getNumRegions() == 0) {
+    if (op->getNumRegions() == 0 || isa<linalg::PadTensorOp>(op)) {
       b.clone(*op, bvm);
       continue;
     }
     // TODO: support more cases as they appear.
     auto forOp = dyn_cast<scf::ForOp>(op);
     assert(forOp && "Expected scf::ForOp when hoisting pad ops");
     // Unused loop, just skip it.
     if (!packingLoops.contains(forOp))
LogicalResult mlir::linalg::hoistPaddingOnTensors(PadTensorOp &padTensorOp,
 
   // Stack step 2. create SubTensorInsertOp at the top of the stack.
   // offsets = [clonedLoopIvs, 0 .. 0].
   SmallVector<OpFoldResult> offsets(clonedLoopIvs.begin(), clonedLoopIvs.end());
   offsets.append(paddedRank, b.getIndexAttr(0));
   // sizes = [1 .. 1, paddedShape].
   SmallVector<OpFoldResult> sizes(nLoops, b.getIndexAttr(1));
   for (int64_t sz : paddedTensorType.getShape()) {
-    // TODO: go grab dims when necessary, for now SimplePadOp returns a static
+    // TODO: go grab dims when necessary, for now PadTensorOp returns a static
     // tensor.
     assert(!ShapedType::isDynamic(sz) && "padded tensor needs static sizes");
     sizes.push_back(b.getIndexAttr(sz));
   }
   // strides = [1 .. 1].
   SmallVector<OpFoldResult> strides(nLoops + paddedRank, b.getIndexAttr(1));
 
   Value inserted =
-      b.create<SubTensorInsertOp>(loc, bvm.lookup(simplePadOp.result()),
+      b.create<SubTensorInsertOp>(loc, bvm.lookup(padTensorOp.result()),
                                   packedTensor, offsets, sizes, strides);
 
   // Stack step 3. iteratively pop the stack and propagate the yield.
   Value valueToYield = inserted;
   for (Value iv : llvm::reverse(clonedLoopIvs)) {
     auto forOp = scf::getForInductionVarOwner(iv);
     b.setInsertionPointToEnd(&forOp.getRegion().front());
     b.create<scf::YieldOp>(loc, valueToYield);
     valueToYield = forOp.getResult(0);
   }
 
   // Now the packed tensor is ready, replace the original padding op by a
   // 1x..x1 SubTensor [originalLoopIvs, 0 .. 0][1 .. 1, paddedShape][1 .. 1].
-  b.setInsertionPoint(simplePadOp);
+  b.setInsertionPoint(padTensorOp);
   SmallVector<Value> originalLoopIvs =
       llvm::to_vector<4>(llvm::map_range(packingLoops, [](Operation *loop) {
         return cast<scf::ForOp>(loop).getInductionVar();
       }));
   // offsets = [originalLoopIvs, 0 .. 0].
   offsets.assign(originalLoopIvs.begin(), originalLoopIvs.end());
   offsets.append(paddedRank, b.getIndexAttr(0));
   // sizes = [1 .. 1, paddedShape] (definedabove).
   // strides = [1 .. 1] (defined above)
   packedTensor =
       scf::getForInductionVarOwner(clonedLoopIvs.front())->getResult(0);
-  simplePadOp.replaceAllUsesWith(
-      b.create<SubTensorOp>(loc, simplePadOp.getResultType(), packedTensor,
+  padTensorOp.replaceAllUsesWith(
+      b.create<SubTensorOp>(loc, padTensorOp.getResultType(), packedTensor,
                             offsets, sizes, strides)
           ->getResult(0));
 
-  Operation *toErase = simplePadOp;
+  Operation *toErase = padTensorOp;
 
-  // Make the newly cloned `simplePadOp` available to the caller.
-  simplePadOp =
-      cast<SimplePadOp>(bvm.lookup(simplePadOp.result()).getDefiningOp());
+  // Make the newly cloned `padTensorOp` available to the caller.
+  padTensorOp =
+      cast<PadTensorOp>(bvm.lookup(padTensorOp.result()).getDefiningOp());
 
   toErase->erase();
 
   return success();
 }

mlir/test/Dialect/Linalg/hoist-padding.mlir

func @matmul_tensors(
   %2 = dim %arg1, %c1 : tensor<?x?xf32>
 
   //      CHECK: scf.for
   //      CHECK:   linalg.init_tensor [%{{.*}}, 2, 4] : tensor<?x2x4xf32>
   // 1-D loop
   //      CHECK:   %[[A:.*]] = scf.for
   //  CHECK-NOT:     scf.for
   //      CHECK:     subtensor %{{.*}} [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
-  //      CHECK:     linalg.simple_pad %{{.*}} : tensor<?x?xf32> to tensor<2x4xf32> pad f32
+  //      CHECK:     linalg.pad_tensor %{{.*}}
+  //      CHECK:       : tensor<?x?xf32> to tensor<2x4xf32>
   //      CHECK:     subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, 0, 0]
   // CHECK-SAME:       [1, 2, 4] [1, 1, 1] : tensor<2x4xf32> into tensor<?x2x4xf32>
   // 2-D loop
   //      CHECK:   linalg.init_tensor [%{{.*}}, %{{.*}}, 4, 3] : tensor<?x?x4x3xf32>
   //      CHECK:   %[[B:.*]] = scf.for
   //      CHECK:     scf.for
   //  CHECK-NOT:       scf.for
   //      CHECK:       subtensor %{{.*}} [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
-  //      CHECK:       linalg.simple_pad %{{.*}} : tensor<?x?xf32> to tensor<4x3xf32> pad f32
+  //      CHECK:       linalg.pad_tensor %{{.*}}
+  //      CHECK:         : tensor<?x?xf32> to tensor<4x3xf32>
   //      CHECK:       subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0, 0]
   // CHECK-SAME:         [1, 1, 4, 3] [1, 1, 1, 1] : tensor<4x3xf32> into tensor<?x?x4x3xf32>
   // 2-D loop
   //      CHECK:   scf.for %[[J:[0-9a-zA-Z]+]]
   //      CHECK:     scf.for %[[K:[0-9a-zA-Z]+]]
   //  CHECK-NOT:       scf.for
   //      CHECK:       %[[stA:.*]] = subtensor %[[A]][%[[K]], 0, 0] [1, 2, 4] [1, 1, 1] :
   // CHECK-SAME:         tensor<?x2x4xf32> to tensor<2x4xf32>
   //      CHECK:       %[[stB:.*]] = subtensor %[[B]][%[[K]], %[[J]], 0, 0] [1, 1, 4, 3] [1, 1, 1, 1] :
   // CHECK-SAME:         tensor<?x?x4x3xf32> to tensor<4x3xf32>
-  //      CHECK:       %[[stC:.*]] = linalg.simple_pad %{{.*}} pad %{{.*}} :
-  // CHECK-SAME:         tensor<?x?xf32> to tensor<2x3xf32> pad f32
+  //      CHECK:       %[[stC:.*]] = linalg.pad_tensor %{{.*}}
+  //      CHECK:         : tensor<?x?xf32> to tensor<2x3xf32>
   //      CHECK:       linalg.matmul ins(%[[stA]], %[[stB]] : tensor<2x4xf32>, tensor<4x3xf32>)
   // CHECK-SAME:         outs(%[[stC]] : tensor<2x3xf32>) -> tensor<2x3xf32>
   %3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
     %4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
       %5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
         %6 = dim %arg0, %c0 : tensor<?x?xf32>
         %7 = affine.min #map0(%arg3)[%6]
         %8 = dim %arg0, %c1 : tensor<?x?xf32>
         %9 = affine.min #map1(%arg7)[%8]
         %10 = subtensor %arg0[%arg3, %arg7] [%7, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
         %11 = dim %arg1, %c0 : tensor<?x?xf32>
         %12 = affine.min #map1(%arg7)[%11]
         %13 = dim %arg1, %c1 : tensor<?x?xf32>
         %14 = affine.min #map2(%arg5)[%13]
         %15 = subtensor %arg1[%arg7, %arg5] [%12, %14] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
         %16 = dim %arg8, %c0 : tensor<?x?xf32>
         %17 = affine.min #map3(%16, %arg3)
         %18 = dim %arg8, %c1 : tensor<?x?xf32>
         %19 = affine.min #map4(%18, %arg5)
         %20 = subtensor %arg8[%arg3, %arg5] [%17, %19] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
-        %21 = linalg.simple_pad %10 pad %cst : tensor<?x?xf32> to tensor<2x4xf32> pad f32
-        %22 = linalg.simple_pad %15 pad %cst : tensor<?x?xf32> to tensor<4x3xf32> pad f32
-        %23 = linalg.simple_pad %20 pad %cst : tensor<?x?xf32> to tensor<2x3xf32> pad f32
-        %24 = linalg.matmul ins(%21, %22 : tensor<2x4xf32>, tensor<4x3xf32>) outs(%23 : tensor<2x3xf32>) -> tensor<2x3xf32>
-        %25 = subtensor %24[0, 0] [%7, %14] [1, 1] : tensor<2x3xf32> to tensor<?x?xf32>
-        %26 = subtensor_insert %25 into %arg8[%arg3, %arg5] [%17, %19] [%c1, %c1] : tensor<?x?xf32> into tensor<?x?xf32>
-        scf.yield %26 : tensor<?x?xf32>
+        %21 = subi %c2, %7 : index
+        %22 = subi %c4, %9 : index
+        %23 = linalg.pad_tensor %10 low[%c0, %c0] high[%21, %22] {
+        ^bb0(%arg9: index, %arg10: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x?xf32> to tensor<2x4xf32>
+        %24 = subi %c4, %12 : index
+        %25 = subi %c3, %14 : index
+        %26 = linalg.pad_tensor %15 low[%c0, %c0] high[%24, %25] {
+        ^bb0(%arg9: index, %arg10: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x?xf32> to tensor<4x3xf32>
+        %27 = subi %c2, %17 : index
+        %28 = subi %c3, %19 : index
+        %29 = linalg.pad_tensor %20 low[%c0, %c0] high[%27, %28] {
+        ^bb0(%arg9: index, %arg10: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x?xf32> to tensor<2x3xf32>
+        %30 = linalg.matmul ins(%23, %26 : tensor<2x4xf32>, tensor<4x3xf32>) outs(%29 : tensor<2x3xf32>) -> tensor<2x3xf32>
+        %31 = subtensor %30[0, 0] [%7, %14] [1, 1] : tensor<2x3xf32> to tensor<?x?xf32>
+        %32 = subtensor_insert %31 into %arg8[%arg3, %arg5] [%17, %19] [%c1, %c1] : tensor<?x?xf32> into tensor<?x?xf32>
+        scf.yield %32 : tensor<?x?xf32>
       }
       scf.yield %5 : tensor<?x?xf32>
     }
     scf.yield %4 : tensor<?x?xf32>
   }
   return %3 : tensor<?x?xf32>
 }

mlir/test/Dialect/Linalg/roundtrip.mlir

 // -----
 
 func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
   %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
   %1 = linalg.fill(%0, %arg2) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 // CHECK: %{{.+}} = linalg.fill(%{{.+}}, %{{.+}}) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>
-
-// -----
-
-// TODO: this op should disappear once pad_tensors is available and connected.
-// CHECK-LABEL: func @simple_pad
-func @simple_pad(%0: tensor<?x4x?xf32>, %pad: f32) {
-//     CHECK:   linalg.simple_pad %{{.+}} pad %{{.+}}: tensor<?x4x?xf32> to tensor<8x4x8xf32>
-  %1 = linalg.simple_pad %0 pad %pad: tensor<?x4x?xf32> to tensor<8x4x8xf32> pad f32
-  return
-}

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getFunction());
   if (testAffineMinSCFCanonicalizationPatterns)
     return applyAffineMinSCFCanonicalizationPatterns(getFunction());
   if (testTileAndPadPattern)
     return applyTileAndPadPattern(getFunction());
   if (testHoistPadding2Levels) {
-    getFunction().walk([](linalg::SimplePadOp simplePadOp) {
-      linalg::hoistPaddingOnTensors(simplePadOp, 2);
+    getFunction().walk([](linalg::PadTensorOp padTensorOp) {
+      linalg::hoistPaddingOnTensors(padTensorOp, 2);
     });
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace test
 } // namespace mlir