[mlir][tensor] Allow tensors with zero dims in tensor.pad/extract_slice swapping

[springerm]

Summary

The tensor.pad(tensor.extract_slice) swapping pattern used to generate "scf.if" and "tensor.generate" ops to avoid tensors that have a dimension that has a static or dynamic size zero. This revision removes that functionality. Tensors with a static dimension zero size may appear in the resulting IR. Such IR can be bufferized and is ignored by vectorization. Ops involving tensors that have zero elements usually fold away at some point.

Depends On: D156137

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

auto max = [&](OpFoldResult v1, OpFoldResult v2) {
     return affine::makeComposedFoldedAffineMax(b, loc, idMap, {v1, v2});
   };
   // Zero index-typed integer.
   OpFoldResult zero = b.getIndexAttr(0);
 
   // Compute new offsets, lengths, low padding, high padding.
   SmallVector<OpFoldResult> newOffsets, newLengths, newStrides;
   SmallVector<OpFoldResult> newLows, newHighs;
-  // Set to true if the original data source is not read at all.
-  bool hasZeroLen = false;
-  // Same as hasZeroLen, but for dynamic dimension sizes. This condition
-  // is true if the original data source turns out to be unused at runtime.
-  Value dynHasZeroLenCond;
 
   int64_t rank = padOp.getSourceType().getRank();
   for (unsigned dim = 0; dim < rank; ++dim) {
     auto low = padOp.getMixedLowPad()[dim];
     bool hasLowPad = !isConstantIntValue(low, 0);
     auto high = padOp.getMixedHighPad()[dim];
     bool hasHighPad = !isConstantIntValue(high, 0);
     auto offset = offsets[dim];
     auto length = sizes[dim];
     auto srcSize = tensor::getMixedSize(b, loc, padOp.getSource(), dim);
 
     // The new amount of low padding is `low - offset`. Except for the case
     // where none of the low padding is read. In that case, the new amount of
-    // low padding is zero.
+    // low padding is zero. Also, do not add more low padding than the size of
+    // result dimension.
     //
     // Optimization: If low = 0, then newLow = 0.
-    OpFoldResult newLow = hasLowPad ? max(zero, sub(low, offset)) : zero;
+    OpFoldResult newLow =
+        hasLowPad ? min(max(zero, sub(low, offset)), length) : zero;
     newLows.push_back(newLow);
 
     // Start reading the data from position `offset - low`. Since the original
     // read may have started in the low padding zone, this value could be
     // negative. Therefore, start reading from:
     //
     // max(offset - low, 0)
     //
for (unsigned dim = 0; dim < rank; ++dim) {
     //
     // Optimization: If low = 0, then the formula can be simplified.
     OpFoldResult endLoc =
         hasLowPad ? min(max(add(sub(offset, low), length), zero), srcSize)
                   : min(add(offset, length), srcSize);
     OpFoldResult newLength = sub(endLoc, newOffset);
     newLengths.push_back(newLength);
 
-    // Check if newLength is zero. In that case, no SubTensorOp should be
-    // executed.
-    if (isConstantIntValue(newLength, 0)) {
-      hasZeroLen = true;
-    } else if (!hasZeroLen) {
-      Value check = b.create<arith::CmpIOp>(
-          loc, arith::CmpIPredicate::eq,
-          getValueOrCreateConstantIndexOp(b, loc, newLength),
-          getValueOrCreateConstantIndexOp(b, loc, zero));
-      dynHasZeroLenCond =
-          dynHasZeroLenCond
-              ? b.create<arith::OrIOp>(loc, check, dynHasZeroLenCond)
-              : check;
-    }
-
     // The amount of high padding is simply the number of elements remaining,
     // so that the result has the same length as the original ExtractSliceOp.
     // As an optimization, if the original high padding is zero, then the new
     // high padding must also be zero.
     OpFoldResult newHigh =
         hasHighPad ? sub(sub(length, newLength), newLow) : zero;
     newHighs.push_back(newHigh);
 
FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
 
   // Insert cast to ensure that types match. (May be folded away.)
   auto castResult = [&](Value val) -> Value {
     if (resultType == val.getType())
       return val;
     return b.create<tensor::CastOp>(loc, resultType, val);
   };
 
-  // In cases where the original data source is unused: Emit a GenerateOp and
-  // do not generate a SliceOp. (The result shape of the SliceOp would
-  // have a dimension of size 0, the semantics of which is unclear.)
-  auto createGenerateOp = [&]() {
-    // Create GenerateOp.
-    auto generateOp = b.create<tensor::GenerateOp>(
-        loc, resultType, dynDims,
-        [&](OpBuilder &builder, Location gLoc, ValueRange indices) {
-          builder.create<tensor::YieldOp>(gLoc, padValue);
-        });
-    return generateOp;
-  };
-
-  // Emit a SliceOp and a PadOp. Should not be used in cases where
-  // the result shape of the new SliceOp has a zero dimension.
+  // Emit a SliceOp and a PadOp.
   auto createPadOfExtractSlice = [&]() {
-    // Create pad(extract_slice(x)).
-    Value newSliceOp = b.create<tensor::ExtractSliceOp>(
-        loc, padOp.getSource(), newOffsets, newLengths, newStrides);
+    Value newSlice =
+        b.create<tensor::ExtractSliceOp>(loc, padOp.getSource(), newOffsets,
+                                         newLengths, newStrides)
+            .getResult();
     auto newPadOp = b.create<PadOp>(
-        loc, Type(), newSliceOp, newLows, newHighs,
+        loc, Type(), newSlice, newLows, newHighs,
         /*nofold=*/padOp.getNofold(),
         getPrunedAttributeList(padOp, PadOp::getAttributeNames()));
 
     // Copy region to new PadOp.
     IRMapping bvm;
     padOp.getRegion().cloneInto(&newPadOp.getRegion(), bvm);
 
     // Cast result and return.
     return newPadOp;
   };
 
-  // Rewrite extract_slice(pad(x)) into a GenerateOp it is statically known that
-  // the original data source x is not used.
-  if (hasZeroLen) {
-    Operation *generateOp = createGenerateOp();
-    return TilingResult{{generateOp}, {castResult(generateOp->getResult(0))}};
-  }
-
-  // If there are dynamic dimensions: Generate an scf.if check to avoid
-  // creating SliceOps with result dimensions of size 0 at runtime.
-  if (generateZeroSliceGuard && dynHasZeroLenCond) {
-    Operation *thenOp;
-    Operation *elseOp;
-    auto result = b.create<scf::IfOp>(
-        loc, dynHasZeroLenCond,
-        /*thenBuilder=*/
-        [&](OpBuilder &b, Location loc) {
-          thenOp = createGenerateOp();
-          b.create<scf::YieldOp>(loc, castResult(thenOp->getResult(0)));
-        },
-        /*elseBuilder=*/
-        [&](OpBuilder &b, Location loc) {
-          elseOp = createPadOfExtractSlice();
-          b.create<scf::YieldOp>(loc, castResult(elseOp->getResult(0)));
-        });
-    return TilingResult{{elseOp}, SmallVector<Value>(result->getResults())};
-  }
-
   Operation *newPadOp = createPadOfExtractSlice();
   return TilingResult{{newPadOp}, {castResult(newPadOp->getResult(0))}};
 }
 
 void mlir::tensor::registerTilingInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
     tensor::PadOp::attachInterface<PadOpTiling>(*ctx);

mlir/test/Dialect/Linalg/subtensor-of-padtensor.mlir

func.func @static_data_only(%arg0 : tensor<4x5xf32>, %pad : f32)
   %1 = tensor.extract_slice %0[1, 2] [2, 1] [1, 1] : tensor<11x13xf32> to tensor<2x1xf32>
   return %1 : tensor<2x1xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @static_high_pad_only
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f32
-//   CHECK-NOT:   tensor.pad
 //   CHECK-NOT:   tensor.extract_slice
-//       CHECK:   %[[RESULT:.*]] = tensor.generate
+//       CHECK:   %[[EMPTY:.*]] = tensor.extract_slice %[[ARG0]][4, 5] [0, 0] [1, 1] : tensor<4x5xf32> to tensor<0x0xf32>
+//       CHECK:   %[[RESULT:.*]] = tensor.pad %[[EMPTY]] low[0, 0] high[2, 4]
 //       CHECK:     tensor.yield %[[PAD]]
 //       CHECK:   return %[[RESULT]] : tensor<2x4xf32>
 func.func @static_high_pad_only(%arg0 : tensor<4x5xf32>, %pad : f32)
     -> tensor<2x4xf32> {
   %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<4x5xf32> to tensor<11x13xf32>
   %1 = tensor.extract_slice %0[4, 5] [2, 4] [1, 1] : tensor<11x13xf32> to tensor<2x4xf32>
   return %1 : tensor<2x4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @static_low_pad_only
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f32
-//   CHECK-NOT:   tensor.pad
 //   CHECK-NOT:   tensor.extract_slice
-//       CHECK:   %[[RESULT:.*]] = tensor.generate
+//       CHECK:   %[[EMPTY:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [0, 0] [1, 1] : tensor<4x5xf32> to tensor<0x0xf32>
+//       CHECK:   %[[RESULT:.*]] = tensor.pad %[[EMPTY]] low[2, 3] high[0, 0]
 //       CHECK:     tensor.yield %[[PAD]]
 //       CHECK:   return %[[RESULT]] : tensor<2x3xf32>
 func.func @static_low_pad_only(%arg0 : tensor<4x5xf32>, %pad : f32)
     -> tensor<2x3xf32> {
   %0 = tensor.pad %arg0 low[3, 7] high[7, 8] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<4x5xf32> to tensor<14x20xf32>
   %1 = tensor.extract_slice %0[1, 3] [2, 3] [1, 1] : tensor<14x20xf32> to tensor<2x3xf32>
   return %1 : tensor<2x3xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @static_low_pad_only_2
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f32
-//   CHECK-NOT:   tensor.pad
 //   CHECK-NOT:   tensor.extract_slice
-//       CHECK:   %[[RESULT:.*]] = tensor.generate
+//       CHECK:   %[[EMPTY:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [0, 0] [1, 1] : tensor<4x5xf32> to tensor<0x0xf32>
+//       CHECK:   %[[RESULT:.*]] = tensor.pad %[[EMPTY]] low[1, 3] high[0, 0]
 //       CHECK:     tensor.yield %[[PAD]]
 //       CHECK:   return %[[RESULT]] : tensor<1x3xf32>
 func.func @static_low_pad_only_2(%arg0 : tensor<4x5xf32>, %pad : f32)
     -> tensor<1x3xf32> {
   %0 = tensor.pad %arg0 low[3, 7] high[7, 8] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<4x5xf32> to tensor<14x20xf32>

 
 // -----
 
 // CHECK-LABEL: @dynamic_high_pad
 //  CHECK-SAME:     %[[ARG0:.*]]: tensor<?x5xf32>
 //   CHECK-NOT:   tensor.pad
 //       CHECK:   %[[C0:.*]] = arith.constant 0 : index
 //       CHECK:   tensor.dim %[[ARG0]], %[[C0]]
-//       CHECK:   %[[RESULT:.*]] = scf.if %{{.*}} -> (tensor<3x4xf32>) {
-//       CHECK:     %[[GEN:.*]] = tensor.generate
-//       CHECK:     scf.yield %[[GEN]]
-//       CHECK:   } else {
-//       CHECK:     %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>
-//       CHECK:     %[[PADTENSOR:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]
-//       CHECK:     scf.yield %[[PADTENSOR]]
-//       CHECK:   }
-//       CHECK:   return %[[RESULT]]
+//       CHECK:   %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>
+//       CHECK:   %[[PADTENSOR:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]
+//       CHECK:   return %[[PADTENSOR]]
 func.func @dynamic_high_pad(%arg0 : tensor<?x5xf32>, %h1: index, %pad : f32) -> tensor<3x4xf32> {
   %0 = tensor.pad %arg0 low[0, 0] high[%h1, 8] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<?x5xf32> to tensor<?x13xf32>
   %1 = tensor.extract_slice %0[2, 4] [3, 4] [1, 1] : tensor<?x13xf32> to tensor<3x4xf32>
   return %1 : tensor<3x4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @dynamic_extract_size
 //  CHECK-SAME:     %[[ARG0:.*]]: tensor<?x5xf32>, %[[ARG1:.*]]: index
 //   CHECK-NOT:   tensor.pad
 //       CHECK:   %[[C0:.*]] = arith.constant 0 : index
 //       CHECK:   tensor.dim %[[ARG0]], %[[C0]]
-//       CHECK:   %[[RESULT:.*]] = scf.if %{{.*}} -> (tensor<?x4xf32>) {
-//       CHECK:     %[[GEN:.*]] = tensor.generate %[[ARG1]]
-//       CHECK:     scf.yield %[[GEN]]
-//       CHECK:   } else {
-//       CHECK:     %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>
-//       CHECK:     %[[PADTENSOR:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]
-//       CHECK:     scf.yield %[[PADTENSOR]]
-//       CHECK:   }
-//       CHECK:   return %[[RESULT]]
+//       CHECK:   %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>
+//       CHECK:   %[[PADTENSOR:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]
+//       CHECK:   return %[[PADTENSOR]]
 func.func @dynamic_extract_size(%arg0 : tensor<?x5xf32>, %s1: index, %pad : f32) -> tensor<?x4xf32> {
   %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<?x5xf32> to tensor<?x13xf32>
   %1 = tensor.extract_slice %0[2, 4] [%s1, 4] [1, 1] : tensor<?x13xf32> to tensor<?x4xf32>
   return %1 : tensor<?x4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @dynamic_zero_low_padding
-//       CHECK:   scf.if
-//       CHECK:     tensor.generate
-//       CHECK:   else
-//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice
-//       CHECK:     tensor.pad %[[SLICE]] low[0, 0]
+//       CHECK:   %[[SLICE:.*]] = tensor.extract_slice
+//       CHECK:   tensor.pad %[[SLICE]] low[0, 0]
 func.func @dynamic_zero_low_padding(%arg0 : tensor<?x?xf32>, %pad : f32,
                                %o1 : index, %o2 : index,
                                %s1 : index, %s2 : index)
     -> tensor<?x?xf32> {
   %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<?x?xf32> to tensor<?x?xf32>
   %1 = tensor.extract_slice %0[%o1, %o2] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @dynamic_zero_high_padding
-//       CHECK:   scf.if
-//       CHECK:     tensor.generate
-//       CHECK:   else
-//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice
-//       CHECK:     tensor.pad %[[SLICE]] low[%{{.*}}, %{{.*}}] high[0, 0]
+//       CHECK:   %[[SLICE:.*]] = tensor.extract_slice
+//       CHECK:   tensor.pad %[[SLICE]] low[%{{.*}}, %{{.*}}] high[0, 0]
 func.func @dynamic_zero_high_padding(%arg0 : tensor<?x?xf32>, %pad : f32,
                                 %o1 : index, %o2 : index,
                                 %s1 : index, %s2 : index)
     -> tensor<?x?xf32> {
   %0 = tensor.pad %arg0 low[7, 8] high[0, 0] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad : f32
     } : tensor<?x?xf32> to tensor<?x?xf32>
   %1 = tensor.extract_slice %0[%o1, %o2] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }

mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir

 #map = affine_map<(d0, d1) -> (d0, d1)>
 //     CHECK: func @pad_generic_static
 // CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:   %[[C16:.*]] = arith.constant 16 : index
 // CHECK-DAG:   %[[C32:.*]] = arith.constant 32 : index
 // CHECK-DAG:   %[[C64:.*]] = arith.constant 64 : index
 // CHECK-DAG:   %[[C128:.*]] = arith.constant 128 : index
 //     CHECK:   scf.for %{{.*}} = %[[C0]] to %[[C64]] step %[[C16]]
-//     CHECK:     %[[CMPI1:.*]] = arith.cmpi eq
 //     CHECK:     scf.for %{{.*}} = %[[C0]] to %[[C128]] step %[[C32]]
-//     CHECK:       %[[CMPI2:.*]] = arith.cmpi eq
-//     CHECK:       %[[HASZERO:.*]] = arith.ori %[[CMPI2]], %[[CMPI1]] : i1
-//     CHECK:       scf.if %[[HASZERO]]
-//     CHECK:         tensor.generate
-//     CHECK:       else
-//     CHECK:         tensor.extract_slice
-//     CHECK:         tensor.pad
+//     CHECK:       tensor.extract_slice
+//     CHECK:       tensor.pad
 //     CHECK:       tensor.extract_slice
 //     CHECK:       tensor.extract_slice
 //     CHECK:       linalg.generic
 //     CHECK:       tensor.insert_slice
 func.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tensor<64x128xf32>) -> tensor<64x128xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c16 = arith.constant 16 : index

mlir/test/Dialect/Linalg/transform-op-tile.mlir

 // -----
 
 // CHECK-LABEL: tile_tensor_pad
 func.func @tile_tensor_pad(
   %arg0 : tensor<?x?xf32>, %cst : f32, %low: index, %high: index)
     -> tensor<20x40xf32>
 {
   // CHECK: scf.forall
-  // CHECK:   scf.if
-  // CHECK:     tensor.generate
-  // CHECK:   else
-  // CHECK:     tensor.pad {{.*}} nofold
+  // CHECK:   tensor.pad {{.*}} nofold
   %0 = tensor.pad %arg0 nofold low[%low, %low] high[%high, %high] {
         ^bb0(%arg9: index, %arg10: index):
           tensor.yield %cst : f32
   } : tensor<?x?xf32> to tensor<20x40xf32>
   return %0 : tensor<20x40xf32>
 }
 
 transform.sequence failures(propagate) {

mlir/test/Dialect/Tensor/tiling.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize -cse -split-input-file | FileCheck %s
 
 //  CHECK-DAG:  #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 8)>
 //  CHECK-DAG:  #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 7)>
 //       CHECK: func @dynamic_pad_tensor_3_4(
 //  CHECK-SAME:     %[[IN:.*]]: tensor<?x?xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[DIM_IN0:.*]] = tensor.dim %[[IN]], %[[C0]]
 //   CHECK-DAG:   %[[DIM_IN1:.*]] = tensor.dim %[[IN]], %[[C1]]
 //   CHECK-DAG:   %[[DIM0:.*]] = affine.apply #[[MAP0]]()[%[[DIM_IN0]]]
 //   CHECK-DAG:   %[[DIM1:.*]] = affine.apply #[[MAP1]]()[%[[DIM_IN1]]]
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM0]] step %[[C2]]
 //       CHECK:     scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:       %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:         tensor.generate
-//       CHECK:       else
-//       CHECK:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
-//       CHECK:         %[[PAD:.*]] = tensor.pad %[[SLICE]]
-//       CHECK:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]]
+//       CHECK:       tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 func.func @dynamic_pad_tensor_3_4(%input_tensor: tensor<?x?xf32>,
                          %pad_value: f32) -> tensor<?x?xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } : tensor<?x?xf32> to tensor<?x?xf32>

 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[DIM_IN1:.*]] = tensor.dim %[[IN]], %[[C1]]
 //   CHECK-DAG:   %[[DIM1:.*]] = affine.apply #[[MAP0]]()[%[[DIM_IN1]]]
 //   CHECK-DAG:   %[[DIM_IN0:.*]] = tensor.dim %[[IN]], %[[C0]]
 //   CHECK-DAG:   %[[DIM0:.*]] = affine.apply #[[MAP1]]()[%[[DIM_IN0]]]
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:     %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:       tensor.generate
-//       CHECK:     else
-//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
-//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[{{.*}}, {{.*}}]
-//       CHECK:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]
+//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]] low[{{.*}}, %{{.*}}] high[{{.*}}, {{.*}}]
+//       CHECK:     tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][0, {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 func.func @dynamic_pad_tensor_0_3(%input_tensor: tensor<?x?xf32>,
                          %pad_value: f32) -> tensor<?x?xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } : tensor<?x?xf32> to tensor<?x?xf32>

 //  CHECK-SAME:     %[[IN:.*]]: tensor<7x9xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[C15:.*]] = arith.constant 15 : index
 //   CHECK-DAG:   %[[C16:.*]] = arith.constant 16 : index
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]
 //       CHECK:     scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:       %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:         tensor.generate
