mlir PadOp Tiling Interface implement does not include generateResultTileValue

[hesse-x]

We found that the implementation of the PadOp TilingInterface:PadOpTiling does not include the generateResultTileValue interface. This has led to unexpected results when we attempt to fuse the pad operation into a containing loop using transform.structured.fuse_into_containing_op.I would like to understand if this is intended behavior or a defect. Additionally, I would like to know the correct approach to fuse the pad operation into a loop.

Here are our test results on llvm-15.x

#map0 = affine_map<(d0, d1) -> (d0, d1)>
func.func @pad(%arg0: tensor<58x1xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> attributes {pad} {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c16 = arith.constant 16 : index
  %c32 = arith.constant 32 : index
  %zero = arith.constant 0.0 : f32
 
  %pad = tensor.pad %arg0 low[4, 60] high[2, 67] {
  ^bb0(%arg4: index, %arg5: index):
    tensor.yield %zero : f32
  } : tensor<58x1xf32> to tensor<64x128xf32>

  %1 = linalg.init_tensor [64, 128] : tensor<64x128xf32>
  %2 = scf.foreach_thread (%arg2, %arg3) in (%c1, %c1) -> (tensor<64x128xf32>) {
    %3 = tensor.extract_slice %arg1[%arg2, %arg3] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
    %4 = tensor.extract_slice %pad[%arg2, %arg3] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
    %5 = linalg.init_tensor [8, 128] : tensor<8x128xf32>
    %6 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
        ins(%3, %4 : tensor<8x128xf32>, tensor<8x128xf32>)
        outs(%5 : tensor<8x128xf32>) {
      ^bb0(%arg4: f32, %arg5: f32, %arg6: f32):
        %7 = arith.addf %arg4, %arg5 : f32
        linalg.yield %7 : f32
    } -> tensor<8x128xf32>
    scf.foreach_thread.perform_concurrently {
      tensor.parallel_insert_slice %6 into %1[%arg2, %arg3] [8, 128] [1, 1] : tensor<8x128xf32> into tensor<64x128xf32>
    }  
  } {thread_dim_mapping = [2, 4]} 
  return %2 : tensor<64x128xf32>
}

transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation):
  transform.genesis.canonicalized_sequence %arg0 failures(propagate) {
  ^bb0(%arg1: !pdl.operation):
    %device_func = transform.genesis.match ops{["func.func"]} attributes {pad} in %arg1
    %foreach_thread_op = transform.genesis.match ops{["scf.foreach_thread"]} in %device_func

    // fuse and tile
    %expand_shape = transform.genesis.match ops{["tensor.pad"]} in %device_func
    transform.structured.fuse_into_containing_op %expand_shape into %foreach_thread_op
  }  
}

#map = affine_map<(d0, d1) -> (d0, d1)>
module {
  func.func @pad(%arg0: tensor<58x1xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> attributes {pad} {
    %c1 = arith.constant 1 : index
    %cst = arith.constant 0.000000e+00 : f32
    %0 = linalg.init_tensor [64, 128] : tensor<64x128xf32>
    %1 = scf.foreach_thread (%arg2, %arg3) in (%c1, %c1) -> (tensor<64x128xf32>) {
      %2 = tensor.extract_slice %arg1[%arg2, %arg3] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %3 = tensor.pad %arg0 low[4, 60] high[2, 67] {
      ^bb0(%arg4: index, %arg5: index):
        tensor.yield %cst : f32
      } : tensor<58x1xf32> to tensor<64x128xf32>
      %4 = tensor.extract_slice %3[%arg2, %arg3] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %5 = linalg.init_tensor [8, 128] : tensor<8x128xf32>
      %6 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%2, %4 : tensor<8x128xf32>, tensor<8x128xf32>) outs(%5 : tensor<8x128xf32>) {
      ^bb0(%arg4: f32, %arg5: f32, %arg6: f32):
        %7 = arith.addf %arg4, %arg5 : f32
        linalg.yield %7 : f32
      } -> tensor<8x128xf32>
      scf.foreach_thread.perform_concurrently {
        tensor.parallel_insert_slice %6 into %0[%arg2, %arg3] [8, 128] [1, 1] : tensor<8x128xf32> into tensor<64x128xf32>
      }
    } {thread_dim_mapping = [2, 4]}
    return %1 : tensor<64x128xf32>
  }
  transform.with_pdl_patterns {
  ^bb0(%arg0: !pdl.operation):
    transform.genesis.canonicalized_sequence %arg0 failures(propagate) {
    ^bb0(%arg1: !pdl.operation):
      %0 = transform.genesis.match ops{["func.func"]} attributes {pad} in %arg1
      %1 = transform.genesis.match ops{["scf.foreach_thread"]} in %0
      %2 = transform.genesis.match ops{["tensor.pad"]} in %0
      %3 = transform.structured.fuse_into_containing_op %2 into %1
    }
  }
}

[llvmbot]

