[--linalg-bufferize] pass couldn't lower tensor.pad into linalg.fill

MLIR
1

Hello all,

I am trying bufferization passes in my tensor->memref lowering experiment. With recent bufferization update, -linalg-bufferize pass couldn’t lower tensor.pad ops to linalg.fill.

// conv2d.mlir
module attributes {torch.debug_module_name = "Conv2D"} {
  func @forward(%arg0: tensor<1x64x64x64xf32>) -> tensor<1x128x32x32xf32> {
    %c2_i64 = arith.constant 2 : i64
    %c1_i64 = arith.constant 1 : i64
    %cst = arith.constant 0.000000e+00 : f32
    %cst_0 = arith.constant dense<1.000000e+00> : tensor<128x64x3x3xf32>
    %cst_1 = arith.constant dense<1.000000e+00> : tensor<128xf32>
    %0 = arith.cmpi eq, %c1_i64, %c1_i64 : i64
    assert %0, "expect groups to be 1"
    %1 = tensor.pad %arg0 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst : f32
    } : tensor<1x64x64x64xf32> to tensor<1x64x66x66xf32>
    %2 = linalg.init_tensor [1, 128, 32, 32] : tensor<1x128x32x32xf32>
    %3 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_1 : tensor<128xf32>) outs(%2 : tensor<1x128x32x32xf32>) {
    ^bb0(%arg1: f32, %arg2: f32):
      linalg.yield %arg1 : f32
    } -> tensor<1x128x32x32xf32>
    %4 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} ins(%1, %cst_0 : tensor<1x64x66x66xf32>, tensor<128x64x3x3xf32>) outs(%3 : tensor<1x128x32x32xf32>) -> tensor<1x128x32x32xf32>
    return %4 : tensor<1x128x32x32xf32>
  }
}

$ mlir-opt --linalg-bufferize --arith-bufferize --func-bufferize --tensor-bufferize --finalizing-bufferize -o conv2d.linalg-memrefs.mlir conv2d.mlir

Generated following error:

conv2d.linalg-memrefs.mlir:13:10: error: failed to legalize operation 'tensor.pad'
    %5 = tensor.pad %0 low[0, 0, 1, 1] high[0, 0, 1, 1] {
         ^
e.conv2d.linalg-memrefs.mlir:13:10: note: see current operation: %6 = "tensor.pad"(%0) ({
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
  "tensor.yield"(%1) : (f32) -> ()
}) {operand_segment_sizes = dense<[1, 0, 0]> : vector<3xi32>, static_high = [0, 0, 1, 1], static_low = [0, 0, 1, 1]} : (tensor<1x64x64x64xf32>) -> tensor<1x64x66x66xf32>
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::PreservedAnalyses::AllAnalysesType)
Makefile:16: recipe for target 'conv2d.linalg-memrefs.mlir' failed
make: *** [conv2d.linalg-memrefs.mlir] Error 1

llvm-project commitHash: e1318078a4e160eb723bcbcfcdcc9a1b618f7067

Is this issue fixed in updated commit or is there any new bufferization pass I need to apply. Thanks!!

2

Hey everyone!

We looked at this again together, and it seems like there isn’t an interface implementation for the BufferizableOpInterface for the PadOp in lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp, which explains why it remains there.

@akafik pointed out that this padding should be replaced by a linalg::FillOp anyways. Looking through the code I found a PadOpTransformationPattern in include/mlir/Dialect/Linalg/Transforms/Transforms.h that seems to do just that. However, it seems that pattern is not available to any pass pipeline except for in the testing code.

We suspect that if this is performed before bufferization, our code would work. We’ll try that for now.

Are there any plans to expose this pass, or is the absence of a bufferization pattern for the PadOp just an oversight?

3

Hello all,

Adding following -test-linalg-transform-patterns=test-transform-pad-tensor command resolve this issue for now.

$ mlir-opt -test-linalg-transform-patterns=test-transform-pad-tensor --linalg-bufferize --arith-bufferize --func-bufferize --tensor-bufferize --finalizing-bufferize -o conv2d.linalg-memrefs.mlir conv2d.mlir

Thanks!!

4

@tpopp added a tensor-to-linalg pass for this purpose I believe.