-//       CHECK:       else
-//       CHECK:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
-//       CHECK:         %[[PAD:.*]] = tensor.pad %[[SLICE]]
-//       CHECK:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]]
+//       CHECK:       tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 func.func @static_pad_tensor_3_4(%input_tensor: tensor<7x9xf32>,
                         %pad_value: f32) -> tensor<15x16xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } : tensor<7x9xf32> to tensor<15x16xf32>

 // -----
 
 // CHECK-LABEL: func @static_pad_tensor_0_3(
 //  CHECK-SAME:     %[[IN:.*]]: tensor<7x9xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[C16:.*]] = arith.constant 16 : index
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:     %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:       tensor.generate
-//       CHECK:     else
-//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]
-//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]
-//       CHECK:     %[[CAST_SWAP_RESULT:.*]] = tensor.cast %[[SWAP_RESULT]] : tensor<?x?xf32> to tensor<15x?xf32>
-//       CHECK:     tensor.insert_slice %[[CAST_SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [15, {{.*}}] [1, 1]
+//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]
+//       CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]
+//       CHECK:     tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][0, {{.*}}] [15, {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 func.func @static_pad_tensor_0_3(%input_tensor: tensor<7x9xf32>,
                         %pad_value: f32) -> tensor<15x16xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } : tensor<7x9xf32> to tensor<15x16xf32>

 // -----
 
 // CHECK-LABEL: func @static_pad_tile_evenly_0_3(
 //  CHECK-SAME:     %[[IN:.*]]: tensor<7x9xf32>, %[[OUT:.*]]: tensor<14x15xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[C15:.*]] = arith.constant 15 : index
 //       CHECK:   %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C15]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:     %[[R2:.*]] = scf.if
