1 //===- PolynomialApproximation.cpp - Approximate math operations ----------===//
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 // This file implements expansion of math operations to fast approximations
10 // that do not rely on any of the library functions.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
15 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
16 #include "mlir/Dialect/Math/IR/Math.h"
17 #include "mlir/Dialect/Math/Transforms/Passes.h"
18 #include "mlir/Dialect/Vector/VectorOps.h"
19 #include "mlir/IR/Builders.h"
20 #include "mlir/IR/ImplicitLocOpBuilder.h"
21 #include "mlir/Transforms/Bufferize.h"
22 #include "mlir/Transforms/DialectConversion.h"
23 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
24 #include <climits>
25
26 using namespace mlir;
27 using namespace mlir::vector;
28
29 using TypePredicate = llvm::function_ref<bool(Type)>;
30
31 // Returns vector width if the element type is matching the predicate (scalars
32 // that do match the predicate have width equal to `1`).
vectorWidth(Type type,TypePredicate pred)33 static Optional<int> vectorWidth(Type type, TypePredicate pred) {
34 // If the type matches the predicate then its width is `1`.
35 if (pred(type))
36 return 1;
37
38 // Otherwise check if the type is a vector type.
39 auto vectorType = type.dyn_cast<VectorType>();
40 if (vectorType && pred(vectorType.getElementType())) {
41 assert(vectorType.getRank() == 1 && "only 1d vectors are supported");
42 return vectorType.getDimSize(0);
43 }
44
45 return llvm::None;
46 }
47
48 // Returns vector width of the type. If the type is a scalar returns `1`.
vectorWidth(Type type)49 static int vectorWidth(Type type) {
50 auto vectorType = type.dyn_cast<VectorType>();
51 return vectorType ? vectorType.getDimSize(0) : 1;
52 }
53
54 // Returns vector element type. If the type is a scalar returns the argument.
elementType(Type type)55 LLVM_ATTRIBUTE_UNUSED static Type elementType(Type type) {
56 auto vectorType = type.dyn_cast<VectorType>();
57 return vectorType ? vectorType.getElementType() : type;
58 }
59
isF32(Type type)60 LLVM_ATTRIBUTE_UNUSED static bool isF32(Type type) { return type.isF32(); }
61
isI32(Type type)62 LLVM_ATTRIBUTE_UNUSED static bool isI32(Type type) {
63 return type.isInteger(32);
64 }
65
66 //----------------------------------------------------------------------------//
67 // Broadcast scalar types and values into vector types and values.
68 //----------------------------------------------------------------------------//
69
70 // Broadcasts scalar type into vector type (iff width is greater then 1).
broadcast(Type type,int width)71 static Type broadcast(Type type, int width) {
72 assert(!type.isa<VectorType>() && "must be scalar type");
73 return width > 1 ? VectorType::get({width}, type) : type;
74 }
75
76 // Broadcasts scalar value into vector (iff width is greater then 1).
broadcast(ImplicitLocOpBuilder & builder,Value value,int width)77 static Value broadcast(ImplicitLocOpBuilder &builder, Value value, int width) {
78 assert(!value.getType().isa<VectorType>() && "must be scalar value");
79 auto type = broadcast(value.getType(), width);
80 return width > 1 ? builder.create<BroadcastOp>(type, value) : value;
81 }
82
83 //----------------------------------------------------------------------------//
84 // Helper functions to create constants.
85 //----------------------------------------------------------------------------//
86
f32Cst(ImplicitLocOpBuilder & builder,float value)87 static Value f32Cst(ImplicitLocOpBuilder &builder, float value) {
88 return builder.create<ConstantOp>(builder.getF32Type(),
89 builder.getF32FloatAttr(value));
90 }
91
i32Cst(ImplicitLocOpBuilder & builder,int32_t value)92 static Value i32Cst(ImplicitLocOpBuilder &builder, int32_t value) {
93 return builder.create<ConstantOp>(builder.getI32Type(),
94 builder.getI32IntegerAttr(value));
95 }
96
f32FromBits(ImplicitLocOpBuilder & builder,uint32_t bits)97 static Value f32FromBits(ImplicitLocOpBuilder &builder, uint32_t bits) {
98 Value i32Value = i32Cst(builder, static_cast<int32_t>(bits));
99 return builder.create<LLVM::BitcastOp>(builder.getF32Type(), i32Value);
100 }
101
102 //----------------------------------------------------------------------------//
103 // Helper functions to build math functions approximations.
104 //----------------------------------------------------------------------------//
105
min(ImplicitLocOpBuilder & builder,Value a,Value b)106 static Value min(ImplicitLocOpBuilder &builder, Value a, Value b) {
107 return builder.create<SelectOp>(
108 builder.create<CmpFOp>(CmpFPredicate::OLT, a, b), a, b);
109 }
110
max(ImplicitLocOpBuilder & builder,Value a,Value b)111 static Value max(ImplicitLocOpBuilder &builder, Value a, Value b) {
112 return builder.create<SelectOp>(
113 builder.create<CmpFOp>(CmpFPredicate::OGT, a, b), a, b);
114 }
115
clamp(ImplicitLocOpBuilder & builder,Value value,Value lowerBound,Value upperBound)116 static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound,
117 Value upperBound) {
118 return max(builder, min(builder, value, upperBound), lowerBound);
119 }
120
121 // Decomposes given floating point value `arg` into a normalized fraction and
122 // an integral power of two (see std::frexp). Returned values have float type.
frexp(ImplicitLocOpBuilder & builder,Value arg,bool is_positive=false)123 static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
124 bool is_positive = false) {
125 assert(isF32(elementType(arg.getType())) && "argument must be f32 type");
126
127 int width = vectorWidth(arg.getType());
128
129 auto bcast = [&](Value value) -> Value {
130 return broadcast(builder, value, width);
131 };
132
133 auto i32 = builder.getIntegerType(32);
134 auto i32Vec = broadcast(i32, width);
135 auto f32Vec = broadcast(builder.getF32Type(), width);
136
137 Value cst126f = f32Cst(builder, 126.0f);
138 Value cstHalf = f32Cst(builder, 0.5f);
139 Value cstInvMantMask = f32FromBits(builder, ~0x7f800000u);
140
141 // Bitcast to i32 for bitwise operations.
142 Value i32Half = builder.create<LLVM::BitcastOp>(i32, cstHalf);
143 Value i32InvMantMask = builder.create<LLVM::BitcastOp>(i32, cstInvMantMask);
144 Value i32Arg = builder.create<LLVM::BitcastOp>(i32Vec, arg);
145
146 // Compute normalized fraction.
147 Value tmp0 = builder.create<LLVM::AndOp>(i32Arg, bcast(i32InvMantMask));
148 Value tmp1 = builder.create<LLVM::OrOp>(tmp0, bcast(i32Half));
149 Value normalizedFraction = builder.create<LLVM::BitcastOp>(f32Vec, tmp1);
150
151 // Compute exponent.
152 Value arg0 = is_positive ? arg : builder.create<AbsFOp>(arg);
153 Value biasedExponentBits = builder.create<UnsignedShiftRightOp>(
154 builder.create<LLVM::BitcastOp>(i32Vec, arg0),
155 bcast(i32Cst(builder, 23)));
156 Value biasedExponent = builder.create<SIToFPOp>(f32Vec, biasedExponentBits);
157 Value exponent = builder.create<SubFOp>(biasedExponent, bcast(cst126f));
158
159 return {normalizedFraction, exponent};
160 }
161
162 // Computes exp2 for an i32 argument.
exp2I32(ImplicitLocOpBuilder & builder,Value arg)163 static Value exp2I32(ImplicitLocOpBuilder &builder, Value arg) {
164 assert(isI32(elementType(arg.getType())) && "argument must be i32 type");
165
166 int width = vectorWidth(arg.getType());
167
168 auto bcast = [&](Value value) -> Value {
169 return broadcast(builder, value, width);
170 };
171
172 auto f32Vec = broadcast(builder.getF32Type(), width);
173 // The exponent of f32 located at 23-bit.
174 auto exponetBitLocation = bcast(i32Cst(builder, 23));
175 // Set the exponent bias to zero.
176 auto bias = bcast(i32Cst(builder, 127));
177
178 Value biasedArg = builder.create<AddIOp>(arg, bias);
179 Value exp2ValueInt =
180 builder.create<ShiftLeftOp>(biasedArg, exponetBitLocation);
181 Value exp2ValueF32 = builder.create<LLVM::BitcastOp>(f32Vec, exp2ValueInt);
182
183 return exp2ValueF32;
184 }
185
186 //----------------------------------------------------------------------------//
187 // TanhOp approximation.
188 //----------------------------------------------------------------------------//
189
190 namespace {
191 struct TanhApproximation : public OpRewritePattern<math::TanhOp> {
192 public:
193 using OpRewritePattern::OpRewritePattern;
194
195 LogicalResult matchAndRewrite(math::TanhOp op,
196 PatternRewriter &rewriter) const final;
197 };
198 } // namespace
199
200 LogicalResult
matchAndRewrite(math::TanhOp op,PatternRewriter & rewriter) const201 TanhApproximation::matchAndRewrite(math::TanhOp op,
202 PatternRewriter &rewriter) const {
203 auto width = vectorWidth(op.operand().getType(), isF32);
204 if (!width.hasValue())
205 return rewriter.notifyMatchFailure(op, "unsupported operand type");
206
207 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
208 auto bcast = [&](Value value) -> Value {
209 return broadcast(builder, value, *width);
210 };
211
212 // Clamp operand into [plusClamp, minusClamp] range.
213 Value minusClamp = bcast(f32Cst(builder, -7.9053111076354980f));
214 Value plusClamp = bcast(f32Cst(builder, 7.90531110763549805f));
215 Value x = clamp(builder, op.operand(), minusClamp, plusClamp);
216
217 // Mask for tiny values that are approximated with `operand`.
218 Value tiny = bcast(f32Cst(builder, 0.0004f));
219 Value tinyMask = builder.create<CmpFOp>(
220 CmpFPredicate::OLT, builder.create<AbsFOp>(op.operand()), tiny);
221
222 // The monomial coefficients of the numerator polynomial (odd).
223 Value alpha1 = bcast(f32Cst(builder, 4.89352455891786e-03f));
224 Value alpha3 = bcast(f32Cst(builder, 6.37261928875436e-04f));
225 Value alpha5 = bcast(f32Cst(builder, 1.48572235717979e-05f));
226 Value alpha7 = bcast(f32Cst(builder, 5.12229709037114e-08f));
227 Value alpha9 = bcast(f32Cst(builder, -8.60467152213735e-11f));
228 Value alpha11 = bcast(f32Cst(builder, 2.00018790482477e-13f));
229 Value alpha13 = bcast(f32Cst(builder, -2.76076847742355e-16f));
230
231 // The monomial coefficients of the denominator polynomial (even).
232 Value beta0 = bcast(f32Cst(builder, 4.89352518554385e-03f));
233 Value beta2 = bcast(f32Cst(builder, 2.26843463243900e-03f));
234 Value beta4 = bcast(f32Cst(builder, 1.18534705686654e-04f));
235 Value beta6 = bcast(f32Cst(builder, 1.19825839466702e-06f));
236
237 // Since the polynomials are odd/even, we need x^2.
238 Value x2 = builder.create<MulFOp>(x, x);
239
240 // Evaluate the numerator polynomial p.
241 Value p = builder.create<FmaFOp>(x2, alpha13, alpha11);
242 p = builder.create<FmaFOp>(x2, p, alpha9);
243 p = builder.create<FmaFOp>(x2, p, alpha7);
244 p = builder.create<FmaFOp>(x2, p, alpha5);
245 p = builder.create<FmaFOp>(x2, p, alpha3);
246 p = builder.create<FmaFOp>(x2, p, alpha1);
247 p = builder.create<MulFOp>(x, p);
248
249 // Evaluate the denominator polynomial q.
250 Value q = builder.create<FmaFOp>(x2, beta6, beta4);
251 q = builder.create<FmaFOp>(x2, q, beta2);
252 q = builder.create<FmaFOp>(x2, q, beta0);
253
254 // Divide the numerator by the denominator.
255 Value res =
256 builder.create<SelectOp>(tinyMask, x, builder.create<DivFOp>(p, q));
257
258 rewriter.replaceOp(op, res);
259
260 return success();
261 }
262
263 #define LN2_VALUE \
264 0.693147180559945309417232121458176568075500134360255254120680009493393621L
265 #define LOG2E_VALUE \
266 1.442695040888963407359924681001892137426645954152985934135449406931109219L
267
268 //----------------------------------------------------------------------------//
269 // LogOp and Log2Op approximation.
270 //----------------------------------------------------------------------------//
271
272 namespace {
273 template <typename Op>
274 struct LogApproximationBase : public OpRewritePattern<Op> {
275 using OpRewritePattern<Op>::OpRewritePattern;
276
277 /// Base 2 if 'base2' is set; natural logarithm (base e) otherwise.
278 LogicalResult logMatchAndRewrite(Op op, PatternRewriter &rewriter,
279 bool base2) const;
280 };
281 } // namespace
282
283 // This approximation comes from Julien Pommier's SSE math library.
284 // Link: http://gruntthepeon.free.fr/ssemath
285 template <typename Op>
286 LogicalResult
logMatchAndRewrite(Op op,PatternRewriter & rewriter,bool base2) const287 LogApproximationBase<Op>::logMatchAndRewrite(Op op, PatternRewriter &rewriter,
288 bool base2) const {
289 auto width = vectorWidth(op.operand().getType(), isF32);
290 if (!width.hasValue())
291 return rewriter.notifyMatchFailure(op, "unsupported operand type");
292
293 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
294 auto bcast = [&](Value value) -> Value {
295 return broadcast(builder, value, *width);
296 };
297
298 Value cstZero = bcast(f32Cst(builder, 0.0f));
299 Value cstOne = bcast(f32Cst(builder, 1.0f));
300 Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
301
302 // The smallest non denormalized float number.
303 Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
304 Value cstMinusInf = bcast(f32FromBits(builder, 0xff800000u));
305 Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u));
306 Value cstNan = bcast(f32FromBits(builder, 0x7fc00000));
307
308 // Polynomial coefficients.
309 Value cstCephesSQRTHF = bcast(f32Cst(builder, 0.707106781186547524f));
310 Value cstCephesLogP0 = bcast(f32Cst(builder, 7.0376836292E-2f));
311 Value cstCephesLogP1 = bcast(f32Cst(builder, -1.1514610310E-1f));
312 Value cstCephesLogP2 = bcast(f32Cst(builder, 1.1676998740E-1f));
313 Value cstCephesLogP3 = bcast(f32Cst(builder, -1.2420140846E-1f));
314 Value cstCephesLogP4 = bcast(f32Cst(builder, +1.4249322787E-1f));
315 Value cstCephesLogP5 = bcast(f32Cst(builder, -1.6668057665E-1f));
316 Value cstCephesLogP6 = bcast(f32Cst(builder, +2.0000714765E-1f));
317 Value cstCephesLogP7 = bcast(f32Cst(builder, -2.4999993993E-1f));
318 Value cstCephesLogP8 = bcast(f32Cst(builder, +3.3333331174E-1f));
319
320 Value x = op.operand();
321
322 // Truncate input values to the minimum positive normal.
323 x = max(builder, x, cstMinNormPos);
324
325 // Extract significant in the range [0.5,1) and exponent.
326 std::pair<Value, Value> pair = frexp(builder, x, /*is_positive=*/true);
327 x = pair.first;
328 Value e = pair.second;
329
330 // Shift the inputs from the range [0.5,1) to [sqrt(1/2), sqrt(2)) and shift
331 // by -1.0. The values are then centered around 0, which improves the
332 // stability of the polynomial evaluation:
333 //
334 // if( x < SQRTHF ) {
335 // e -= 1;
336 // x = x + x - 1.0;
337 // } else { x = x - 1.0; }
338 Value mask = builder.create<CmpFOp>(CmpFPredicate::OLT, x, cstCephesSQRTHF);
339 Value tmp = builder.create<SelectOp>(mask, x, cstZero);
340
341 x = builder.create<SubFOp>(x, cstOne);
342 e = builder.create<SubFOp>(e,
343 builder.create<SelectOp>(mask, cstOne, cstZero));
344 x = builder.create<AddFOp>(x, tmp);
345
346 Value x2 = builder.create<MulFOp>(x, x);
347 Value x3 = builder.create<MulFOp>(x2, x);
348
349 // Evaluate the polynomial approximant of degree 8 in three parts.
350 Value y0, y1, y2;
351 y0 = builder.create<FmaFOp>(cstCephesLogP0, x, cstCephesLogP1);
352 y1 = builder.create<FmaFOp>(cstCephesLogP3, x, cstCephesLogP4);
353 y2 = builder.create<FmaFOp>(cstCephesLogP6, x, cstCephesLogP7);
354 y0 = builder.create<FmaFOp>(y0, x, cstCephesLogP2);
355 y1 = builder.create<FmaFOp>(y1, x, cstCephesLogP5);
356 y2 = builder.create<FmaFOp>(y2, x, cstCephesLogP8);
357 y0 = builder.create<FmaFOp>(y0, x3, y1);
358 y0 = builder.create<FmaFOp>(y0, x3, y2);
359 y0 = builder.create<MulFOp>(y0, x3);
360
361 y0 = builder.create<FmaFOp>(cstNegHalf, x2, y0);
362 x = builder.create<AddFOp>(x, y0);
363
364 if (base2) {
365 Value cstLog2e = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE)));
366 x = builder.create<FmaFOp>(x, cstLog2e, e);
367 } else {
368 Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
369 x = builder.create<FmaFOp>(e, cstLn2, x);
370 }
371
372 Value invalidMask =
373 builder.create<CmpFOp>(CmpFPredicate::ULT, op.operand(), cstZero);
374 Value zeroMask =
375 builder.create<CmpFOp>(CmpFPredicate::OEQ, op.operand(), cstZero);
376 Value posInfMask =
377 builder.create<CmpFOp>(CmpFPredicate::OEQ, op.operand(), cstPosInf);
378
379 // Filter out invalid values:
380 // • x == 0 -> -INF
381 // • x < 0 -> NAN
382 // • x == +INF -> +INF
383 Value aproximation = builder.create<SelectOp>(
384 zeroMask, cstMinusInf,
385 builder.create<SelectOp>(
386 invalidMask, cstNan,
387 builder.create<SelectOp>(posInfMask, cstPosInf, x)));
388
389 rewriter.replaceOp(op, aproximation);
390
391 return success();
392 }
393
394 namespace {
395 struct LogApproximation : public LogApproximationBase<math::LogOp> {
396 using LogApproximationBase::LogApproximationBase;
397
matchAndRewrite__anone93dec420711::LogApproximation398 LogicalResult matchAndRewrite(math::LogOp op,
399 PatternRewriter &rewriter) const final {
400 return logMatchAndRewrite(op, rewriter, /*base2=*/false);
401 }
402 };
403 } // namespace
404
405 namespace {
406 struct Log2Approximation : public LogApproximationBase<math::Log2Op> {
407 using LogApproximationBase::LogApproximationBase;
408
matchAndRewrite__anone93dec420811::Log2Approximation409 LogicalResult matchAndRewrite(math::Log2Op op,
410 PatternRewriter &rewriter) const final {
411 return logMatchAndRewrite(op, rewriter, /*base2=*/true);
412 }
413 };
414 } // namespace
415
416 //----------------------------------------------------------------------------//
417 // Log1p approximation.
418 //----------------------------------------------------------------------------//
419
420 namespace {
421 struct Log1pApproximation : public OpRewritePattern<math::Log1pOp> {
422 public:
423 using OpRewritePattern::OpRewritePattern;
424
425 LogicalResult matchAndRewrite(math::Log1pOp op,
426 PatternRewriter &rewriter) const final;
427 };
428 } // namespace
429
430 // Approximate log(1+x).
431 LogicalResult
matchAndRewrite(math::Log1pOp op,PatternRewriter & rewriter) const432 Log1pApproximation::matchAndRewrite(math::Log1pOp op,
433 PatternRewriter &rewriter) const {
434 auto width = vectorWidth(op.operand().getType(), isF32);
435 if (!width.hasValue())
436 return rewriter.notifyMatchFailure(op, "unsupported operand type");
437
438 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
439 auto bcast = [&](Value value) -> Value {
440 return broadcast(builder, value, *width);
441 };
442
443 // Approximate log(1+x) using the following, due to W. Kahan:
444 // u = x + 1.0;
445 // if (u == 1.0 || u == inf) return x;
446 // return x * log(u) / (u - 1.0);
447 // ^^^^^^^^^^^^^^^^^^^^^^
448 // "logLarge" below.
449 Value cstOne = bcast(f32Cst(builder, 1.0f));
450 Value x = op.operand();
451 Value u = builder.create<AddFOp>(x, cstOne);
452 Value uSmall = builder.create<CmpFOp>(CmpFPredicate::OEQ, u, cstOne);
453 Value logU = builder.create<math::LogOp>(u);
454 Value uInf = builder.create<CmpFOp>(CmpFPredicate::OEQ, u, logU);
455 Value logLarge = builder.create<MulFOp>(
456 x, builder.create<DivFOp>(logU, builder.create<SubFOp>(u, cstOne)));
457 Value approximation = builder.create<SelectOp>(
458 builder.create<LLVM::OrOp>(uSmall, uInf), x, logLarge);
459 rewriter.replaceOp(op, approximation);
460 return success();
461 }
462
463 //----------------------------------------------------------------------------//
464 // Exp approximation.
465 //----------------------------------------------------------------------------//
466
467 namespace {
468
469 struct ExpApproximation : public OpRewritePattern<math::ExpOp> {
470 public:
471 using OpRewritePattern::OpRewritePattern;
472
473 LogicalResult matchAndRewrite(math::ExpOp op,
474 PatternRewriter &rewriter) const final;
475 };
476 } // namespace
477
478 // Approximate exp(x) using its reduced range exp(y) where y is in the range
479 // [0, ln(2)], let y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2), exp(x)
480 // = exp(y) * 2^k. exp(y).
481 LogicalResult
matchAndRewrite(math::ExpOp op,PatternRewriter & rewriter) const482 ExpApproximation::matchAndRewrite(math::ExpOp op,
483 PatternRewriter &rewriter) const {
484 auto width = vectorWidth(op.operand().getType(), isF32);
485 if (!width.hasValue())
486 return rewriter.notifyMatchFailure(op, "unsupported operand type");
487 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
488
489 // TODO: Consider a common pattern rewriter with all methods below to
490 // write the approximations.
491 auto bcast = [&](Value value) -> Value {
492 return broadcast(builder, value, *width);
493 };
494 auto fmla = [&](Value a, Value b, Value c) {
495 return builder.create<FmaFOp>(a, b, c);
496 };
497 auto mul = [&](Value a, Value b) -> Value {
498 return builder.create<MulFOp>(a, b);
499 };
500 auto sub = [&](Value a, Value b) -> Value {
501 return builder.create<SubFOp>(a, b);
502 };
503 auto floor = [&](Value a) { return builder.create<FloorFOp>(a); };
504
505 Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
506 Value cstLog2E = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE)));
507
508 // Polynomial coefficients.
509 Value cstCephesExpP0 = bcast(f32Cst(builder, 1.0));
510 Value cstCephesExpP1 = bcast(f32Cst(builder, 1.0));
511 Value cstCephesExpP2 = bcast(f32Cst(builder, 0.49970514590562437052f));
512 Value cstCephesExpP3 = bcast(f32Cst(builder, 0.16873890085469545053f));
513 Value cstCephesExpP4 = bcast(f32Cst(builder, 0.03668965196652099192f));
514 Value cstCephesExpP5 = bcast(f32Cst(builder, 0.01314350012789660196f));
515
516 Value x = op.operand();
517
518 // Reduced y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2)
519 Value xL2Inv = mul(x, cstLog2E);
520 Value kF32 = floor(xL2Inv);
521 Value kLn2 = mul(kF32, cstLn2);
522 Value y = sub(x, kLn2);
523
524 // Use Estrin's evaluation scheme with 3 independent parts:
525 // P(y)^y : (c0 + c1 y) + (c2 + c3 y) y^2 + (c4 + c5 y) y^4
526 Value y2 = mul(y, y);
527 Value y4 = mul(y2, y2);
528
529 Value q0 = fmla(cstCephesExpP1, y, cstCephesExpP0);
530 Value q1 = fmla(cstCephesExpP3, y, cstCephesExpP2);
531 Value q2 = fmla(cstCephesExpP5, y, cstCephesExpP4);
532 Value expY = fmla(q1, y2, q0);
533 expY = fmla(q2, y4, expY);
534
535 auto i32Vec = broadcast(builder.getI32Type(), *width);
536
537 // exp2(k)
538 Value k = builder.create<FPToSIOp>(kF32, i32Vec);
539 Value exp2KValue = exp2I32(builder, k);
540
541 // exp(x) = exp(y) * exp2(k)
542 expY = mul(expY, exp2KValue);
543
544 // Handle overflow, inf and underflow of exp(x). exp(x) range is [0, inf], its
545 // partitioned as the following:
546 // exp(x) = 0, x <= -inf
547 // exp(x) = underflow (min_float), x <= -88
548 // exp(x) = inf (min_float), x >= 88
549 // Note: |k| = 127 is the value where the 8-bits exponent saturates.
550 Value zerof32Const = bcast(f32Cst(builder, 0));
551 auto constPosInfinity =
552 bcast(f32Cst(builder, std::numeric_limits<float>::infinity()));
553 auto constNegIfinity =
554 bcast(f32Cst(builder, -std::numeric_limits<float>::infinity()));
555 auto underflow = bcast(f32Cst(builder, std::numeric_limits<float>::min()));
556
557 Value kMaxConst = bcast(i32Cst(builder, 127));
558 Value kMaxNegConst = bcast(i32Cst(builder, -127));
559 Value rightBound = builder.create<CmpIOp>(CmpIPredicate::sle, k, kMaxConst);
560 Value leftBound = builder.create<CmpIOp>(CmpIPredicate::sge, k, kMaxNegConst);
561
562 Value isNegInfinityX =
563 builder.create<CmpFOp>(CmpFPredicate::OEQ, x, constNegIfinity);
564 Value isPostiveX =
565 builder.create<CmpFOp>(CmpFPredicate::OGT, x, zerof32Const);
566 Value isComputable = builder.create<AndOp>(rightBound, leftBound);
567
568 expY = builder.create<SelectOp>(
569 isComputable, expY,
570 builder.create<SelectOp>(
571 isPostiveX, constPosInfinity,
572 builder.create<SelectOp>(isNegInfinityX, zerof32Const, underflow)));
573
574 rewriter.replaceOp(op, expY);
575
576 return success();
577 }
578
579 //----------------------------------------------------------------------------//
580 // ExpM1 approximation.
581 //----------------------------------------------------------------------------//
582
583 namespace {
584
585 struct ExpM1Approximation : public OpRewritePattern<math::ExpM1Op> {
586 public:
587 using OpRewritePattern::OpRewritePattern;
588
589 LogicalResult matchAndRewrite(math::ExpM1Op op,
590 PatternRewriter &rewriter) const final;
591 };
592 } // namespace
593
594 LogicalResult
matchAndRewrite(math::ExpM1Op op,PatternRewriter & rewriter) const595 ExpM1Approximation::matchAndRewrite(math::ExpM1Op op,
596 PatternRewriter &rewriter) const {
597 auto width = vectorWidth(op.operand().getType(), isF32);
598 if (!width.hasValue())
599 return rewriter.notifyMatchFailure(op, "unsupported operand type");
600
601 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
602 auto bcast = [&](Value value) -> Value {
603 return broadcast(builder, value, *width);
604 };
605
606 // expm1(x) = exp(x) - 1 = u - 1.
607 // We have to handle it carefully when x is near 0, i.e. u ~= 1,
608 // and when the input is ~= -inf, i.e. u - 1 ~= -1.
609 Value cstOne = bcast(f32Cst(builder, 1.0f));
610 Value cstNegOne = bcast(f32Cst(builder, -1.0f));
611 Value x = op.operand();
612 Value u = builder.create<math::ExpOp>(x);
613 Value uEqOne = builder.create<CmpFOp>(CmpFPredicate::OEQ, u, cstOne);
614 Value uMinusOne = builder.create<SubFOp>(u, cstOne);
615 Value uMinusOneEqNegOne =
616 builder.create<CmpFOp>(CmpFPredicate::OEQ, uMinusOne, cstNegOne);
617 // logU = log(u) ~= x
618 Value logU = builder.create<math::LogOp>(u);
619
620 // Detect exp(x) = +inf; written this way to avoid having to form +inf.
621 Value isInf = builder.create<CmpFOp>(CmpFPredicate::OEQ, logU, u);
622
623 // (u - 1) * (x / ~x)
624 Value expm1 =
625 builder.create<MulFOp>(uMinusOne, builder.create<DivFOp>(x, logU));
626 expm1 = builder.create<SelectOp>(isInf, u, expm1);
627 Value approximation = builder.create<SelectOp>(
628 uEqOne, x, builder.create<SelectOp>(uMinusOneEqNegOne, cstNegOne, expm1));
629 rewriter.replaceOp(op, approximation);
630 return success();
631 }
632
633 //----------------------------------------------------------------------------//
634 // Sin and Cos approximation.
635 //----------------------------------------------------------------------------//
636
637 namespace {
638
639 template <bool isSine, typename OpTy>
640 struct SinAndCosApproximation : public OpRewritePattern<OpTy> {
641 public:
642 using OpRewritePattern<OpTy>::OpRewritePattern;
643
644 LogicalResult matchAndRewrite(OpTy op, PatternRewriter &rewriter) const final;
645 };
646 } // namespace
647
648 #define TWO_OVER_PI \
649 0.6366197723675813430755350534900574481378385829618257949906693762L
650 #define PI_OVER_2 \
651 1.5707963267948966192313216916397514420985846996875529104874722961L
652
653 // Approximates sin(x) or cos(x) by finding the best approximation polynomial in
654 // the reduced range [0, pi/2] for both sin(x) and cos(x). Then given y in the
655 // reduced range sin(x) will be computed as sin(y), -sin(y), cos(y) or -cos(y).
656 template <bool isSine, typename OpTy>
matchAndRewrite(OpTy op,PatternRewriter & rewriter) const657 LogicalResult SinAndCosApproximation<isSine, OpTy>::matchAndRewrite(
658 OpTy op, PatternRewriter &rewriter) const {
659 static_assert(
660 llvm::is_one_of<OpTy, math::SinOp, math::CosOp>::value,
661 "SinAndCosApproximation pattern expects math::SinOp or math::CosOp");
662 auto width = vectorWidth(op.operand().getType(), isF32);
663 if (!width.hasValue())
664 return rewriter.notifyMatchFailure(op, "unsupported operand type");
665
666 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
667 auto bcast = [&](Value value) -> Value {
668 return broadcast(builder, value, *width);
669 };
670 auto mul = [&](Value a, Value b) -> Value {
671 return builder.create<MulFOp>(a, b);
672 };
673 auto sub = [&](Value a, Value b) -> Value {
674 return builder.create<SubFOp>(a, b);
675 };
676 auto floor = [&](Value a) { return builder.create<FloorFOp>(a); };
677
678 auto i32Vec = broadcast(builder.getI32Type(), *width);
679 auto fPToSingedInteger = [&](Value a) -> Value {
680 return builder.create<FPToSIOp>(a, i32Vec);
681 };
682
683 auto modulo4 = [&](Value a) -> Value {
684 return builder.create<AndOp>(a, bcast(i32Cst(builder, 3)));
685 };
686
687 auto isEqualTo = [&](Value a, Value b) -> Value {
688 return builder.create<CmpIOp>(CmpIPredicate::eq, a, b);
689 };
690
691 auto isGreaterThan = [&](Value a, Value b) -> Value {
692 return builder.create<CmpIOp>(CmpIPredicate::sgt, a, b);
693 };
694
695 auto select = [&](Value cond, Value t, Value f) -> Value {
696 return builder.create<SelectOp>(cond, t, f);
697 };
698
699 auto fmla = [&](Value a, Value b, Value c) {
700 return builder.create<FmaFOp>(a, b, c);
701 };
702
703 auto bitwiseOr = [&](Value a, Value b) { return builder.create<OrOp>(a, b); };
704
705 Value twoOverPi = bcast(f32Cst(builder, TWO_OVER_PI));
706 Value piOverTwo = bcast(f32Cst(builder, PI_OVER_2));
707
708 Value x = op.operand();
709
710 Value k = floor(mul(x, twoOverPi));
711
712 Value y = sub(x, mul(k, piOverTwo));
713
714 Value cstOne = bcast(f32Cst(builder, 1.0));
715 Value cstNegativeOne = bcast(f32Cst(builder, -1.0));
716
717 Value cstSC2 = bcast(f32Cst(builder, -0.16666667163372039794921875f));
718 Value cstSC4 = bcast(f32Cst(builder, 8.333347737789154052734375e-3f));
719 Value cstSC6 = bcast(f32Cst(builder, -1.9842604524455964565277099609375e-4f));
720 Value cstSC8 =
721 bcast(f32Cst(builder, 2.760012648650445044040679931640625e-6f));
722 Value cstSC10 =
723 bcast(f32Cst(builder, -2.50293279435709337121807038784027099609375e-8f));
724
725 Value cstCC2 = bcast(f32Cst(builder, -0.5f));
726 Value cstCC4 = bcast(f32Cst(builder, 4.166664183139801025390625e-2f));
727 Value cstCC6 = bcast(f32Cst(builder, -1.388833043165504932403564453125e-3f));
728 Value cstCC8 = bcast(f32Cst(builder, 2.47562347794882953166961669921875e-5f));
729 Value cstCC10 =
730 bcast(f32Cst(builder, -2.59630184018533327616751194000244140625e-7f));
731
732 Value kMod4 = modulo4(fPToSingedInteger(k));
733
734 Value kR0 = isEqualTo(kMod4, bcast(i32Cst(builder, 0)));
735 Value kR1 = isEqualTo(kMod4, bcast(i32Cst(builder, 1)));
736 Value kR2 = isEqualTo(kMod4, bcast(i32Cst(builder, 2)));
737 Value kR3 = isEqualTo(kMod4, bcast(i32Cst(builder, 3)));
738
739 Value sinuseCos = isSine ? bitwiseOr(kR1, kR3) : bitwiseOr(kR0, kR2);
740 Value negativeRange = isSine ? isGreaterThan(kMod4, bcast(i32Cst(builder, 1)))
741 : bitwiseOr(kR1, kR2);
742
743 Value y2 = mul(y, y);
744
745 Value base = select(sinuseCos, cstOne, y);
746 Value cstC2 = select(sinuseCos, cstCC2, cstSC2);
747 Value cstC4 = select(sinuseCos, cstCC4, cstSC4);
748 Value cstC6 = select(sinuseCos, cstCC6, cstSC6);
749 Value cstC8 = select(sinuseCos, cstCC8, cstSC8);
750 Value cstC10 = select(sinuseCos, cstCC10, cstSC10);
751
752 Value v1 = fmla(y2, cstC10, cstC8);
753 Value v2 = fmla(y2, v1, cstC6);
754 Value v3 = fmla(y2, v2, cstC4);
755 Value v4 = fmla(y2, v3, cstC2);
756 Value v5 = fmla(y2, v4, cstOne);
757 Value v6 = mul(base, v5);
758
759 Value approximation = select(negativeRange, mul(cstNegativeOne, v6), v6);
760
761 rewriter.replaceOp(op, approximation);
762
763 return success();
764 }
765
766 //----------------------------------------------------------------------------//
767
populateMathPolynomialApproximationPatterns(RewritePatternSet & patterns)768 void mlir::populateMathPolynomialApproximationPatterns(
769 RewritePatternSet &patterns) {
770 patterns.add<TanhApproximation, LogApproximation, Log2Approximation,
771 Log1pApproximation, ExpApproximation, ExpM1Approximation,
772 SinAndCosApproximation<true, math::SinOp>,
773 SinAndCosApproximation<false, math::CosOp>>(
774 patterns.getContext());
775 }
776