[mlir][tosa] Fix dynamic shape inference in conv2d

The comment in the code correctly states the equation for the shape inference as follows:

```
H = ((IH+pad_top+pad_bottom-(dilation_y*(KH-1)+1))/stride_y)+1
```

However the final operation is generated as `-` instead of `+`. I believe `+`
is indeed correct. For example if we have an image with dimension 6 and kernel
of dimension 3 (assuming padding is 0 and stride and dilation are both 1) we
are expecting 4 elements in the output (computed for image elements `(0, 1, 2)
x kernel`, `(1, 2, 3) x kernel`, `(2, 3, 4) x kernel` and `(3, 4, 5) x kernel`.
However currently only 2 elements are produced in the output.

Reviewed By: NatashaKnk

Differential Revision: https://reviews.llvm.org/D133208

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 42bca1e..9ed3262 100644 (file)
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -93,7 +93,7 @@ getConvOutputDim(Location loc, Value initDim, Attribute padBeforeAttr,
   Value subtract = builder.create<arith::SubIOp>(paddedAfter, addOne);
   Value stride = reifyConstantDim(strideAttr, builder);
   Value divide = builder.create<arith::DivUIOp>(subtract, stride);
-  return builder.create<arith::SubIOp>(divide, one);
+  return builder.create<arith::AddIOp>(divide, one);
 }
 
 // Creates a vector of the dynamic output dims for Conv2D and Depthwise_Conv2D

diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index db81b20..d956e82 100644 (file)
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -405,7 +405,7 @@ func.func @conv2d_dyn_w_h(%input: tensor<1x?x?x27xf32>, %weights: tensor<28x3x3x
   // CHECK: %[[SUBTRACTED:.+]] = arith.subi %[[ADD_PAD_1]], %[[ADD_ONE]] : index
   // CHECK: %[[STRIDE_H:.+]] = arith.constant 1 : index
   // CHECK: %[[DIVIDED:.+]] = arith.divui %[[SUBTRACTED]], %[[STRIDE_H]] : index
-  // CHECK: %[[H_OUT:.+]] = arith.subi %[[DIVIDED]], %[[ONE]] : index
+  // CHECK: %[[H_OUT:.+]] = arith.addi %[[DIVIDED]], %[[ONE]] : index
 
   // Computing output width
   // CHECK: %[[C2:.+]] = arith.constant 2
@@ -424,7 +424,7 @@ func.func @conv2d_dyn_w_h(%input: tensor<1x?x?x27xf32>, %weights: tensor<28x3x3x
   // CHECK: %[[SUBTRACTED_0:.+]] = arith.subi %[[ADD_PAD_3]], %[[ADD_ONE_0]] : index
   // CHECK: %[[STRIDE_W:.+]] = arith.constant 1 : index
   // CHECK: %[[DIVIDED_0:.+]] = arith.divui %[[SUBTRACTED_0]], %[[STRIDE_W]] : index
-  // CHECK: %[[W_OUT:.+]] = arith.subi %[[DIVIDED_0]], %[[ONE_0]] : index
+  // CHECK: %[[W_OUT:.+]] = arith.addi %[[DIVIDED_0]], %[[ONE_0]] : index
 
   // Running convolution
   // CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]>