Skip to content

[TOSA] Add transposed conv support #4360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+140 −23
Open

[TOSA] Add transposed conv support #4360

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 120 additions & 13 deletions lib/Conversion/TorchToTosa/TorchToTosa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2304,9 +2304,6 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
if (!matchPattern(op.getTransposed(), m_TorchConstantBool(&transposed)))
return rewriter.notifyMatchFailure(
op, "Unimplemented: non-constant value for transposed not supported");
if (transposed)
return rewriter.notifyMatchFailure(
op, "Unimplemented: transposed convolution not supported");

auto input = adaptor.getInput();
auto weight = adaptor.getWeight();
Expand Down Expand Up @@ -2338,12 +2335,17 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
auto bias = adaptor.getBias();

if (isa<Torch::NoneType>(bias.getType())) {
auto bias_result = tosa::getConvBiasForNoneType(op, rewriter, inputElemTy,
outputElemTy, weightShape);
if (failed(bias_result))
SmallVector<int64_t, 4> biasWeightShape =
transposed ? SmallVector<int64_t, 4>{weightShape[1], weightShape[0],
Copy link
Member

@sahas3 sahas3 Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you add a comment about why the biasShape is different for tranposed case?

weightShape[2], weightShape[3]}
: weightShape;

auto biasResult = tosa::getConvBiasForNoneType(
op, rewriter, inputElemTy, outputElemTy, biasWeightShape);
if (failed(biasResult))
return rewriter.notifyMatchFailure(
op, "Failed to create bias tensor for none type.");
bias = bias_result.value();
bias = biasResult.value();
} else {
if (!isa<RankedTensorType>(bias.getType()))
return rewriter.notifyMatchFailure(
Expand All @@ -2370,8 +2372,8 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
m_TorchListOfConstantInts(padding_2d)))
return rewriter.notifyMatchFailure(op,
"non-const padding list unsupported");
// TOSA uses 4D padding {top, bottom, left, right} while Torch defines 2D
// padding {height, width}. The Torch OFM computation uses 2*pad in each
// TOSA uses 4D padding {top, bottom, left, right} while PyTorch defines 2D
// padding {height, width}. The PyTorch OFM computation uses 2*pad in each
// spatial direction, implying the same top=bottom=height and left=right=width
// values for TOSA.
SmallVector<int64_t> padding(
Expand All @@ -2388,11 +2390,19 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
return rewriter.notifyMatchFailure(
op, "failed to get accumulator type for convolution ops");

// Weight layout reference:
// Conv : PyTorch OIHW -> TOSA OHWI
// Depthwise : PyTorch OIHW* -> TOSA HWIM
// (PyTorch depthwise uses out_ch=in_ch*depth_multiplier)
// Grouped : PyTorch O(I/G)HW -> N/A
// Transposed : PyTorch IOHW -> TOSA OHWI
// TOSA works in NHWC and takes OHWI (conv) / HWIM (depthwise conv) weights.
// Perform the necessary transformations.
SmallVector<int32_t> nchwToNhwcDims({0, 2, 3, 1});
SmallVector<int64_t> transposedInputShape(
{inputShape[0], inputShape[2], inputShape[3], inputShape[1]});
SmallVector<int32_t> nhwcToNchwDims({0, 3, 1, 2});
SmallVector<int64_t, 4> transposedInputShape;
for (int32_t dim : nchwToNhwcDims)
transposedInputShape.push_back(inputShape[dim]);
auto transposedInputType = RankedTensorType::get(
makeShapeLLVMCompatible(transposedInputShape), inputElemTy);
auto transposedInput =
Expand All @@ -2403,6 +2413,104 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
rewriter.getDenseI32ArrayAttr(nchwToNhwcDims))
.getResult();

