feat: basic support for ConvTranspose lowering to tosa.

src/Conversion/ONNXToTOSA/Math/Conv2D.cpp

@@ -216,12 +216,112 @@ class ONNXConvOpLoweringToTOSA : public ConversionPattern {
   int64_t groupedConvThreshold;
 };
 
+class ONNXConvTransposeOpLoweringToTOSA
+    : public OpConversionPattern<ONNXConvTransposeOp> {
+public:
+  ONNXConvTransposeOpLoweringToTOSA(MLIRContext *ctx)
+      : OpConversionPattern<ONNXConvTransposeOp>(ctx) {}
+
+  using OpAdaptor = typename ONNXConvTransposeOp::Adaptor;
+  LogicalResult matchAndRewrite(ONNXConvTransposeOp convOp, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+
+    auto getDynamicShapeTypeTensor4D = [](Type elementType) {
+      return mlir::RankedTensorType::get(
+          llvm::SmallVector<int64_t, 4>(4, ShapedType::kDynamic), elementType);
+    };
+
+    auto loc = convOp->getLoc();
+    TosaBuilder tosaBuilder(rewriter, loc);
+
+    auto input = adaptor.getX();
+    auto weights = adaptor.getW();
+    auto bias = adaptor.getB();
+
+    auto inputType = input.getType().cast<TensorType>();
+    auto inputShape = inputType.getShape();
+    auto weightType = weights.getType().cast<TensorType>();
+    auto weightShape = weightType.getShape();
+    auto outputType = convOp.getResult().getType().cast<TensorType>();
+    auto outputShape = outputType.getShape();
+
+    if (inputShape.size() != 4)
+      return rewriter.notifyMatchFailure(
+          convOp, "only 2d tensor support is implemented");
+
+    if (adaptor.getGroup() != 1)
+      return rewriter.notifyMatchFailure(
+          convOp, "grouped ConvTranspose not supported");
+
+    auto dilations = adaptor.getDilations();
+    if (dilations.has_value() && llvm::any_of(*dilations, [](Attribute attr) {
+          return attr.cast<IntegerAttr>().getValue() != 1;
+        }))
+      return rewriter.notifyMatchFailure(
+          convOp, "dilation more than 1 is not supported");
+
+    // Get shapehelper for autopad attributes
+    IndexExprBuilderForTosa createTosaIE(rewriter, loc);
+    ONNXConvTransposeOpShapeHelper shapeHelper(
+        convOp, adaptor.getOperands(), &createTosaIE);
+    shapeHelper.computeShapeAndAssertOnFailure();
+
+    if (!IndexExpr::isLiteral(shapeHelper.pads))
+      return rewriter.notifyMatchFailure(convOp, "pads is not a literal.");
+
+    // Convert input [N,IC,IH,IW] -> [N,IH,IW,IC]
+    Value newInput = tosaBuilder.transpose(input, {0, 2, 3, 1});
+
+    // Convert weights [OC,IC,KH,KW] -> [OC,KH,KW,IC]
+    Value newWeight = tosaBuilder.transpose(weights, {0, 2, 3, 1});
+
+    if (bias.getType().isa<NoneType>()) {
+      DenseElementsAttr newBiasAttr = DenseElementsAttr::get(
+          RankedTensorType::get({weightShape[0]}, rewriter.getF32Type()),
+          {0.0F});
+      bias = rewriter.create<mlir::tosa::ConstOp>(
+          loc, newBiasAttr.getType(), newBiasAttr);
+    }
+
+    llvm::SmallVector<int64_t, 4> tosaLayoutOutShape{
+        outputShape[0], outputShape[2], outputShape[3], outputShape[1]};
+    DenseI64ArrayAttr strides =
+        rewriter.getDenseI64ArrayAttr(shapeHelper.strides);
+
+    llvm::SmallVector<int64_t, 4> pads;
+    IndexExpr::getLiteral(shapeHelper.pads, pads);
+
+    Value padding = tosa::buildOnnxToTosaPaddingConstOp(
+        rewriter, pads, loc, {0, 0}, {0, 0});
+    auto constTosaTensor =
+        tosaBuilder.getSplattedConst(0.0, inputType.getElementType());
+
+    auto padOp = tosa::CreateOpAndInfer<mlir::tosa::PadOp>(rewriter, loc,
+        getDynamicShapeTypeTensor4D(inputType.getElementType()), newInput,
+        padding, constTosaTensor);
+
+    Value transposeConv2d =
+        tosa::CreateOpAndInfer<mlir::tosa::TransposeConv2DOp>(rewriter, loc,
+            getDynamicShapeTypeTensor4D(outputType.getElementType()), padOp,
+            newWeight, bias,
+            rewriter.getDenseI64ArrayAttr({0, shapeHelper.outputPadding[0], 0,
+                shapeHelper.outputPadding[1]}),
+            strides, rewriter.getDenseI64ArrayAttr(tosaLayoutOutShape));
+
+    // Convert output [N,OH,OW,OC] -> [N,OC,OH,OW]
+    Value newOutput = tosaBuilder.transpose(transposeConv2d, {0, 3, 1, 2});
+
+    rewriter.replaceOp(convOp, {newOutput});
+    return success();
+  }
+};
 } // namespace
 
 void populateLoweringONNXConvOpToTOSAPattern(ConversionTarget &target,
     RewritePatternSet &patterns, TypeConverter &typeConverter, MLIRContext *ctx,
     int64_t groupedConvThreshold) {
   patterns.insert<ONNXConvOpLoweringToTOSA>(ctx, groupedConvThreshold);
+  patterns.insert<ONNXConvTransposeOpLoweringToTOSA>(ctx);
 }
 
 } // namespace onnx_mlir

