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

[Anastasia]

Summary

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.

[NatashaKnk]

Good catch, thanks!

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

getConvOutputDim(Location loc, Value initDim, Attribute padBeforeAttr,
   Value subOne = builder.create<arith::SubIOp>(kernelDim, one);
   Value dilation = reifyConstantDim(dilationAttr, builder);
   Value dilated = builder.create<arith::MulIOp>(dilation, subOne);
   Value addOne = builder.create<arith::AddIOp>(dilated, one);
 
   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
 static SmallVector<Value> inferDynamicDimsForConv(
     Location loc, Value input, Value weight, ShapedType resultTy,
     ArrayAttr padAttr, ArrayAttr strideAttr, ArrayAttr dilationAttr,
     int64_t weightHDim, int64_t weightWDim, OpBuilder &rewriter) {
   ShapedType inputTy = input.getType().cast<ShapedType>();

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir

func.func @conv2d_dyn_w_h(%input: tensor<1x?x?x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
   // CHECK: %[[ADD_PAD_1:.+]] = arith.addi %[[ADD_PAD_0]], %[[PAD_1]] : index
   // CHECK: %[[SUB_ONE:.+]] = arith.subi %[[KH]], %[[ONE]] : index
   // CHECK: %[[DIL_H:.+]] = arith.constant 2 : index
   // CHECK: %[[DILATED:.+]] = arith.muli %[[DIL_H]], %[[SUB_ONE]] : index
   // CHECK: %[[ADD_ONE:.+]] = arith.addi %[[DILATED]], %[[ONE]] : index
   // 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
   // CHECK: %[[W:.+]] = tensor.dim %arg0, %[[C2]]
   // CHECK: %[[C2_0:.+]] = arith.constant 2
   // CHECK: %[[KW:.+]] = tensor.dim %arg1, %[[C2_0]]
   // CHECK: %[[ONE_0:.+]] = arith.constant 1 : index
   // CHECK: %[[PAD_2:.+]] = arith.constant 0 : index
   // CHECK: %[[ADD_PAD_2:.+]] = arith.addi %[[W]], %[[PAD_2]] : index
   // CHECK: %[[PAD_3:.+]] = arith.constant 0 : index
   // CHECK: %[[ADD_PAD_3:.+]] = arith.addi %[[ADD_PAD_2]], %[[PAD_3]] : index
   // CHECK: %[[SUB_ONE_0:.+]] = arith.subi %[[KW]], %[[ONE_0]] : index
   // CHECK: %[[DIL_W:.+]] = arith.constant 1 : index
   // CHECK: %[[DILATED_0:.+]] = arith.muli %[[DIL_W]], %[[SUB_ONE_0]] : index
   // CHECK: %[[ADD_ONE_0:.+]] = arith.addi %[[DILATED_0]], %[[ONE_0]] : index
   // 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]>
   // CHECK: %[[WEIGHT:.+]] = "tosa.transpose"(%arg1, %[[PERM]])
   // CHECK: %[[M_IN:.+]] = linalg.init_tensor [1, %[[H_OUT]], %[[W_OUT]], 28]
   // CHECK: %[[CST:.+]] = arith.constant 0
   // CHECK: %[[FILL:.+]] = linalg.fill
   // CHECK: %[[B_IN:.+]] = linalg.init_tensor [1, %[[H_OUT]], %[[W_OUT]], 28]