[mlir][tosa] Add dilation to tosa.transpose_conv2d lowering

[rsuderman]

Summary

Dilation only requires increasing the padding on the left/right side of the
input, and including dilation in the convolution. This implementation still
lacks support for strided convolutions.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

matchAndRewrite(tosa::TransposeConv2DOp op, ArrayRef<Value> args,
     llvm::SmallVector<int64_t> pad;
     llvm::SmallVector<int64_t> stride;
     llvm::SmallVector<int64_t> dilation;
 
     getValuesFromIntArrayAttribute(op.out_pad().cast<ArrayAttr>(), pad);
     getValuesFromIntArrayAttribute(op.stride().cast<ArrayAttr>(), stride);
     getValuesFromIntArrayAttribute(op.dilation().cast<ArrayAttr>(), dilation);
 
-    // We have not solved for stride / dilation yet. Dilation should be
-    // straight forward but stride is more complicated. Linalg work is likely
-    // required for efficient implementation.
-    if (llvm::any_of(stride, [](int64_t v) { return v != 1; }))
-      return failure();
-    if (llvm::any_of(dilation, [](int64_t v) { return v != 1; }))
-      return failure();
-
+    // If striding is all 1 we can modify padding and reverse the kernel along
+    // the x/y direction to make it a regular convolution. This is much simpler
+    // then handling striding....
+    if (llvm::all_of(stride, [](int64_t v) { return v == 1; })) {
       if (!inputTy.hasStaticShape() || !weightTy.hasStaticShape() ||
           !biasTy.hasStaticShape() || !resultTy.hasStaticShape())
         return failure();
 
-    int64_t inputHeight = inputTy.getDimSize(1);
-    int64_t inputWidth = inputTy.getDimSize(2);
-    int64_t kernelHeight = weightTy.getDimSize(1);
-    int64_t kernelWidth = weightTy.getDimSize(2);
-    int64_t outputHeight = resultTy.getDimSize(1);
-    int64_t outputWidth = resultTy.getDimSize(2);
-
-    int64_t requiredInputHeight = outputHeight + kernelHeight - 1;
-    int64_t requiredInputWidth = outputWidth + kernelWidth - 1;
-
-    llvm::SmallVector<int64_t> newPad(4, 0);
-    newPad[0] = kernelHeight - 1 - pad[0];
-    newPad[2] = kernelWidth - 1 - pad[1];
-
-    newPad[1] = requiredInputHeight - newPad[0] - inputHeight;
-    newPad[3] = requiredInputWidth - newPad[2] - inputWidth;
+      int64_t kernelHeight = (weightTy.getDimSize(1) - 1) * dilation[0] + 1;
+      int64_t kernelWidth = (weightTy.getDimSize(2) - 1) * dilation[1] + 1;
+      int64_t requiredInputHeight = resultTy.getDimSize(1) + kernelHeight - 1;
+      int64_t requiredInputWidth = resultTy.getDimSize(2) + kernelWidth - 1;
+
+      llvm::SmallVector<int64_t> convPad(4, 0);
+      convPad[0] = kernelHeight - 1 - pad[0];
+      convPad[2] = kernelWidth - 1 - pad[1];
+      convPad[1] = requiredInputHeight - convPad[0] - inputTy.getDimSize(1);
+      convPad[3] = requiredInputWidth - convPad[2] - inputTy.getDimSize(2);
 
       auto reverse1 = rewriter.create<tosa::ReverseOp>(
           loc, weightTy, weight, rewriter.getI64IntegerAttr(1));
       auto reverse2 = rewriter.create<tosa::ReverseOp>(
           loc, weightTy, reverse1, rewriter.getI64IntegerAttr(2));
 
       Value conv2d;
       if (op.quantization_info().hasValue()) {
         conv2d = rewriter.create<tosa::Conv2DOp>(
             loc, resultTy, input, reverse2, bias,
-          rewriter.getI64ArrayAttr(newPad), rewriter.getI64ArrayAttr(stride),
+            rewriter.getI64ArrayAttr(convPad), rewriter.getI64ArrayAttr(stride),
             rewriter.getI64ArrayAttr(dilation),
             op.quantization_info().getValue());
       } else {
         conv2d = rewriter.create<tosa::Conv2DOp>(
             loc, resultTy, input, reverse2, bias,
-          rewriter.getI64ArrayAttr(newPad), rewriter.getI64ArrayAttr(stride),
+            rewriter.getI64ArrayAttr(convPad), rewriter.getI64ArrayAttr(stride),
             rewriter.getI64ArrayAttr(dilation));
       }
 
       rewriter.replaceOp(op, conv2d);
       return success();
     }
+
+    return failure();
+  }
 };
 
 class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
 public:
   using OpConversionPattern<tosa::MatMulOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tosa::MatMulOp op, ArrayRef<Value> args,
                   ConversionPatternRewriter &rewriter) const final {

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

 // CHECK-LABEL: @transpose_conv
 func @transpose_conv(%arg0 : tensor<1x12x12x2xf32>, %arg1 : tensor<4x3x3x2xf32>, %arg2 : tensor<4xf32>) -> () {
   // CHECK: [[PAD:%.+]] = linalg.pad_tensor %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0]
   // CHECK: linalg.conv_2d_input_nhwc_filter_ohwi_poly {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins([[PAD]], {{%.+}} : tensor<1x16x16x2xf32>, tensor<4x3x3x2xf32>)
   %0 = "tosa.transpose_conv2d"(%arg0, %arg1, %arg2) {dilation = [1, 1], out_pad = [0, 0], out_shape = [1, 14, 14, 4], stride = [1, 1]} : (tensor<1x12x12x2xf32>, tensor<4x3x3x2xf32>, tensor<4xf32>) -> tensor<1x14x14x4xf32>
   return
 }
 
+// -----
+
+// CHECK-LABEL: @transpose_conv_dilated
+func @transpose_conv_dilated(%arg0 : tensor<1x12x12x2xf32>, %arg1 : tensor<4x3x3x2xf32>, %arg2 : tensor<4xf32>) -> () {
+  // CHECK: [[PAD:%.+]] = linalg.pad_tensor %arg0 low[0, 4, 4, 0] high[0, 4, 4, 0]
+  // CHECK: linalg.conv_2d_input_nhwc_filter_ohwi_poly {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins([[PAD]], {{%.+}} : tensor<1x20x20x2xf32>, tensor<4x3x3x2xf32>)
+  %0 = "tosa.transpose_conv2d"(%arg0, %arg1, %arg2) {dilation = [2, 2], out_pad = [0, 0], out_shape = [1, 16, 16, 4], stride = [1, 1]} : (tensor<1x12x12x2xf32>, tensor<4x3x3x2xf32>, tensor<4xf32>) -> tensor<1x16x16x4xf32>
+  return
+}
+
 
 // -----
 
 // CHECK-LABEL: @resize_nearest
 func @resize_nearest(%input: tensor<1x2x2x1xf32>) -> () {
   // CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1]
   // CHECK: %[[GENERIC:.+]] = linalg.generic
   // CHECK: %[[IDX0:.+]] = linalg.index 0