Skip to content

Commit 0029fc2

Browse files
TessilTessil
authored
[TorchToTosa] Add direct aten.addmm legalization (#4712)
Lower addmm directly over decomposition to retain the rank-3 TOSA matmul result until the final reshape, avoiding unnecessary intermediates.
1 parent d622902 commit 0029fc2

2 files changed

Lines changed: 257 additions & 4 deletions

File tree

‎lib/Conversion/TorchToTosa/TorchToTosa.cpp‎

Lines changed: 142 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2369,10 +2369,19 @@ class ConvertAtenMatmulBaseOp : public TorchToTosaOpConversionPattern<AtenOpT> {
23692369
const TypeConverter *typeConverter) const override {
23702370
if constexpr (!std::is_same_v<AtenOpT, AtenMatmulOp> &&
23712371
!std::is_same_v<AtenOpT, AtenMmOp> &&
2372-
!std::is_same_v<AtenOpT, AtenBmmOp>) {
2372+
!std::is_same_v<AtenOpT, AtenBmmOp> &&
2373+
!std::is_same_v<AtenOpT, AtenAddmmOp>) {
23732374
return false;
23742375
} else {
2375-
auto lhs = adaptor.getSelf();
2376+
Value lhs;
2377+
if constexpr (std::is_same_v<AtenOpT, AtenAddmmOp>) {
2378+
lhs = adaptor.getMat1();
2379+
auto biasTy = dyn_cast<RankedTensorType>(adaptor.getSelf().getType());
2380+
if (biasTy && mlir::tosa::typeHasZeroDim(biasTy))
2381+
return false;
2382+
} else {
2383+
lhs = adaptor.getSelf();
2384+
}
23762385
Value rhs;
23772386
if constexpr (std::is_same_v<AtenOpT, AtenMatmulOp>)
23782387
rhs = adaptor.getOther();
@@ -2393,10 +2402,13 @@ class ConvertAtenMatmulBaseOp : public TorchToTosaOpConversionPattern<AtenOpT> {
23932402
}
23942403
}
23952404

2405+
// When keepRank3Result is true, return the native rank-3 TOSA matmul result;
2406+
// the caller is responsible for reshaping it to the operation's result shape.
23962407
LogicalResult performMatmul(AtenOpT op, OpAdaptor adaptor,
23972408
ConversionPatternRewriter &rewriter, Value &lhs,
23982409
Value &rhs, Value &lhsZp, Value &rhsZp,
2399-
Value &output) const {
2410+
Value &output,
2411+
bool keepRank3Result = false) const {
24002412

24012413
auto lhsTy = cast<RankedTensorType>(lhs.getType());
24022414
auto rhsTy = cast<RankedTensorType>(rhs.getType());
@@ -2842,7 +2854,8 @@ class ConvertAtenMatmulBaseOp : public TorchToTosaOpConversionPattern<AtenOpT> {
28422854
// Perform the reshape to output shape. This is always required unless max
28432855
// input rank=3 and there was no broadcasting, in which case the tosa.matmul
28442856
// output itself is correctly shaped.
2845-
bool performOpReshape = !(maxInputRank == 3 && !performBatchDimBroadcast);
2857+
bool performOpReshape =
2858+
!(maxInputRank == 3 && !performBatchDimBroadcast) && !keepRank3Result;
28462859

28472860
if (performOpReshape) {
28482861
// Since the output shape may be unknown, we construct it
@@ -3101,6 +3114,127 @@ class ConvertAtenMmOp : public ConvertAtenMatmulBaseOp<AtenOpT> {
31013114
}
31023115
};
31033116

3117+
// Lowers statically shaped floating-point addmm while retaining the rank-3
3118+
// matmul result so that the bias add can be done before the reshape.
3119+
class ConvertAtenAddmmOp : public ConvertAtenMatmulBaseOp<AtenAddmmOp> {
3120+
static bool isStaticRanked(RankedTensorType type, int64_t rank) {
3121+
return type && type.hasStaticShape() && type.getRank() == rank;
3122+
}
3123+
3124+
static std::optional<double> getConstantScalar(Value value) {
3125+
double floatValue;
3126+
if (matchPattern(value, m_TorchConstantFloat(&floatValue))) {
3127+
return floatValue;
3128+
}
3129+
3130+
int64_t intValue;
3131+
if (matchPattern(value, m_TorchConstantInt(&intValue))) {
3132+
return static_cast<double>(intValue);
3133+
}
3134+
return std::nullopt;
3135+
}
3136+
3137+
static FailureOr<Value> scaleTensor(AtenAddmmOp op, Value tensor,
3138+
Value scalar, double scalarValue,
3139+
ConversionPatternRewriter &rewriter) {
3140+
if (scalarValue == 1.0) {
3141+
return tensor;
3142+
}
3143+
3144+
auto tensorTy = cast<RankedTensorType>(tensor.getType());
3145+
SmallVector<int64_t> scalarShape(tensorTy.getRank(), 1);
3146+
Value scalarTensor;
3147+
if (failed(torchScalarToTosaTensor(rewriter, op, scalar, scalarTensor,
3148+
tensorTy.getElementType(),
3149+
scalarShape))) {
3150+
return failure();
3151+
}
3152+
return tosa::createMulOpAndCast(rewriter, op, tensorTy, tensor,
3153+
scalarTensor, /*shift=*/0)
3154+
.getResult();
3155+
}
3156+
3157+
public:
3158+
using ConvertAtenMatmulBaseOp<AtenAddmmOp>::ConvertAtenMatmulBaseOp;
3159+
using OpAdaptor = AtenAddmmOp::Adaptor;
3160+
3161+
LogicalResult
3162+
matchAndRewriteImpl(AtenAddmmOp op, OpAdaptor adaptor,
3163+
ConversionPatternRewriter &rewriter) const override {
3164+
Value lhs = adaptor.getMat1();
3165+
Value rhs = adaptor.getMat2();
3166+
Value bias = adaptor.getSelf();
3167+
auto lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
3168+
auto rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
3169+
auto biasTy = dyn_cast<RankedTensorType>(bias.getType());
3170+
auto resultTy = dyn_cast<RankedTensorType>(
3171+
this->getTypeConverter()->convertType(op.getType()));
3172+
3173+
std::optional<double> alpha = getConstantScalar(op.getAlpha());
3174+
std::optional<double> beta = getConstantScalar(op.getBeta());
3175+
if (!alpha || !beta) {
3176+
return rewriter.notifyMatchFailure(op,
3177+
"requires constant alpha and beta");
3178+
}
3179+
if (!isStaticRanked(lhsTy, 2) || !isStaticRanked(rhsTy, 2) ||
3180+
!isStaticRanked(resultTy, 2)) {
3181+
return rewriter.notifyMatchFailure(
3182+
op, "requires static rank-2 matrices and result");
3183+
}
3184+
if (!biasTy || !biasTy.hasStaticShape() || biasTy.getRank() > 2) {
3185+
return rewriter.notifyMatchFailure(
3186+
op, "requires static broadcastable bias of rank at most 2");
3187+
}
3188+
if (!isa<FloatType>(lhsTy.getElementType()) ||
3189+
!isa<FloatType>(rhsTy.getElementType()) ||
3190+
!isa<FloatType>(biasTy.getElementType()) ||
3191+
!isa<FloatType>(resultTy.getElementType())) {
3192+
return rewriter.notifyMatchFailure(op, "requires floating-point tensors");
3193+
}
3194+
3195+
Value lhsZp, rhsZp, matmul;
3196+
if (failed(this->performMatmul(op, adaptor, rewriter, lhs, rhs, lhsZp,
3197+
rhsZp, matmul,
3198+
/*keepRank3Result=*/true))) {
3199+
return rewriter.notifyMatchFailure(op, "failed to lower addmm matmul");
3200+
}
3201+
3202+
auto matmulTy = cast<RankedTensorType>(matmul.getType());
3203+
matmul = tosa::tosaCastTensorToType(
3204+
rewriter, matmul, matmulTy.clone(resultTy.getElementType()))
3205+
.value();
3206+
3207+
FailureOr<Value> scaledMatmul =
3208+
scaleTensor(op, matmul, op.getAlpha(), *alpha, rewriter);
3209+
if (failed(scaledMatmul)) {
3210+
return rewriter.notifyMatchFailure(op, "failed to apply addmm alpha");
3211+
}
3212+
3213+
Value result = *scaledMatmul;
3214+
if (*beta != 0.0) {
3215+
FailureOr<Value> scaledBias =
3216+
scaleTensor(op, bias, op.getBeta(), *beta, rewriter);
3217+
if (failed(scaledBias)) {
3218+
return rewriter.notifyMatchFailure(op, "failed to apply addmm beta");
3219+
}
3220+
bias = *scaledBias;
3221+
3222+
if (failed(tosa::EqualizeRanks(rewriter, op.getLoc(), result, bias))) {
3223+
return rewriter.notifyMatchFailure(
3224+
op, "failed to broadcast bias to addmm result");
3225+
}
3226+
result = tosa::AddOp::create(rewriter, op.getLoc(), result.getType(),
3227+
result, bias)
3228+
.getResult();
3229+
}
3230+
3231+
rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
3232+
op, resultTy, result,
3233+
tosa::getTosaConstShape(rewriter, op.getLoc(), resultTy.getShape()));
3234+
return success();
3235+
}
3236+
};
3237+
31043238
// Implements handling of aten.linear op.
31053239
template <typename AtenOpT>
31063240
class ConvertAtenLinearOp : public ConvertAtenMatmulBaseOp<AtenOpT> {
@@ -11878,6 +12012,10 @@ std::set<StringRef> populateTorchToTosaConversionPatternsAndIllegalOps(
1187812012
INSERT_MM_ATENOP_PATTERN(AtenBmmOp);
1187912013
#undef INSERT_MM_ATENOP_PATTERN
1188012014

12015+
illegalOps.insert(AtenAddmmOp::getOperationName());
12016+
patterns.addWithLabel<ConvertAtenAddmmOp>(AtenAddmmOp::getOperationName(),
12017+
typeConverter, context);
12018+
1188112019
#define INSERT_LINEAR_ATENOP_PATTERN(AtenOp) \
1188212020
illegalOps.insert(AtenOp::getOperationName()); \
1188312021
patterns.addWithLabel<ConvertAtenLinearOp<AtenOp>>( \

‎test/Conversion/TorchToTosa/basic.mlir‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5520,6 +5520,121 @@ func.func @torch.aten.mm$f32(%arg0: !torch.vtensor<[1,22],f32>, %arg1: !torch.vt
55205520
return %0 : !torch.vtensor<[1,10],f32>
55215521
}
55225522

5523+
// -----
5524+
// CHECK-LABEL: func.func @torch.aten.addmm$f32
5525+
// CHECK: %[[MATMUL:.*]] = tosa.matmul
5526+
// CHECK-SAME: -> tensor<1x6x4xf32>
5527+
// CHECK: %[[ADD:.*]] = tosa.add
5528+
// CHECK-SAME: (tensor<1x6x4xf32>, tensor<1x1x4xf32>) -> tensor<1x6x4xf32>
5529+
// CHECK: %[[RESULT:.*]] = tosa.reshape %[[ADD]]
5530+
// CHECK-SAME: -> tensor<6x4xf32>
5531+
// CHECK-NOT: torch.aten.addmm
5532+
func.func @torch.aten.addmm$f32(%bias: !torch.vtensor<[4],f32>, %mat1: !torch.vtensor<[6,8],f32>, %mat2: !torch.vtensor<[8,4],f32>) -> !torch.vtensor<[6,4],f32> {
5533+
%one = torch.constant.int 1
5534+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %one, %one : !torch.vtensor<[4],f32>, !torch.vtensor<[6,8],f32>, !torch.vtensor<[8,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[6,4],f32>
5535+
return %0 : !torch.vtensor<[6,4],f32>
5536+
}
5537+
5538+
// -----
5539+
// CHECK-LABEL: func.func @torch.aten.addmm$scaled_f32
5540+
// CHECK: %[[MATMUL:.*]] = tosa.matmul
5541+
// CHECK-SAME: -> tensor<1x6x4xf32>
5542+
// CHECK-DAG: %[[BETA:.*]] = tosa.mul {{.*}} -> tensor<4xf32>
5543+
// CHECK-DAG: %[[ALPHA:.*]] = tosa.mul %[[MATMUL]]{{.*}} -> tensor<1x6x4xf32>
5544+
// CHECK: %[[ADD:.*]] = tosa.add
5545+
// CHECK-SAME: -> tensor<1x6x4xf32>
5546+
// CHECK: tosa.reshape %[[ADD]]
5547+
// CHECK-SAME: -> tensor<6x4xf32>
5548+
// CHECK-NOT: torch.aten.addmm
5549+
func.func @torch.aten.addmm$scaled_f32(%bias: !torch.vtensor<[4],f32>, %mat1: !torch.vtensor<[6,8],f32>, %mat2: !torch.vtensor<[8,4],f32>) -> !torch.vtensor<[6,4],f32> {
5550+
%two = torch.constant.int 2
5551+
%three = torch.constant.int 3
5552+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %three, %two : !torch.vtensor<[4],f32>, !torch.vtensor<[6,8],f32>, !torch.vtensor<[8,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[6,4],f32>
5553+
return %0 : !torch.vtensor<[6,4],f32>
5554+
}
5555+
5556+
// -----
5557+
// CHECK-LABEL: func.func @torch.aten.addmm$float_scalars_scalar_bias
5558+
// CHECK: %[[MATMUL:.*]] = tosa.matmul
5559+
// CHECK-SAME: -> tensor<1x6x4xf32>
5560+
// CHECK: tosa.mul %[[MATMUL]]
5561+
// CHECK-SAME: -> tensor<1x6x4xf32>
5562+
// CHECK: tosa.mul
5563+
// CHECK-SAME: -> tensor<f32>
5564+
// CHECK: %[[ADD:.*]] = tosa.add
5565+
// CHECK-SAME: -> tensor<1x6x4xf32>
5566+
// CHECK: tosa.reshape %[[ADD]]
5567+
// CHECK-SAME: -> tensor<6x4xf32>
5568+
// CHECK-NOT: torch.aten.addmm
5569+
func.func @torch.aten.addmm$float_scalars_scalar_bias(%bias: !torch.vtensor<[],f32>, %mat1: !torch.vtensor<[6,8],f32>, %mat2: !torch.vtensor<[8,4],f32>) -> !torch.vtensor<[6,4],f32> {
5570+
%half = torch.constant.float 5.000000e-01
5571+
%two = torch.constant.float 2.000000e+00
5572+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %two, %half : !torch.vtensor<[],f32>, !torch.vtensor<[6,8],f32>, !torch.vtensor<[8,4],f32>, !torch.float, !torch.float -> !torch.vtensor<[6,4],f32>
5573+
return %0 : !torch.vtensor<[6,4],f32>
5574+
}
5575+
5576+
// -----
5577+
// CHECK-LABEL: func.func @torch.aten.addmm$beta_zero_f32
5578+
// CHECK: %[[MATMUL:.*]] = tosa.matmul
5579+
// CHECK-SAME: -> tensor<1x6x4xf32>
5580+
// CHECK-NOT: tosa.mul
5581+
// CHECK-NOT: tosa.add
5582+
// CHECK: tosa.reshape %[[MATMUL]]
5583+
// CHECK-SAME: -> tensor<6x4xf32>
5584+
// CHECK-NOT: torch.aten.addmm
5585+
func.func @torch.aten.addmm$beta_zero_f32(%bias: !torch.vtensor<[4],f32>, %mat1: !torch.vtensor<[6,8],f32>, %mat2: !torch.vtensor<[8,4],f32>) -> !torch.vtensor<[6,4],f32> {
5586+
%zero = torch.constant.int 0
5587+
%one = torch.constant.int 1
5588+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %zero, %one : !torch.vtensor<[4],f32>, !torch.vtensor<[6,8],f32>, !torch.vtensor<[8,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[6,4],f32>
5589+
return %0 : !torch.vtensor<[6,4],f32>
5590+
}
5591+
5592+
// -----
5593+
// CHECK-LABEL: func.func @torch.aten.addmm$f16
5594+
// CHECK: %[[MATMUL:.*]] = tosa.matmul
5595+
// CHECK-SAME: -> tensor<1x6x4xf32>
5596+
// CHECK: %[[CAST:.*]] = tosa.cast %[[MATMUL]]
5597+
// CHECK-SAME: -> tensor<1x6x4xf16>
5598+
// CHECK: %[[ADD:.*]] = tosa.add
5599+
// CHECK-SAME: -> tensor<1x6x4xf16>
5600+
// CHECK: tosa.reshape %[[ADD]]
5601+
// CHECK-SAME: -> tensor<6x4xf16>
5602+
// CHECK-NOT: torch.aten.addmm
5603+
func.func @torch.aten.addmm$f16(%bias: !torch.vtensor<[4],f16>, %mat1: !torch.vtensor<[6,8],f16>, %mat2: !torch.vtensor<[8,4],f16>) -> !torch.vtensor<[6,4],f16> {
5604+
%one = torch.constant.int 1
5605+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %one, %one : !torch.vtensor<[4],f16>, !torch.vtensor<[6,8],f16>, !torch.vtensor<[8,4],f16>, !torch.int, !torch.int -> !torch.vtensor<[6,4],f16>
5606+
return %0 : !torch.vtensor<[6,4],f16>
5607+
}
5608+
5609+
// -----
5610+
// CHECK-LABEL: func.func @torch.aten.addmm$rank2_bias
5611+
// CHECK: %[[MATMUL:.*]] = tosa.matmul
5612+
// CHECK-SAME: -> tensor<1x6x4xf32>
5613+
// CHECK: %[[ADD:.*]] = tosa.add
5614+
// CHECK-SAME: -> tensor<1x6x4xf32>
5615+
// CHECK: tosa.reshape %[[ADD]]
5616+
// CHECK-SAME: -> tensor<6x4xf32>
5617+
// CHECK-NOT: torch.aten.addmm
5618+
func.func @torch.aten.addmm$rank2_bias(%bias: !torch.vtensor<[6,4],f32>, %mat1: !torch.vtensor<[6,8],f32>, %mat2: !torch.vtensor<[8,4],f32>) -> !torch.vtensor<[6,4],f32> {
5619+
%one = torch.constant.int 1
5620+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %one, %one : !torch.vtensor<[6,4],f32>, !torch.vtensor<[6,8],f32>, !torch.vtensor<[8,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[6,4],f32>
5621+
return %0 : !torch.vtensor<[6,4],f32>
5622+
}
5623+
5624+
// -----
5625+
// CHECK-LABEL: func.func @torch.aten.addmm$zero_k
5626+
// CHECK-NOT: tosa.matmul
5627+
// CHECK: %[[ZERO:.*]] = "tosa.const"()
5628+
// CHECK-SAME: tensor<6x4xf32>
5629+
// CHECK: %[[ADD:.*]] = tosa.add
5630+
// CHECK-SAME: (tensor<6x4xf32>, tensor<1x4xf32>) -> tensor<6x4xf32>
5631+
// CHECK-NOT: torch.aten.addmm
5632+
func.func @torch.aten.addmm$zero_k(%bias: !torch.vtensor<[4],f32>, %mat1: !torch.vtensor<[6,0],f32>, %mat2: !torch.vtensor<[0,4],f32>) -> !torch.vtensor<[6,4],f32> {
5633+
%one = torch.constant.int 1
5634+
%0 = torch.aten.addmm %bias, %mat1, %mat2, %one, %one : !torch.vtensor<[4],f32>, !torch.vtensor<[6,0],f32>, !torch.vtensor<[0,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[6,4],f32>
5635+
return %0 : !torch.vtensor<[6,4],f32>
5636+
}
5637+
55235638
// -----
55245639
// CHECK-LABEL: func.func @torch.aten.mm$si8
55255640
// CHECK: tosa.matmul

0 commit comments

Comments
 (0)