MLIR  14.0.0git
PolynomialApproximation.cpp
Go to the documentation of this file.
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 <climits>
15 #include <cstddef>
16 
24 #include "mlir/IR/Builders.h"
26 #include "mlir/IR/TypeUtilities.h"
29 #include "llvm/ADT/ArrayRef.h"
30 
31 using namespace mlir;
32 using namespace mlir::math;
33 using namespace mlir::vector;
34 
35 // Returns vector shape if the type is a vector. Returns an empty shape if it is
36 // not a vector.
38  auto vectorType = type.dyn_cast<VectorType>();
39  return vectorType ? vectorType.getShape() : ArrayRef<int64_t>();
40 }
41 
43  return vectorShape(value.getType());
44 }
45 
46 //----------------------------------------------------------------------------//
47 // Broadcast scalar types and values into vector types and values.
48 //----------------------------------------------------------------------------//
49 
50 // Broadcasts scalar type into vector type (iff shape is non-scalar).
51 static Type broadcast(Type type, ArrayRef<int64_t> shape) {
52  assert(!type.isa<VectorType>() && "must be scalar type");
53  return !shape.empty() ? VectorType::get(shape, type) : type;
54 }
55 
56 // Broadcasts scalar value into vector (iff shape is non-scalar).
58  ArrayRef<int64_t> shape) {
59  assert(!value.getType().isa<VectorType>() && "must be scalar value");
60  auto type = broadcast(value.getType(), shape);
61  return !shape.empty() ? builder.create<BroadcastOp>(type, value) : value;
62 }
63 
64 //----------------------------------------------------------------------------//
65 // Helper function to handle n-D vectors with 1-D operations.
66 //----------------------------------------------------------------------------//
67 
68 // Expands and unrolls n-D vector operands into multiple fixed size 1-D vectors
69 // and calls the compute function with 1-D vector operands. Stitches back all
70 // results into the original n-D vector result.
71 //
72 // Examples: vectorWidth = 8
73 // - vector<4x8xf32> unrolled 4 times
74 // - vector<16xf32> expanded to vector<2x8xf32> and unrolled 2 times
75 // - vector<4x16xf32> expanded to vector<4x2x8xf32> and unrolled 4*2 times
76 //
77 // Some math approximations rely on ISA-specific operations that only accept
78 // fixed size 1-D vectors (e.g. AVX expects vectors of width 8).
79 //
80 // It is the caller's responsibility to verify that the inner dimension is
81 // divisible by the vectorWidth, and that all operands have the same vector
82 // shape.
83 static Value
85  ValueRange operands, int64_t vectorWidth,
87  assert(!operands.empty() && "operands must be not empty");
88  assert(vectorWidth > 0 && "vector width must be larger than 0");
89 
90  VectorType inputType = operands[0].getType().cast<VectorType>();
91  ArrayRef<int64_t> inputShape = inputType.getShape();
92 
93  // If input shape matches target vector width, we can just call the
94  // user-provided compute function with the operands.
95  if (inputShape == llvm::makeArrayRef(vectorWidth))
96  return compute(operands);
97 
98  // Check if the inner dimension has to be expanded, or we can directly iterate
99  // over the outer dimensions of the vector.
100  int64_t innerDim = inputShape.back();
101  int64_t expansionDim = innerDim / vectorWidth;
102  assert((innerDim % vectorWidth == 0) && "invalid inner dimension size");
103 
104  // Maybe expand operands to the higher rank vector shape that we'll use to
105  // iterate over and extract one dimensional vectors.
106  SmallVector<int64_t> expandedShape(inputShape.begin(), inputShape.end());
107  SmallVector<Value> expandedOperands(operands);
108 
109  if (expansionDim > 1) {
110  // Expand shape from [..., innerDim] to [..., expansionDim, vectorWidth].
111  expandedShape.insert(expandedShape.end() - 1, expansionDim);
112  expandedShape.back() = vectorWidth;
113 
114  for (unsigned i = 0; i < operands.size(); ++i) {
115  auto operand = operands[i];
116  auto eltType = operand.getType().cast<VectorType>().getElementType();
117  auto expandedType = VectorType::get(expandedShape, eltType);
118  expandedOperands[i] =
119  builder.create<vector::ShapeCastOp>(expandedType, operand);
120  }
121  }
122 
123  // Iterate over all outer dimensions of the compute shape vector type.
124  auto iterationDims = ArrayRef<int64_t>(expandedShape).drop_back();
125  int64_t maxLinearIndex = computeMaxLinearIndex(iterationDims);
126 
127  SmallVector<int64_t> ones(iterationDims.size(), 1);
128  auto strides = computeStrides(iterationDims, ones);
129 
130  // Compute results for each one dimensional vector.
131  SmallVector<Value> results(maxLinearIndex);
132 
133  for (int64_t i = 0; i < maxLinearIndex; ++i) {
134  auto offsets = delinearize(strides, i);
135 
136  SmallVector<Value> extracted(expandedOperands.size());
137  for (const auto &tuple : llvm::enumerate(expandedOperands))
138  extracted[tuple.index()] =
139  builder.create<vector::ExtractOp>(tuple.value(), offsets);
140 
141  results[i] = compute(extracted);
142  }
143 
144  // Stitch results together into one large vector.
145  Type resultEltType = results[0].getType().cast<VectorType>().getElementType();
146  Type resultExpandedType = VectorType::get(expandedShape, resultEltType);
147  Value result = builder.create<ConstantOp>(
148  resultExpandedType, builder.getZeroAttr(resultExpandedType));
149 
150  for (int64_t i = 0; i < maxLinearIndex; ++i)
151  result = builder.create<vector::InsertOp>(results[i], result,
152  delinearize(strides, i));
153 
154  // Reshape back to the original vector shape.
155  return builder.create<vector::ShapeCastOp>(
156  VectorType::get(inputShape, resultEltType), result);
157 }
158 
159 //----------------------------------------------------------------------------//
160 // Helper functions to create constants.
161 //----------------------------------------------------------------------------//
162 
163 static Value f32Cst(ImplicitLocOpBuilder &builder, float value) {
164  return builder.create<arith::ConstantOp>(builder.getF32FloatAttr(value));
165 }
166 
167 static Value i32Cst(ImplicitLocOpBuilder &builder, int32_t value) {
168  return builder.create<arith::ConstantOp>(builder.getI32IntegerAttr(value));
169 }
170 
171 static Value f32FromBits(ImplicitLocOpBuilder &builder, uint32_t bits) {
172  Value i32Value = i32Cst(builder, static_cast<int32_t>(bits));
173  return builder.create<arith::BitcastOp>(builder.getF32Type(), i32Value);
174 }
175 
176 //----------------------------------------------------------------------------//
177 // Helper functions to build math functions approximations.
178 //----------------------------------------------------------------------------//
179 
180 static Value min(ImplicitLocOpBuilder &builder, Value a, Value b) {
181  return builder.create<SelectOp>(
182  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, a, b), a, b);
183 }
184 
185 static Value max(ImplicitLocOpBuilder &builder, Value a, Value b) {
186  return builder.create<SelectOp>(
187  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, a, b), a, b);
188 }
189 
190 static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound,
191  Value upperBound) {
192  return max(builder, min(builder, value, upperBound), lowerBound);
193 }
194 
195 // Decomposes given floating point value `arg` into a normalized fraction and
196 // an integral power of two (see std::frexp). Returned values have float type.
197 static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
198  bool isPositive = false) {
199  assert(getElementTypeOrSelf(arg).isF32() && "arg must be f32 type");
200  ArrayRef<int64_t> shape = vectorShape(arg);
201 
202  auto bcast = [&](Value value) -> Value {
203  return broadcast(builder, value, shape);
204  };
205 
206  auto i32 = builder.getIntegerType(32);
207  auto i32Vec = broadcast(i32, shape);
208  auto f32Vec = broadcast(builder.getF32Type(), shape);
209 
210  Value cst126f = f32Cst(builder, 126.0f);
211  Value cstHalf = f32Cst(builder, 0.5f);
212  Value cstInvMantMask = f32FromBits(builder, ~0x7f800000u);
213 
214  // Bitcast to i32 for bitwise operations.
215  Value i32Half = builder.create<arith::BitcastOp>(i32, cstHalf);
216  Value i32InvMantMask = builder.create<arith::BitcastOp>(i32, cstInvMantMask);
217  Value i32Arg = builder.create<arith::BitcastOp>(i32Vec, arg);
218 
219  // Compute normalized fraction.
220  Value tmp0 = builder.create<arith::AndIOp>(i32Arg, bcast(i32InvMantMask));
221  Value tmp1 = builder.create<arith::OrIOp>(tmp0, bcast(i32Half));
222  Value normalizedFraction = builder.create<arith::BitcastOp>(f32Vec, tmp1);
223 
224  // Compute exponent.
225  Value arg0 = isPositive ? arg : builder.create<math::AbsOp>(arg);
226  Value biasedExponentBits = builder.create<arith::ShRUIOp>(
227  builder.create<arith::BitcastOp>(i32Vec, arg0),
228  bcast(i32Cst(builder, 23)));
229  Value biasedExponent =
230  builder.create<arith::SIToFPOp>(f32Vec, biasedExponentBits);
231  Value exponent =
232  builder.create<arith::SubFOp>(biasedExponent, bcast(cst126f));
233 
234  return {normalizedFraction, exponent};
235 }
236 
237 // Computes exp2 for an i32 argument.
238 static Value exp2I32(ImplicitLocOpBuilder &builder, Value arg) {
239  assert(getElementTypeOrSelf(arg).isInteger(32) && "arg must be i32 type");
240  ArrayRef<int64_t> shape = vectorShape(arg);
241 
242  auto bcast = [&](Value value) -> Value {
243  return broadcast(builder, value, shape);
244  };
245 
246  auto f32Vec = broadcast(builder.getF32Type(), shape);
247  // The exponent of f32 located at 23-bit.
248  auto exponetBitLocation = bcast(i32Cst(builder, 23));
249  // Set the exponent bias to zero.
250  auto bias = bcast(i32Cst(builder, 127));
251 
252  Value biasedArg = builder.create<arith::AddIOp>(arg, bias);
253  Value exp2ValueInt =
254  builder.create<arith::ShLIOp>(biasedArg, exponetBitLocation);
255  Value exp2ValueF32 = builder.create<arith::BitcastOp>(f32Vec, exp2ValueInt);
256 
257  return exp2ValueF32;
258 }
259 
260 namespace {
261 Value makePolynomialCalculation(ImplicitLocOpBuilder &builder,
262  llvm::ArrayRef<Value> coeffs, Value x) {
263  assert(getElementTypeOrSelf(x).isF32() && "x must be f32 type");
264  ArrayRef<int64_t> shape = vectorShape(x);
265 
266  if (coeffs.empty())
267  return broadcast(builder, f32Cst(builder, 0.0f), shape);
268 
269  if (coeffs.size() == 1)
270  return coeffs[0];
271 
272  Value res = builder.create<math::FmaOp>(x, coeffs[coeffs.size() - 1],
273  coeffs[coeffs.size() - 2]);
274  for (auto i = ptrdiff_t(coeffs.size()) - 3; i >= 0; --i) {
275  res = builder.create<math::FmaOp>(x, res, coeffs[i]);
276  }
277  return res;
278 }
279 } // namespace
280 
281 //----------------------------------------------------------------------------//
282 // AtanOp approximation.
283 //----------------------------------------------------------------------------//
284 
285 namespace {
286 struct AtanApproximation : public OpRewritePattern<math::AtanOp> {
287 public:
289 
290  LogicalResult matchAndRewrite(math::AtanOp op,
291  PatternRewriter &rewriter) const final;
292 };
293 } // namespace
294 
296 AtanApproximation::matchAndRewrite(math::AtanOp op,
297  PatternRewriter &rewriter) const {
298  auto operand = op.getOperand();
299  if (!getElementTypeOrSelf(operand).isF32())
300  return rewriter.notifyMatchFailure(op, "unsupported operand type");
301 
302  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
303 
304  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
305  auto one = broadcast(builder, f32Cst(builder, 1.0f), shape);
306 
307  // Remap the problem over [0.0, 1.0] by looking at the absolute value and the
308  // handling symmetry.
309  Value abs = builder.create<math::AbsOp>(operand);
310  Value reciprocal = builder.create<arith::DivFOp>(one, abs);
311  Value compare =
312  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, abs, reciprocal);
313  Value x = builder.create<SelectOp>(compare, abs, reciprocal);
314 
315  // Perform the Taylor series approximation for atan over the range
316  // [-1.0, 1.0].
317  auto n1 = broadcast(builder, f32Cst(builder, 0.14418283), shape);
318  auto n2 = broadcast(builder, f32Cst(builder, -0.34999234), shape);
319  auto n3 = broadcast(builder, f32Cst(builder, -0.01067831), shape);
320  auto n4 = broadcast(builder, f32Cst(builder, 1.00209986), shape);
321 
322  Value p = builder.create<math::FmaOp>(x, n1, n2);
323  p = builder.create<math::FmaOp>(x, p, n3);
324  p = builder.create<math::FmaOp>(x, p, n4);
325  p = builder.create<arith::MulFOp>(x, p);
326 
327  // Remap the solution for over [0.0, 1.0] to [0.0, inf]
328  auto half_pi = broadcast(builder, f32Cst(builder, 1.57079632679f), shape);
329  Value sub = builder.create<arith::SubFOp>(half_pi, p);
330  Value select = builder.create<SelectOp>(compare, p, sub);
331 
332  // Correct for signing of the input.
333  rewriter.replaceOpWithNewOp<math::CopySignOp>(op, select, operand);
334  return success();
335 }
336 
337 //----------------------------------------------------------------------------//
338 // AtanOp approximation.
339 //----------------------------------------------------------------------------//
340 
341 namespace {
342 struct Atan2Approximation : public OpRewritePattern<math::Atan2Op> {
343 public:
345 
346  LogicalResult matchAndRewrite(math::Atan2Op op,
347  PatternRewriter &rewriter) const final;
348 };
349 } // namespace
350 
352 Atan2Approximation::matchAndRewrite(math::Atan2Op op,
353  PatternRewriter &rewriter) const {
354  auto y = op.getOperand(0);
355  auto x = op.getOperand(1);
356  if (!getElementTypeOrSelf(x).isF32())
357  return rewriter.notifyMatchFailure(op, "unsupported operand type");
358 
359  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
360  ArrayRef<int64_t> shape = vectorShape(op.getResult());
361 
362  // Compute atan in the valid range.
363  auto div = builder.create<arith::DivFOp>(y, x);
364  auto atan = builder.create<math::AtanOp>(div);
365 
366  // Determine what the atan would be for a 180 degree rotation.
367  auto zero = broadcast(builder, f32Cst(builder, 0.0f), shape);
368  auto pi = broadcast(builder, f32Cst(builder, 3.14159265359f), shape);
369  auto add_pi = builder.create<arith::AddFOp>(atan, pi);
370  auto sub_pi = builder.create<arith::SubFOp>(atan, pi);
371  auto atan_gt =
372  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, atan, zero);
373  auto flipped_atan = builder.create<SelectOp>(atan_gt, sub_pi, add_pi);
374 
375  // Determine whether to directly use atan or use the 180 degree flip
376  auto x_gt = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, x, zero);
377  Value result = builder.create<SelectOp>(x_gt, atan, flipped_atan);
378 
379  // Handle x = 0, y > 0
380  Value x_zero =
381  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, x, zero);
382  Value y_gt =
383  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, y, zero);
384  Value is_half_pi = builder.create<arith::AndIOp>(x_zero, y_gt);
385  auto half_pi = broadcast(builder, f32Cst(builder, 1.57079632679f), shape);
386  result = builder.create<SelectOp>(is_half_pi, half_pi, result);
387 
388  // Handle x = 0, y < 0
389  Value y_lt =
390  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, y, zero);
391  Value is_negative_half_pi_pi = builder.create<arith::AndIOp>(x_zero, y_lt);
392  auto negative_half_pi_pi =
393  broadcast(builder, f32Cst(builder, -1.57079632679), shape);
394  result = builder.create<SelectOp>(is_negative_half_pi_pi, negative_half_pi_pi,
395  result);
396 
397  // Handle x = 0, y = 0;
398  Value y_zero =
399  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, y, zero);
400  Value is_nan = builder.create<arith::AndIOp>(x_zero, y_zero);
401  Value cst_nan = broadcast(builder, f32FromBits(builder, 0x7fc00000), shape);
402  result = builder.create<SelectOp>(is_nan, cst_nan, result);
403 
404  rewriter.replaceOp(op, result);
405  return success();
406 }
407 
408 //----------------------------------------------------------------------------//
409 // TanhOp approximation.
410 //----------------------------------------------------------------------------//
411 
412 namespace {
413 struct TanhApproximation : public OpRewritePattern<math::TanhOp> {
414 public:
416 
417  LogicalResult matchAndRewrite(math::TanhOp op,
418  PatternRewriter &rewriter) const final;
419 };
420 } // namespace
421 
423 TanhApproximation::matchAndRewrite(math::TanhOp op,
424  PatternRewriter &rewriter) const {
425  if (!getElementTypeOrSelf(op.getOperand()).isF32())
426  return rewriter.notifyMatchFailure(op, "unsupported operand type");
427 
428  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
429 
430  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
431  auto bcast = [&](Value value) -> Value {
432  return broadcast(builder, value, shape);
433  };
434 
435  // Clamp operand into [plusClamp, minusClamp] range.
436  Value minusClamp = bcast(f32Cst(builder, -7.99881172180175781f));
437  Value plusClamp = bcast(f32Cst(builder, 7.99881172180175781f));
438  Value x = clamp(builder, op.getOperand(), minusClamp, plusClamp);
439 
440  // Mask for tiny values that are approximated with `operand`.
441  Value tiny = bcast(f32Cst(builder, 0.0004f));
442  Value tinyMask = builder.create<arith::CmpFOp>(
443  arith::CmpFPredicate::OLT, builder.create<math::AbsOp>(op.getOperand()),
444  tiny);
445 
446  // The monomial coefficients of the numerator polynomial (odd).
447  Value alpha1 = bcast(f32Cst(builder, 4.89352455891786e-03f));
448  Value alpha3 = bcast(f32Cst(builder, 6.37261928875436e-04f));
449  Value alpha5 = bcast(f32Cst(builder, 1.48572235717979e-05f));
450  Value alpha7 = bcast(f32Cst(builder, 5.12229709037114e-08f));
451  Value alpha9 = bcast(f32Cst(builder, -8.60467152213735e-11f));
452  Value alpha11 = bcast(f32Cst(builder, 2.00018790482477e-13f));
453  Value alpha13 = bcast(f32Cst(builder, -2.76076847742355e-16f));
454 
455  // The monomial coefficients of the denominator polynomial (even).
456  Value beta0 = bcast(f32Cst(builder, 4.89352518554385e-03f));
457  Value beta2 = bcast(f32Cst(builder, 2.26843463243900e-03f));
458  Value beta4 = bcast(f32Cst(builder, 1.18534705686654e-04f));
459  Value beta6 = bcast(f32Cst(builder, 1.19825839466702e-06f));
460 
461  // Since the polynomials are odd/even, we need x^2.
462  Value x2 = builder.create<arith::MulFOp>(x, x);
463 
464  // Evaluate the numerator polynomial p.
465  Value p = builder.create<math::FmaOp>(x2, alpha13, alpha11);
466  p = builder.create<math::FmaOp>(x2, p, alpha9);
467  p = builder.create<math::FmaOp>(x2, p, alpha7);
468  p = builder.create<math::FmaOp>(x2, p, alpha5);
469  p = builder.create<math::FmaOp>(x2, p, alpha3);
470  p = builder.create<math::FmaOp>(x2, p, alpha1);
471  p = builder.create<arith::MulFOp>(x, p);
472 
473  // Evaluate the denominator polynomial q.
474  Value q = builder.create<math::FmaOp>(x2, beta6, beta4);
475  q = builder.create<math::FmaOp>(x2, q, beta2);
476  q = builder.create<math::FmaOp>(x2, q, beta0);
477 
478  // Divide the numerator by the denominator.
479  Value res = builder.create<SelectOp>(tinyMask, x,
480  builder.create<arith::DivFOp>(p, q));
481 
482  rewriter.replaceOp(op, res);
483 
484  return success();
485 }
486 
487 #define LN2_VALUE \
488  0.693147180559945309417232121458176568075500134360255254120680009493393621L
489 #define LOG2E_VALUE \
490  1.442695040888963407359924681001892137426645954152985934135449406931109219L
491 
492 //----------------------------------------------------------------------------//
493 // LogOp and Log2Op approximation.
494 //----------------------------------------------------------------------------//
495 
496 namespace {
497 template <typename Op>
498 struct LogApproximationBase : public OpRewritePattern<Op> {
500 
501  /// Base 2 if 'base2' is set; natural logarithm (base e) otherwise.
502  LogicalResult logMatchAndRewrite(Op op, PatternRewriter &rewriter,
503  bool base2) const;
504 };
505 } // namespace
506 
507 // This approximation comes from Julien Pommier's SSE math library.
508 // Link: http://gruntthepeon.free.fr/ssemath
509 template <typename Op>
511 LogApproximationBase<Op>::logMatchAndRewrite(Op op, PatternRewriter &rewriter,
512  bool base2) const {
513  if (!getElementTypeOrSelf(op.getOperand()).isF32())
514  return rewriter.notifyMatchFailure(op, "unsupported operand type");
515 
516  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
517 
518  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
519  auto bcast = [&](Value value) -> Value {
520  return broadcast(builder, value, shape);
521  };
522 
523  Value cstZero = bcast(f32Cst(builder, 0.0f));
524  Value cstOne = bcast(f32Cst(builder, 1.0f));
525  Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
526 
527  // The smallest non denormalized float number.
528  Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
529  Value cstMinusInf = bcast(f32FromBits(builder, 0xff800000u));
530  Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u));
531  Value cstNan = bcast(f32FromBits(builder, 0x7fc00000));
532 
533  // Polynomial coefficients.
534  Value cstCephesSQRTHF = bcast(f32Cst(builder, 0.707106781186547524f));
535  Value cstCephesLogP0 = bcast(f32Cst(builder, 7.0376836292E-2f));
536  Value cstCephesLogP1 = bcast(f32Cst(builder, -1.1514610310E-1f));
537  Value cstCephesLogP2 = bcast(f32Cst(builder, 1.1676998740E-1f));
538  Value cstCephesLogP3 = bcast(f32Cst(builder, -1.2420140846E-1f));
539  Value cstCephesLogP4 = bcast(f32Cst(builder, +1.4249322787E-1f));
540  Value cstCephesLogP5 = bcast(f32Cst(builder, -1.6668057665E-1f));
541  Value cstCephesLogP6 = bcast(f32Cst(builder, +2.0000714765E-1f));
542  Value cstCephesLogP7 = bcast(f32Cst(builder, -2.4999993993E-1f));
543  Value cstCephesLogP8 = bcast(f32Cst(builder, +3.3333331174E-1f));
544 
545  Value x = op.getOperand();
546 
547  // Truncate input values to the minimum positive normal.
548  x = max(builder, x, cstMinNormPos);
549 
550  // Extract significant in the range [0.5,1) and exponent.
551  std::pair<Value, Value> pair = frexp(builder, x, /*isPositive=*/true);
552  x = pair.first;
553  Value e = pair.second;
554 
555  // Shift the inputs from the range [0.5,1) to [sqrt(1/2), sqrt(2)) and shift
556  // by -1.0. The values are then centered around 0, which improves the
557  // stability of the polynomial evaluation:
558  //
559  // if( x < SQRTHF ) {
560  // e -= 1;
561  // x = x + x - 1.0;
562  // } else { x = x - 1.0; }
563  Value mask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, x,
564  cstCephesSQRTHF);
565  Value tmp = builder.create<SelectOp>(mask, x, cstZero);
566 
567  x = builder.create<arith::SubFOp>(x, cstOne);
568  e = builder.create<arith::SubFOp>(
569  e, builder.create<SelectOp>(mask, cstOne, cstZero));
570  x = builder.create<arith::AddFOp>(x, tmp);
571 
572  Value x2 = builder.create<arith::MulFOp>(x, x);
573  Value x3 = builder.create<arith::MulFOp>(x2, x);
574 
575  // Evaluate the polynomial approximant of degree 8 in three parts.
576  Value y0, y1, y2;
577  y0 = builder.create<math::FmaOp>(cstCephesLogP0, x, cstCephesLogP1);
578  y1 = builder.create<math::FmaOp>(cstCephesLogP3, x, cstCephesLogP4);
579  y2 = builder.create<math::FmaOp>(cstCephesLogP6, x, cstCephesLogP7);
580  y0 = builder.create<math::FmaOp>(y0, x, cstCephesLogP2);
581  y1 = builder.create<math::FmaOp>(y1, x, cstCephesLogP5);
582  y2 = builder.create<math::FmaOp>(y2, x, cstCephesLogP8);
583  y0 = builder.create<math::FmaOp>(y0, x3, y1);
584  y0 = builder.create<math::FmaOp>(y0, x3, y2);
585  y0 = builder.create<arith::MulFOp>(y0, x3);
586 
587  y0 = builder.create<math::FmaOp>(cstNegHalf, x2, y0);
588  x = builder.create<arith::AddFOp>(x, y0);
589 
590  if (base2) {
591  Value cstLog2e = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE)));
592  x = builder.create<math::FmaOp>(x, cstLog2e, e);
593  } else {
594  Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
595  x = builder.create<math::FmaOp>(e, cstLn2, x);
596  }
597 
598  Value invalidMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::ULT,
599  op.getOperand(), cstZero);
600  Value zeroMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ,
601  op.getOperand(), cstZero);
602  Value posInfMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ,
603  op.getOperand(), cstPosInf);
604 
605  // Filter out invalid values:
606  // • x == 0 -> -INF
607  // • x < 0 -> NAN
608  // • x == +INF -> +INF
609  Value aproximation = builder.create<SelectOp>(
610  zeroMask, cstMinusInf,
611  builder.create<SelectOp>(
612  invalidMask, cstNan,
613  builder.create<SelectOp>(posInfMask, cstPosInf, x)));
614 
615  rewriter.replaceOp(op, aproximation);
616 
617  return success();
618 }
619 
620 namespace {
621 struct LogApproximation : public LogApproximationBase<math::LogOp> {
622  using LogApproximationBase::LogApproximationBase;
623 
624  LogicalResult matchAndRewrite(math::LogOp op,
625  PatternRewriter &rewriter) const final {
626  return logMatchAndRewrite(op, rewriter, /*base2=*/false);
627  }
628 };
629 } // namespace
630 
631 namespace {
632 struct Log2Approximation : public LogApproximationBase<math::Log2Op> {
633  using LogApproximationBase::LogApproximationBase;
634 
635  LogicalResult matchAndRewrite(math::Log2Op op,
636  PatternRewriter &rewriter) const final {
637  return logMatchAndRewrite(op, rewriter, /*base2=*/true);
638  }
639 };
640 } // namespace
641 
642 //----------------------------------------------------------------------------//
643 // Log1p approximation.
644 //----------------------------------------------------------------------------//
645 
646 namespace {
647 struct Log1pApproximation : public OpRewritePattern<math::Log1pOp> {
648 public:
650 
651  LogicalResult matchAndRewrite(math::Log1pOp op,
652  PatternRewriter &rewriter) const final;
653 };
654 } // namespace
655 
656 // Approximate log(1+x).
658 Log1pApproximation::matchAndRewrite(math::Log1pOp op,
659  PatternRewriter &rewriter) const {
660  if (!getElementTypeOrSelf(op.getOperand()).isF32())
661  return rewriter.notifyMatchFailure(op, "unsupported operand type");
662 
663  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
664 
665  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
666  auto bcast = [&](Value value) -> Value {
667  return broadcast(builder, value, shape);
668  };
669 
670  // Approximate log(1+x) using the following, due to W. Kahan:
671  // u = x + 1.0;
672  // if (u == 1.0 || u == inf) return x;
673  // return x * log(u) / (u - 1.0);
674  // ^^^^^^^^^^^^^^^^^^^^^^
675  // "logLarge" below.
676  Value cstOne = bcast(f32Cst(builder, 1.0f));
677  Value x = op.getOperand();
678  Value u = builder.create<arith::AddFOp>(x, cstOne);
679  Value uSmall =
680  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, cstOne);
681  Value logU = builder.create<math::LogOp>(u);
682  Value uInf =
683  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, logU);
684  Value logLarge = builder.create<arith::MulFOp>(
685  x, builder.create<arith::DivFOp>(
686  logU, builder.create<arith::SubFOp>(u, cstOne)));
687  Value approximation = builder.create<SelectOp>(
688  builder.create<arith::OrIOp>(uSmall, uInf), x, logLarge);
689  rewriter.replaceOp(op, approximation);
690  return success();
691 }
692 
693 //----------------------------------------------------------------------------//
694 // Erf approximation.
695 //----------------------------------------------------------------------------//
696 
697 // Approximates erf(x) with
698 // a - P(x)/Q(x)
699 // where P and Q are polynomials of degree 4.
700 // Different coefficients are chosen based on the value of x.
701 // The approximation error is ~2.5e-07.
702 // Boost's minimax tool that utilizes the Remez method was used to find the
703 // coefficients.
706  PatternRewriter &rewriter) const {
707  if (!getElementTypeOrSelf(op.getOperand()).isF32())
708  return rewriter.notifyMatchFailure(op, "unsupported operand type");
709 
710  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
711 
712  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
713  auto bcast = [&](Value value) -> Value {
714  return broadcast(builder, value, shape);
715  };
716 
717  const int intervalsCount = 3;
718  const int polyDegree = 4;
719 
720  Value zero = bcast(f32Cst(builder, 0));
721  Value one = bcast(f32Cst(builder, 1));
722  Value pp[intervalsCount][polyDegree + 1];
723  pp[0][0] = bcast(f32Cst(builder, +0.00000000000000000e+00f));
724  pp[0][1] = bcast(f32Cst(builder, +1.12837916222975858e+00f));
725  pp[0][2] = bcast(f32Cst(builder, -5.23018562988006470e-01f));
726  pp[0][3] = bcast(f32Cst(builder, +2.09741709609267072e-01f));
727  pp[0][4] = bcast(f32Cst(builder, +2.58146801602987875e-02f));
728  pp[1][0] = bcast(f32Cst(builder, +0.00000000000000000e+00f));
729  pp[1][1] = bcast(f32Cst(builder, +1.12750687816789140e+00f));
730  pp[1][2] = bcast(f32Cst(builder, -3.64721408487825775e-01f));
731  pp[1][3] = bcast(f32Cst(builder, +1.18407396425136952e-01f));
732  pp[1][4] = bcast(f32Cst(builder, +3.70645533056476558e-02f));
733  pp[2][0] = bcast(f32Cst(builder, -3.30093071049483172e-03f));
734  pp[2][1] = bcast(f32Cst(builder, +3.51961938357697011e-03f));
735  pp[2][2] = bcast(f32Cst(builder, -1.41373622814988039e-03f));
736  pp[2][3] = bcast(f32Cst(builder, +2.53447094961941348e-04f));
737  pp[2][4] = bcast(f32Cst(builder, -1.71048029455037401e-05f));
738 
739  Value qq[intervalsCount][polyDegree + 1];
740  qq[0][0] = bcast(f32Cst(builder, +1.000000000000000000e+00f));
741  qq[0][1] = bcast(f32Cst(builder, -4.635138185962547255e-01f));
742  qq[0][2] = bcast(f32Cst(builder, +5.192301327279782447e-01f));
743  qq[0][3] = bcast(f32Cst(builder, -1.318089722204810087e-01f));
744  qq[0][4] = bcast(f32Cst(builder, +7.397964654672315005e-02f));
745  qq[1][0] = bcast(f32Cst(builder, +1.00000000000000000e+00f));
746  qq[1][1] = bcast(f32Cst(builder, -3.27607011824493086e-01f));
747  qq[1][2] = bcast(f32Cst(builder, +4.48369090658821977e-01f));
748  qq[1][3] = bcast(f32Cst(builder, -8.83462621207857930e-02f));
749  qq[1][4] = bcast(f32Cst(builder, +5.72442770283176093e-02f));
750  qq[2][0] = bcast(f32Cst(builder, +1.00000000000000000e+00f));
751  qq[2][1] = bcast(f32Cst(builder, -2.06069165953913769e+00f));
752  qq[2][2] = bcast(f32Cst(builder, +1.62705939945477759e+00f));
753  qq[2][3] = bcast(f32Cst(builder, -5.83389859211130017e-01f));
754  qq[2][4] = bcast(f32Cst(builder, +8.21908939856640930e-02f));
755 
756  Value offsets[intervalsCount];
757  offsets[0] = bcast(f32Cst(builder, 0.0f));
758  offsets[1] = bcast(f32Cst(builder, 0.0f));
759  offsets[2] = bcast(f32Cst(builder, 1.0f));
760 
761  Value bounds[intervalsCount];
762  bounds[0] = bcast(f32Cst(builder, 0.8f));
763  bounds[1] = bcast(f32Cst(builder, 2.0f));
764  bounds[2] = bcast(f32Cst(builder, 3.75f));
765 
766  Value isNegativeArg = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT,
767  op.getOperand(), zero);
768  Value negArg = builder.create<arith::NegFOp>(op.getOperand());
769  Value x = builder.create<SelectOp>(isNegativeArg, negArg, op.getOperand());
770 
771  Value offset = offsets[0];
772  Value p[polyDegree + 1];
773  Value q[polyDegree + 1];
774  for (int i = 0; i <= polyDegree; ++i) {
775  p[i] = pp[0][i];
776  q[i] = qq[0][i];
777  }
778 
779  // TODO: maybe use vector stacking to reduce the number of selects.
780  Value isLessThanBound[intervalsCount];
781  for (int j = 0; j < intervalsCount - 1; ++j) {
782  isLessThanBound[j] =
783  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, x, bounds[j]);
784  for (int i = 0; i <= polyDegree; ++i) {
785  p[i] = builder.create<SelectOp>(isLessThanBound[j], p[i], pp[j + 1][i]);
786  q[i] = builder.create<SelectOp>(isLessThanBound[j], q[i], qq[j + 1][i]);
787  }
788  offset =
789  builder.create<SelectOp>(isLessThanBound[j], offset, offsets[j + 1]);
790  }
791  isLessThanBound[intervalsCount - 1] = builder.create<arith::CmpFOp>(
792  arith::CmpFPredicate::ULT, x, bounds[intervalsCount - 1]);
793 
794  Value pPoly = makePolynomialCalculation(builder, p, x);
795  Value qPoly = makePolynomialCalculation(builder, q, x);
796  Value rationalPoly = builder.create<arith::DivFOp>(pPoly, qPoly);
797  Value formula = builder.create<arith::AddFOp>(offset, rationalPoly);
798  formula = builder.create<SelectOp>(isLessThanBound[intervalsCount - 1],
799  formula, one);
800 
801  // erf is odd function: erf(x) = -erf(-x).
802  Value negFormula = builder.create<arith::NegFOp>(formula);
803  Value res = builder.create<SelectOp>(isNegativeArg, negFormula, formula);
804 
805  rewriter.replaceOp(op, res);
806 
807  return success();
808 }
809 
810 //----------------------------------------------------------------------------//
811 // Exp approximation.
812 //----------------------------------------------------------------------------//
813 
814 namespace {
815 
816 struct ExpApproximation : public OpRewritePattern<math::ExpOp> {
817 public:
819 
820  LogicalResult matchAndRewrite(math::ExpOp op,
821  PatternRewriter &rewriter) const final;
822 };
823 } // namespace
824 
825 // Approximate exp(x) using its reduced range exp(y) where y is in the range
826 // [0, ln(2)], let y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2), exp(x)
827 // = exp(y) * 2^k. exp(y).
829 ExpApproximation::matchAndRewrite(math::ExpOp op,
830  PatternRewriter &rewriter) const {
831  if (!getElementTypeOrSelf(op.getOperand()).isF32())
832  return rewriter.notifyMatchFailure(op, "unsupported operand type");
833 
834  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
835 
836  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
837 
838  // TODO: Consider a common pattern rewriter with all methods below to
839  // write the approximations.
840  auto bcast = [&](Value value) -> Value {
841  return broadcast(builder, value, shape);
842  };
843  auto fmla = [&](Value a, Value b, Value c) {
844  return builder.create<math::FmaOp>(a, b, c);
845  };
846  auto mul = [&](Value a, Value b) -> Value {
847  return builder.create<arith::MulFOp>(a, b);
848  };
849  auto sub = [&](Value a, Value b) -> Value {
850  return builder.create<arith::SubFOp>(a, b);
851  };
852  auto floor = [&](Value a) { return builder.create<math::FloorOp>(a); };
853 
854  Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
855  Value cstLog2E = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE)));
856 
857  // Polynomial coefficients.
858  Value cstCephesExpP0 = bcast(f32Cst(builder, 1.0));
859  Value cstCephesExpP1 = bcast(f32Cst(builder, 1.0));
860  Value cstCephesExpP2 = bcast(f32Cst(builder, 0.49970514590562437052f));
861  Value cstCephesExpP3 = bcast(f32Cst(builder, 0.16873890085469545053f));
862  Value cstCephesExpP4 = bcast(f32Cst(builder, 0.03668965196652099192f));
863  Value cstCephesExpP5 = bcast(f32Cst(builder, 0.01314350012789660196f));
864 
865  Value x = op.getOperand();
866 
867  // Reduced y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2)
868  Value xL2Inv = mul(x, cstLog2E);
869  Value kF32 = floor(xL2Inv);
870  Value kLn2 = mul(kF32, cstLn2);
871  Value y = sub(x, kLn2);
872 
873  // Use Estrin's evaluation scheme with 3 independent parts:
874  // P(y)^y : (c0 + c1 y) + (c2 + c3 y) y^2 + (c4 + c5 y) y^4
875  Value y2 = mul(y, y);
876  Value y4 = mul(y2, y2);
877 
878  Value q0 = fmla(cstCephesExpP1, y, cstCephesExpP0);
879  Value q1 = fmla(cstCephesExpP3, y, cstCephesExpP2);
880  Value q2 = fmla(cstCephesExpP5, y, cstCephesExpP4);
881  Value expY = fmla(q1, y2, q0);
882  expY = fmla(q2, y4, expY);
883 
884  auto i32Vec = broadcast(builder.getI32Type(), shape);
885 
886  // exp2(k)
887  Value k = builder.create<arith::FPToSIOp>(kF32, i32Vec);
888  Value exp2KValue = exp2I32(builder, k);
889 
890  // exp(x) = exp(y) * exp2(k)
891  expY = mul(expY, exp2KValue);
892 
893  // Handle overflow, inf and underflow of exp(x). exp(x) range is [0, inf], its
894  // partitioned as the following:
895  // exp(x) = 0, x <= -inf
896  // exp(x) = underflow (min_float), x <= -88
897  // exp(x) = inf (min_float), x >= 88
898  // Note: |k| = 127 is the value where the 8-bits exponent saturates.
899  Value zerof32Const = bcast(f32Cst(builder, 0));
900  auto constPosInfinity =
901  bcast(f32Cst(builder, std::numeric_limits<float>::infinity()));
902  auto constNegIfinity =
903  bcast(f32Cst(builder, -std::numeric_limits<float>::infinity()));
904  auto underflow = bcast(f32Cst(builder, std::numeric_limits<float>::min()));
905 
906  Value kMaxConst = bcast(i32Cst(builder, 127));
907  Value kMaxNegConst = bcast(i32Cst(builder, -127));
908  Value rightBound =
909  builder.create<arith::CmpIOp>(arith::CmpIPredicate::sle, k, kMaxConst);
910  Value leftBound =
911  builder.create<arith::CmpIOp>(arith::CmpIPredicate::sge, k, kMaxNegConst);
912 
913  Value isNegInfinityX = builder.create<arith::CmpFOp>(
914  arith::CmpFPredicate::OEQ, x, constNegIfinity);
915  Value isPosInfinityX = builder.create<arith::CmpFOp>(
916  arith::CmpFPredicate::OEQ, x, constPosInfinity);
917  Value isPostiveX =
918  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, x, zerof32Const);
919  Value isComputable = builder.create<arith::AndIOp>(rightBound, leftBound);
920 
921  expY = builder.create<SelectOp>(
922  isNegInfinityX, zerof32Const,
923  builder.create<SelectOp>(
924  isPosInfinityX, constPosInfinity,
925  builder.create<SelectOp>(isComputable, expY,
926  builder.create<SelectOp>(isPostiveX,
927  constPosInfinity,
928  underflow))));
929 
930  rewriter.replaceOp(op, expY);
931 
932  return success();
933 }
934 
935 //----------------------------------------------------------------------------//
936 // ExpM1 approximation.
937 //----------------------------------------------------------------------------//
938 
939 namespace {
940 
941 struct ExpM1Approximation : public OpRewritePattern<math::ExpM1Op> {
942 public:
944 
945  LogicalResult matchAndRewrite(math::ExpM1Op op,
946  PatternRewriter &rewriter) const final;
947 };
948 } // namespace
949 
951 ExpM1Approximation::matchAndRewrite(math::ExpM1Op op,
952  PatternRewriter &rewriter) const {
953  if (!getElementTypeOrSelf(op.getOperand()).isF32())
954  return rewriter.notifyMatchFailure(op, "unsupported operand type");
955 
956  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
957 
958  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
959  auto bcast = [&](Value value) -> Value {
960  return broadcast(builder, value, shape);
961  };
962 
963  // expm1(x) = exp(x) - 1 = u - 1.
964  // We have to handle it carefully when x is near 0, i.e. u ~= 1,
965  // and when the input is ~= -inf, i.e. u - 1 ~= -1.
966  Value cstOne = bcast(f32Cst(builder, 1.0f));
967  Value cstNegOne = bcast(f32Cst(builder, -1.0f));
968  Value x = op.getOperand();
969  Value u = builder.create<math::ExpOp>(x);
970  Value uEqOne =
971  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, cstOne);
972  Value uMinusOne = builder.create<arith::SubFOp>(u, cstOne);
973  Value uMinusOneEqNegOne = builder.create<arith::CmpFOp>(
974  arith::CmpFPredicate::OEQ, uMinusOne, cstNegOne);
975  // logU = log(u) ~= x
976  Value logU = builder.create<math::LogOp>(u);
977 
978  // Detect exp(x) = +inf; written this way to avoid having to form +inf.
979  Value isInf =
980  builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, logU, u);
981 
982  // (u - 1) * (x / ~x)
983  Value expm1 = builder.create<arith::MulFOp>(
984  uMinusOne, builder.create<arith::DivFOp>(x, logU));
985  expm1 = builder.create<SelectOp>(isInf, u, expm1);
986  Value approximation = builder.create<SelectOp>(
987  uEqOne, x, builder.create<SelectOp>(uMinusOneEqNegOne, cstNegOne, expm1));
988  rewriter.replaceOp(op, approximation);
989  return success();
990 }
991 
992 //----------------------------------------------------------------------------//
993 // Sin and Cos approximation.
994 //----------------------------------------------------------------------------//
995 
996 namespace {
997 
998 template <bool isSine, typename OpTy>
999 struct SinAndCosApproximation : public OpRewritePattern<OpTy> {
1000 public:
1002 
1003  LogicalResult matchAndRewrite(OpTy op, PatternRewriter &rewriter) const final;
1004 };
1005 } // namespace
1006 
1007 #define TWO_OVER_PI \
1008  0.6366197723675813430755350534900574481378385829618257949906693762L
1009 #define PI_OVER_2 \
1010  1.5707963267948966192313216916397514420985846996875529104874722961L
1011 
1012 // Approximates sin(x) or cos(x) by finding the best approximation polynomial in
1013 // the reduced range [0, pi/2] for both sin(x) and cos(x). Then given y in the
1014 // reduced range sin(x) will be computed as sin(y), -sin(y), cos(y) or -cos(y).
1015 template <bool isSine, typename OpTy>
1016 LogicalResult SinAndCosApproximation<isSine, OpTy>::matchAndRewrite(
1017  OpTy op, PatternRewriter &rewriter) const {
1018  static_assert(
1020  "SinAndCosApproximation pattern expects math::SinOp or math::CosOp");
1021 
1022  if (!getElementTypeOrSelf(op.getOperand()).isF32())
1023  return rewriter.notifyMatchFailure(op, "unsupported operand type");
1024 
1025  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
1026 
1027  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
1028  auto bcast = [&](Value value) -> Value {
1029  return broadcast(builder, value, shape);
1030  };
1031  auto mul = [&](Value a, Value b) -> Value {
1032  return builder.create<arith::MulFOp>(a, b);
1033  };
1034  auto sub = [&](Value a, Value b) -> Value {
1035  return builder.create<arith::SubFOp>(a, b);
1036  };
1037  auto floor = [&](Value a) { return builder.create<math::FloorOp>(a); };
1038 
1039  auto i32Vec = broadcast(builder.getI32Type(), shape);
1040  auto fPToSingedInteger = [&](Value a) -> Value {
1041  return builder.create<arith::FPToSIOp>(a, i32Vec);
1042  };
1043 
1044  auto modulo4 = [&](Value a) -> Value {
1045  return builder.create<arith::AndIOp>(a, bcast(i32Cst(builder, 3)));
1046  };
1047 
1048  auto isEqualTo = [&](Value a, Value b) -> Value {
1049  return builder.create<arith::CmpIOp>(arith::CmpIPredicate::eq, a, b);
1050  };
1051 
1052  auto isGreaterThan = [&](Value a, Value b) -> Value {
1053  return builder.create<arith::CmpIOp>(arith::CmpIPredicate::sgt, a, b);
1054  };
1055 
1056  auto select = [&](Value cond, Value t, Value f) -> Value {
1057  return builder.create<SelectOp>(cond, t, f);
1058  };
1059 
1060  auto fmla = [&](Value a, Value b, Value c) {
1061  return builder.create<math::FmaOp>(a, b, c);
1062  };
1063 
1064  auto bitwiseOr = [&](Value a, Value b) {
1065  return builder.create<arith::OrIOp>(a, b);
1066  };
1067 
1068  Value twoOverPi = bcast(f32Cst(builder, TWO_OVER_PI));
1069  Value piOverTwo = bcast(f32Cst(builder, PI_OVER_2));
1070 
1071  Value x = op.getOperand();
1072 
1073  Value k = floor(mul(x, twoOverPi));
1074 
1075  Value y = sub(x, mul(k, piOverTwo));
1076 
1077  Value cstOne = bcast(f32Cst(builder, 1.0));
1078  Value cstNegativeOne = bcast(f32Cst(builder, -1.0));
1079 
1080  Value cstSC2 = bcast(f32Cst(builder, -0.16666667163372039794921875f));
1081  Value cstSC4 = bcast(f32Cst(builder, 8.333347737789154052734375e-3f));
1082  Value cstSC6 = bcast(f32Cst(builder, -1.9842604524455964565277099609375e-4f));
1083  Value cstSC8 =
1084  bcast(f32Cst(builder, 2.760012648650445044040679931640625e-6f));
1085  Value cstSC10 =
1086  bcast(f32Cst(builder, -2.50293279435709337121807038784027099609375e-8f));
1087 
1088  Value cstCC2 = bcast(f32Cst(builder, -0.5f));
1089  Value cstCC4 = bcast(f32Cst(builder, 4.166664183139801025390625e-2f));
1090  Value cstCC6 = bcast(f32Cst(builder, -1.388833043165504932403564453125e-3f));
1091  Value cstCC8 = bcast(f32Cst(builder, 2.47562347794882953166961669921875e-5f));
1092  Value cstCC10 =
1093  bcast(f32Cst(builder, -2.59630184018533327616751194000244140625e-7f));
1094 
1095  Value kMod4 = modulo4(fPToSingedInteger(k));
1096 
1097  Value kR0 = isEqualTo(kMod4, bcast(i32Cst(builder, 0)));
1098  Value kR1 = isEqualTo(kMod4, bcast(i32Cst(builder, 1)));
1099  Value kR2 = isEqualTo(kMod4, bcast(i32Cst(builder, 2)));
1100  Value kR3 = isEqualTo(kMod4, bcast(i32Cst(builder, 3)));
1101 
1102  Value sinuseCos = isSine ? bitwiseOr(kR1, kR3) : bitwiseOr(kR0, kR2);
1103  Value negativeRange = isSine ? isGreaterThan(kMod4, bcast(i32Cst(builder, 1)))
1104  : bitwiseOr(kR1, kR2);
1105 
1106  Value y2 = mul(y, y);
1107 
1108  Value base = select(sinuseCos, cstOne, y);
1109  Value cstC2 = select(sinuseCos, cstCC2, cstSC2);
1110  Value cstC4 = select(sinuseCos, cstCC4, cstSC4);
1111  Value cstC6 = select(sinuseCos, cstCC6, cstSC6);
1112  Value cstC8 = select(sinuseCos, cstCC8, cstSC8);
1113  Value cstC10 = select(sinuseCos, cstCC10, cstSC10);
1114 
1115  Value v1 = fmla(y2, cstC10, cstC8);
1116  Value v2 = fmla(y2, v1, cstC6);
1117  Value v3 = fmla(y2, v2, cstC4);
1118  Value v4 = fmla(y2, v3, cstC2);
1119  Value v5 = fmla(y2, v4, cstOne);
1120  Value v6 = mul(base, v5);
1121 
1122  Value approximation = select(negativeRange, mul(cstNegativeOne, v6), v6);
1123 
1124  rewriter.replaceOp(op, approximation);
1125 
1126  return success();
1127 }
1128 
1129 //----------------------------------------------------------------------------//
1130 // Rsqrt approximation.
1131 //----------------------------------------------------------------------------//
1132 
1133 namespace {
1134 struct RsqrtApproximation : public OpRewritePattern<math::RsqrtOp> {
1136 
1137  LogicalResult matchAndRewrite(math::RsqrtOp op,
1138  PatternRewriter &rewriter) const final;
1139 };
1140 } // namespace
1141 
1143 RsqrtApproximation::matchAndRewrite(math::RsqrtOp op,
1144  PatternRewriter &rewriter) const {
1145  if (!getElementTypeOrSelf(op.getOperand()).isF32())
1146  return rewriter.notifyMatchFailure(op, "unsupported operand type");
1147 
1148  ArrayRef<int64_t> shape = vectorShape(op.getOperand());
1149 
1150  // Only support already-vectorized rsqrt's.
1151  if (shape.empty() || shape.back() % 8 != 0)
1152  return rewriter.notifyMatchFailure(op, "unsupported operand type");
1153 
1154  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
1155  auto bcast = [&](Value value) -> Value {
1156  return broadcast(builder, value, shape);
1157  };
1158 
1159  Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u));
1160  Value cstOnePointFive = bcast(f32Cst(builder, 1.5f));
1161  Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
1162  Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
1163 
1164  Value negHalf = builder.create<arith::MulFOp>(op.getOperand(), cstNegHalf);
1165 
1166  // Select only the inverse sqrt of positive normals (denormals are
1167  // flushed to zero).
1168  Value ltMinMask = builder.create<arith::CmpFOp>(
1169  arith::CmpFPredicate::OLT, op.getOperand(), cstMinNormPos);
1170  Value infMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ,
1171  op.getOperand(), cstPosInf);
1172  Value notNormalFiniteMask = builder.create<arith::OrIOp>(ltMinMask, infMask);
1173 
1174  // Compute an approximate result.
1176  builder, op->getOperands(), 8, [&builder](ValueRange operands) -> Value {
1177  return builder.create<x86vector::RsqrtOp>(operands);
1178  });
1179 
1180  // Do a single step of Newton-Raphson iteration to improve the approximation.
1181  // This uses the formula y_{n+1} = y_n * (1.5 - y_n * (0.5 * x) * y_n).
1182  // It is essential to evaluate the inner term like this because forming
1183  // y_n^2 may over- or underflow.
1184  Value inner = builder.create<arith::MulFOp>(negHalf, yApprox);
1185  Value fma = builder.create<math::FmaOp>(yApprox, inner, cstOnePointFive);
1186  Value yNewton = builder.create<arith::MulFOp>(yApprox, fma);
1187 
1188  // Select the result of the Newton-Raphson step for positive normal arguments.
1189  // For other arguments, choose the output of the intrinsic. This will
1190  // return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(x) = +inf if
1191  // x is zero or a positive denormalized float (equivalent to flushing positive
1192  // denormalized inputs to zero).
1193  Value res = builder.create<SelectOp>(notNormalFiniteMask, yApprox, yNewton);
1194  rewriter.replaceOp(op, res);
1195 
1196  return success();
1197 }
1198 
1199 //----------------------------------------------------------------------------//
1200 
1202  RewritePatternSet &patterns,
1204  patterns.add<AtanApproximation, Atan2Approximation, TanhApproximation,
1205  LogApproximation, Log2Approximation, Log1pApproximation,
1206  ErfPolynomialApproximation, ExpApproximation, ExpM1Approximation,
1207  SinAndCosApproximation<true, math::SinOp>,
1208  SinAndCosApproximation<false, math::CosOp>>(
1209  patterns.getContext());
1210  if (options.enableAvx2)
1211  patterns.add<RsqrtApproximation>(patterns.getContext());
1212 }
static Value f32Cst(ImplicitLocOpBuilder &builder, float value)
Include the generated interface declarations.
SmallVector< int64_t, 4 > computeStrides(ArrayRef< int64_t > shape, ArrayRef< int64_t > sizes)
Given the shape and sizes of a vector, returns the corresponding strides for each dimension...
Definition: VectorUtils.cpp:52
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:881
static Value min(ImplicitLocOpBuilder &builder, Value a, Value b)
Attribute getZeroAttr(Type type)
Definition: Builders.cpp:264
static Value i32Cst(ImplicitLocOpBuilder &builder, int32_t value)
int64_t floor(Fraction f)
Definition: Fraction.h:55
static Type getElementType(Type type, ArrayRef< int32_t > indices, function_ref< InFlightDiagnostic(StringRef)> emitErrorFn)
Walks the given type hierarchy with the given indices, potentially down to component granularity...
Definition: SPIRVOps.cpp:639
FloatType getF32Type()
Definition: Builders.cpp:40
#define PI_OVER_2
static ArrayRef< int64_t > vectorShape(Type type)
static constexpr const bool value
IntegerAttr getI32IntegerAttr(int32_t value)
Definition: Builders.cpp:148
void populateMathPolynomialApproximationPatterns(RewritePatternSet &patterns, const MathPolynomialApproximationOptions &options={})
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
#define TWO_OVER_PI
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
#define LOG2E_VALUE
U dyn_cast() const
Definition: Types.h:244
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:206
static Value exp2I32(ImplicitLocOpBuilder &builder, Value arg)
IntegerType getIntegerType(unsigned width)
Definition: Builders.cpp:58
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:360
int64_t computeMaxLinearIndex(ArrayRef< int64_t > basis)
Return the number of elements of basis, 0 if empty.
Definition: VectorUtils.cpp:45
#define LN2_VALUE
static Value f32FromBits(ImplicitLocOpBuilder &builder, uint32_t bits)
auto getType() const
Eliminates identifier at the specified position using Fourier-Motzkin variable elimination.
static std::pair< Value, Value > frexp(ImplicitLocOpBuilder &builder, Value arg, bool isPositive=false)
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:72
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:84
static llvm::ManagedStatic< PassManagerOptions > options
SmallVector< int64_t, 4 > delinearize(ArrayRef< int64_t > strides, int64_t linearIndex)
Given the strides together with a linear index in the dimension space, returns the vector-space offse...
Definition: VectorUtils.cpp:75
static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound, Value upperBound)
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:355
OpTy create(Args &&...args)
Create an operation of specific op type at the current insertion point and location.
OpTy replaceOpWithNewOp(Operation *op, Args &&... args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:741
Type getType() const
Return the type of this value.
Definition: Value.h:117
Location getLoc()
The source location the operation was defined or derived from.
Definition: OpDefinition.h:124
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
Definition: PatternMatch.h:930
LogicalResult matchAndRewrite(math::ErfOp op, PatternRewriter &rewriter) const final
ImplicitLocOpBuilder maintains a &#39;current location&#39;, allowing use of the create<> method without spec...
static VectorType vectorType(CodeGen &codegen, Type etp)
Constructs vector type.
int compare(Fraction x, Fraction y)
Three-way comparison between two fractions.
Definition: Fraction.h:46
This provides public APIs that all operations should have.
static Type broadcast(Type type, ArrayRef< int64_t > shape)
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Operation *op, CallbackT &&reasonCallback)
Used to notify the rewriter that the IR failed to be rewritten because of a match failure...
Definition: PatternMatch.h:802
bool isa() const
Definition: Types.h:234
FloatAttr getF32FloatAttr(float value)
Definition: Builders.cpp:185
static Value handleMultidimensionalVectors(ImplicitLocOpBuilder &builder, ValueRange operands, int64_t vectorWidth, llvm::function_ref< Value(ValueRange)> compute)
This class provides an abstraction over the different types of ranges over Values.
MLIRContext * getContext() const
Definition: PatternMatch.h:906
IntegerType getI32Type()
Definition: Builders.cpp:54
static Value max(ImplicitLocOpBuilder &builder, Value a, Value b)