if (transposed) {
if (groups != 1)
return rewriter.notifyMatchFailure(
op, "Unimplemented: grouped transposed convolution not supported by "
"TOSA");
if (dilation[0] != 1 || dilation[1] != 1)
return rewriter.notifyMatchFailure(
op, "Unimplemented: dilated transposed convolution not supported by "
"TOSA");

Comment on lines +2418 to +2425
Copy link
Member

@sahas3 sahas3 Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think these notify failures need to happen before any IR rewrites take place, otherwise the pattern rewriter ends up in a recursive loop. For example, on line 2410 we'd have already introduced tosa.transpose and then we'll bail out from here for a failure case.

SmallVector<int32_t> iohwToOhwi({1, 2, 3, 0});
SmallVector<int64_t, 4> ohwiWeightShape;
for (int32_t dim : iohwToOhwi)
ohwiWeightShape.push_back(weightShape[dim]);
auto ohwiWeightType = RankedTensorType::get(
makeShapeLLVMCompatible(ohwiWeightShape), weightElemTy);
Value transformedWeight =
rewriter
.create<tosa::TransposeOp>(
op->getLoc(), getTypeConverter()->convertType(ohwiWeightType),
weight, rewriter.getDenseI32ArrayAttr(iohwToOhwi))
.getResult();

// TOSA 'out_pad' is a 4D array {top,bottom,left,right}.
// Map from PyTorch's (padding, output_padding):
// out_pad_total(H/W) = output_padding(H/W) - 2*padding(H/W)
// Negative values are allowed and will be handled by the TOSA
// decomposition.
SmallVector<int64_t, 2> outPadding2D;
if (!matchPattern(adaptor.getOutputPadding(),
m_TorchListOfConstantInts(outPadding2D)))
return rewriter.notifyMatchFailure(
op, "non-const output_padding list unsupported for transposed conv");

int64_t outPadH = outPadding2D[0] - 2 * padding_2d[0];
int64_t outPadW = outPadding2D[1] - 2 * padding_2d[1];
int64_t outPadTop = outPadH / 2;
int64_t outPadBottom = outPadH - outPadTop;
int64_t outPadLeft = outPadW / 2;
int64_t outPadRight = outPadW - outPadLeft;
SmallVector<int64_t, 4> outPad(
{outPadTop, outPadBottom, outPadLeft, outPadRight});

// Result type is NHWC (we'll transpose back).
auto outNCHW = makeShapeTorchCompatible(outputTy.getShape());
SmallVector<int64_t, 4> outNHWC;
for (int32_t dim : nchwToNhwcDims)
outNHWC.push_back(outNCHW[dim]);
auto transConvOpTy =
RankedTensorType::get(makeShapeLLVMCompatible(outNHWC), biasElemTy);

// Zero-points.
auto zps = tosa::createZPsAsConst(rewriter, input, weight);
Value inputZp = zps.first ? zps.first
: tosa::createZeroPointTensor(
rewriter, op->getLoc(), inputElemTy, 0)
.value();
Value weightZp = zps.second ? zps.second
: tosa::createZeroPointTensor(
rewriter, op->getLoc(), weightElemTy, 0)
.value();

Value convTOut =
rewriter
.create<tosa::TransposeConv2DOp>(
op->getLoc(), getTypeConverter()->convertType(transConvOpTy),
transposedInput, transformedWeight, bias, inputZp, weightZp,
rewriter.getDenseI64ArrayAttr(outPad),
rewriter.getDenseI64ArrayAttr(stride), accType)
.getResult();

SmallVector<int64_t, 4> transposedOutputShape;
for (int32_t dim : nhwcToNchwDims)
transposedOutputShape.push_back(outNHWC[dim]);
auto transposedOutputType = RankedTensorType::get(
makeShapeLLVMCompatible(transposedOutputShape), biasElemTy);
Value transposedOutput =
rewriter
.create<tosa::TransposeOp>(
op->getLoc(),
getTypeConverter()->convertType(transposedOutputType), convTOut,
rewriter.getDenseI32ArrayAttr(nhwcToNchwDims))
.getResult();

// Quantized rescale.
Value rescaledResult = transposedOutput;
if (isa<quant::QuantizedType>(inputElemTy)) {
rescaledResult = tosa::buildRescaleOpConvOutput(
Copy link
Member

@sahas3 sahas3 Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ConvTranspose2DQInt8_basic e2e test likely triggers this code path -- any idea why that is still failing? If that's not the correct test, are there any e2e tests that triggers this code?

Copy link
Contributor Author

@Lallapallooza Lallapallooza Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The failing XFAIL stops at torch.aten._make_per_tensor_quantized_tensor, which Torch->TOSA doesn't legalize today, so this isn't coming from the new transpose-conv path.
I'd like to land this patch as-is and follow up with a dedicated change that adds that conversion (plus quantized coverage).

rewriter, op, transposedOutput, inputTy, weightTy, outputTy);
}

// Final cast to requested output type.
rewriter.replaceOp(
op, {tosa::tosaCastTensorToType(rewriter, rescaledResult, outputTy)
.value()});
return success();
}

