[Mlir-commits] [mlir] [BugFix] : Fix padding calculation for tosa.transpose_conv2d decomposition. (PR #68167)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

<details>
<summary>Changes</summary>

This change fixes the padding computation for tosa.transpose_conv2d decomposition to tosa.conv2d for unit stride case. Prior to this change the following MLIR code:

`%0 = tosa.transpose_conv2d %arg0, %cst_0, %cst {out_pad = array<i64: 1, 1, 1, 1>, out_shape = array<i64: 1, 13, 10, 2>, stride = array<i64: 1, 1>} : (tensor<1x13x10x3xf32>, tensor<2x3x3x3xf32>, tensor<2xf32>) -> tensor<1x13x10x2xf32>`

produces the following decomposition:

```
%0 = tosa.reverse %cst_0 {axis = 1 : i32} : (tensor<2x3x3x3xf32>) -> tensor<2x3x3x3xf32>
%1 = tosa.reverse %0 {axis = 2 : i32} : (tensor<2x3x3x3xf32>) -> tensor<2x3x3x3xf32>
%2 = tosa.conv2d %arg0, %1, %cst {dilation = array<i64: 1, 1>, pad = array<i64: 3, 3, 3, 3>, stride = array<i64: 1, 1>} : (tensor<1x13x10x3xf32>, tensor<2x3x3x3xf32>, tensor<2xf32>) -> tensor<1x13x10x2xf32>

```
Note that the for the padding of (3, 3, 3, 3) the output shape of tosa.conv2d for the given input will not be 1x13x10x2. The correct padding value to obtain the shape of 1x13x10x2 as the output shape is (1, 1, 1, 1) instead. This change fixes the padding computation logic to produce the padding as (1, 1, 1, 1) for the decomposition.

---
Full diff: https://github.com/llvm/llvm-project/pull/68167.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp (+4-4) 
- (modified) mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir (+1-1) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp
index 8d937217d706554..cde2f3d6f94b31e 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp
@@ -106,10 +106,10 @@ class TransposeConvNonStridedConverter
     int64_t kernelWidth = weightTy.getDimSize(2);
 
     llvm::SmallVector<int64_t> convPad(4, 0);
-    convPad[0] = kernelHeight - 1 + pad[0];
-    convPad[1] = kernelHeight - 1 + pad[1];
-    convPad[2] = kernelWidth - 1 + pad[2];
-    convPad[3] = kernelWidth - 1 + pad[3];
+    convPad[0] = kernelHeight - 1 - pad[0];
+    convPad[1] = kernelHeight - 1 - pad[1];
+    convPad[2] = kernelWidth - 1 - pad[2];
+    convPad[3] = kernelWidth - 1 - pad[3];
 
     auto reverse1 = rewriter.create<tosa::ReverseOp>(
         loc, weightTy, weight, /* axis = */ rewriter.getI32IntegerAttr(1));
diff --git a/mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir b/mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir
index 5f9fe68c5789254..9fef5a0e09d4be3 100644
--- a/mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir
@@ -29,7 +29,7 @@ func.func @transpose_conv2d_quantized_padded(%arg0: tensor<2x16x14x3xi8>, %arg1:
   // CHECK-DAG: %[[REV0:.+]] = tosa.reverse %0 {axis = 2 : i32}
   // CHECK-DAG: %[[REV1:.+]] = tosa.reverse %arg1 {axis = 1 : i32}
   // CHECK: tosa.conv2d %arg0, %1, %arg2
-  // CHECK-SAME: dilation = array<i64: 1, 1>, pad = array<i64: 3, 4, 8, 9>,
+  // CHECK-SAME: dilation = array<i64: 1, 1>, pad = array<i64: 1, 0, 2, 1>,
   // CHECK-SAME: quantization_info = #tosa.conv_quant<input_zp = -22, weight_zp = 42>, stride = array<i64: 1, 1>}
   %0 = tosa.transpose_conv2d %arg0, %arg1, %arg2 {
     out_pad = array<i64: 1, 2, 3, 4>,

``````````

</details>


https://github.com/llvm/llvm-project/pull/68167