1 //===- PolynomialApproximation.cpp - Approximate math operations ----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements expansion of math operations to fast approximations
10 // that do not rely on any of the library functions.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
15 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
16 #include "mlir/Dialect/Math/IR/Math.h"
17 #include "mlir/Dialect/Math/Transforms/Passes.h"
18 #include "mlir/Dialect/Vector/VectorOps.h"
19 #include "mlir/IR/Builders.h"
20 #include "mlir/IR/ImplicitLocOpBuilder.h"
21 #include "mlir/Transforms/DialectConversion.h"
22 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
23 
24 using namespace mlir;
25 using namespace mlir::vector;
26 
27 using TypePredicate = llvm::function_ref<bool(Type)>;
28 
29 static bool isF32(Type type) { return type.isF32(); }
30 
31 // Returns vector width if the element type is matching the predicate (scalars
32 // that do match the predicate have width equal to `1`).
33 static Optional<int> vectorWidth(Type type, TypePredicate pred) {
34   // If the type matches the predicate then its width is `1`.
35   if (pred(type))
36     return 1;
37 
38   // Otherwise check if the type is a vector type.
39   auto vectorType = type.dyn_cast<VectorType>();
40   if (vectorType && pred(vectorType.getElementType())) {
41     assert(vectorType.getRank() == 1 && "only 1d vectors are supported");
42     return vectorType.getDimSize(0);
43   }
44 
45   return llvm::None;
46 }
47 
48 // Returns vector width of the type. If the type is a scalar returns `1`.
49 static int vectorWidth(Type type) {
50   auto vectorType = type.dyn_cast<VectorType>();
51   return vectorType ? vectorType.getDimSize(0) : 1;
52 }
53 
54 // Returns vector element type. If the type is a scalar returns the argument.
55 static Type elementType(Type type) {
56   auto vectorType = type.dyn_cast<VectorType>();
57   return vectorType ? vectorType.getElementType() : type;
58 }
59 
60 //----------------------------------------------------------------------------//
61 // Broadcast scalar types and values into vector types and values.
62 //----------------------------------------------------------------------------//
63 
64 // Broadcasts scalar type into vector type (iff width is greater then 1).
65 static Type broadcast(Type type, int width) {
66   assert(!type.isa<VectorType>() && "must be scalar type");
67   return width > 1 ? VectorType::get({width}, type) : type;
68 }
69 
70 // Broadcasts scalar value into vector (iff width is greater then 1).
71 static Value broadcast(ImplicitLocOpBuilder &builder, Value value, int width) {
72   assert(!value.getType().isa<VectorType>() && "must be scalar value");
73   auto type = broadcast(value.getType(), width);
74   return width > 1 ? builder.create<BroadcastOp>(type, value) : value;
75 }
76 
77 //----------------------------------------------------------------------------//
78 // Helper functions to create constants.
79 //----------------------------------------------------------------------------//
80 
81 static Value f32Cst(ImplicitLocOpBuilder &builder, float value) {
82   return builder.create<ConstantOp>(builder.getF32Type(),
83                                     builder.getF32FloatAttr(value));
84 }
85 
86 static Value i32Cst(ImplicitLocOpBuilder &builder, int32_t value) {
87   return builder.create<ConstantOp>(builder.getI32Type(),
88                                     builder.getI32IntegerAttr(value));
89 }
90 
91 static Value f32FromBits(ImplicitLocOpBuilder &builder, uint32_t bits) {
92   Value i32Value = i32Cst(builder, static_cast<int32_t>(bits));
93   return builder.create<LLVM::BitcastOp>(builder.getF32Type(), i32Value);
94 }
95 
96 //----------------------------------------------------------------------------//
97 // Helper functions to build math functions approximations.
98 //----------------------------------------------------------------------------//
99 
100 static Value min(ImplicitLocOpBuilder &builder, Value a, Value b) {
101   return builder.create<SelectOp>(
102       builder.create<CmpFOp>(CmpFPredicate::OLT, a, b), a, b);
103 }
104 
105 static Value max(ImplicitLocOpBuilder &builder, Value a, Value b) {
106   return builder.create<SelectOp>(
107       builder.create<CmpFOp>(CmpFPredicate::OGT, a, b), a, b);
108 }
109 
110 static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound,
111                    Value upperBound) {
112   return max(builder, min(builder, value, upperBound), lowerBound);
113 }
114 
115 // Decomposes given floating point value `arg` into a normalized fraction and
116 // an integral power of two (see std::frexp). Returned values have float type.
117 static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
118                                      bool is_positive = false) {
119   assert(isF32(elementType(arg.getType())) && "argument must be f32 type");
120 
121   int width = vectorWidth(arg.getType());
122 
123   auto bcast = [&](Value value) -> Value {
124     return broadcast(builder, value, width);
125   };
126 
127   auto i32 = builder.getIntegerType(32);
128   auto i32Vec = broadcast(i32, width);
129   auto f32Vec = broadcast(builder.getF32Type(), width);
130 
131   Value cst126f = f32Cst(builder, 126.0f);
132   Value cstHalf = f32Cst(builder, 0.5f);
133   Value cstInvMantMask = f32FromBits(builder, ~0x7f800000u);
134 
135   // Bitcast to i32 for bitwise operations.
136   Value i32Half = builder.create<LLVM::BitcastOp>(i32, cstHalf);
137   Value i32InvMantMask = builder.create<LLVM::BitcastOp>(i32, cstInvMantMask);
138   Value i32Arg = builder.create<LLVM::BitcastOp>(i32Vec, arg);
139 
140   // Compute normalized fraction.
141   Value tmp0 = builder.create<LLVM::AndOp>(i32Arg, bcast(i32InvMantMask));
142   Value tmp1 = builder.create<LLVM::OrOp>(tmp0, bcast(i32Half));
143   Value normalizedFraction = builder.create<LLVM::BitcastOp>(f32Vec, tmp1);
144 
145   // Compute exponent.
146   Value arg0 = is_positive ? arg : builder.create<AbsFOp>(arg);
147   Value biasedExponentBits = builder.create<UnsignedShiftRightOp>(
148       builder.create<LLVM::BitcastOp>(i32Vec, arg0),
149       bcast(i32Cst(builder, 23)));
150   Value biasedExponent = builder.create<SIToFPOp>(f32Vec, biasedExponentBits);
151   Value exponent = builder.create<SubFOp>(biasedExponent, bcast(cst126f));
152 
153   return {normalizedFraction, exponent};
154 }
155 
156 //----------------------------------------------------------------------------//
157 // TanhOp approximation.
158 //----------------------------------------------------------------------------//
159 
160 namespace {
161 struct TanhApproximation : public OpRewritePattern<math::TanhOp> {
162 public:
163   using OpRewritePattern::OpRewritePattern;
164 
165   LogicalResult matchAndRewrite(math::TanhOp op,
166                                 PatternRewriter &rewriter) const final;
167 };
168 } // namespace
169 
170 LogicalResult
171 TanhApproximation::matchAndRewrite(math::TanhOp op,
172                                    PatternRewriter &rewriter) const {
173   auto width = vectorWidth(op.operand().getType(), isF32);
174   if (!width.hasValue())
175     return rewriter.notifyMatchFailure(op, "unsupported operand type");
176 
177   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
178   auto bcast = [&](Value value) -> Value {
179     return broadcast(builder, value, *width);
180   };
181 
182   // Clamp operand into [plusClamp, minusClamp] range.
183   Value minusClamp = bcast(f32Cst(builder, -7.9053111076354980f));
184   Value plusClamp = bcast(f32Cst(builder, 7.90531110763549805f));
185   Value x = clamp(builder, op.operand(), minusClamp, plusClamp);
186 
187   // Mask for tiny values that are approximated with `operand`.
188   Value tiny = bcast(f32Cst(builder, 0.0004f));
189   Value tinyMask = builder.create<CmpFOp>(
190       CmpFPredicate::OLT, builder.create<AbsFOp>(op.operand()), tiny);
191 
192   // The monomial coefficients of the numerator polynomial (odd).
193   Value alpha1 = bcast(f32Cst(builder, 4.89352455891786e-03f));
194   Value alpha3 = bcast(f32Cst(builder, 6.37261928875436e-04f));
195   Value alpha5 = bcast(f32Cst(builder, 1.48572235717979e-05f));
196   Value alpha7 = bcast(f32Cst(builder, 5.12229709037114e-08f));
197   Value alpha9 = bcast(f32Cst(builder, -8.60467152213735e-11f));
198   Value alpha11 = bcast(f32Cst(builder, 2.00018790482477e-13f));
199   Value alpha13 = bcast(f32Cst(builder, -2.76076847742355e-16f));
200 
201   // The monomial coefficients of the denominator polynomial (even).
202   Value beta0 = bcast(f32Cst(builder, 4.89352518554385e-03f));
203   Value beta2 = bcast(f32Cst(builder, 2.26843463243900e-03f));
204   Value beta4 = bcast(f32Cst(builder, 1.18534705686654e-04f));
205   Value beta6 = bcast(f32Cst(builder, 1.19825839466702e-06f));
206 
207   // Since the polynomials are odd/even, we need x^2.
208   Value x2 = builder.create<MulFOp>(x, x);
209 
210   // Evaluate the numerator polynomial p.
211   Value p = builder.create<FmaFOp>(x2, alpha13, alpha11);
212   p = builder.create<FmaFOp>(x2, p, alpha9);
213   p = builder.create<FmaFOp>(x2, p, alpha7);
214   p = builder.create<FmaFOp>(x2, p, alpha5);
215   p = builder.create<FmaFOp>(x2, p, alpha3);
216   p = builder.create<FmaFOp>(x2, p, alpha1);
217   p = builder.create<MulFOp>(x, p);
218 
219   // Evaluate the denominator polynomial q.
220   Value q = builder.create<FmaFOp>(x2, beta6, beta4);
221   q = builder.create<FmaFOp>(x2, q, beta2);
222   q = builder.create<FmaFOp>(x2, q, beta0);
223 
224   // Divide the numerator by the denominator.
225   Value res =
226       builder.create<SelectOp>(tinyMask, x, builder.create<DivFOp>(p, q));
227 
228   rewriter.replaceOp(op, res);
229 
230   return success();
231 }
232 
233 //----------------------------------------------------------------------------//
234 // LogOp approximation.
235 //----------------------------------------------------------------------------//
236 
237 namespace {
238 
239 // This approximations comes from the Julien Pommier's SSE math library.
240 // Link: http://gruntthepeon.free.fr/ssemath
241 struct LogApproximation : public OpRewritePattern<math::LogOp> {
242 public:
243   using OpRewritePattern::OpRewritePattern;
244 
245   LogicalResult matchAndRewrite(math::LogOp op,
246                                 PatternRewriter &rewriter) const final;
247 };
248 } // namespace
249 
250 #define LN2_VALUE                                                              \
251   0.693147180559945309417232121458176568075500134360255254120680009493393621L
252 
253 LogicalResult
254 LogApproximation::matchAndRewrite(math::LogOp op,
255                                   PatternRewriter &rewriter) const {
256   auto width = vectorWidth(op.operand().getType(), isF32);
257   if (!width.hasValue())
258     return rewriter.notifyMatchFailure(op, "unsupported operand type");
259 
260   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
261   auto bcast = [&](Value value) -> Value {
262     return broadcast(builder, value, *width);
263   };
264 
265   Value cstZero = bcast(f32Cst(builder, 0.0f));
266   Value cstOne = bcast(f32Cst(builder, 1.0f));
267   Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
268 
269   // The smallest non denormalized float number.
270   Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
271   Value cstMinusInf = bcast(f32FromBits(builder, 0xff800000u));
272   Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u));
273   Value cstNan = bcast(f32FromBits(builder, 0x7fc00000));
274 
275   // Polynomial coefficients.
276   Value cstCephesSQRTHF = bcast(f32Cst(builder, 0.707106781186547524f));
277   Value cstCephesLogP0 = bcast(f32Cst(builder, 7.0376836292E-2f));
278   Value cstCephesLogP1 = bcast(f32Cst(builder, -1.1514610310E-1f));
279   Value cstCephesLogP2 = bcast(f32Cst(builder, 1.1676998740E-1f));
280   Value cstCephesLogP3 = bcast(f32Cst(builder, -1.2420140846E-1f));
281   Value cstCephesLogP4 = bcast(f32Cst(builder, +1.4249322787E-1f));
282   Value cstCephesLogP5 = bcast(f32Cst(builder, -1.6668057665E-1f));
283   Value cstCephesLogP6 = bcast(f32Cst(builder, +2.0000714765E-1f));
284   Value cstCephesLogP7 = bcast(f32Cst(builder, -2.4999993993E-1f));
285   Value cstCephesLogP8 = bcast(f32Cst(builder, +3.3333331174E-1f));
286 
287   Value x = op.operand();
288 
289   // Truncate input values to the minimum positive normal.
290   x = max(builder, x, cstMinNormPos);
291 
292   // Extract significant in the range [0.5,1) and exponent.
293   std::pair<Value, Value> pair = frexp(builder, x, /*is_positive=*/true);
294   x = pair.first;
295   Value e = pair.second;
296 
297   // Shift the inputs from the range [0.5,1) to [sqrt(1/2), sqrt(2)) and shift
298   // by -1.0. The values are then centered around 0, which improves the
299   // stability of the polynomial evaluation:
300   //
301   //   if( x < SQRTHF ) {
302   //     e -= 1;
303   //     x = x + x - 1.0;
304   //   } else { x = x - 1.0; }
305   Value mask = builder.create<CmpFOp>(CmpFPredicate::OLT, x, cstCephesSQRTHF);
306   Value tmp = builder.create<SelectOp>(mask, x, cstZero);
307 
308   x = builder.create<SubFOp>(x, cstOne);
309   e = builder.create<SubFOp>(e,
310                              builder.create<SelectOp>(mask, cstOne, cstZero));
311   x = builder.create<AddFOp>(x, tmp);
312 
313   Value x2 = builder.create<MulFOp>(x, x);
314   Value x3 = builder.create<MulFOp>(x2, x);
315 
316   // Evaluate the polynomial approximant of degree 8 in three parts.
317   Value y0, y1, y2;
318   y0 = builder.create<FmaFOp>(cstCephesLogP0, x, cstCephesLogP1);
319   y1 = builder.create<FmaFOp>(cstCephesLogP3, x, cstCephesLogP4);
320   y2 = builder.create<FmaFOp>(cstCephesLogP6, x, cstCephesLogP7);
321   y0 = builder.create<FmaFOp>(y0, x, cstCephesLogP2);
322   y1 = builder.create<FmaFOp>(y1, x, cstCephesLogP5);
323   y2 = builder.create<FmaFOp>(y2, x, cstCephesLogP8);
324   y0 = builder.create<FmaFOp>(y0, x3, y1);
325   y0 = builder.create<FmaFOp>(y0, x3, y2);
326   y0 = builder.create<MulFOp>(y0, x3);
327 
328   y0 = builder.create<FmaFOp>(cstNegHalf, x2, y0);
329   x = builder.create<AddFOp>(x, y0);
330 
331   Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
332   x = builder.create<FmaFOp>(e, cstLn2, x);
333 
334   Value invalidMask =
335       builder.create<CmpFOp>(CmpFPredicate::ULT, op.operand(), cstZero);
336   Value zeroMask =
337       builder.create<CmpFOp>(CmpFPredicate::OEQ, op.operand(), cstZero);
338   Value posInfMask =
339       builder.create<CmpFOp>(CmpFPredicate::OEQ, op.operand(), cstPosInf);
340 
341   // Filter out invalid values:
342   //  • x == 0     -> -INF
343   //  • x < 0      ->  NAN
344   //  • x == +INF  -> +INF
345   Value aproximation = builder.create<SelectOp>(
346       zeroMask, cstMinusInf,
347       builder.create<SelectOp>(
348           invalidMask, cstNan,
349           builder.create<SelectOp>(posInfMask, cstPosInf, x)));
350 
351   rewriter.replaceOp(op, aproximation);
352 
353   return success();
354 }
355 
356 //----------------------------------------------------------------------------//
357 
358 void mlir::populateMathPolynomialApproximationPatterns(
359     OwningRewritePatternList &patterns, MLIRContext *ctx) {
360   patterns.insert<TanhApproximation, LogApproximation>(ctx);
361 }
362