[MLIR][PadTilingInterface] Incorrect input padding for convolutions

[yzhang93]

Description:

When padding convolution-style ops using affine indexing maps of the form (dX + dY), the current logic in computePaddedShape incorrectly computes the required input padding. Treating the affine map result as the padded shape causes over-padding of the input tensor.

Example:

For the following conv_2d_nhwc_fhwc op, if I only pad f and c dimensions with multiples of {0, 0, 0, 32, 0, 0, 32}, it also pads the convolved dimensions and generates:

%padded_2 = tensor.pad %padded low[0, 0, 0, 0] high[0, 1, 1, 1] {
  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
    tensor.yield %cst_1 : bf16
  } : tensor<16x26x19x287xbf16> to tensor<16x27x20x288xbf16>
  %padded_4 = tensor.pad %4 low[0, 0, 0, 0] high[1, 0, 0, 1] {
  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
    tensor.yield %cst_3 : bf16
  } : tensor<287x3x3x287xbf16> to tensor<288x3x3x288xbf16>
  %padded_6 = tensor.pad %6 low[0, 0, 0, 0] high[0, 0, 0, 1] {
  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
    tensor.yield %cst_5 : f32
  } : tensor<16x24x17x287xf32> to tensor<16x24x17x288xf32>
  %7 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%padded_2, %padded_4 : tensor<16x27x20x288xbf16>, tensor<288x3x3x288xbf16>) outs(%padded_6 : tensor<16x24x17x288xf32>) {
  ^bb0(%in: bf16, %in_7: bf16, %out: f32):
    %10 = arith.extf %in : bf16 to f32
    %11 = arith.extf %in_7 : bf16 to f32
    %12 = arith.mulf %10, %11 : f32
    %13 = arith.addf %out, %12 : f32
    linalg.yield %13 : f32
  } -> tensor<16x24x17x288xf32>

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Vivian Zhang (yzhang93)

**Description:**

When padding convolution-style ops using affine indexing maps of the form (dX + dY), the current logic in computePaddedShape incorrectly computes the required input padding. Specifically, directly using the affine map result (dX_size + dY_size) without considering the actual extent of kernel coverage leads to incorrect padding for convolved dimensions.

Example:

For the following conv_2d_nhwc_fhwc op, if I only pad f and c dimensions with multiples of {0, 0, 0, 32, 0, 0, 32}, it also pads the convolved dimensions and generates:

%padded_2 = tensor.pad %padded low[0, 0, 0, 0] high[0, 1, 1, 1] {
  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
    tensor.yield %cst_1 : bf16
  } : tensor<16x26x19x287xbf16> to tensor<16x27x20x288xbf16>
  %padded_4 = tensor.pad %4 low[0, 0, 0, 0] high[1, 0, 0, 1] {
  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
    tensor.yield %cst_3 : bf16
  } : tensor<287x3x3x287xbf16> to tensor<288x3x3x288xbf16>
  %padded_6 = tensor.pad %6 low[0, 0, 0, 0] high[0, 0, 0, 1] {
  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
    tensor.yield %cst_5 : f32
  } : tensor<16x24x17x287xf32> to tensor<16x24x17x288xf32>
  %7 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%padded_2, %padded_4 : tensor<16x27x20x288xbf16>, tensor<288x3x3x288xbf16>) outs(%padded_6 : tensor<16x24x17x288xf32>) {
  ^bb0(%in: bf16, %in_7: bf16, %out: f32):
    %10 = arith.extf %in : bf16 to f32
    %11 = arith.extf %in_7 : bf16 to f32
    %12 = arith.mulf %10, %11 : f32
    %13 = arith.addf %out, %12 : f32
    linalg.yield %13 : f32
  } -> tensor<16x24x17x288xf32>

[MaheshRavishankar]

@yzhang93 could you give me a way to repro this?

[yzhang93]

@yzhang93 could you give me a way to repro this?

You can try this simple example:

// RUN: mlir-opt --transform-interpreter -split-input-file --verify-diagnostics %s | FileCheck %s

func.func @conv_2d_nhwc_fhwc(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<16x3x3x4xf32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {
    %0 = linalg.conv_2d_nhwc_fhwc
      {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
       ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<16x3x3x4xf32>)
      outs(%arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32>
    return %0 : tensor<1x14x14x16xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [0, 0, 0, 0, 0, 0, 16] pad_to_multiple_of {
      padding_values = [0.0 : f32, 0.0 : f32, 0.0 : f32, 0.0 : f32, 0.0 : f32, 0.0 : f32, 0.0 : f32]
    } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    transform.yield
  }
}