[mlir][linalg] Fold linalg.pad_tensor if src type == result type

[springerm]

Summary

Fold PadTensorOp to source if source type and result type have static shape and are equal.

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

let builders = [
     // Build a PadTensorOp with mixed static and dynamic entries and custom
     // result type. If the type passed is nullptr, it is inferred.
     OpBuilder<(ins "Type":$resultType, "Value":$source,
       "ArrayRef<OpFoldResult>":$low, "ArrayRef<OpFoldResult>":$high,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 def Linalg_RangeOp :
     Linalg_Op<"range", [NoSideEffect]>,
     Arguments<(ins Index:$min, Index:$max, Index:$step)>,
     Results<(outs Range)> {
   let summary = "Create a `range` type value, used to create `view`s";
   let description = [{

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

if (matchPattern(padValue, m_Constant()))
     return padValue;
   // Check if yield value is defined inside the PadTensorOp block.
   if (padValue.getParentBlock() == &getRegion().front())
     return {};
   // Else: Yield value defined outside of the PadTensorOp block.
   return padValue;
 }
 
+OpFoldResult PadTensorOp::fold(ArrayRef<Attribute>) {
+  if (getResultType().hasStaticShape() && getResultType() == getSourceType())
+    return source();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 Optional<SmallVector<ReassociationIndices>>
 mlir::linalg::getReassociationIndicesForReshape(ShapedType sourceType,
                                                 ShapedType targetType) {
   // Make the sourceType greater rank than the targetType. If they are same

mlir/test/Dialect/Linalg/canonicalize.mlir

%3 = linalg.pad_tensor %arg2 low[%c0, %c0] high[%high, %high] {
         ^bb0(%arg9: index, %arg10: index):  // no predecessors
           linalg.yield %cst : f32
   } : tensor<?x?xf32> to tensor<2x4xf32>
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @pad_tensor_same_static_shape(
+//  CHECK-SAME:   %[[ARG0:.*]]: tensor<5x6xf32>
+//   CHECK-NOT:   linalg.pad_tensor
+//       CHECK:   return %[[ARG0]]
+func @pad_tensor_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
+    -> tensor<5x6xf32> {
+  %cst = constant 0.000000e+00 : f32
+  %0 = linalg.pad_tensor %arg0 low[%a, 0] high[0, %a] {
+        ^bb0(%arg1: index, %arg2: index):
+          linalg.yield %cst : f32
+  } : tensor<5x6xf32> to tensor<5x6xf32>
+  return %0 : tensor<5x6xf32>
+}
+
+// -----
+
 func @dim_reshape_expansion(%arg0 : tensor<6x5x?xf32>) -> (index, index, index)
 {
   %c1 = constant 1 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2], [3, 4, 5]]
       : tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
   %1 = memref.dim %0, %c1 : tensor<2x3x5x4x?x7xf32>