SmallVector<int64_t> transformedWeightShape;
RankedTensorType transformedWeightType;
Value transformedWeight;
Expand Down Expand Up @@ -2487,7 +2595,7 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
if (remainderHDim != 0) {
if (remainderHDim > padding[1]) {
SmallVector<int64_t> startHSlice(inputTy.getRank(), 0);
SmallVector<int64_t> sizeHSlice(transposedInputShape);
SmallVector<int64_t, 4> sizeHSlice(transposedInputShape);
// TOSA uses NHWC, so we will slice dim 1 for Height value
sizeHSlice[1] = inputHDim - (remainderHDim - padding[1]);
transposedInput = tosa::CreateOpAndInfer<tosa::SliceOp>(
Expand Down Expand Up @@ -2583,7 +2691,6 @@ LogicalResult ConvertAtenOp<AtenConvolutionOp>::matchAndRewrite(
llvm_unreachable("Unhandled convolution type");
}

SmallVector<int32_t> nhwcToNchwDims({0, 3, 1, 2});
SmallVector<int64_t> transposedOutputShape(
{outputShape[0], outputShape[3], outputShape[1], outputShape[2]});
auto transposedOutputType = RankedTensorType::get(
Expand Down
6 changes: 0 additions & 6 deletions projects/pt1/e2e_testing/xfail_sets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3588,7 +3588,6 @@
"AvgPool3dCountIncludePadFalseWithoutPadding_basic",
"Conv_Transpose1dModule_basic",
"Conv_Transpose1dStaticModule_basic",
"Conv_Transpose2dStaticModule_basic",
"Conv_Transpose3dModule_basic",
"Conv_Transpose3dStaticModule_basic",
"IndexPutWithNoneAndBroadcastModule_basic",
Expand Down Expand Up @@ -3713,16 +3712,11 @@
"Conv3dWithValidPaddingModule_basic",
"ConvTbcModule_basic",
"ConvTranspose2DQInt8_basic",
"Conv_Transpose2dModule_basic",
"ConvolutionBackwardModule2DPadded_basic",
"ConvolutionBackwardModule2DStatic_basic",
"ConvolutionBackwardModule2DStrided_basic",
"ConvolutionBackwardModule2D_basic",
"ConvolutionModule2DGroups_basic",
"ConvolutionModule2DTransposeNonUnitOutputPadding_basic",
"ConvolutionModule2DTransposeStridedStatic_basic",
"ConvolutionModule2DTransposeStrided_basic",
"ConvolutionModule2DTranspose_basic",
"ConvolutionModule2DGroupedTranspose_basic",
"ConvolutionModule3DGroups_basic",
"ConvolutionModule3DGroupsStrided_basic",
Expand Down
2 changes: 2 additions & 0 deletions projects/pt1/python/torch_mlir_e2e_test/tosa_backends/linalg_on_tensors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
# that depend on TOSA as well as TOSA-to-Standard.
"tosa-to-arith",
"tosa-to-scf",
# Required for transposed convolution support (decomposes to conv ops).
"tosa-optional-decompositions",
# Named ops must be legalized prior to general tosa-to-linalg
"tosa-to-linalg-named",
# TOSA-to-LinAlg may generate tosa.const() ops, so we want to lower them
Expand Down
22 changes: 18 additions & 4 deletions test/Conversion/TorchToTosa/conv2d_transpose.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,23 @@
// RUN: torch-mlir-opt <%s -convert-torch-to-tosa -split-input-file -verify-diagnostics
// RUN: torch-mlir-opt <%s -convert-torch-to-tosa -split-input-file | FileCheck %s

// The following test ensures that a tranposed convolution op is not
// lowered in the torch-to-tosa conversion pass.
// The lowering now legalizes transpose convolutions into the TOSA dialect.
// Verify that we emit tosa.transpose_conv2d with the expected reshapes/
// permutations.

// CHECK-LABEL: func.func @forward
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[1,64,1,100],f32>) -> !torch.vtensor<[1,64,2,200],f32> {
// CHECK: %[[IN_TENSOR:.*]] = torch_c.to_builtin_tensor %[[INPUT]] : !torch.vtensor<[1,64,1,100],f32> -> tensor<1x64x1x100xf32>
// CHECK: %[[WEIGHT:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<64x64x3x3xf32>}> : () -> tensor<64x64x3x3xf32>
// CHECK: %[[BIAS:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<64xf32>}> : () -> tensor<64xf32>
// CHECK: %[[TRANS_IN:.*]] = tosa.transpose %[[IN_TENSOR]] {perms = array<i32: 0, 2, 3, 1>} : (tensor<1x64x1x100xf32>) -> tensor<1x1x100x64xf32>
// CHECK: %[[W_OHWI:.*]] = tosa.transpose %[[WEIGHT]] {perms = array<i32: 1, 2, 3, 0>} : (tensor<64x64x3x3xf32>) -> tensor<64x3x3x64xf32>
// CHECK: %[[ZP0:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32>
// CHECK: %[[ZP1:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32>
// CHECK: %[[TCONV:.*]] = tosa.transpose_conv2d %[[TRANS_IN]], %[[W_OHWI]], %[[BIAS]], %[[ZP0]], %[[ZP1]] {acc_type = f32, out_pad = array<i64: 0, -1, 0, -1>, stride = array<i64: 2, 2>} : (tensor<1x1x100x64xf32>, tensor<64x3x3x64xf32>, tensor<64xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x2x200x64xf32>
// CHECK: %[[TRANS_OUT:.*]] = tosa.transpose %[[TCONV]] {perms = array<i32: 0, 3, 1, 2>} : (tensor<1x2x200x64xf32>) -> tensor<1x64x2x200xf32>
// CHECK: %[[RESULT:.*]] = torch_c.from_builtin_tensor %[[TRANS_OUT]] : tensor<1x64x2x200xf32> -> !torch.vtensor<[1,64,2,200],f32>
// CHECK: return %[[RESULT]] : !torch.vtensor<[1,64,2,200],f32>
// CHECK: }
func.func @forward(%input: !torch.vtensor<[1,64,1,100],f32>) -> !torch.vtensor<[1,64,2,200],f32> {
%true = torch.constant.bool true
%int1 = torch.constant.int 1
Expand All @@ -11,7 +26,6 @@ func.func @forward(%input: !torch.vtensor<[1,64,1,100],f32>) -> !torch.vtensor<[
%bias = torch.vtensor.literal(dense<0.0> : tensor<64xf32>) : !torch.vtensor<[64],f32>
%stride = torch.prim.ListConstruct %int2, %int2 : (!torch.int, !torch.int) -> !torch.list<int>
%int1x1 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<int>
// expected-error@+1 {{failed to legalize operation 'torch.aten.convolution' that was explicitly marked illegal}}
%output = torch.aten.convolution %input, %weight, %bias, %stride, %int1x1, %int1x1, %true, %int1x1, %int1 : !torch.vtensor<[1,64,1,100],f32>, !torch.vtensor<[64,64,3,3],f32>, !torch.vtensor<[64],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.list<int>, !torch.int -> !torch.vtensor<[1,64,2,200],f32>
return %output : !torch.vtensor<[1,64,2,200],f32>
}
Loading