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