Skip to content

Commit 7db2a11

Browse files
hanhanWhanhanW
authored
[TorchOnnxToTorch] Support ConvTranspose output_shape for VALID (#4567)
Support onnx.ConvTranspose output_shape lowering when auto_pad=VALID by inferring output_padding from the requested shape with zero padding. Also reject the ambiguous output_shape + output_padding combination instead of silently overriding the explicit attribute. Add regression coverage for the VALID + output_shape case and a diagnostics test for the mixed-attribute rejection. --------- Signed-off-by: hanhanW <hanhan0912@gmail.com>
1 parent 90fc215 commit 7db2a11

3 files changed

Lines changed: 177 additions & 7 deletions

File tree

‎lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp‎

Lines changed: 73 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1875,23 +1875,89 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
18751875
binder.op, "Expected input type having sizes");
18761876
}
18771877
ArrayRef<int64_t> inputShape = inputTensorType.getSizes();
1878+
if (binder.s64IntegerArrayAttr(outputShape, "output_shape", {}))
1879+
return failure();
1880+
if (!outputShape.empty() && outputShape.size() != rank - 2) {
1881+
return rewriter.notifyMatchFailure(
1882+
binder.op,
1883+
"output_shape list size does not match the number of axes");
1884+
}
1885+
1886+
auto inferOutputPaddingFromOutputShape =
1887+
[&]() -> FailureOr<SmallVector<int64_t>> {
1888+
if (padding.size() != rank - 2 && padding.size() != 2 * (rank - 2)) {
1889+
return rewriter.notifyMatchFailure(
1890+
binder.op,
1891+
"padding list size does not match the number of axes");
1892+
}
1893+
bool isPerAxisPadded = padding.size() == rank - 2;
1894+
SmallVector<int64_t> inferredOutputPadding;
1895+
inferredOutputPadding.reserve(rank - 2);
1896+
for (unsigned i = 0; i < rank - 2; i++) {
1897+
// ONNX pads are laid out as [x1_begin, ..., xN_begin, x1_end,
1898+
// ..., xN_end] when fully specified, or as a per-axis symmetric
1899+
// value when half-sized.
1900+
int64_t totalPadding = isPerAxisPadded
1901+
? 2 * padding[i]
1902+
: padding[i] + padding[i + rank - 2];
1903+
int64_t inferredDim = strides[i] * (inputShape[2 + i] - 1) -
1904+
totalPadding +
1905+
((kernelShape[i] - 1) * dilations[i] + 1);
1906+
int64_t inferredOutputPaddingValue = outputShape[i] - inferredDim;
1907+
if (inferredOutputPaddingValue < 0) {
1908+
return rewriter.notifyMatchFailure(
1909+
binder.op,
1910+
"output_shape would require a negative output_padding");
1911+
}
1912+
if (inferredOutputPaddingValue >= strides[i]) {
1913+
return rewriter.notifyMatchFailure(
1914+
binder.op,
1915+
"output_shape would require output_padding >= stride, "
1916+
"which violates the ONNX ConvTranspose specification");
1917+
}
1918+
inferredOutputPadding.push_back(inferredOutputPaddingValue);
1919+
}
1920+
return inferredOutputPadding;
1921+
};
1922+
1923+
auto applyOutputPaddingFromOutputShape = [&]() -> LogicalResult {
1924+
FailureOr<SmallVector<int64_t>> inferredOutputPadding =
1925+
inferOutputPaddingFromOutputShape();
1926+
if (failed(inferredOutputPadding))
1927+
return failure();
1928+
if (outputPadding != defaultOutputPadding &&
1929+
outputPadding != *inferredOutputPadding) {
1930+
return rewriter.notifyMatchFailure(
1931+
binder.op, "output_shape and output_padding imply different "
1932+
"output_padding values");
1933+
}
1934+
outputPadding = *inferredOutputPadding;
1935+
return success();
1936+
};
18781937

