Skip to content

Commit 654142c

Browse files
tatwaichongFranklandJack
tatwaichong
authored and
FranklandJack
committed
[mlir][tosa] Make TOSA MUL's Shift an Input
The TOSA-v1.0 specification makes the shift attribute of the MUL (Hammard product) operator an input. Move the `shift` parameter of the MUL operator in the MILR TOSA dialect from an attribute to an input and update any lit tests appropriately. Expand the verifier of the `tosa::MulOp` operation to check the various constraints defined in the TOSA-v1.0 specification. Specifically, ensure that all input operands (excluding the optional shift) are of the same rank. This means that broadcasting tests which previously checked rank-0 tensors would be broadcast are no longer valid and are removed. Signed-off-by: Jack Frankland <[email protected]>
1 parent 6c5941b commit 654142c

File tree

14 files changed

+211
-83
lines changed

14 files changed

+211
-83
lines changed

‎mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -805,7 +805,7 @@ def Tosa_MulOp : Tosa_ElementwiseOp<"mul", [
805805
let arguments = (ins
806806
Tosa_Tensor:$input1,
807807
Tosa_Tensor:$input2,
808-
I8Attr:$shift
808+
Optional<TosaTensorRankOf<[Tosa_Int8], [0]>>:$shift
809809
);
810810

811811
let results = (outs

‎mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp‎

Lines changed: 56 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -90,43 +90,59 @@ static Value createLinalgBodyCalculationForElementwiseOp(
9090
}
9191

9292
// tosa::MulOp
93-
if (isa<tosa::MulOp>(op) && isa<FloatType>(elementTy))
94-
return rewriter.create<arith::MulFOp>(loc, resultTypes, args);
95-
96-
if (isa<tosa::MulOp>(op) && isa<IntegerType>(elementTy)) {
97-
Value a = args[0];
98-
Value b = args[1];
99-
auto shift =
100-
cast<IntegerAttr>(op->getAttr("shift")).getValue().getSExtValue();
101-
if (shift > 0) {
102-
auto shiftConst =
103-
rewriter.create<arith::ConstantIntOp>(loc, shift, /*bitwidth=*/8);
104-
if (!a.getType().isInteger(32))
105-
a = rewriter.create<arith::ExtSIOp>(loc, rewriter.getI32Type(), a);
106-
107-
if (!b.getType().isInteger(32))
108-
b = rewriter.create<arith::ExtSIOp>(loc, rewriter.getI32Type(), b);
109-
110-
auto result = rewriter.create<tosa::ApplyScaleOp>(
111-
loc, rewriter.getI32Type(), a, b, shiftConst,
112-
rewriter.getBoolAttr(false));
113-
114-
if (elementTy.isInteger(32))
115-
return result;
116-
117-
return rewriter.create<arith::TruncIOp>(loc, elementTy, result);
93+
if (isa<tosa::MulOp>(op)) {
94+
auto shift_val = cast<tosa::MulOp>(op).getShift();
95+
if (!elementTy.isInteger(32) && shift_val.getImpl()) {
96+
(void)rewriter.notifyMatchFailure(
97+
op, "Cannot have shift value for non i32 output");
98+
return nullptr;
99+
};
100+
101+
if (isa<FloatType>(elementTy)) {
102+
return rewriter.create<arith::MulFOp>(loc, resultTypes, args[0], args[1]);
118103
}
119104

120-
int aWidth = a.getType().getIntOrFloatBitWidth();
121-
int bWidth = b.getType().getIntOrFloatBitWidth();
122-
int cWidth = resultTypes[0].getIntOrFloatBitWidth();
105+
if (isa<IntegerType>(elementTy)) {
106+
int32_t shift = 0;
107+
ElementsAttr shift_elem;
108+
if (shift_val.getImpl() &&
109+
matchPattern(shift_val, m_Constant(&shift_elem))) {
110+
// Explicit shift is set.
111+
shift = shift_elem.getValues<IntegerAttr>()[0].getInt();
112+
}
113+
114+
Value a = args[0];
115+
Value b = args[1];
116+
if (shift > 0) {
117+
auto shiftConst =
118+
rewriter.create<arith::ConstantIntOp>(loc, shift, /*bitwidth=*/8);
119+
if (!a.getType().isInteger(32))
120+
a = rewriter.create<arith::ExtSIOp>(loc, rewriter.getI32Type(), a);
121+
122+
if (!b.getType().isInteger(32))
123+
b = rewriter.create<arith::ExtSIOp>(loc, rewriter.getI32Type(), b);
124+
125+
auto result = rewriter.create<tosa::ApplyScaleOp>(
126+
loc, rewriter.getI32Type(), a, b, shiftConst,
127+
rewriter.getBoolAttr(false));
123128

124-
if (aWidth < cWidth)
125-
a = rewriter.create<arith::ExtSIOp>(loc, resultTypes[0], a);
126-
if (bWidth < cWidth)
127-
b = rewriter.create<arith::ExtSIOp>(loc, resultTypes[0], b);
129+
if (elementTy.isInteger(32))
130+
return result;
128131

129-
return rewriter.create<arith::MulIOp>(loc, resultTypes, a, b);
132+
return rewriter.create<arith::TruncIOp>(loc, elementTy, result);
133+
}
134+
135+
int aWidth = a.getType().getIntOrFloatBitWidth();
136+
int bWidth = b.getType().getIntOrFloatBitWidth();
137+
int cWidth = resultTypes[0].getIntOrFloatBitWidth();
138+
139+
if (aWidth < cWidth)
140+
a = rewriter.create<arith::ExtSIOp>(loc, resultTypes[0], a);
141+
if (bWidth < cWidth)
142+
b = rewriter.create<arith::ExtSIOp>(loc, resultTypes[0], b);
143+
144+
return rewriter.create<arith::MulIOp>(loc, resultTypes, a, b);
145+
}
130146
}
131147

132148
// tosa::NegateOp
@@ -934,7 +950,13 @@ elementwiseMatchAndRewriteHelper(Operation *operation, ValueRange operands,
934950
auto loc = operation->getLoc();
935951
auto rank =
936952
cast<RankedTensorType>(operation->getResultTypes().front()).getRank();
937-
auto expandedOperands = expandInputRanks(rewriter, loc, operands, rank);
953+
// For the mul op we need to avoid expanding the rank of the optional shift
954+
// input.
955+
auto operandsToExpand =
956+
isa<tosa::MulOp>(operation) ? operands.take_front(2) : operands;
957+
958+
auto expandedOperands =
959+
expandInputRanks(rewriter, loc, operandsToExpand, rank);
938960
auto [targetShape, masterOperands] =
939961
computeTargetShape(rewriter, loc, indexPool, expandedOperands);
940962
auto broadcastOperands = broadcastDynamicDimensions(

‎mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -614,7 +614,18 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
614614
auto rhsAttr =
615615
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
616616

617-
const int64_t shift = llvm::isa<IntegerType>(resultETy) ? getShift() : 0;
617+
// Result right shift on i32_t data type only. For simplification, synthesize
618+
// a zero shift for other date type.
619+
int32_t shift = 0;
620+
if (resultETy.isInteger(32)) {
621+
ElementsAttr shift_elem;
622+
if (getShift().getImpl()) {
623+
if (!matchPattern(getShift(), m_Constant(&shift_elem)))
624+
// cannot be folded when the shift value is unknown.
625+
return {};
626+
shift = shift_elem.getValues<IntegerAttr>()[0].getInt();
627+
}
628+
}
618629

619630
if (rhsTy == resultTy) {
620631
if (isSplatZero(resultETy, lhsAttr))
@@ -629,7 +640,7 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
629640
return lhs;
630641
}
631642

632-
return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, getShift());
643+
return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, shift);
633644
}
634645

635646
OpFoldResult SubOp::fold(FoldAdaptor adaptor) {

‎mlir/lib/Dialect/Tosa/IR/TosaOps.cpp‎

Lines changed: 76 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -916,9 +916,82 @@ LogicalResult tosa::SliceOp::verify() {
916916
}
917917

918918
LogicalResult tosa::MulOp::verify() {
919-
Type elementTy = getInput1().getType().getElementType();
920-
if (isa<FloatType>(elementTy) && getShift() != 0)
921-
return emitOpError() << "require shift to be 0 for float type";
919+
auto resElemType = getElementTypeOrSelf(getOutput());
920+
921+
// Verify if the element type amoung operands and result match tosa
922+
// specification.
923+
if (auto resIntType = dyn_cast<IntegerType>(resElemType)) {
924+
IntegerType lhsIntType =
925+
cast<IntegerType>(getElementTypeOrSelf(getInput1()));
926+
IntegerType rhsIntType =
927+
cast<IntegerType>(getElementTypeOrSelf(getInput2()));
928+
if (lhsIntType != rhsIntType)
929+
return emitOpError("requires the same element type for all operands");
930+
931+
// Though the spec requires the element type of result to be i32, a more
932+
// relaxed way is provided at dialect level for easier cooperating with
933+
// other dialects.
934+
if (lhsIntType.getWidth() > resIntType.getWidth())
935+
return emitOpError("invalid data type size for operands or result");
936+
937+
} else {
938+
// For other supported type, the spec requires requires the same element
939+
// type for all operands (excludes `shift` operand) and results.
940+
for (int i = 0; i < 2; ++i) {
941+
if (getElementTypeOrSelf(getOperand(i)) != resElemType)
942+
return emitOpError(
943+
"requires the same element type for all operands and results");
944+
}
945+
}
946+
947+
// Check if the shift value apply to non-i32 output type as that is not
948+
// allowed in the spec.
949+
if (!(llvm::isa<IntegerType>(resElemType) && resElemType.isInteger(32)))
950+
if (getShift().getImpl())
951+
return emitOpError("right shift output only on i32 data type");
952+
953+
// Verify the op has same ranks for all main operands (excludes extra operands
954+
// such as shift of mul op, so this is the only difference with the built-in
955+
// `SameOperandsAndResultRank` trait) and results types, if known.
956+
957+
// delegate function that returns true if type is a shaped type with known
958+
// rank
959+
auto hasRank = [](const Type type) {
960+
if (auto shaped_type = dyn_cast<ShapedType>(type))
961+
return shaped_type.hasRank();
962+
963+
return false;
964+
};
965+
966+
auto rankedOperandTypes =
967+
llvm::to_vector(llvm::make_filter_range(getOperandTypes(), hasRank));
968+
969+
auto rankedResultTypes =
970+
llvm::make_filter_range(getOperation()->getResultTypes(), hasRank);
971+
972+
// If all operands and results are unranked, then no further verification.
973+
if (rankedOperandTypes.empty() && rankedResultTypes.empty())
974+
return success();
975+
976+
// delegate function that returns rank of shaped type with known rank
977+
auto getRank = [](const Type type) {
978+
return cast<ShapedType>(type).getRank();
979+
};
980+
981+
auto rank = !rankedOperandTypes.empty() ? getRank(*rankedOperandTypes.begin())
982+
: getRank(*rankedResultTypes.begin());
983+
984+
for (size_t i = 0; i < 2; ++i) {
985+
if (rank != getRank(rankedOperandTypes[i])) {
986+
return emitOpError("operands don't have matching ranks");
987+
}
988+
}
989+
990+
for (const auto type : rankedResultTypes) {
991+
if (rank != getRank(type)) {
992+
return emitOpError("result type has different rank than operands");
993+
}
994+
}
922995

923996
return success();
924997
}

‎mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeDepthwise.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ struct DepthwiseConv2DIsMul : public OpRewritePattern<tosa::DepthwiseConv2DOp> {
137137

138138
Value mulValue = rewriter
139139
.create<tosa::MulOp>(op.getLoc(), mulShapeType, input,
140-
weight, /*shift=*/0)
140+
weight, Value{} /* zero_shift */)
141141
.getResult();
142142

143143
// Reshape output to [N, H, W, C * M].

‎mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ struct ConvertTosaOp<tosa::MulOp> : public OpRewritePattern<tosa::MulOp> {
113113

114114
Value input1 = tosaBinaryOp.getInput1();
115115
Value input2 = tosaBinaryOp.getInput2();
116-
int32_t shift = tosaBinaryOp.getShift();
116+
Value shift = tosaBinaryOp.getShift();
117117
Value output = tosaBinaryOp.getResult();
118118
auto outputType = dyn_cast<RankedTensorType>(output.getType());
119119
if (!outputType)

‎mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -474,7 +474,7 @@ func.func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
474474

475475
// CHECK: linalg.generic
476476
// CHECK: arith.mulf
477-
%4 = tosa.mul %0, %1 {shift = 0 : i8} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
477+
%4 = tosa.mul %0, %1 : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
478478

479479
// CHECK: linalg.generic
480480
// CHECK: arith.negf
@@ -620,7 +620,7 @@ func.func @test_simple_i16(%arg0: tensor<1xi16>) -> () {
620620
// CHECK: arith.extsi
621621
// CHECK: arith.extsi
622622
// CHECK: arith.muli
623-
%0 = tosa.mul %arg0, %arg0 {shift = 0 : i8} : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi32>
623+
%0 = tosa.mul %arg0, %arg0 : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi32>
624624

625625
return
626626
}
@@ -648,12 +648,14 @@ func.func @test_simple_i32(%arg0: tensor<1xi32>, %unsigned: tensor<1xui32>, %uns
648648

649649
// CHECK: linalg.generic
650650
// CHECK: arith.muli
651-
%2 = tosa.mul %arg0, %arg0 {shift = 0 : i8} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
651+
%shift1 = "tosa.const"() <{value = dense<0> : tensor<i8>}> : () -> tensor<i8>
652+
%2 = tosa.mul %arg0, %arg0, %shift1 : (tensor<1xi32>, tensor<1xi32>, tensor<i8>) -> tensor<1xi32>
652653

653654
// CHECK: linalg.generic
654655
// CHECK: arith.constant 2
655656
// CHECK: apply_scale
656-
%3 = tosa.mul %arg0, %arg0 {shift = 2 : i8} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
657+
%shift2 = "tosa.const"() <{value = dense<2> : tensor<i8>}> : () -> tensor<i8>
658+
%3 = tosa.mul %arg0, %arg0, %shift2: (tensor<1xi32>, tensor<1xi32>, tensor<i8>) -> tensor<1xi32>
657659

658660
// CHECK: linalg.generic
659661
// CHECK: arith.divsi

‎mlir/test/Dialect/Tosa/broadcast.mlir‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -169,15 +169,6 @@ func.func @test_broadcast20(%arg0: tensor<3x3x4x1xf32>, %arg1: tensor<4x5xf32>)
169169
return %0 : tensor<3x3x4x5xf32>
170170
}
171171

172-
// -----
173-
// CHECK-LABEL: broadcast_mul
174-
func.func @test_broadcast_mul(%arg0: tensor<15x14xi32>, %arg1: tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32> {
175-
// CHECK-DAG: %[[VAR0:.*]] = tosa.reshape %arg0 {new_shape = array<i64: 1, 1, 15, 14>}
176-
// CHECK: %[[VAR1:.*]] = tosa.mul %[[VAR0]], %arg1
177-
%0 = tosa.mul %arg0, %arg1 {shift = 1 : i8 } : (tensor<15x14xi32>, tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32>
178-
return %0 : tensor<17x16x15x14xi32>
179-
}
180-
181172
// -----
182173
// CHECK-LABEL: broadcast_arithmetic_right_shift
183174
func.func @test_broadcast_arithmetic_right_shift(%arg0: tensor<15x14xi32>, %arg1: tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32> {

‎mlir/test/Dialect/Tosa/canonicalize.mlir‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,7 @@ func.func @mul_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
277277
// CHECK: return %arg0
278278
// CHECK-NOT: tosa.mul
279279
%ones = "tosa.const"() {value = dense<1.0> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
280-
%1 = tosa.mul %arg0, %ones {shift = 0 : i8} : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
280+
%1 = tosa.mul %arg0, %ones : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
281281
return %1 : tensor<2x3xf32>
282282
}
283283

@@ -288,7 +288,7 @@ func.func @mul_bcast_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
288288
// CHECK: return %arg0
289289
// CHECK-NOT: tosa.mul
290290
%ones = "tosa.const"() {value = dense<1.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
291-
%1 = tosa.mul %ones, %arg0 {shift = 0 : i8} : (tensor<1x1xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
291+
%1 = tosa.mul %ones, %arg0 : (tensor<1x1xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
292292
return %1 : tensor<2x3xf32>
293293
}
294294

@@ -299,7 +299,20 @@ func.func @mul_one_int(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
299299
// CHECK: return %arg0
300300
// CHECK-NOT: tosa.mul
301301
%ones = "tosa.const"() {value = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
302-
%1 = tosa.mul %arg0, %ones {shift = 0 : i8} : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
302+
%1 = tosa.mul %arg0, %ones : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
303+
return %1 : tensor<2x3xi32>
304+
}
305+
306+
// -----
307+
308+
// CHECK-LABEL: @mul_one_int_and_shift
309+
func.func @mul_one_int_and_shift(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
310+
// CHECK-DAG: %[[VAL_1:.*]] = "tosa.const"() <{value = dense<1> : tensor<2x3xi32>}>
311+
// CHECK-DAG: %[[VAL_2:.*]] = "tosa.const"() <{value = dense<31> : tensor<i8>}>
312+
// CHECK: %[[VAL_3:.*]] = tosa.mul %arg0, %[[VAL_1]], %[[VAL_2]] : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<i8>)
313+
%ones = "tosa.const"() {value = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
314+
%shift = "tosa.const"() <{value = dense<31> : tensor<i8>}> : () -> tensor<i8>
315+
%1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<i8>) -> tensor<2x3xi32>
303316
return %1 : tensor<2x3xi32>
304317
}
305318

@@ -310,11 +323,11 @@ func.func @mul_zero_broadcast(%arg0: tensor<2x3xf32>) -> (tensor<2x3xf32>, tenso
310323
// CHECK: %[[ZERO:.*]] = "tosa.const"() <{value = dense<0.000000e+00> : tensor<2x3xf32>}
311324
// CHECK-NOT: tosa.mul
312325
%zeros = "tosa.const"() {value = dense<0.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
313-
%1 = tosa.mul %arg0, %zeros {shift = 0 : i8} : (tensor<2x3xf32>, tensor<1x1xf32>) -> tensor<2x3xf32>
326+
%1 = tosa.mul %arg0, %zeros : (tensor<2x3xf32>, tensor<1x1xf32>) -> tensor<2x3xf32>
314327

315328
// CHECK-NOT: tosa.mul
316329
// CHECK: return %[[ZERO]], %[[ZERO]]
317-
%2 = tosa.mul %zeros, %arg0 {shift = 0 : i8} : (tensor<1x1xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
330+
%2 = tosa.mul %zeros, %arg0 : (tensor<1x1xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
318331
return %1, %2 : tensor<2x3xf32>, tensor<2x3xf32>
319332
}
320333

@@ -869,7 +882,7 @@ func.func @mul_quant_nofold() -> tensor<1x!quant.uniform<i8:f32, 3.0757404601899
869882
// CHECK: tosa.mul
870883
%0 = "tosa.const"() {value = dense<0> : tensor<1xi8>} : () -> tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>
871884
%1 = "tosa.const"() {value = dense<1> : tensor<1xi8>} : () -> tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>
872-
%2 = tosa.mul %0, %1 { shift = 0 : i8} : (tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>, tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>) -> tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>
885+
%2 = tosa.mul %0, %1 : (tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>, tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>) -> tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>
873886
return %2 : tensor<1x!quant.uniform<i8:f32, 3.0757404601899907E-5:-128>>
874887
}
875888

0 commit comments

Comments
 (0)