@llvm/issue-subscribers-mlir

[EugeneZelenko]

Could you please try 16 or main branch?

[hesse-x]

I test on main branch

#map0 = affine_map<(d0, d1) -> (d0, d1)>
func.func @expand0(%arg0: tensor<58x1xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> attributes {pad} {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c8 = arith.constant 8 : index
  %c64 = arith.constant 64 : index
  %c128 = arith.constant 128 : index
  %zero = arith.constant 0.0 : f32
 
  %pad = tensor.pad %arg0 low[4, 60] high[2, 67] {
  ^bb0(%arg4: index, %arg5: index):
    tensor.yield %zero : f32
  } : tensor<58x1xf32> to tensor<64x128xf32>

  %1 = tensor.empty() : tensor<64x128xf32>
  %for0 = scf.for %arg2 = %c0 to %c64 step %c8 iter_args(%arg4 = %1) -> tensor<64x128xf32> {
    %for1 = scf.for %arg3 = %c0 to %c128 step %c128 iter_args(%arg5 = %arg4) -> tensor<64x128xf32> {
      %3 = tensor.extract_slice %arg5[%arg2, %arg3] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %4 = tensor.extract_slice %pad[%arg2, %arg3] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %5 = tensor.empty() : tensor<8x128xf32>
      %6 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
          ins(%3, %4 : tensor<8x128xf32>, tensor<8x128xf32>)
          outs(%5 : tensor<8x128xf32>) {
        ^bb0(%arg7: f32, %arg6: f32, %arg8: f32):
          %7 = arith.addf %arg7, %arg6 : f32
          linalg.yield %7 : f32
      } -> tensor<8x128xf32>
      %insert = tensor.insert_slice %6 into %1[%arg2, %arg3] [8, 128] [1, 1] : tensor<8x128xf32> into tensor<64x128xf32>
      scf.yield %insert : tensor<64x128xf32>
    }
    scf.yield %for1 : tensor<64x128xf32>
  }
  return %for0 : tensor<64x128xf32>
}

transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
  %device_func = transform.structured.match ops{["func.func"]} attributes {pad} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
  %foreach_thread_op = transform.structured.match ops{["scf.for"]} in %device_func : (!transform.op<"func.func">) -> !transform.op<"scf.for">
  %for0, %for1 = transform.split_handle %foreach_thread_op : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)

  // fuse and tile
  %expand_shape = transform.structured.match ops{["tensor.pad"]} in %device_func : (!transform.op<"func.func">) -> !transform.op<"tensor.pad">
  transform.structured.fuse_into_containing_op %expand_shape into %for1 : (!transform.op<"tensor.pad">, !transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
}

Its result is same as 15.0

#map = affine_map<(d0, d1) -> (d0, d1)>
module {
  func.func @expand0(%arg0: tensor<58x1xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> attributes {pad} {
    %c0 = arith.constant 0 : index
    %c8 = arith.constant 8 : index
    %c64 = arith.constant 64 : index
    %cst = arith.constant 0.000000e+00 : f32
    %0 = tensor.empty() : tensor<64x128xf32>
    %1 = scf.for %arg2 = %c0 to %c64 step %c8 iter_args(%arg3 = %0) -> (tensor<64x128xf32>) {
      %extracted_slice = tensor.extract_slice %arg3[%arg2, 0] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %padded = tensor.pad %arg0 low[4, 60] high[2, 67] {
      ^bb0(%arg4: index, %arg5: index):
        tensor.yield %cst : f32
      } : tensor<58x1xf32> to tensor<64x128xf32>
      %extracted_slice_0 = tensor.extract_slice %padded[%arg2, 0] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %2 = tensor.empty() : tensor<8x128xf32>
      %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice, %extracted_slice_0 : tensor<8x128xf32>, tensor<8x128xf32>) outs(%2 : tensor<8x128xf32>) {
      ^bb0(%in: f32, %in_1: f32, %out: f32):
        %4 = arith.addf %in, %in_1 : f32
        linalg.yield %4 : f32
      } -> tensor<8x128xf32>
      %inserted_slice = tensor.insert_slice %3 into %0[%arg2, 0] [8, 128] [1, 1] : tensor<8x128xf32> into tensor<64x128xf32>
      scf.yield %inserted_slice : tensor<64x128xf32>
    }
    return %1 : tensor<64x128xf32>
  }
  transform.sequence  failures(propagate) {
  ^bb0(%arg0: !transform.any_op):
    %0 = transform.structured.match ops{["func.func"]} attributes {pad} in %arg0 : (!transform.any_op) -> !transform.op<"func.func">
    %1 = transform.structured.match ops{["scf.for"]} in %0 : (!transform.op<"func.func">) -> !transform.op<"scf.for">
    %2:2 = split_handle %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
    %3 = transform.structured.match ops{["tensor.pad"]} in %0 : (!transform.op<"func.func">) -> !transform.op<"tensor.pad">
    %fused_op, %new_containing_op = transform.structured.fuse_into_containing_op %3 into %2#1 : (!transform.op<"tensor.pad">, !transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
  }
}