18791938
if (autoPad == "VALID") {
18801939
// Zero padding.
18811940
padding = defaultPadding;
1941+
if (!outputShape.empty()) {
1942+
if (failed(applyOutputPaddingFromOutputShape()))
1943+
return failure();
1944+
}
18821945
} else if (autoPad == "NOTSET") {
18831946
// Explicit padding; read pads with defaults.
18841947
if (binder.s64IntegerArrayAttr(padding, "pads", defaultPadding))
18851948
return failure();
1949+
if (!outputShape.empty()) {
1950+
if (failed(applyOutputPaddingFromOutputShape()))
1951+
return failure();
1952+
}
18861953
} else { // autopad == SAME_UPPER or SAME_LOWER
1887-
// Auto-padding; output_shape defaults to input_shape * strides.
1888-
SmallVector<int64_t> defaultOutputShape;
1889-
for (unsigned i = 0; i < rank - 2; i++) {
1890-
defaultOutputShape.push_back(inputShape[2 + i] * strides[i]);
1954+
// Auto-padding. When output_shape is not specified, default it to
1955+
// input_shape * strides.
1956+
if (outputShape.empty()) {
1957+
for (unsigned i = 0; i < rank - 2; i++) {
1958+
outputShape.push_back(inputShape[2 + i] * strides[i]);
1959+
}
18911960
}
1892-
if (binder.s64IntegerArrayAttr(outputShape, "output_shape",
1893-
defaultOutputShape))
1894-
return failure();
18951961
SmallVector<int64_t> paddingEnd;
18961962
for (unsigned i = 0; i < rank - 2; i++) {
18971963
int64_t totalPadding =

‎test/Conversion/TorchOnnxToTorch/convtranspose_diagnostics.mlir‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// RUN: torch-mlir-opt <%s -split-input-file -verify-diagnostics -convert-torch-onnx-to-torch
2+
3+
func.func @test_convtranspose_output_shape_with_conflicting_output_padding(
4+
%arg0: !torch.vtensor<[1,1,3,3],f32>,
5+
%arg1: !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32>
6+
attributes {torch.onnx_meta.ir_version = 10 : si64,
7+
torch.onnx_meta.opset_version = 22 : si64,
8+
torch.onnx_meta.producer_name = "backend-test",
9+
torch.onnx_meta.producer_version = ""} {
10+
// expected-error @below {{failed to legalize operation 'torch.operator' that was explicitly marked illegal}}
11+
%0 = torch.operator "onnx.ConvTranspose"(%arg0, %arg1) {
12+
torch.onnx.output_padding = [0 : si64, 1 : si64],
13+
torch.onnx.output_shape = [10 : si64, 8 : si64],
14+
torch.onnx.strides = [3 : si64, 2 : si64]
15+
} : (!torch.vtensor<[1,1,3,3],f32>, !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32>
16+
return %0 : !torch.vtensor<[1,2,10,8],f32>
17+
}

‎test/Conversion/TorchOnnxToTorch/simple_ops_a_to_f.mlir‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1296,6 +1296,93 @@ func.func @test_convtranspose(%arg0: !torch.vtensor<[1,1,3,3],f32>, %arg1: !torc
12961296

12971297
// -----
12981298

1299+
// CHECK-LABEL: @test_convtranspose_output_shape_autopad_valid
1300+
func.func @test_convtranspose_output_shape_autopad_valid(%arg0: !torch.vtensor<[1,1,3,3],f32>, %arg1: !torch.vtensor<[1,2,4,4],f32>) -> !torch.vtensor<[1,2,9,9],f32> attributes {torch.onnx_meta.ir_version = 10 : si64, torch.onnx_meta.opset_version = 22 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
1301+
// CHECK: %[[C0:.*]] = torch.constant.int 0
1302+
// CHECK: %[[C0_0:.*]] = torch.constant.int 0
1303+
// CHECK: %[[C1:.*]] = torch.constant.int 1
1304+
// CHECK: %[[C1_0:.*]] = torch.constant.int 1
1305+
// CHECK: %[[C2:.*]] = torch.constant.int 2
1306+
// CHECK: %[[C2_0:.*]] = torch.constant.int 2
1307+
// CHECK: %[[C1_1:.*]] = torch.constant.int 1
1308+
// CHECK: %[[C1_2:.*]] = torch.constant.int 1
1309+
// CHECK: %[[PADDING:.*]] = torch.prim.ListConstruct %[[C0]], %[[C0_0]] : (!torch.int, !torch.int) -> !torch.list<int>
1310+
// CHECK: %[[DILATIONS:.*]] = torch.prim.ListConstruct %[[C1]], %[[C1_0]] : (!torch.int, !torch.int) -> !torch.list<int>
1311+
// CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[C2]], %[[C2_0]] : (!torch.int, !torch.int) -> !torch.list<int>
1312+
// CHECK: %[[OUTPUT_PADDING:.*]] = torch.prim.ListConstruct %[[C1_1]], %[[C1_2]] : (!torch.int, !torch.int) -> !torch.list<int>
1313+
// CHECK: torch.aten.convolution %arg0, %arg1, {{.*}}, %[[STRIDE]], %[[PADDING]], %[[DILATIONS]], {{.*}}, %[[OUTPUT_PADDING]], {{.*}} -> !torch.vtensor<[1,2,9,9],f32>
1314+
%0 = torch.operator "onnx.ConvTranspose"(%arg0, %arg1) {torch.onnx.auto_pad="VALID", torch.onnx.output_shape = [9 : si64, 9 : si64], torch.onnx.strides = [2 : si64, 2 : si64]} : (!torch.vtensor<[1,1,3,3],f32>, !torch.vtensor<[1,2,4,4],f32>) -> !torch.vtensor<[1,2,9,9],f32>
1315+
return %0 : !torch.vtensor<[1,2,9,9],f32>
1316+
}
1317+
1318+
// -----
1319+
1320+
// CHECK-LABEL: @test_convtranspose_output_shape
1321+
func.func @test_convtranspose_output_shape(%arg0: !torch.vtensor<[1,1,3,3],f32>, %arg1: !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32> attributes {torch.onnx_meta.ir_version = 10 : si64, torch.onnx_meta.opset_version = 22 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
1322+
// CHECK: %[[C0:.*]] = torch.constant.int 0
1323+
// CHECK: %[[C0_0:.*]] = torch.constant.int 0
1324+
// CHECK: %[[C1:.*]] = torch.constant.int 1
1325+
// CHECK: %[[C1_0:.*]] = torch.constant.int 1
1326+
// CHECK: %[[C3:.*]] = torch.constant.int 3
1327+
// CHECK: %[[C2:.*]] = torch.constant.int 2
1328+
// CHECK: %[[C1_1:.*]] = torch.constant.int 1
1329+
// CHECK: %[[C1_2:.*]] = torch.constant.int 1
1330+
// CHECK: %[[PADDING:.*]] = torch.prim.ListConstruct %[[C0]], %[[C0_0]] : (!torch.int, !torch.int) -> !torch.list<int>
1331+
// CHECK: %[[DILATIONS:.*]] = torch.prim.ListConstruct %[[C1]], %[[C1_0]] : (!torch.int, !torch.int) -> !torch.list<int>
1332+
// CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[C3]], %[[C2]] : (!torch.int, !torch.int) -> !torch.list<int>
1333+
// CHECK: %[[OUTPUT_PADDING:.*]] = torch.prim.ListConstruct %[[C1_1]], %[[C1_2]] : (!torch.int, !torch.int) -> !torch.list<int>
1334+
// CHECK: %[[TRANSPOSED:.*]] = torch.constant.bool true
1335+
// CHECK: %[[BIAS:.*]] = torch.constant.none
1336+
// CHECK: %[[GROUPS:.*]] = torch.constant.int 1
1337+
// CHECK: torch.aten.convolution %arg0, %arg1, %[[BIAS]], %[[STRIDE]], %[[PADDING]], %[[DILATIONS]], %[[TRANSPOSED]], %[[OUTPUT_PADDING]], %[[GROUPS]] : !torch.vtensor<[1,1,3,3],f32>, !torch.vtensor<[1,2,3,3],f32>, !torch.none, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.list<int>, !torch.int -> !torch.vtensor<[1,2,10,8],f32>
1338+
%0 = torch.operator "onnx.ConvTranspose"(%arg0, %arg1) {torch.onnx.output_shape = [10 : si64, 8 : si64], torch.onnx.strides = [3 : si64, 2 : si64]} : (!torch.vtensor<[1,1,3,3],f32>, !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32>
1339+
return %0 : !torch.vtensor<[1,2,10,8],f32>
1340+
}
1341+
1342+
// -----
1343+
1344+
// CHECK-LABEL: @test_convtranspose_output_shape_with_output_padding
1345+
func.func @test_convtranspose_output_shape_with_output_padding(%arg0: !torch.vtensor<[1,1,3,3],f32>, %arg1: !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32> attributes {torch.onnx_meta.ir_version = 10 : si64, torch.onnx_meta.opset_version = 22 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
1346+
// CHECK: %[[PADV0:.*]] = torch.constant.int 0
1347+
// CHECK: %[[PADV1:.*]] = torch.constant.int 0
1348+
// CHECK: %[[DILV0:.*]] = torch.constant.int 1
1349+
// CHECK: %[[DILV1:.*]] = torch.constant.int 1
1350+
// CHECK: %[[STRV0:.*]] = torch.constant.int 3
1351+
// CHECK: %[[STRV1:.*]] = torch.constant.int 2
1352+
// CHECK: %[[OPADV0:.*]] = torch.constant.int 1
1353+
// CHECK: %[[OPADV1:.*]] = torch.constant.int 1
1354+
// CHECK: %[[PADDING:.*]] = torch.prim.ListConstruct %[[PADV0]], %[[PADV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1355+
// CHECK: %[[DILATIONS:.*]] = torch.prim.ListConstruct %[[DILV0]], %[[DILV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1356+
// CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[STRV0]], %[[STRV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1357+
// CHECK: %[[OUTPUT_PADDING:.*]] = torch.prim.ListConstruct %[[OPADV0]], %[[OPADV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1358+
// CHECK: torch.aten.convolution %arg0, %arg1, {{.*}}, %[[STRIDE]], %[[PADDING]], %[[DILATIONS]], {{.*}}, %[[OUTPUT_PADDING]], {{.*}} -> !torch.vtensor<[1,2,10,8],f32>
1359+
%0 = torch.operator "onnx.ConvTranspose"(%arg0, %arg1) {torch.onnx.output_padding = [1 : si64, 1 : si64], torch.onnx.output_shape = [10 : si64, 8 : si64], torch.onnx.strides = [3 : si64, 2 : si64]} : (!torch.vtensor<[1,1,3,3],f32>, !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32>
1360+
return %0 : !torch.vtensor<[1,2,10,8],f32>
1361+
}
1362+
1363+
// -----
1364+
1365+
// CHECK-LABEL: @test_convtranspose_output_shape_with_pads
1366+
func.func @test_convtranspose_output_shape_with_pads(%arg0: !torch.vtensor<[1,1,3,3],f32>, %arg1: !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,8,5],f32> attributes {torch.onnx_meta.ir_version = 10 : si64, torch.onnx_meta.opset_version = 22 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
1367+
// CHECK: %[[PADV0:.*]] = torch.constant.int 1
1368+
// CHECK: %[[PADV1:.*]] = torch.constant.int 1
1369+
// CHECK: %[[DILV0:.*]] = torch.constant.int 1
1370+
// CHECK: %[[DILV1:.*]] = torch.constant.int 1
1371+
// CHECK: %[[STRV0:.*]] = torch.constant.int 3
1372+
// CHECK: %[[STRV1:.*]] = torch.constant.int 2
1373+
// CHECK: %[[OPADV0:.*]] = torch.constant.int 1
1374+
// CHECK: %[[OPADV1:.*]] = torch.constant.int 0
1375+
// CHECK: %[[PADDING:.*]] = torch.prim.ListConstruct %[[PADV0]], %[[PADV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1376+
// CHECK: %[[DILATIONS:.*]] = torch.prim.ListConstruct %[[DILV0]], %[[DILV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1377+
// CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[STRV0]], %[[STRV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1378+
// CHECK: %[[OUTPUT_PADDING:.*]] = torch.prim.ListConstruct %[[OPADV0]], %[[OPADV1]] : (!torch.int, !torch.int) -> !torch.list<int>
1379+
// CHECK: torch.aten.convolution %arg0, %arg1, {{.*}}, %[[STRIDE]], %[[PADDING]], %[[DILATIONS]], {{.*}}, %[[OUTPUT_PADDING]], {{.*}} -> !torch.vtensor<[1,2,8,5],f32>
1380+
%0 = torch.operator "onnx.ConvTranspose"(%arg0, %arg1) {torch.onnx.output_shape = [8 : si64, 5 : si64], torch.onnx.pads = [1 : si64, 1 : si64, 1 : si64, 1 : si64], torch.onnx.strides = [3 : si64, 2 : si64]} : (!torch.vtensor<[1,1,3,3],f32>, !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,8,5],f32>
1381+
return %0 : !torch.vtensor<[1,2,8,5],f32>
1382+
}
1383+
1384+
// -----
1385+
12991386
// CHECK-LABEL: @test_convtranspose_pad
13001387
func.func @test_convtranspose_pad(%arg0: !torch.vtensor<[1,1,3,3],f32>, %arg1: !torch.vtensor<[1,2,3,3],f32>) -> !torch.vtensor<[1,2,10,8],f32> attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 11 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
13011388
// CHECK: %[[C0:.*]] = torch.constant.int 0

0 commit comments

Comments
 (0)