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