Skip to content

Commit 2073784

Browse files
sahas3sahas3
authored
[TorchToTosa] Use i48 accumulator type for any i16 inputs to conv (#4583)
TOSA conv2d verifier requires the accumulator type to be i48 whenever the input element type is i16, regardless of weight type https://github.com/llvm/llvm-project/blob/4bc1cd92e41f682a5aa49fd5ca10a2002d8470ba/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp#L836. The previous helper only matched `i16-input + i8-weight + i48-output` and silently fell through to `acc_type = outputType` for other valid configurations (e.g., quantized a16/w16 models with i32 output), producing IR rejected by the verifier.
1 parent aad9cee commit 2073784

4 files changed

Lines changed: 81 additions & 2 deletions

File tree

‎lib/Conversion/TorchToTosa/TosaLegalizeUtils.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -511,8 +511,7 @@ LogicalResult getConvOpsAccType(PatternRewriter &rewriter,
511511
(weightElemTy.isInteger(8) || weightElemTy.isInteger(4)) &&
512512
outputElemTy.isInteger(32)) {
513513
accType = mlir::TypeAttr::get(rewriter.getIntegerType(32));
514-
} else if (inputElemTy.isInteger(16) && weightElemTy.isInteger(8) &&
515-
outputElemTy.isInteger(48)) {
514+
} else if (inputElemTy.isInteger(16)) {
516515
accType = mlir::TypeAttr::get(rewriter.getIntegerType(48));
517516
} else if ((isa<Float8E4M3Type>(inputElemTy) &&
518517
isa<Float8E4M3Type>(weightElemTy) && outputElemTy.isF16()) ||

‎projects/pt1/e2e_testing/xfail_sets.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -694,6 +694,7 @@
694694
"ContainsIntList_False",
695695
"ContainsIntList_True",
696696
"Conv2dFP16NoBiasModule_basic",
697+
"Conv2dQInt16Module_basic",
697698
"Conv2dQInt8Module_basic",
698699
"Conv2dQInt8Module_depthwise",
699700
"Conv2dQInt8Module_grouped",
@@ -2945,6 +2946,7 @@
29452946
"Conv2dModule_basic",
29462947
"Conv2dNoPaddingModule_basic",
29472948
"Conv2dFP16NoBiasModule_basic",
2949+
"Conv2dQInt16Module_basic",
29482950
"Conv2dQInt8Module_basic",
29492951
"Conv2dQInt8Module_depthwise",
29502952
"Conv2dQInt8Module_grouped",

‎projects/pt1/python/torch_mlir_e2e_test/test_suite/conv.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1835,6 +1835,54 @@ def Conv2dQInt8Module_not_depthwise(module, tu: TestUtils):
18351835
module.forward(inputVec, weight, bias)
18361836

18371837

1838+
class Conv2dQInt16Module(torch.nn.Module):
1839+
def __init__(self):
1840+
super().__init__()
1841+
1842+
@export
1843+
@annotate_args(
1844+
[
1845+
None,
1846+
([-1, -1, -1, -1], torch.int16, True),
1847+
([-1, -1, -1, -1], torch.int16, True),
1848+
([-1], torch.int32, True),
1849+
]
1850+
)
1851+
def forward(self, inputVec, weight, bias):
1852+
inputVec = torch.ops.quantized_decomposed.dequantize_per_tensor.default(
1853+
inputVec, 0.01, 0, -(2**15), 2**15 - 1, torch.int16
1854+
)
1855+
weight = torch.ops.quantized_decomposed.dequantize_per_tensor.default(
1856+
weight, 0.01, 0, -(2**15), 2**15 - 1, torch.int16
1857+
)
1858+
bias = torch.ops.quantized_decomposed.dequantize_per_tensor.default(
1859+
bias, 1e-4, 0, -(2**31), 2**31 - 1, torch.int32
1860+
)
1861+
1862+
conv = torch.ops.aten.conv2d(
1863+
inputVec,
1864+
weight,
1865+
bias=bias,
1866+
stride=[1, 1],
1867+
padding=[0, 0],
1868+
dilation=[1, 1],
1869+
groups=1,
1870+
)
1871+
1872+
# Use int32 to avoid overflows.
1873+
return torch.ops.quantized_decomposed.quantize_per_tensor.default(
1874+
conv, 1e-4, 0, -(2**31), 2**31 - 1, torch.int32
1875+
)
1876+
1877+
1878+
@register_test_case(module_factory=lambda: Conv2dQInt16Module())
1879+
def Conv2dQInt16Module_basic(module, tu: TestUtils):
1880+
inputVec = tu.randint(2, 4, 7, 8, low=-128, high=127).to(torch.int16)
1881+
weight = tu.randint(3, 4, 3, 2, low=-128, high=127).to(torch.int16)
1882+
bias = tu.randint(3, low=-1000, high=1000).to(torch.int32)
1883+
module.forward(inputVec, weight, bias)
1884+
1885+
18381886
class ConvTranspose2DQInt8Module(torch.nn.Module):
18391887

18401888
def __init__(self):

‎test/Conversion/TorchToTosa/quantization.mlir‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,3 +139,33 @@ func.func @quantized_conv(%arg0: !torch.vtensor<[?,4,7,8],si8>, %arg1: !torch.vt
139139
%11 = torch.aten.dequantize.tensor %10 : !torch.vtensor<[?,3,5,7],!torch.qint32> -> !torch.vtensor<[?,3,5,7],f32>
140140
return %11 : !torch.vtensor<[?,3,5,7],f32>
141141
}
142+
143+
// -----
144+
145+
// CHECK-LABEL: func.func @quantized_conv_i16(
146+
// CHECK: %[[CONV:.*]] = tosa.conv2d
147+
// CHECK-SAME: {acc_type = i48,
148+
func.func @quantized_conv_i16(%arg0: !torch.vtensor<[?,4,7,8],si16>, %arg1: !torch.vtensor<[3,4,3,2],si16>, %arg2: !torch.vtensor<[?],f32>) -> !torch.vtensor<[?,3,5,7],f32> {
149+
%false = torch.constant.bool false
150+
%int1 = torch.constant.int 1
151+
%int0 = torch.constant.int 0
152+
%float1.000000e-04 = torch.constant.float 1.000000e-04
153+
%int3 = torch.constant.int 3
154+
%int7 = torch.constant.int 7
155+
%float1.000000e-02 = torch.constant.float 1.000000e-02
156+
%int14 = torch.constant.int 14
157+
%0 = torch.aten.quantize_per_tensor %arg2, %float1.000000e-04, %int0, %int14 : !torch.vtensor<[?],f32>, !torch.float, !torch.int, !torch.int -> !torch.vtensor<[?],!torch.qint32>
158+
%1 = torch.aten.dequantize.self %0 : !torch.vtensor<[?],!torch.qint32> -> !torch.vtensor<[?],f32>
159+
%2 = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<int>
160+
%3 = torch.prim.ListConstruct %int0, %int0 : (!torch.int, !torch.int) -> !torch.list<int>
161+
%4 = torch.prim.ListConstruct : () -> !torch.list<int>
162+
// TOSA spec requires zero-point = 0 for non-int8 integer convs.
163+
%5 = torch.aten._make_per_tensor_quantized_tensor %arg0, %float1.000000e-02, %int0 : !torch.vtensor<[?,4,7,8],si16>, !torch.float, !torch.int -> !torch.vtensor<[?,4,7,8],!torch.qint16>
164+
%6 = torch.aten._make_per_tensor_quantized_tensor %arg1, %float1.000000e-02, %int0 : !torch.vtensor<[3,4,3,2],si16>, !torch.float, !torch.int -> !torch.vtensor<[3,4,3,2],!torch.qint16>
165+
%7 = torch.aten.quantize_per_tensor %1, %float1.000000e-04, %int0, %int14 : !torch.vtensor<[?],f32>, !torch.float, !torch.int, !torch.int -> !torch.vtensor<[?],!torch.qint32>
166+
%8 = torch.aten.int_repr %7 : !torch.vtensor<[?],!torch.qint32> -> !torch.vtensor<[?],si32>
167+
%9 = torch.aten.convolution %5, %6, %8, %2, %3, %2, %false, %4, %int1 : !torch.vtensor<[?,4,7,8],!torch.qint16>, !torch.vtensor<[3,4,3,2],!torch.qint16>, !torch.vtensor<[?],si32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.list<int>, !torch.int -> !torch.vtensor<[?,3,5,7],si32>
168+
%10 = torch.aten._make_per_tensor_quantized_tensor %9, %float1.000000e-04, %int0 : !torch.vtensor<[?,3,5,7],si32>, !torch.float, !torch.int -> !torch.vtensor<[?,3,5,7],!torch.qint32>
169+
%11 = torch.aten.dequantize.tensor %10 : !torch.vtensor<[?,3,5,7],!torch.qint32> -> !torch.vtensor<[?,3,5,7],f32>
170+
return %11 : !torch.vtensor<[?,3,5,7],f32>
171+
}

0 commit comments

Comments
 (0)