-//       CHECK:       %[[GEN:.*]] = tensor.generate
-//       CHECK:       %[[cast_0:.*]] = tensor.cast %[[GEN]] : tensor<14x3xf32> to tensor<?x?xf32>
-//       CHECK:       scf.yield %[[cast_0]] : tensor<?x?xf32>
-//       CHECK:     else
-//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %arg0[0, %{{.*}}] [7, %{{.*}}] [1, 1] : tensor<7x9xf32> to tensor<7x?xf32>
-//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]] low[0, 0] high[7, %{{.*}}]
-//       CHECK:       %[[cast_1:.*]] = tensor.cast %[[PAD]] : tensor<14x?xf32> to tensor<?x?xf32>
-//       CHECK:       scf.yield %[[cast_1]] : tensor<?x?xf32>
-//       CHECK:     %[[cast:.*]] = tensor.cast %[[R2]] : tensor<?x?xf32> to tensor<14x3xf32>
-//       CHECK:     %[[R3:.*]] = tensor.insert_slice %[[cast]] into %[[INNER_OUT]][0, %[[IV]]] [14, 3] [1, 1] : tensor<14x3xf32> into tensor<14x15xf32>
+//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice %arg0[0, %{{.*}}] [7, %{{.*}}] [1, 1] : tensor<7x9xf32> to tensor<7x?xf32>
+//       CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]] low[0, 0] high[7, %{{.*}}]
+//       CHECK:     %[[R3:.*]] = tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][0, %[[IV]]] [14, 3] [1, 1] : tensor<14x3xf32> into tensor<14x15xf32>
 //       CHECK:     scf.yield %[[R3]] : tensor<14x15xf32>
 //       CHECK:   return %[[RESULT]] : tensor<14x15xf32>
 
 func.func @static_pad_tile_evenly_0_3(%input_tensor: tensor<7x9xf32>,
                              %output_tensor: tensor<14x15xf32>,
                              %pad_value: f32) -> tensor<14x15xf32> {
   %0 = tensor.pad %input_tensor low[0, 0] high[7, 6] {
     ^bb0(%arg1: index, %arg2: index):

mlir/test/Interfaces/TilingInterface/tile-pad-using-interface.mlir

 //   CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
 //   CHECK-DAG:   %[[C3:.+]] = arith.constant 3 : index
 //       CHECK:   %[[DIM_IN0:.+]] = tensor.dim %[[IN]], %[[C0]]
 //       CHECK:   %[[DIM0:.+]] = affine.apply #[[MAP0]]()[%[[DIM_IN0]]]
 //       CHECK:   %[[DIM_IN1:.+]] = tensor.dim %[[IN]], %[[C1]]
 //       CHECK:   %[[DIM1:.+]] = affine.apply #[[MAP1]]()[%[[DIM_IN1]]]
 //       CHECK:   %[[RESULT:[a-zA-Z0-9]+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] = %[[C0]] to %[[DIM0]] step %[[C2]]
 //       CHECK:     scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:       %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:         tensor.generate
-//       CHECK:       else
-//       CHECK:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
-//       CHECK:         %[[PAD:.*]] = tensor.pad %[[SLICE]]
-//       CHECK:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]]
+//       CHECK:       tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
 func.func @dynamic_2d_pad_tensor_inner_tiling(%input_tensor: tensor<?x?xf32>,
                          %pad_value: f32) -> tensor<?x?xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } {__internal_linalg_transform__ = "pad_inner_tiling"}: tensor<?x?xf32> to tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
 //   CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 8)>
 //   CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 7)>
 //       CHECK: func @dynamic_2d_pad_tensor_inner_tiling(
 //  CHECK-SAME:     %[[IN:.*]]: tensor<?x?xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //       CHECK:   %[[DIM_IN0:.*]] = tensor.dim %[[IN]], %[[C0]]
 //       CHECK:   %[[DIM0:.*]] = affine.apply #[[MAP0]]()[%[[DIM_IN0]]]
 //       CHECK:   %[[DIM_IN1:.*]] = tensor.dim %[[IN]], %[[C1]]
 //       CHECK:   %[[DIM1:.*]] = affine.apply #[[MAP1]]()[%[[DIM_IN1]]]
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:     %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:       tensor.generate
-//       CHECK:     else
-//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
-//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[{{.*}}, {{.*}}]
-//       CHECK:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][%[[C0]], {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]
+//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]] low[{{.*}}, %{{.*}}] high[{{.*}}, {{.*}}]
+//       CHECK:     tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][%[[C0]], {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
 func.func @static_pad_tensor(%input_tensor: tensor<7x9xf32>,
                         %pad_value: f32) -> tensor<15x16xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } {__internal_linalg_transform__ = "pad_2dtiling"} : tensor<7x9xf32> to tensor<15x16xf32>
   return %0 : tensor<15x16xf32>
 }
 // CHECK-LABEL: func @static_pad_tensor(
 //  CHECK-SAME:     %[[IN:.*]]: tensor<7x9xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[C15:.*]] = arith.constant 15 : index
 //   CHECK-DAG:   %[[C16:.*]] = arith.constant 16 : index
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]
 //       CHECK:     scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:       %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:         tensor.generate
