[DispatchCreation] Skip Pad decomposition within dispatches. (iree-org#20373)

If a pad is already in a dispatch it shouldn't be decomposed. That is
effectively user assertion to not decompose the pad.

Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>

Author: MaheshRavishankar

compiler/src/iree/compiler/DispatchCreation/TensorPadToTensorInsertSlice.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
+#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -41,6 +42,9 @@ struct TensorPadOpConversion : public OpRewritePattern<tensor::PadOp> {
 
   LogicalResult matchAndRewrite(tensor::PadOp padTensorOp,
                                 PatternRewriter &rewriter) const override {
+    if (!IREE::Flow::isNonNullAndOutsideDispatch(padTensorOp)) {
+      return failure();
+    }
     // Check that the region is just a yield operation which is returning a
     // scalar that is not one of the arguments of the linalg operation.
     Region &region = padTensorOp.getRegion();

compiler/src/iree/compiler/DispatchCreation/test/tensor_pad_to_tensor_insert_slice.mlir

@@ -77,7 +77,7 @@ util.func public @_main(%arg0: tensor<1x33x33x480xf32>, %arg1: tensor<3x3x480x1x
 // CHECK-NOT: tensor.pad
 // SKIP: tensor.pad
 
-// ----
+// -----
 
 #encoding = #iree_encoding.testing_encoding<>
 util.func public @dispatch_dispatch_0_generic_512x1024_f32(
@@ -105,3 +105,24 @@ util.func public @dispatch_dispatch_0_generic_512x1024_f32(
 // CHECK:    %[[PAD:.+]] = tensor.pad %[[LOAD]] low
 // CHECK:    %[[ENCODE:.+]] = iree_encoding.set_encoding %[[PAD]] : tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING]]>
 // CHECK:    flow.dispatch.tensor.store %[[ENCODE]],
+
+// -----
+
+// Do not break up pad within dispatches.
+
+util.func @pad_within_dispatch(%arg0 : tensor<500x1000xf32>) -> tensor<512x1024xf32> {
+  %cst = arith.constant 0.0 : f32
+  %0 = flow.dispatch.region -> (tensor<512x1024xf32>) {
+    %1 = tensor.pad %arg0 low [0, 0] high[12, 24] {
+    ^bb0(%arg1 : index, %arg2 : index):
+      tensor.yield %cst : f32
+    } : tensor<500x1000xf32> to tensor<512x1024xf32>
+    flow.return %1 : tensor<512x1024xf32>
+  }
+  util.return %0 : tensor<512x1024xf32>
+}
+// CHECK-LABEL: func public @pad_within_dispatch
+//       CHECK:   %[[DISPATCH:.+]] = flow.dispatch.region
+//       CHECK:     %[[PAD:.+]] = tensor.pad
+//       CHECK:     flow.return %[[PAD]]
+//       CHECK:   return %[[DISPATCH]]