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

[Sayan Saha]

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

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.

>From 3299e2f7ad38c320fa97a4a0f6665fa13ea8dd2c Mon Sep 17 00:00:00 2001
From: Sayan Saha <sayans at mathworks.com>
Date: Tue, 3 Oct 2023 18:43:12 -0400
Subject: [PATCH] [BugFix] : Fix padding calculation for tosa.transpose_conv2d
 decomposition.

---
 .../Tosa/Transforms/TosaDecomposeTransposeConv.cpp        | 8 ++++----
 mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

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>,