test/mlir/conversion/onnx_to_tosa/Math/Conv.mlir

@@ -188,4 +188,18 @@ func.func @test_onnx_conv2d_dyn_shapes_with_shape_inference(%arg0: tensor<5x3x25
   return %0 : tensor<?x?x?x?xf32>
 // CHECK-LABEL:  func.func @test_onnx_conv2d_dyn_shapes_with_shape_inference
 // CHECK: tosa.conv
+}
+
+func.func @test_onnx_convtranspose(%arg0: tensor<32x128x64x64xf32>, %arg1 : tensor<128x64x2x2xf32>, %arg2: tensor<64xf32>) ->  tensor<32x64x128x128xf32> {
+  %0 = "onnx.ConvTranspose"(%arg0, %arg1, %arg2) {auto_pad = "NOTSET", dilations = [1, 1], group = 1 : si64, kernel_shape = [2, 2], onnx_node_name = "/module_37/ConvTranspose", pads = [0, 0, 0, 0], strides = [2, 2]} : (tensor<32x128x64x64xf32>, tensor<128x64x2x2xf32>, tensor<64xf32>) -> tensor<32x64x128x128xf32>
+  return %0 : tensor<32x64x128x128xf32>
+// CHECK-LABEL:  func.func @test_onnx_convtranspose
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<32x128x64x64xf32>, [[PARAM_1_:%.+]]: tensor<128x64x2x2xf32>, [[PARAM_2_:%.+]]: tensor<64xf32>) -> tensor<32x64x128x128xf32> {
+// CHECK:           [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}> : () -> tensor<4xi32>
+// CHECK:           [[VAR_1_:%.+]] = "tosa.transpose"([[PARAM_0_]], [[VAR_0_]]) : (tensor<32x128x64x64xf32>, tensor<4xi32>) -> tensor<32x64x64x128xf32>
+// CHECK:           [[VAR_2_:%.+]] = "tosa.transpose"([[PARAM_1_]], [[VAR_0_]]) : (tensor<128x64x2x2xf32>, tensor<4xi32>) -> tensor<128x2x2x64xf32>
+// CHECK:           [[VAR_3_:%.+]] = "tosa.transpose_conv2d"([[VAR_1_]], [[VAR_2_]], [[PARAM_2_]]) <{out_pad = array<i64: 0, 0, 0, 0>, out_shape = array<i64: 32, 128, 128, 64>, stride = array<i64: 2, 2>}> : (tensor<32x64x64x128xf32>, tensor<128x2x2x64xf32>, tensor<64xf32>) -> tensor<32x128x128x64xf32>
+// CHECK:           [[VAR_4_:%.+]] = "tosa.const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> : () -> tensor<4xi32>
+// CHECK:           [[VAR_5_:%.+]] = "tosa.transpose"([[VAR_3_]], [[VAR_4_]]) : (tensor<32x128x128x64xf32>, tensor<4xi32>) -> tensor<32x64x128x128xf32>
+// CHECK:           return [[VAR_5_]] : tensor<32x64x128x128xf32>
 }