1//===- TosaToArith.cpp - Lowering Tosa to Arith Dialect -------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// These rewriters lower from the Tosa to the Arith dialect.
10//
11//===----------------------------------------------------------------------===//
12
13#include "mlir/Conversion/TosaToArith/TosaToArith.h"
14#include "mlir/Dialect/Arith/IR/Arith.h"
15#include "mlir/Dialect/Tosa/IR/TosaOps.h"
16#include "mlir/IR/PatternMatch.h"
17#include "mlir/IR/TypeUtilities.h"
18
19using namespace mlir;
20using namespace tosa;
21
22namespace {
23
24class ConstOpConverter : public OpRewritePattern<tosa::ConstOp> {
25public:
26 using OpRewritePattern<tosa::ConstOp>::OpRewritePattern;
27
28 LogicalResult matchAndRewrite(tosa::ConstOp op,
29 PatternRewriter &rewriter) const final {
30 rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, args: op.getValues());
31 return success();
32 }
33};
34
35Type matchContainerType(Type element, Type container) {
36 if (auto shapedTy = dyn_cast<ShapedType>(Val&: container))
37 return shapedTy.clone(elementType: element);
38
39 return element;
40}
41
42TypedAttr getConstantAttr(Type type, int64_t value, PatternRewriter &rewriter) {
43 if (auto shapedTy = dyn_cast<ShapedType>(Val&: type)) {
44 Type eTy = shapedTy.getElementType();
45 APInt valueInt(eTy.getIntOrFloatBitWidth(), value, /*isSigned=*/true);
46 return DenseIntElementsAttr::get(type: shapedTy, arg&: valueInt);
47 }
48
49 return rewriter.getIntegerAttr(type, value);
50}
51
52Value getConstantValue(Location loc, Type type, int64_t value,
53 PatternRewriter &rewriter) {
54 return rewriter.create<arith::ConstantOp>(
55 location: loc, args: getConstantAttr(type, value, rewriter));
56}
57
58// This converts the TOSA ApplyScale operator to a set of arithmetic ops,
59// using 64-bit operations to perform the necessary multiply, bias, and shift.
60class ApplyScaleGenericOpConverter
61 : public OpRewritePattern<tosa::ApplyScaleOp> {
62public:
63 using OpRewritePattern<tosa::ApplyScaleOp>::OpRewritePattern;
64
65 LogicalResult matchAndRewrite(tosa::ApplyScaleOp op,
66 PatternRewriter &rewriter) const final {
67 StringRef roundingMode = op.getRoundingMode();
68 if (roundingMode != "DOUBLE_ROUND" && roundingMode != "SINGLE_ROUND") {
69 return failure();
70 }
71
72 Location loc = op.getLoc();
73 Value value = op.getValue();
74 Value multiplier32 = op.getMultiplier();
75
76 Type resultTy = op.getType();
77 Type valueTy = value.getType();
78 Type i32Ty = matchContainerType(element: rewriter.getI32Type(), container: resultTy);
79 Type i64Ty = matchContainerType(element: rewriter.getI64Type(), container: resultTy);
80
81 Value zero = getConstantValue(loc, type: valueTy, value: 0, rewriter);
82 Value one64 = getConstantValue(loc, type: i64Ty, value: 1, rewriter);
83 Value thirtyOne32 = getConstantValue(loc, type: i32Ty, value: 31, rewriter);
84
85 Value shift32 = rewriter.create<arith::ExtUIOp>(location: loc, args&: i32Ty, args: op.getShift());
86
87 // Compute the multiplication in 64-bits then select the high / low parts.
88 Value value64 = value;
89 if (getElementTypeOrSelf(type: valueTy) != rewriter.getI64Type())
90 value64 = rewriter.create<arith::ExtSIOp>(location: loc, args&: i64Ty, args&: value);
91 Value multiplier64 =
92 rewriter.create<arith::ExtSIOp>(location: loc, args&: i64Ty, args&: multiplier32);
93 Value multiply64 =
94 rewriter.create<arith::MulIOp>(location: loc, args&: value64, args&: multiplier64);
95
96 // Apply normal rounding.
97 Value shift64 = rewriter.create<arith::ExtUIOp>(location: loc, args&: i64Ty, args&: shift32);
98 Value round = rewriter.create<arith::ShLIOp>(location: loc, args&: one64, args&: shift64);
99 round = rewriter.create<arith::ShRUIOp>(location: loc, args&: round, args&: one64);
100 multiply64 = rewriter.create<arith::AddIOp>(location: loc, args&: multiply64, args&: round);
101
102 // Apply double rounding if necessary.
103 if (op.getRoundingMode() == "DOUBLE_ROUND") {
104 int64_t roundInt = 1 << 30;
105 Value roundUp = getConstantValue(loc, type: i64Ty, value: roundInt, rewriter);
106 Value roundDown = getConstantValue(loc, type: i64Ty, value: -roundInt, rewriter);
107 Value positive = rewriter.create<arith::CmpIOp>(
108 location: loc, args: arith::CmpIPredicate::sge, args&: value, args&: zero);
109 Value dir =
110 rewriter.create<arith::SelectOp>(location: loc, args&: positive, args&: roundUp, args&: roundDown);
111 Value val = rewriter.create<arith::AddIOp>(location: loc, args&: dir, args&: multiply64);
112 Value valid = rewriter.create<arith::CmpIOp>(
113 location: loc, args: arith::CmpIPredicate::sgt, args&: shift32, args&: thirtyOne32);
114 multiply64 =
115 rewriter.create<arith::SelectOp>(location: loc, args&: valid, args&: val, args&: multiply64);
116 }
117
118 Value result64 = rewriter.create<arith::ShRSIOp>(location: loc, args&: multiply64, args&: shift64);
119 Value result32 = rewriter.create<arith::TruncIOp>(location: loc, args&: i32Ty, args&: result64);
120
121 rewriter.replaceOp(op, newValues: result32);
122 return success();
123 }
124};
125
126class ApplyScale32BitOpConverter : public OpRewritePattern<tosa::ApplyScaleOp> {
127public:
128 using OpRewritePattern<tosa::ApplyScaleOp>::OpRewritePattern;
129
130 LogicalResult matchAndRewrite(tosa::ApplyScaleOp op,
131 PatternRewriter &rewriter) const final {
132 StringRef roundingMode = op.getRoundingMode();
133 if (roundingMode != "DOUBLE_ROUND" && roundingMode != "SINGLE_ROUND") {
134 return failure();
135 }
136
137 Location loc = op.getLoc();
138
139 Type resultTy = op.getType();
140 Type i32Ty = matchContainerType(element: rewriter.getI32Type(), container: resultTy);
141
142 Value value = op.getValue();
143 if (getElementTypeOrSelf(type: value.getType()).getIntOrFloatBitWidth() > 32) {
144 return failure();
145 }
146
147 Value value32 = op.getValue();
148 Value multiplier32 = op.getMultiplier();
149 Value shift32 = rewriter.create<arith::ExtUIOp>(location: loc, args&: i32Ty, args: op.getShift());
150
151 // Constants used during the scaling operation.
152 Value zero32 = getConstantValue(loc, type: i32Ty, value: 0, rewriter);
153 Value one32 = getConstantValue(loc, type: i32Ty, value: 1, rewriter);
154 Value two32 = getConstantValue(loc, type: i32Ty, value: 2, rewriter);
155 Value thirty32 = getConstantValue(loc, type: i32Ty, value: 30, rewriter);
156 Value thirtyTwo32 = getConstantValue(loc, type: i32Ty, value: 32, rewriter);
157
158 // Compute the multiplication in 64-bits then select the high / low parts.
159 // Grab out the high/low of the computation
160 auto value64 =
161 rewriter.create<arith::MulSIExtendedOp>(location: loc, args&: value32, args&: multiplier32);
162 Value low32 = value64.getLow();
163 Value high32 = value64.getHigh();
164
165 // Determine the direction and amount to shift the high bits.
166 Value shiftOver32 = rewriter.create<arith::CmpIOp>(
167 location: loc, args: arith::CmpIPredicate::sge, args&: shift32, args&: thirtyTwo32);
168 Value roundHighBits = rewriter.create<arith::CmpIOp>(
169 location: loc, args: arith::CmpIPredicate::sgt, args&: shift32, args&: thirtyTwo32);
170
171 Value shiftHighL =
172 rewriter.create<arith::SubIOp>(location: loc, args&: thirtyTwo32, args&: shift32);
173 Value shiftHighR =
174 rewriter.create<arith::SubIOp>(location: loc, args&: shift32, args&: thirtyTwo32);
175
176 shiftHighL =
177 rewriter.create<arith::SelectOp>(location: loc, args&: shiftOver32, args&: zero32, args&: shiftHighL);
178 shiftHighR =
179 rewriter.create<arith::SelectOp>(location: loc, args&: shiftOver32, args&: shiftHighR, args&: zero32);
180
181 // Conditionally perform our double round.
182 if (op.getRoundingMode() == "DOUBLE_ROUND") {
183 Value negOne32 = getConstantValue(loc, type: i32Ty, value: -1, rewriter);
184 Value valuePositive = rewriter.create<arith::CmpIOp>(
185 location: loc, args: arith::CmpIPredicate::sge, args&: value32, args&: zero32);
186
187 Value roundDir =
188 rewriter.create<arith::SelectOp>(location: loc, args&: valuePositive, args&: one32, args&: negOne32);
189 roundDir =
190 rewriter.create<arith::SelectOp>(location: loc, args&: shiftOver32, args&: roundDir, args&: zero32);
191
192 Value shiftLow = rewriter.create<arith::ShRUIOp>(location: loc, args&: low32, args&: thirty32);
193 Value rounded = rewriter.create<arith::AddIOp>(location: loc, args&: shiftLow, args&: roundDir);
194 Value carry = rewriter.create<arith::ShRSIOp>(location: loc, args&: rounded, args&: two32);
195
196 Value shiftRound =
197 rewriter.create<arith::ShLIOp>(location: loc, args&: roundDir, args&: thirty32);
198
199 low32 = rewriter.create<arith::AddIOp>(location: loc, args&: low32, args&: shiftRound);
200 high32 = rewriter.create<arith::AddIOp>(location: loc, args&: high32, args&: carry);
201 }
202
203 // Conditionally apply rounding in the low bits.
204 {
205 Value shiftSubOne = rewriter.create<arith::SubIOp>(location: loc, args&: shift32, args&: one32);
206 Value roundBit = rewriter.create<arith::ShLIOp>(location: loc, args&: one32, args&: shiftSubOne);
207 roundBit = rewriter.create<arith::SelectOp>(location: loc, args&: roundHighBits, args&: zero32,
208 args&: roundBit);
209
210 Value newLow32 = rewriter.create<arith::AddIOp>(location: loc, args&: low32, args&: roundBit);
211 Value wasRounded = rewriter.create<arith::CmpIOp>(
212 location: loc, args: arith::CmpIPredicate::ugt, args&: low32, args&: newLow32);
213 low32 = newLow32;
214
215 Value rounded32 = rewriter.create<arith::ExtUIOp>(location: loc, args&: i32Ty, args&: wasRounded);
216 high32 = rewriter.create<arith::AddIOp>(location: loc, args&: high32, args&: rounded32);
217 }
218
219 // Conditionally apply rounding in the high bits.
220 {
221 Value shiftSubOne =
222 rewriter.create<arith::SubIOp>(location: loc, args&: shiftHighR, args&: one32);
223 Value roundBit = rewriter.create<arith::ShLIOp>(location: loc, args&: one32, args&: shiftSubOne);
224 roundBit = rewriter.create<arith::SelectOp>(location: loc, args&: roundHighBits, args&: roundBit,
225 args&: zero32);
226 high32 = rewriter.create<arith::AddIOp>(location: loc, args&: high32, args&: roundBit);
227 }
228
229 // Combine the correct high/low bits into the final rescale result.
230 high32 = rewriter.create<arith::ShLIOp>(location: loc, args&: high32, args&: shiftHighL);
231 high32 = rewriter.create<arith::ShRSIOp>(location: loc, args&: high32, args&: shiftHighR);
232 low32 = rewriter.create<arith::ShRUIOp>(location: loc, args&: low32, args&: shift32);
233 low32 = rewriter.create<arith::SelectOp>(location: loc, args&: shiftOver32, args&: zero32, args&: low32);
234
235 // Apply the rounding behavior and shift to the final alignment.
236 Value result = rewriter.create<arith::AddIOp>(location: loc, args&: low32, args&: high32);
237
238 // Truncate if necessary.
239 if (!getElementTypeOrSelf(type: resultTy).isInteger(width: 32)) {
240 result = rewriter.create<arith::TruncIOp>(location: loc, args&: resultTy, args&: result);
241 }
242
243 rewriter.replaceOp(op, newValues: result);
244 return success();
245 }
246};
247
248} // namespace
249
250void mlir::tosa::populateTosaToArithConversionPatterns(
251 RewritePatternSet *patterns) {
252 patterns->add<ConstOpConverter>(arg: patterns->getContext());
253}
254
255void mlir::tosa::populateTosaRescaleToArithConversionPatterns(
256 RewritePatternSet *patterns, bool include32Bit) {
257 patterns->add<ApplyScaleGenericOpConverter>(arg: patterns->getContext(), args: 100);
258 if (include32Bit) {
259 patterns->add<ApplyScale32BitOpConverter>(arg: patterns->getContext(), args: 200);
260 }
261}
262

source code of mlir/lib/Conversion/TosaToArith/TosaToArith.cpp