After applying the above patch, I was able to obtain the expected results as follows. Therefore, I would like to confirm whether the absence of the implementation of generateResultTileValue was intentional to discourage the fusion of pad in this manner, or if it is a bug.

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
index 56cfdfcf0b..9676ea6bd8 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
@@ -67,6 +67,18 @@ struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
     resultSizes.assign(sizes.begin(), sizes.end());
     return success();
   }
+ 
+  FailureOr<TilingResult> generateResultTileValue(Operation *op, OpBuilder &b,
+                                           unsigned resultNumber,
+                                           ArrayRef<OpFoldResult> offsets,
+                                           ArrayRef<OpFoldResult> sizes) const {
+    FailureOr<TilingResult> tilingResult =
+        getTiledImplementation(op, b, offsets, sizes);
+    if (failed(tilingResult))
+      return failure();
+    return tilingResult.value();
+  }
+
 };

#map = affine_map<(d0) -> (-d0 + 4, 0)>
#map1 = affine_map<(d0) -> (0, d0 - 4)>
#map2 = affine_map<(d0) -> (58, d0)>
#map3 = affine_map<(d0) -> (0, d0 + 4)>
#map4 = affine_map<(d0, d1) -> (d0 - d1)>
#map5 = affine_map<(d0, d1, d2) -> (-d0 - d1 + d2 + 8)>
#map6 = affine_map<(d0, d1) -> (d0, d1)>
module {
  func.func @expand0(%arg0: tensor<58x1xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> attributes {pad} {
    %c67 = arith.constant 67 : index
    %c60 = arith.constant 60 : index
    %c0 = arith.constant 0 : index
    %c8 = arith.constant 8 : index
    %c64 = arith.constant 64 : index
    %cst = arith.constant 0.000000e+00 : f32
    %0 = tensor.empty() : tensor<64x128xf32>
    %1 = scf.for %arg2 = %c0 to %c64 step %c8 iter_args(%arg3 = %0) -> (tensor<64x128xf32>) {
      %extracted_slice = tensor.extract_slice %arg3[%arg2, 0] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
      %2 = affine.max #map(%arg2)
      %3 = affine.max #map1(%arg2)
      %4 = affine.min #map2(%3)
      %5 = affine.max #map3(%arg2)
      %6 = affine.min #map2(%5)
      %7 = affine.apply #map4(%6, %4)
      %8 = arith.cmpi eq, %7, %c0 : index
      %9 = affine.apply #map5(%2, %6, %4)
      %10 = scf.if %8 -> (tensor<8x128xf32>) {
        %generated = tensor.generate  {
        ^bb0(%arg4: index, %arg5: index):
          tensor.yield %cst : f32
        } : tensor<8x128xf32>
        scf.yield %generated : tensor<8x128xf32>
      } else {
        %extracted_slice_0 = tensor.extract_slice %arg0[%4, 0] [%7, 1] [1, 1] : tensor<58x1xf32> to tensor<?x1xf32>
        %padded = tensor.pad %extracted_slice_0 low[%2, %c60] high[%9, %c67] {
        ^bb0(%arg4: index, %arg5: index):
          tensor.yield %cst : f32
        } : tensor<?x1xf32> to tensor<8x128xf32>
        scf.yield %padded : tensor<8x128xf32>
      }
      %11 = tensor.empty() : tensor<8x128xf32>
      %12 = linalg.generic {indexing_maps = [#map6, #map6, #map6], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice, %10 : tensor<8x128xf32>, tensor<8x128xf32>) outs(%11 : tensor<8x128xf32>) {
      ^bb0(%in: f32, %in_0: f32, %out: f32):
        %13 = arith.addf %in, %in_0 : f32
        linalg.yield %13 : f32
      } -> tensor<8x128xf32>
      %inserted_slice = tensor.insert_slice %12 into %0[%arg2, 0] [8, 128] [1, 1] : tensor<8x128xf32> into tensor<64x128xf32>
      scf.yield %inserted_slice : tensor<64x128xf32>
    }
    return %1 : tensor<64x128xf32>
  }
  transform.sequence  failures(propagate) {
  ^bb0(%arg0: !transform.any_op):
    %0 = transform.structured.match ops{["func.func"]} attributes {pad} in %arg0 : (!transform.any_op) -> !transform.op<"func.func">
    %1 = transform.structured.match ops{["scf.for"]} in %0 : (!transform.op<"func.func">) -> !transform.op<"scf.for">
    %2:2 = split_handle %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
    %3 = transform.structured.match ops{["tensor.pad"]} in %0 : (!transform.op<"func.func">) -> !transform.op<"tensor.pad">
    %fused_op, %new_containing_op = transform.structured.fuse_into_containing_op %3 into %2#1 : (!transform.op<"tensor.pad">, !transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
  }
}