-//       CHECK:       else
-//       CHECK:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
-//       CHECK:         %[[PAD:.*]] = tensor.pad %[[SLICE]]
-//       CHECK:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]]
+//       CHECK:       tensor.insert_slice %[[PAD]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
 func.func @static_pad_tensor_inner_tiling(%input_tensor: tensor<7x9xf32>,
                         %pad_value: f32) -> tensor<15x16xf32> {
   %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %pad_value : f32
     } {__internal_linalg_transform__ = "pad_inner_tiling"} : tensor<7x9xf32> to tensor<15x16xf32>
   return %0 : tensor<15x16xf32>
 }
 // CHECK-LABEL: func @static_pad_tensor_inner_tiling(
 //  CHECK-SAME:     %[[IN:.*]]: tensor<7x9xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[C15:.*]] = arith.constant 15 : index
 //   CHECK-DAG:   %[[C16:.*]] = arith.constant 16 : index
 //       CHECK:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
-//       CHECK:     %[[SWAP_RESULT:.*]] = scf.if
-//       CHECK:       tensor.generate
-//       CHECK:     else
-//       CHECK:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]
-//       CHECK:       %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]
-//       CHECK:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][%[[C0]], {{.*}}] [%[[C15]], {{.*}}] [1, 1]
+//       CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]
+//       CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]
+//       CHECK:     %[[CAST:.*]] = tensor.cast %[[PAD]]
+//       CHECK:     tensor.insert_slice %[[CAST]] into %[[INNER_OUT]][%[[C0]], {{.*}}] [%[[C15]], {{.*}}] [1, 1]
 //       CHECK:   return %[[RESULT]]
 
 /// Rest of the tests only check that they dont fail.
 
 // -----
 
 func.func @dynamic_2d_pad_tensor_outer_tiling(%input_tensor: tensor<?x?xf32>,
                          %pad_value: f32) -> tensor<?x?xf32> {