Skip to content

Commit 41b3cd4

Browse files
catcor01catcor01
authored
[TorchToTosa] Lower zero-K matmul to zero const (#4595)
## Summary - allow Torch-to-TOSA matmul lowering to handle static zero contraction when the result shape is non-empty - lower those cases to a typed zero `tosa.const` instead of rejecting the zero-sized input - keep zero-sized output tensors rejected as unsupported by TOSA
1 parent 1e2ef9f commit 41b3cd4

4 files changed

Lines changed: 220 additions & 2 deletions

File tree

‎lib/Conversion/TorchToTosa/TorchToTosa.cpp‎

Lines changed: 72 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,11 +278,16 @@ class TorchToTosaOpConversionPattern : public OpConversionPattern<AtenOpT> {
278278
LogicalResult
279279
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
280280
ConversionPatternRewriter &rewriter) const final {
281+
const TypeConverter *typeConverter = this->getTypeConverter();
282+
bool canHandleZeroDimInputOperands =
283+
canHandleZeroDimInputs(op, adaptor, typeConverter);
284+
281285
// Pre-check: all tensor operands and outputs must have no zero-sized
282286
// dimensions.
283287
for (auto v : adaptor.getOperands()) {
284288
auto rankedInputType = dyn_cast<RankedTensorType>(v.getType());
285-
if (rankedInputType && mlir::tosa::typeHasZeroDim(rankedInputType)) {
289+
if (rankedInputType && mlir::tosa::typeHasZeroDim(rankedInputType) &&
290+
!canHandleZeroDimInputOperands) {
286291
return rewriter.notifyMatchFailure(
287292
op,
288293
"TOSA lowering does not support input tensors with a zero-sized "
@@ -292,7 +297,6 @@ class TorchToTosaOpConversionPattern : public OpConversionPattern<AtenOpT> {
292297

293298
// not all adaptors have results, instead get the result from the op
294299
// directly
295-
const TypeConverter *typeConverter = this->getTypeConverter();
296300
for (auto res : op->getResults()) {
297301
auto rankedOutputType =
298302
dyn_cast<RankedTensorType>(typeConverter->convertType(res.getType()));
@@ -307,6 +311,12 @@ class TorchToTosaOpConversionPattern : public OpConversionPattern<AtenOpT> {
307311
}
308312

309313
protected:
314+
virtual bool
315+
canHandleZeroDimInputs(AtenOpT op, OpAdaptor adaptor,
316+
const TypeConverter *typeConverter) const {
317+
return false;
318+
}
319+
310320
virtual LogicalResult
311321
matchAndRewriteImpl(AtenOpT op, OpAdaptor adaptor,
312322
ConversionPatternRewriter &rewriter) const = 0;
@@ -1795,6 +1805,36 @@ class ConvertAtenMatmulBaseOp : public TorchToTosaOpConversionPattern<AtenOpT> {
17951805
op,
17961806
"Unimplemented matrix multiplication variant input parsing function");
17971807
}
1808+
1809+
bool
1810+
canHandleZeroDimInputs(AtenOpT op, OpAdaptor adaptor,
1811+
const TypeConverter *typeConverter) const override {
1812+
if constexpr (!std::is_same_v<AtenOpT, AtenMatmulOp> &&
1813+
!std::is_same_v<AtenOpT, AtenMmOp> &&
1814+
!std::is_same_v<AtenOpT, AtenBmmOp>) {
1815+
return false;
1816+
} else {
1817+
auto lhs = adaptor.getSelf();
1818+
Value rhs;
1819+
if constexpr (std::is_same_v<AtenOpT, AtenMatmulOp>)
1820+
rhs = adaptor.getOther();
1821+
else
1822+
rhs = adaptor.getMat2();
1823+
1824+
auto lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
1825+
auto rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
1826+
auto resultTy =
1827+
dyn_cast<RankedTensorType>(typeConverter->convertType(op.getType()));
1828+
if (!lhsTy || !rhsTy || !resultTy)
1829+
return false;
1830+
if (!resultTy.hasStaticShape())
1831+
return false;
1832+
if (mlir::tosa::typeHasZeroDim(resultTy))
1833+
return false;
1834+
return hasStaticZeroContraction(lhsTy, rhsTy);
1835+
}
1836+
}
1837+
17981838
LogicalResult performMatmul(AtenOpT op, OpAdaptor adaptor,
17991839
ConversionPatternRewriter &rewriter, Value &lhs,
18001840
Value &rhs, Value &lhsZp, Value &rhsZp,
@@ -1816,6 +1856,20 @@ class ConvertAtenMatmulBaseOp : public TorchToTosaOpConversionPattern<AtenOpT> {
18161856
return rewriter.notifyMatchFailure(op,
18171857
"Matmul: input datatypes mismatched");
18181858

1859+
auto resultTy = dyn_cast<RankedTensorType>(
1860+
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
1861+
op.getType()));
1862+
if (resultTy && resultTy.hasStaticShape() &&
1863+
!mlir::tosa::typeHasZeroDim(resultTy) &&
1864+
hasStaticZeroContraction(lhsTy, rhsTy)) {
1865+
auto zeroOutput = tosa::getZerosLikeTensor(rewriter, op, resultTy);
1866+
if (!zeroOutput)
1867+
return rewriter.notifyMatchFailure(
1868+
op, "failed to materialize zero-contraction matmul result");
1869+
output = *zeroOutput;
1870+
return success();
1871+
}
1872+
18191873
if (!lhsZp) {
18201874
// Initialize zero constant values as zero-points, if the op operands
18211875
// aren't quantized types
@@ -2352,6 +2406,22 @@ class ConvertAtenMatmulBaseOp : public TorchToTosaOpConversionPattern<AtenOpT> {
23522406

23532407
return success();
23542408
}
2409+
2410+
private:
2411+
static bool hasStaticZeroContraction(RankedTensorType lhsTy,
2412+
RankedTensorType rhsTy) {
2413+
auto lhsShape = makeShapeTorchCompatible(lhsTy.getShape());
2414+
auto rhsShape = makeShapeTorchCompatible(rhsTy.getShape());
2415+
if (lhsShape.empty() || rhsShape.empty())
2416+
return false;
2417+
2418+
int64_t lhsK = lhsShape.back();
2419+
int64_t rhsK =
2420+
rhsShape.size() == 1 ? rhsShape.back() : rhsShape[rhsShape.size() - 2];
2421+
return lhsK == 0 && rhsK == 0;
2422+
}
2423+
2424+
public:
23552425
// The default version just reads two inputs, computes output and returns it.
23562426
// Other versions may add a bias, apply GEMM-style alpha/beta scaling etc.
23572427
virtual LogicalResult

‎projects/pt1/e2e_testing/xfail_sets.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
"SliceCopyStartGreaterThanDimSize_Module_basic",
9191
# unimplemented: for conversion to byte or char type dstOriginalDtype has to be passed to convertScalarToDtype
9292
"AtenMmInt8Types_basic",
93+
"AtenMmInt8ZeroK_basic",
9394
# Hanging tests:
9495
"ConvolutionBackwardModule2DDilated_basic",
9596
"ConvolutionBackwardModule2DStridedPaddedDilatedGrouped_basic",

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,30 @@ def Matmul_2d(module, tu: TestUtils):
8484
# ==============================================================================
8585

8686

87+
class MatmulZeroK(torch.nn.Module):
88+
def __init__(self):
89+
super().__init__()
90+
91+
@export
92+
@annotate_args(
93+
[
94+
None,
95+
([5, 0], torch.float32, True),
96+
([0, 10], torch.float32, True),
97+
]
98+
)
99+
def forward(self, lhs, rhs):
100+
return torch.matmul(lhs, rhs)
101+
102+
103+
@register_test_case(module_factory=lambda: MatmulZeroK())
104+
def MatmulZeroK_basic(module, tu: TestUtils):
105+
module.forward(torch.empty(5, 0), torch.empty(0, 10))
106+
107+
108+
# ==============================================================================
109+
110+
87111
class MatmulVecMat(torch.nn.Module):
88112
def __init__(self):
89113
super().__init__()
@@ -342,6 +366,48 @@ def AtenMmFloatTypes_basic(module, tu: TestUtils):
342366
# ==============================================================================
343367

344368

369+
class AtenMmZeroK(torch.nn.Module):
370+
@export
371+
@annotate_args(
372+
[
373+
None,
374+
([5, 0], torch.float32, True),
375+
([0, 10], torch.float32, True),
376+
]
377+
)
378+
def forward(self, a, b):
379+
return torch.ops.aten.mm(a, b)
380+
381+
382+
@register_test_case(module_factory=lambda: AtenMmZeroK())
383+
def AtenMmZeroK_basic(module, tu: TestUtils):
384+
module.forward(torch.empty(5, 0), torch.empty(0, 10))
385+
386+
387+
# ==============================================================================
388+
389+
390+
class AtenBmmZeroK(torch.nn.Module):
391+
@export
392+
@annotate_args(
393+
[
394+
None,
395+
([2, 5, 0], torch.float32, True),
396+
([2, 0, 10], torch.float32, True),
397+
]
398+
)
399+
def forward(self, a, b):
400+
return torch.ops.aten.bmm(a, b)
401+
402+
403+
@register_test_case(module_factory=lambda: AtenBmmZeroK())
404+
def AtenBmmZeroK_basic(module, tu: TestUtils):
405+
module.forward(torch.empty(2, 5, 0), torch.empty(2, 0, 10))
406+
407+
408+
# ==============================================================================
409+
410+
345411
class AtenMmIntTypes(torch.nn.Module):
346412
@export
347413
@annotate_args(
@@ -387,6 +453,30 @@ def AtenMmInt8Types_basic(module, tu: TestUtils):
387453
# ==============================================================================
388454

389455

456+
class AtenMmInt8ZeroK(torch.nn.Module):
457+
@export
458+
@annotate_args(
459+
[
460+
None,
461+
([3, 0], torch.int8, True),
462+
([0, 3], torch.int8, True),
463+
]
464+
)
465+
def forward(self, a, b):
466+
return torch.ops.aten.mm(a, b)
467+
468+
469+
@register_test_case(module_factory=lambda: AtenMmInt8ZeroK())
470+
def AtenMmInt8ZeroK_basic(module, tu: TestUtils):
471+
module.forward(
472+
torch.empty(3, 0, dtype=torch.int8),
473+
torch.empty(0, 3, dtype=torch.int8),
474+
)
475+
476+
477+
# ==============================================================================
478+
479+
390480
class AtenMmF16Types(torch.nn.Module):
391481
@export
392482
@annotate_args(

‎test/Conversion/TorchToTosa/basic.mlir‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4877,6 +4877,63 @@ func.func @torch.aten.mm$bf16(%arg0: !torch.vtensor<[1,22],bf16>, %arg1: !torch.
48774877
return %4 : !torch.vtensor<[1,10],bf16>
48784878
}
48794879

4880+
// -----
4881+
// CHECK-LABEL: func.func @torch.aten.mm$zero_k_f32(
4882+
// CHECK-SAME: %[[LHS:.*]]: !torch.vtensor<[5,0],f32>,
4883+
// CHECK-SAME: %[[RHS:.*]]: !torch.vtensor<[0,10],f32>) -> !torch.vtensor<[5,10],f32> {
4884+
// CHECK-NOT: tosa.matmul
4885+
// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<5x10xf32>}> : () -> tensor<5x10xf32>
4886+
// CHECK: %[[RES:.*]] = torch_c.from_builtin_tensor %[[ZERO]] : tensor<5x10xf32> -> !torch.vtensor<[5,10],f32>
4887+
// CHECK: return %[[RES]]
4888+
func.func @torch.aten.mm$zero_k_f32(%arg0: !torch.vtensor<[5,0],f32>, %arg1: !torch.vtensor<[0,10],f32>) -> !torch.vtensor<[5,10],f32> {
4889+
%0 = torch.aten.mm %arg0, %arg1 : !torch.vtensor<[5,0],f32>, !torch.vtensor<[0,10],f32> -> !torch.vtensor<[5,10],f32>
4890+
return %0 : !torch.vtensor<[5,10],f32>
4891+
}
4892+
4893+
// -----
4894+
// CHECK-LABEL: func.func @torch.aten.matmul$zero_k_f32(
4895+
// CHECK-SAME: %[[LHS:.*]]: !torch.vtensor<[5,0],f32>,
4896+
// CHECK-SAME: %[[RHS:.*]]: !torch.vtensor<[0,10],f32>) -> !torch.vtensor<[5,10],f32> {
4897+
// CHECK-NOT: tosa.matmul
4898+
// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<5x10xf32>}> : () -> tensor<5x10xf32>
4899+
// CHECK: %[[RES:.*]] = torch_c.from_builtin_tensor %[[ZERO]] : tensor<5x10xf32> -> !torch.vtensor<[5,10],f32>
4900+
// CHECK: return %[[RES]]
4901+
func.func @torch.aten.matmul$zero_k_f32(%arg0: !torch.vtensor<[5,0],f32>, %arg1: !torch.vtensor<[0,10],f32>) -> !torch.vtensor<[5,10],f32> {
4902+
%0 = torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[5,0],f32>, !torch.vtensor<[0,10],f32> -> !torch.vtensor<[5,10],f32>
4903+
return %0 : !torch.vtensor<[5,10],f32>
4904+
}
4905+
4906+
// -----
4907+
// CHECK-LABEL: func.func @torch.aten.bmm$zero_k_f32(
4908+
// CHECK-SAME: %[[LHS:.*]]: !torch.vtensor<[2,5,0],f32>,
4909+
// CHECK-SAME: %[[RHS:.*]]: !torch.vtensor<[2,0,10],f32>) -> !torch.vtensor<[2,5,10],f32> {
4910+
// CHECK-NOT: tosa.matmul
4911+
// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<2x5x10xf32>}> : () -> tensor<2x5x10xf32>
4912+
// CHECK: %[[RES:.*]] = torch_c.from_builtin_tensor %[[ZERO]] : tensor<2x5x10xf32> -> !torch.vtensor<[2,5,10],f32>
4913+
// CHECK: return %[[RES]]
4914+
func.func @torch.aten.bmm$zero_k_f32(%arg0: !torch.vtensor<[2,5,0],f32>, %arg1: !torch.vtensor<[2,0,10],f32>) -> !torch.vtensor<[2,5,10],f32> {
4915+
%0 = torch.aten.bmm %arg0, %arg1 : !torch.vtensor<[2,5,0],f32>, !torch.vtensor<[2,0,10],f32> -> !torch.vtensor<[2,5,10],f32>
4916+
return %0 : !torch.vtensor<[2,5,10],f32>
4917+
}
4918+
4919+
// -----
4920+
module {
4921+
func.func @torch.aten.mm$zero_output_rejected(%arg0: !torch.vtensor<[0,5],f32>, %arg1: !torch.vtensor<[5,10],f32>) -> !torch.vtensor<[0,10],f32> {
4922+
// expected-error @below {{failed to legalize operation 'torch.aten.mm' that was explicitly marked illegal}}
4923+
%0 = torch.aten.mm %arg0, %arg1 : !torch.vtensor<[0,5],f32>, !torch.vtensor<[5,10],f32> -> !torch.vtensor<[0,10],f32>
4924+
return %0 : !torch.vtensor<[0,10],f32>
4925+
}
4926+
}
4927+
4928+
// -----
4929+
module {
4930+
func.func @torch.aten.mm$zero_k_dynamic_result_rejected(%arg0: !torch.vtensor<[?,0],f32>, %arg1: !torch.vtensor<[0,10],f32>) -> !torch.vtensor<[?,10],f32> {
4931+
// expected-error @below {{failed to legalize operation 'torch.aten.mm' that was explicitly marked illegal}}
4932+
%0 = torch.aten.mm %arg0, %arg1 : !torch.vtensor<[?,0],f32>, !torch.vtensor<[0,10],f32> -> !torch.vtensor<[?,10],f32>
4933+
return %0 : !torch.vtensor<[?,10],f32>
4934+
}
4935+
}
4936+
48804937
// -----
48814938
// CHECK-LABEL: func.func @torch.aten.matmul$broadcast(
48824939
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[10,3,4],f32>,

0 commit comments

Comments
 (0)