| //===- PolynomialApproximation.cpp - Approximate math operations ----------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file implements expansion of math operations to fast approximations |
| // that do not rely on any of the library functions. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include <climits> |
| #include <cstddef> |
| |
| #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" |
| #include "mlir/Dialect/Math/IR/Math.h" |
| #include "mlir/Dialect/Math/Transforms/Approximation.h" |
| #include "mlir/Dialect/Math/Transforms/Passes.h" |
| #include "mlir/Dialect/Vector/VectorOps.h" |
| #include "mlir/Dialect/Vector/VectorUtils.h" |
| #include "mlir/Dialect/X86Vector/X86VectorDialect.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/ImplicitLocOpBuilder.h" |
| #include "mlir/IR/TypeUtilities.h" |
| #include "mlir/Transforms/Bufferize.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| #include "llvm/ADT/ArrayRef.h" |
| |
| using namespace mlir; |
| using namespace mlir::math; |
| using namespace mlir::vector; |
| |
| // Returns vector shape if the type is a vector. Returns an empty shape if it is |
| // not a vector. |
| static ArrayRef<int64_t> vectorShape(Type type) { |
| auto vectorType = type.dyn_cast<VectorType>(); |
| return vectorType ? vectorType.getShape() : ArrayRef<int64_t>(); |
| } |
| |
| static ArrayRef<int64_t> vectorShape(Value value) { |
| return vectorShape(value.getType()); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Broadcast scalar types and values into vector types and values. |
| //----------------------------------------------------------------------------// |
| |
| // Broadcasts scalar type into vector type (iff shape is non-scalar). |
| static Type broadcast(Type type, ArrayRef<int64_t> shape) { |
| assert(!type.isa<VectorType>() && "must be scalar type"); |
| return !shape.empty() ? VectorType::get(shape, type) : type; |
| } |
| |
| // Broadcasts scalar value into vector (iff shape is non-scalar). |
| static Value broadcast(ImplicitLocOpBuilder &builder, Value value, |
| ArrayRef<int64_t> shape) { |
| assert(!value.getType().isa<VectorType>() && "must be scalar value"); |
| auto type = broadcast(value.getType(), shape); |
| return !shape.empty() ? builder.create<BroadcastOp>(type, value) : value; |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Helper function to handle n-D vectors with 1-D operations. |
| //----------------------------------------------------------------------------// |
| |
| // Expands and unrolls n-D vector operands into multiple fixed size 1-D vectors |
| // and calls the compute function with 1-D vector operands. Stitches back all |
| // results into the original n-D vector result. |
| // |
| // Examples: vectorWidth = 8 |
| // - vector<4x8xf32> unrolled 4 times |
| // - vector<16xf32> expanded to vector<2x8xf32> and unrolled 2 times |
| // - vector<4x16xf32> expanded to vector<4x2x8xf32> and unrolled 4*2 times |
| // |
| // Some math approximations rely on ISA-specific operations that only accept |
| // fixed size 1-D vectors (e.g. AVX expects vectors of width 8). |
| // |
| // It is the caller's responsibility to verify that the inner dimension is |
| // divisible by the vectorWidth, and that all operands have the same vector |
| // shape. |
| static Value |
| handleMultidimensionalVectors(ImplicitLocOpBuilder &builder, |
| ValueRange operands, int64_t vectorWidth, |
| std::function<Value(ValueRange)> compute) { |
| assert(!operands.empty() && "operands must be not empty"); |
| assert(vectorWidth > 0 && "vector width must be larger than 0"); |
| |
| VectorType inputType = operands[0].getType().cast<VectorType>(); |
| ArrayRef<int64_t> inputShape = inputType.getShape(); |
| |
| // If input shape matches target vector width, we can just call the |
| // user-provided compute function with the operands. |
| if (inputShape == llvm::makeArrayRef(vectorWidth)) |
| return compute(operands); |
| |
| // Check if the inner dimension has to be expanded, or we can directly iterate |
| // over the outer dimensions of the vector. |
| int64_t innerDim = inputShape.back(); |
| int64_t expansionDim = innerDim / vectorWidth; |
| assert((innerDim % vectorWidth == 0) && "invalid inner dimension size"); |
| |
| // Maybe expand operands to the higher rank vector shape that we'll use to |
| // iterate over and extract one dimensional vectors. |
| SmallVector<int64_t> expandedShape(inputShape.begin(), inputShape.end()); |
| SmallVector<Value> expandedOperands(operands); |
| |
| if (expansionDim > 1) { |
| // Expand shape from [..., innerDim] to [..., expansionDim, vectorWidth]. |
| expandedShape.insert(expandedShape.end() - 1, expansionDim); |
| expandedShape.back() = vectorWidth; |
| |
| for (unsigned i = 0; i < operands.size(); ++i) { |
| auto operand = operands[i]; |
| auto eltType = operand.getType().cast<VectorType>().getElementType(); |
| auto expandedType = VectorType::get(expandedShape, eltType); |
| expandedOperands[i] = |
| builder.create<vector::ShapeCastOp>(expandedType, operand); |
| } |
| } |
| |
| // Iterate over all outer dimensions of the compute shape vector type. |
| auto iterationDims = ArrayRef<int64_t>(expandedShape).drop_back(); |
| int64_t maxLinearIndex = computeMaxLinearIndex(iterationDims); |
| |
| SmallVector<int64_t> ones(iterationDims.size(), 1); |
| auto strides = computeStrides(iterationDims, ones); |
| |
| // Compute results for each one dimensional vector. |
| SmallVector<Value> results(maxLinearIndex); |
| |
| for (int64_t i = 0; i < maxLinearIndex; ++i) { |
| auto offsets = delinearize(strides, i); |
| |
| SmallVector<Value> extracted(expandedOperands.size()); |
| for (auto tuple : llvm::enumerate(expandedOperands)) |
| extracted[tuple.index()] = |
| builder.create<vector::ExtractOp>(tuple.value(), offsets); |
| |
| results[i] = compute(extracted); |
| } |
| |
| // Stitch results together into one large vector. |
| Type resultEltType = results[0].getType().cast<VectorType>().getElementType(); |
| Type resultExpandedType = VectorType::get(expandedShape, resultEltType); |
| Value result = builder.create<ConstantOp>( |
| resultExpandedType, builder.getZeroAttr(resultExpandedType)); |
| |
| for (int64_t i = 0; i < maxLinearIndex; ++i) |
| result = builder.create<vector::InsertOp>(results[i], result, |
| delinearize(strides, i)); |
| |
| // Reshape back to the original vector shape. |
| return builder.create<vector::ShapeCastOp>( |
| VectorType::get(inputShape, resultEltType), result); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Helper functions to create constants. |
| //----------------------------------------------------------------------------// |
| |
| static Value f32Cst(ImplicitLocOpBuilder &builder, float value) { |
| return builder.create<arith::ConstantOp>(builder.getF32FloatAttr(value)); |
| } |
| |
| static Value i32Cst(ImplicitLocOpBuilder &builder, int32_t value) { |
| return builder.create<arith::ConstantOp>(builder.getI32IntegerAttr(value)); |
| } |
| |
| static Value f32FromBits(ImplicitLocOpBuilder &builder, uint32_t bits) { |
| Value i32Value = i32Cst(builder, static_cast<int32_t>(bits)); |
| return builder.create<arith::BitcastOp>(builder.getF32Type(), i32Value); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Helper functions to build math functions approximations. |
| //----------------------------------------------------------------------------// |
| |
| static Value min(ImplicitLocOpBuilder &builder, Value a, Value b) { |
| return builder.create<SelectOp>( |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, a, b), a, b); |
| } |
| |
| static Value max(ImplicitLocOpBuilder &builder, Value a, Value b) { |
| return builder.create<SelectOp>( |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, a, b), a, b); |
| } |
| |
| static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound, |
| Value upperBound) { |
| return max(builder, min(builder, value, upperBound), lowerBound); |
| } |
| |
| // Decomposes given floating point value `arg` into a normalized fraction and |
| // an integral power of two (see std::frexp). Returned values have float type. |
| static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg, |
| bool is_positive = false) { |
| assert(getElementTypeOrSelf(arg).isF32() && "arg must be f32 type"); |
| ArrayRef<int64_t> shape = vectorShape(arg); |
| |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| auto i32 = builder.getIntegerType(32); |
| auto i32Vec = broadcast(i32, shape); |
| auto f32Vec = broadcast(builder.getF32Type(), shape); |
| |
| Value cst126f = f32Cst(builder, 126.0f); |
| Value cstHalf = f32Cst(builder, 0.5f); |
| Value cstInvMantMask = f32FromBits(builder, ~0x7f800000u); |
| |
| // Bitcast to i32 for bitwise operations. |
| Value i32Half = builder.create<arith::BitcastOp>(i32, cstHalf); |
| Value i32InvMantMask = builder.create<arith::BitcastOp>(i32, cstInvMantMask); |
| Value i32Arg = builder.create<arith::BitcastOp>(i32Vec, arg); |
| |
| // Compute normalized fraction. |
| Value tmp0 = builder.create<arith::AndIOp>(i32Arg, bcast(i32InvMantMask)); |
| Value tmp1 = builder.create<arith::OrIOp>(tmp0, bcast(i32Half)); |
| Value normalizedFraction = builder.create<arith::BitcastOp>(f32Vec, tmp1); |
| |
| // Compute exponent. |
| Value arg0 = is_positive ? arg : builder.create<math::AbsOp>(arg); |
| Value biasedExponentBits = builder.create<arith::ShRUIOp>( |
| builder.create<arith::BitcastOp>(i32Vec, arg0), |
| bcast(i32Cst(builder, 23))); |
| Value biasedExponent = |
| builder.create<arith::SIToFPOp>(f32Vec, biasedExponentBits); |
| Value exponent = |
| builder.create<arith::SubFOp>(biasedExponent, bcast(cst126f)); |
| |
| return {normalizedFraction, exponent}; |
| } |
| |
| // Computes exp2 for an i32 argument. |
| static Value exp2I32(ImplicitLocOpBuilder &builder, Value arg) { |
| assert(getElementTypeOrSelf(arg).isInteger(32) && "arg must be i32 type"); |
| ArrayRef<int64_t> shape = vectorShape(arg); |
| |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| auto f32Vec = broadcast(builder.getF32Type(), shape); |
| // The exponent of f32 located at 23-bit. |
| auto exponetBitLocation = bcast(i32Cst(builder, 23)); |
| // Set the exponent bias to zero. |
| auto bias = bcast(i32Cst(builder, 127)); |
| |
| Value biasedArg = builder.create<arith::AddIOp>(arg, bias); |
| Value exp2ValueInt = |
| builder.create<arith::ShLIOp>(biasedArg, exponetBitLocation); |
| Value exp2ValueF32 = builder.create<arith::BitcastOp>(f32Vec, exp2ValueInt); |
| |
| return exp2ValueF32; |
| } |
| |
| namespace { |
| Value makePolynomialCalculation(ImplicitLocOpBuilder &builder, |
| llvm::ArrayRef<Value> coeffs, Value x) { |
| assert(getElementTypeOrSelf(x).isF32() && "x must be f32 type"); |
| ArrayRef<int64_t> shape = vectorShape(x); |
| |
| if (coeffs.empty()) |
| return broadcast(builder, f32Cst(builder, 0.0f), shape); |
| |
| if (coeffs.size() == 1) |
| return coeffs[0]; |
| |
| Value res = builder.create<math::FmaOp>(x, coeffs[coeffs.size() - 1], |
| coeffs[coeffs.size() - 2]); |
| for (auto i = ptrdiff_t(coeffs.size()) - 3; i >= 0; --i) { |
| res = builder.create<math::FmaOp>(x, res, coeffs[i]); |
| } |
| return res; |
| } |
| } // namespace |
| |
| //----------------------------------------------------------------------------// |
| // TanhOp approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| struct TanhApproximation : public OpRewritePattern<math::TanhOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(math::TanhOp op, |
| PatternRewriter &rewriter) const final; |
| }; |
| } // namespace |
| |
| LogicalResult |
| TanhApproximation::matchAndRewrite(math::TanhOp op, |
| PatternRewriter &rewriter) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| // Clamp operand into [plusClamp, minusClamp] range. |
| Value minusClamp = bcast(f32Cst(builder, -7.99881172180175781f)); |
| Value plusClamp = bcast(f32Cst(builder, 7.99881172180175781f)); |
| Value x = clamp(builder, op.operand(), minusClamp, plusClamp); |
| |
| // Mask for tiny values that are approximated with `operand`. |
| Value tiny = bcast(f32Cst(builder, 0.0004f)); |
| Value tinyMask = builder.create<arith::CmpFOp>( |
| arith::CmpFPredicate::OLT, builder.create<math::AbsOp>(op.operand()), |
| tiny); |
| |
| // The monomial coefficients of the numerator polynomial (odd). |
| Value alpha1 = bcast(f32Cst(builder, 4.89352455891786e-03f)); |
| Value alpha3 = bcast(f32Cst(builder, 6.37261928875436e-04f)); |
| Value alpha5 = bcast(f32Cst(builder, 1.48572235717979e-05f)); |
| Value alpha7 = bcast(f32Cst(builder, 5.12229709037114e-08f)); |
| Value alpha9 = bcast(f32Cst(builder, -8.60467152213735e-11f)); |
| Value alpha11 = bcast(f32Cst(builder, 2.00018790482477e-13f)); |
| Value alpha13 = bcast(f32Cst(builder, -2.76076847742355e-16f)); |
| |
| // The monomial coefficients of the denominator polynomial (even). |
| Value beta0 = bcast(f32Cst(builder, 4.89352518554385e-03f)); |
| Value beta2 = bcast(f32Cst(builder, 2.26843463243900e-03f)); |
| Value beta4 = bcast(f32Cst(builder, 1.18534705686654e-04f)); |
| Value beta6 = bcast(f32Cst(builder, 1.19825839466702e-06f)); |
| |
| // Since the polynomials are odd/even, we need x^2. |
| Value x2 = builder.create<arith::MulFOp>(x, x); |
| |
| // Evaluate the numerator polynomial p. |
| Value p = builder.create<math::FmaOp>(x2, alpha13, alpha11); |
| p = builder.create<math::FmaOp>(x2, p, alpha9); |
| p = builder.create<math::FmaOp>(x2, p, alpha7); |
| p = builder.create<math::FmaOp>(x2, p, alpha5); |
| p = builder.create<math::FmaOp>(x2, p, alpha3); |
| p = builder.create<math::FmaOp>(x2, p, alpha1); |
| p = builder.create<arith::MulFOp>(x, p); |
| |
| // Evaluate the denominator polynomial q. |
| Value q = builder.create<math::FmaOp>(x2, beta6, beta4); |
| q = builder.create<math::FmaOp>(x2, q, beta2); |
| q = builder.create<math::FmaOp>(x2, q, beta0); |
| |
| // Divide the numerator by the denominator. |
| Value res = builder.create<SelectOp>(tinyMask, x, |
| builder.create<arith::DivFOp>(p, q)); |
| |
| rewriter.replaceOp(op, res); |
| |
| return success(); |
| } |
| |
| #define LN2_VALUE \ |
| 0.693147180559945309417232121458176568075500134360255254120680009493393621L |
| #define LOG2E_VALUE \ |
| 1.442695040888963407359924681001892137426645954152985934135449406931109219L |
| |
| //----------------------------------------------------------------------------// |
| // LogOp and Log2Op approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| template <typename Op> |
| struct LogApproximationBase : public OpRewritePattern<Op> { |
| using OpRewritePattern<Op>::OpRewritePattern; |
| |
| /// Base 2 if 'base2' is set; natural logarithm (base e) otherwise. |
| LogicalResult logMatchAndRewrite(Op op, PatternRewriter &rewriter, |
| bool base2) const; |
| }; |
| } // namespace |
| |
| // This approximation comes from Julien Pommier's SSE math library. |
| // Link: http://gruntthepeon.free.fr/ssemath |
| template <typename Op> |
| LogicalResult |
| LogApproximationBase<Op>::logMatchAndRewrite(Op op, PatternRewriter &rewriter, |
| bool base2) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| Value cstZero = bcast(f32Cst(builder, 0.0f)); |
| Value cstOne = bcast(f32Cst(builder, 1.0f)); |
| Value cstNegHalf = bcast(f32Cst(builder, -0.5f)); |
| |
| // The smallest non denormalized float number. |
| Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u)); |
| Value cstMinusInf = bcast(f32FromBits(builder, 0xff800000u)); |
| Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u)); |
| Value cstNan = bcast(f32FromBits(builder, 0x7fc00000)); |
| |
| // Polynomial coefficients. |
| Value cstCephesSQRTHF = bcast(f32Cst(builder, 0.707106781186547524f)); |
| Value cstCephesLogP0 = bcast(f32Cst(builder, 7.0376836292E-2f)); |
| Value cstCephesLogP1 = bcast(f32Cst(builder, -1.1514610310E-1f)); |
| Value cstCephesLogP2 = bcast(f32Cst(builder, 1.1676998740E-1f)); |
| Value cstCephesLogP3 = bcast(f32Cst(builder, -1.2420140846E-1f)); |
| Value cstCephesLogP4 = bcast(f32Cst(builder, +1.4249322787E-1f)); |
| Value cstCephesLogP5 = bcast(f32Cst(builder, -1.6668057665E-1f)); |
| Value cstCephesLogP6 = bcast(f32Cst(builder, +2.0000714765E-1f)); |
| Value cstCephesLogP7 = bcast(f32Cst(builder, -2.4999993993E-1f)); |
| Value cstCephesLogP8 = bcast(f32Cst(builder, +3.3333331174E-1f)); |
| |
| Value x = op.operand(); |
| |
| // Truncate input values to the minimum positive normal. |
| x = max(builder, x, cstMinNormPos); |
| |
| // Extract significant in the range [0.5,1) and exponent. |
| std::pair<Value, Value> pair = frexp(builder, x, /*is_positive=*/true); |
| x = pair.first; |
| Value e = pair.second; |
| |
| // Shift the inputs from the range [0.5,1) to [sqrt(1/2), sqrt(2)) and shift |
| // by -1.0. The values are then centered around 0, which improves the |
| // stability of the polynomial evaluation: |
| // |
| // if( x < SQRTHF ) { |
| // e -= 1; |
| // x = x + x - 1.0; |
| // } else { x = x - 1.0; } |
| Value mask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, x, |
| cstCephesSQRTHF); |
| Value tmp = builder.create<SelectOp>(mask, x, cstZero); |
| |
| x = builder.create<arith::SubFOp>(x, cstOne); |
| e = builder.create<arith::SubFOp>( |
| e, builder.create<SelectOp>(mask, cstOne, cstZero)); |
| x = builder.create<arith::AddFOp>(x, tmp); |
| |
| Value x2 = builder.create<arith::MulFOp>(x, x); |
| Value x3 = builder.create<arith::MulFOp>(x2, x); |
| |
| // Evaluate the polynomial approximant of degree 8 in three parts. |
| Value y0, y1, y2; |
| y0 = builder.create<math::FmaOp>(cstCephesLogP0, x, cstCephesLogP1); |
| y1 = builder.create<math::FmaOp>(cstCephesLogP3, x, cstCephesLogP4); |
| y2 = builder.create<math::FmaOp>(cstCephesLogP6, x, cstCephesLogP7); |
| y0 = builder.create<math::FmaOp>(y0, x, cstCephesLogP2); |
| y1 = builder.create<math::FmaOp>(y1, x, cstCephesLogP5); |
| y2 = builder.create<math::FmaOp>(y2, x, cstCephesLogP8); |
| y0 = builder.create<math::FmaOp>(y0, x3, y1); |
| y0 = builder.create<math::FmaOp>(y0, x3, y2); |
| y0 = builder.create<arith::MulFOp>(y0, x3); |
| |
| y0 = builder.create<math::FmaOp>(cstNegHalf, x2, y0); |
| x = builder.create<arith::AddFOp>(x, y0); |
| |
| if (base2) { |
| Value cstLog2e = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE))); |
| x = builder.create<math::FmaOp>(x, cstLog2e, e); |
| } else { |
| Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE))); |
| x = builder.create<math::FmaOp>(e, cstLn2, x); |
| } |
| |
| Value invalidMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::ULT, |
| op.operand(), cstZero); |
| Value zeroMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, |
| op.operand(), cstZero); |
| Value posInfMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, |
| op.operand(), cstPosInf); |
| |
| // Filter out invalid values: |
| // • x == 0 -> -INF |
| // • x < 0 -> NAN |
| // • x == +INF -> +INF |
| Value aproximation = builder.create<SelectOp>( |
| zeroMask, cstMinusInf, |
| builder.create<SelectOp>( |
| invalidMask, cstNan, |
| builder.create<SelectOp>(posInfMask, cstPosInf, x))); |
| |
| rewriter.replaceOp(op, aproximation); |
| |
| return success(); |
| } |
| |
| namespace { |
| struct LogApproximation : public LogApproximationBase<math::LogOp> { |
| using LogApproximationBase::LogApproximationBase; |
| |
| LogicalResult matchAndRewrite(math::LogOp op, |
| PatternRewriter &rewriter) const final { |
| return logMatchAndRewrite(op, rewriter, /*base2=*/false); |
| } |
| }; |
| } // namespace |
| |
| namespace { |
| struct Log2Approximation : public LogApproximationBase<math::Log2Op> { |
| using LogApproximationBase::LogApproximationBase; |
| |
| LogicalResult matchAndRewrite(math::Log2Op op, |
| PatternRewriter &rewriter) const final { |
| return logMatchAndRewrite(op, rewriter, /*base2=*/true); |
| } |
| }; |
| } // namespace |
| |
| //----------------------------------------------------------------------------// |
| // Log1p approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| struct Log1pApproximation : public OpRewritePattern<math::Log1pOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(math::Log1pOp op, |
| PatternRewriter &rewriter) const final; |
| }; |
| } // namespace |
| |
| // Approximate log(1+x). |
| LogicalResult |
| Log1pApproximation::matchAndRewrite(math::Log1pOp op, |
| PatternRewriter &rewriter) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| // Approximate log(1+x) using the following, due to W. Kahan: |
| // u = x + 1.0; |
| // if (u == 1.0 || u == inf) return x; |
| // return x * log(u) / (u - 1.0); |
| // ^^^^^^^^^^^^^^^^^^^^^^ |
| // "logLarge" below. |
| Value cstOne = bcast(f32Cst(builder, 1.0f)); |
| Value x = op.operand(); |
| Value u = builder.create<arith::AddFOp>(x, cstOne); |
| Value uSmall = |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, cstOne); |
| Value logU = builder.create<math::LogOp>(u); |
| Value uInf = |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, logU); |
| Value logLarge = builder.create<arith::MulFOp>( |
| x, builder.create<arith::DivFOp>( |
| logU, builder.create<arith::SubFOp>(u, cstOne))); |
| Value approximation = builder.create<SelectOp>( |
| builder.create<arith::OrIOp>(uSmall, uInf), x, logLarge); |
| rewriter.replaceOp(op, approximation); |
| return success(); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Erf approximation. |
| //----------------------------------------------------------------------------// |
| |
| // Approximates erf(x) with |
| // a - P(x)/Q(x) |
| // where P and Q are polynomials of degree 4. |
| // Different coefficients are chosen based on the value of x. |
| // The approximation error is ~2.5e-07. |
| // Boost's minimax tool that utilizes the Remez method was used to find the |
| // coefficients. |
| LogicalResult |
| ErfPolynomialApproximation::matchAndRewrite(math::ErfOp op, |
| PatternRewriter &rewriter) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| const int intervalsCount = 3; |
| const int polyDegree = 4; |
| |
| Value zero = bcast(f32Cst(builder, 0)); |
| Value one = bcast(f32Cst(builder, 1)); |
| Value pp[intervalsCount][polyDegree + 1]; |
| pp[0][0] = bcast(f32Cst(builder, +0.00000000000000000e+00f)); |
| pp[0][1] = bcast(f32Cst(builder, +1.12837916222975858e+00f)); |
| pp[0][2] = bcast(f32Cst(builder, -5.23018562988006470e-01f)); |
| pp[0][3] = bcast(f32Cst(builder, +2.09741709609267072e-01f)); |
| pp[0][4] = bcast(f32Cst(builder, +2.58146801602987875e-02f)); |
| pp[1][0] = bcast(f32Cst(builder, +0.00000000000000000e+00f)); |
| pp[1][1] = bcast(f32Cst(builder, +1.12750687816789140e+00f)); |
| pp[1][2] = bcast(f32Cst(builder, -3.64721408487825775e-01f)); |
| pp[1][3] = bcast(f32Cst(builder, +1.18407396425136952e-01f)); |
| pp[1][4] = bcast(f32Cst(builder, +3.70645533056476558e-02f)); |
| pp[2][0] = bcast(f32Cst(builder, -3.30093071049483172e-03f)); |
| pp[2][1] = bcast(f32Cst(builder, +3.51961938357697011e-03f)); |
| pp[2][2] = bcast(f32Cst(builder, -1.41373622814988039e-03f)); |
| pp[2][3] = bcast(f32Cst(builder, +2.53447094961941348e-04f)); |
| pp[2][4] = bcast(f32Cst(builder, -1.71048029455037401e-05f)); |
| |
| Value qq[intervalsCount][polyDegree + 1]; |
| qq[0][0] = bcast(f32Cst(builder, +1.000000000000000000e+00f)); |
| qq[0][1] = bcast(f32Cst(builder, -4.635138185962547255e-01f)); |
| qq[0][2] = bcast(f32Cst(builder, +5.192301327279782447e-01f)); |
| qq[0][3] = bcast(f32Cst(builder, -1.318089722204810087e-01f)); |
| qq[0][4] = bcast(f32Cst(builder, +7.397964654672315005e-02f)); |
| qq[1][0] = bcast(f32Cst(builder, +1.00000000000000000e+00f)); |
| qq[1][1] = bcast(f32Cst(builder, -3.27607011824493086e-01f)); |
| qq[1][2] = bcast(f32Cst(builder, +4.48369090658821977e-01f)); |
| qq[1][3] = bcast(f32Cst(builder, -8.83462621207857930e-02f)); |
| qq[1][4] = bcast(f32Cst(builder, +5.72442770283176093e-02f)); |
| qq[2][0] = bcast(f32Cst(builder, +1.00000000000000000e+00f)); |
| qq[2][1] = bcast(f32Cst(builder, -2.06069165953913769e+00f)); |
| qq[2][2] = bcast(f32Cst(builder, +1.62705939945477759e+00f)); |
| qq[2][3] = bcast(f32Cst(builder, -5.83389859211130017e-01f)); |
| qq[2][4] = bcast(f32Cst(builder, +8.21908939856640930e-02f)); |
| |
| Value offsets[intervalsCount]; |
| offsets[0] = bcast(f32Cst(builder, 0.0f)); |
| offsets[1] = bcast(f32Cst(builder, 0.0f)); |
| offsets[2] = bcast(f32Cst(builder, 1.0f)); |
| |
| Value bounds[intervalsCount]; |
| bounds[0] = bcast(f32Cst(builder, 0.8f)); |
| bounds[1] = bcast(f32Cst(builder, 2.0f)); |
| bounds[2] = bcast(f32Cst(builder, 3.75f)); |
| |
| Value isNegativeArg = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, |
| op.operand(), zero); |
| Value negArg = builder.create<arith::NegFOp>(op.operand()); |
| Value x = builder.create<SelectOp>(isNegativeArg, negArg, op.operand()); |
| |
| Value offset = offsets[0]; |
| Value p[polyDegree + 1]; |
| Value q[polyDegree + 1]; |
| for (int i = 0; i <= polyDegree; ++i) { |
| p[i] = pp[0][i]; |
| q[i] = qq[0][i]; |
| } |
| |
| // TODO: maybe use vector stacking to reduce the number of selects. |
| Value isLessThanBound[intervalsCount]; |
| for (int j = 0; j < intervalsCount - 1; ++j) { |
| isLessThanBound[j] = |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, x, bounds[j]); |
| for (int i = 0; i <= polyDegree; ++i) { |
| p[i] = builder.create<SelectOp>(isLessThanBound[j], p[i], pp[j + 1][i]); |
| q[i] = builder.create<SelectOp>(isLessThanBound[j], q[i], qq[j + 1][i]); |
| } |
| offset = |
| builder.create<SelectOp>(isLessThanBound[j], offset, offsets[j + 1]); |
| } |
| isLessThanBound[intervalsCount - 1] = builder.create<arith::CmpFOp>( |
| arith::CmpFPredicate::ULT, x, bounds[intervalsCount - 1]); |
| |
| Value pPoly = makePolynomialCalculation(builder, p, x); |
| Value qPoly = makePolynomialCalculation(builder, q, x); |
| Value rationalPoly = builder.create<arith::DivFOp>(pPoly, qPoly); |
| Value formula = builder.create<arith::AddFOp>(offset, rationalPoly); |
| formula = builder.create<SelectOp>(isLessThanBound[intervalsCount - 1], |
| formula, one); |
| |
| // erf is odd function: erf(x) = -erf(-x). |
| Value negFormula = builder.create<arith::NegFOp>(formula); |
| Value res = builder.create<SelectOp>(isNegativeArg, negFormula, formula); |
| |
| rewriter.replaceOp(op, res); |
| |
| return success(); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Exp approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| |
| struct ExpApproximation : public OpRewritePattern<math::ExpOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(math::ExpOp op, |
| PatternRewriter &rewriter) const final; |
| }; |
| } // namespace |
| |
| // Approximate exp(x) using its reduced range exp(y) where y is in the range |
| // [0, ln(2)], let y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2), exp(x) |
| // = exp(y) * 2^k. exp(y). |
| LogicalResult |
| ExpApproximation::matchAndRewrite(math::ExpOp op, |
| PatternRewriter &rewriter) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| |
| // TODO: Consider a common pattern rewriter with all methods below to |
| // write the approximations. |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| auto fmla = [&](Value a, Value b, Value c) { |
| return builder.create<math::FmaOp>(a, b, c); |
| }; |
| auto mul = [&](Value a, Value b) -> Value { |
| return builder.create<arith::MulFOp>(a, b); |
| }; |
| auto sub = [&](Value a, Value b) -> Value { |
| return builder.create<arith::SubFOp>(a, b); |
| }; |
| auto floor = [&](Value a) { return builder.create<math::FloorOp>(a); }; |
| |
| Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE))); |
| Value cstLog2E = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE))); |
| |
| // Polynomial coefficients. |
| Value cstCephesExpP0 = bcast(f32Cst(builder, 1.0)); |
| Value cstCephesExpP1 = bcast(f32Cst(builder, 1.0)); |
| Value cstCephesExpP2 = bcast(f32Cst(builder, 0.49970514590562437052f)); |
| Value cstCephesExpP3 = bcast(f32Cst(builder, 0.16873890085469545053f)); |
| Value cstCephesExpP4 = bcast(f32Cst(builder, 0.03668965196652099192f)); |
| Value cstCephesExpP5 = bcast(f32Cst(builder, 0.01314350012789660196f)); |
| |
| Value x = op.operand(); |
| |
| // Reduced y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2) |
| Value xL2Inv = mul(x, cstLog2E); |
| Value kF32 = floor(xL2Inv); |
| Value kLn2 = mul(kF32, cstLn2); |
| Value y = sub(x, kLn2); |
| |
| // Use Estrin's evaluation scheme with 3 independent parts: |
| // P(y)^y : (c0 + c1 y) + (c2 + c3 y) y^2 + (c4 + c5 y) y^4 |
| Value y2 = mul(y, y); |
| Value y4 = mul(y2, y2); |
| |
| Value q0 = fmla(cstCephesExpP1, y, cstCephesExpP0); |
| Value q1 = fmla(cstCephesExpP3, y, cstCephesExpP2); |
| Value q2 = fmla(cstCephesExpP5, y, cstCephesExpP4); |
| Value expY = fmla(q1, y2, q0); |
| expY = fmla(q2, y4, expY); |
| |
| auto i32Vec = broadcast(builder.getI32Type(), shape); |
| |
| // exp2(k) |
| Value k = builder.create<arith::FPToSIOp>(kF32, i32Vec); |
| Value exp2KValue = exp2I32(builder, k); |
| |
| // exp(x) = exp(y) * exp2(k) |
| expY = mul(expY, exp2KValue); |
| |
| // Handle overflow, inf and underflow of exp(x). exp(x) range is [0, inf], its |
| // partitioned as the following: |
| // exp(x) = 0, x <= -inf |
| // exp(x) = underflow (min_float), x <= -88 |
| // exp(x) = inf (min_float), x >= 88 |
| // Note: |k| = 127 is the value where the 8-bits exponent saturates. |
| Value zerof32Const = bcast(f32Cst(builder, 0)); |
| auto constPosInfinity = |
| bcast(f32Cst(builder, std::numeric_limits<float>::infinity())); |
| auto constNegIfinity = |
| bcast(f32Cst(builder, -std::numeric_limits<float>::infinity())); |
| auto underflow = bcast(f32Cst(builder, std::numeric_limits<float>::min())); |
| |
| Value kMaxConst = bcast(i32Cst(builder, 127)); |
| Value kMaxNegConst = bcast(i32Cst(builder, -127)); |
| Value rightBound = |
| builder.create<arith::CmpIOp>(arith::CmpIPredicate::sle, k, kMaxConst); |
| Value leftBound = |
| builder.create<arith::CmpIOp>(arith::CmpIPredicate::sge, k, kMaxNegConst); |
| |
| Value isNegInfinityX = builder.create<arith::CmpFOp>( |
| arith::CmpFPredicate::OEQ, x, constNegIfinity); |
| Value isPosInfinityX = builder.create<arith::CmpFOp>( |
| arith::CmpFPredicate::OEQ, x, constPosInfinity); |
| Value isPostiveX = |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, x, zerof32Const); |
| Value isComputable = builder.create<arith::AndIOp>(rightBound, leftBound); |
| |
| expY = builder.create<SelectOp>( |
| isNegInfinityX, zerof32Const, |
| builder.create<SelectOp>( |
| isPosInfinityX, constPosInfinity, |
| builder.create<SelectOp>(isComputable, expY, |
| builder.create<SelectOp>(isPostiveX, |
| constPosInfinity, |
| underflow)))); |
| |
| rewriter.replaceOp(op, expY); |
| |
| return success(); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // ExpM1 approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| |
| struct ExpM1Approximation : public OpRewritePattern<math::ExpM1Op> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(math::ExpM1Op op, |
| PatternRewriter &rewriter) const final; |
| }; |
| } // namespace |
| |
| LogicalResult |
| ExpM1Approximation::matchAndRewrite(math::ExpM1Op op, |
| PatternRewriter &rewriter) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| // expm1(x) = exp(x) - 1 = u - 1. |
| // We have to handle it carefully when x is near 0, i.e. u ~= 1, |
| // and when the input is ~= -inf, i.e. u - 1 ~= -1. |
| Value cstOne = bcast(f32Cst(builder, 1.0f)); |
| Value cstNegOne = bcast(f32Cst(builder, -1.0f)); |
| Value x = op.operand(); |
| Value u = builder.create<math::ExpOp>(x); |
| Value uEqOne = |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, cstOne); |
| Value uMinusOne = builder.create<arith::SubFOp>(u, cstOne); |
| Value uMinusOneEqNegOne = builder.create<arith::CmpFOp>( |
| arith::CmpFPredicate::OEQ, uMinusOne, cstNegOne); |
| // logU = log(u) ~= x |
| Value logU = builder.create<math::LogOp>(u); |
| |
| // Detect exp(x) = +inf; written this way to avoid having to form +inf. |
| Value isInf = |
| builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, logU, u); |
| |
| // (u - 1) * (x / ~x) |
| Value expm1 = builder.create<arith::MulFOp>( |
| uMinusOne, builder.create<arith::DivFOp>(x, logU)); |
| expm1 = builder.create<SelectOp>(isInf, u, expm1); |
| Value approximation = builder.create<SelectOp>( |
| uEqOne, x, builder.create<SelectOp>(uMinusOneEqNegOne, cstNegOne, expm1)); |
| rewriter.replaceOp(op, approximation); |
| return success(); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Sin and Cos approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| |
| template <bool isSine, typename OpTy> |
| struct SinAndCosApproximation : public OpRewritePattern<OpTy> { |
| public: |
| using OpRewritePattern<OpTy>::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(OpTy op, PatternRewriter &rewriter) const final; |
| }; |
| } // namespace |
| |
| #define TWO_OVER_PI \ |
| 0.6366197723675813430755350534900574481378385829618257949906693762L |
| #define PI_OVER_2 \ |
| 1.5707963267948966192313216916397514420985846996875529104874722961L |
| |
| // Approximates sin(x) or cos(x) by finding the best approximation polynomial in |
| // the reduced range [0, pi/2] for both sin(x) and cos(x). Then given y in the |
| // reduced range sin(x) will be computed as sin(y), -sin(y), cos(y) or -cos(y). |
| template <bool isSine, typename OpTy> |
| LogicalResult SinAndCosApproximation<isSine, OpTy>::matchAndRewrite( |
| OpTy op, PatternRewriter &rewriter) const { |
| static_assert( |
| llvm::is_one_of<OpTy, math::SinOp, math::CosOp>::value, |
| "SinAndCosApproximation pattern expects math::SinOp or math::CosOp"); |
| |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| auto mul = [&](Value a, Value b) -> Value { |
| return builder.create<arith::MulFOp>(a, b); |
| }; |
| auto sub = [&](Value a, Value b) -> Value { |
| return builder.create<arith::SubFOp>(a, b); |
| }; |
| auto floor = [&](Value a) { return builder.create<math::FloorOp>(a); }; |
| |
| auto i32Vec = broadcast(builder.getI32Type(), shape); |
| auto fPToSingedInteger = [&](Value a) -> Value { |
| return builder.create<arith::FPToSIOp>(a, i32Vec); |
| }; |
| |
| auto modulo4 = [&](Value a) -> Value { |
| return builder.create<arith::AndIOp>(a, bcast(i32Cst(builder, 3))); |
| }; |
| |
| auto isEqualTo = [&](Value a, Value b) -> Value { |
| return builder.create<arith::CmpIOp>(arith::CmpIPredicate::eq, a, b); |
| }; |
| |
| auto isGreaterThan = [&](Value a, Value b) -> Value { |
| return builder.create<arith::CmpIOp>(arith::CmpIPredicate::sgt, a, b); |
| }; |
| |
| auto select = [&](Value cond, Value t, Value f) -> Value { |
| return builder.create<SelectOp>(cond, t, f); |
| }; |
| |
| auto fmla = [&](Value a, Value b, Value c) { |
| return builder.create<math::FmaOp>(a, b, c); |
| }; |
| |
| auto bitwiseOr = [&](Value a, Value b) { |
| return builder.create<arith::OrIOp>(a, b); |
| }; |
| |
| Value twoOverPi = bcast(f32Cst(builder, TWO_OVER_PI)); |
| Value piOverTwo = bcast(f32Cst(builder, PI_OVER_2)); |
| |
| Value x = op.operand(); |
| |
| Value k = floor(mul(x, twoOverPi)); |
| |
| Value y = sub(x, mul(k, piOverTwo)); |
| |
| Value cstOne = bcast(f32Cst(builder, 1.0)); |
| Value cstNegativeOne = bcast(f32Cst(builder, -1.0)); |
| |
| Value cstSC2 = bcast(f32Cst(builder, -0.16666667163372039794921875f)); |
| Value cstSC4 = bcast(f32Cst(builder, 8.333347737789154052734375e-3f)); |
| Value cstSC6 = bcast(f32Cst(builder, -1.9842604524455964565277099609375e-4f)); |
| Value cstSC8 = |
| bcast(f32Cst(builder, 2.760012648650445044040679931640625e-6f)); |
| Value cstSC10 = |
| bcast(f32Cst(builder, -2.50293279435709337121807038784027099609375e-8f)); |
| |
| Value cstCC2 = bcast(f32Cst(builder, -0.5f)); |
| Value cstCC4 = bcast(f32Cst(builder, 4.166664183139801025390625e-2f)); |
| Value cstCC6 = bcast(f32Cst(builder, -1.388833043165504932403564453125e-3f)); |
| Value cstCC8 = bcast(f32Cst(builder, 2.47562347794882953166961669921875e-5f)); |
| Value cstCC10 = |
| bcast(f32Cst(builder, -2.59630184018533327616751194000244140625e-7f)); |
| |
| Value kMod4 = modulo4(fPToSingedInteger(k)); |
| |
| Value kR0 = isEqualTo(kMod4, bcast(i32Cst(builder, 0))); |
| Value kR1 = isEqualTo(kMod4, bcast(i32Cst(builder, 1))); |
| Value kR2 = isEqualTo(kMod4, bcast(i32Cst(builder, 2))); |
| Value kR3 = isEqualTo(kMod4, bcast(i32Cst(builder, 3))); |
| |
| Value sinuseCos = isSine ? bitwiseOr(kR1, kR3) : bitwiseOr(kR0, kR2); |
| Value negativeRange = isSine ? isGreaterThan(kMod4, bcast(i32Cst(builder, 1))) |
| : bitwiseOr(kR1, kR2); |
| |
| Value y2 = mul(y, y); |
| |
| Value base = select(sinuseCos, cstOne, y); |
| Value cstC2 = select(sinuseCos, cstCC2, cstSC2); |
| Value cstC4 = select(sinuseCos, cstCC4, cstSC4); |
| Value cstC6 = select(sinuseCos, cstCC6, cstSC6); |
| Value cstC8 = select(sinuseCos, cstCC8, cstSC8); |
| Value cstC10 = select(sinuseCos, cstCC10, cstSC10); |
| |
| Value v1 = fmla(y2, cstC10, cstC8); |
| Value v2 = fmla(y2, v1, cstC6); |
| Value v3 = fmla(y2, v2, cstC4); |
| Value v4 = fmla(y2, v3, cstC2); |
| Value v5 = fmla(y2, v4, cstOne); |
| Value v6 = mul(base, v5); |
| |
| Value approximation = select(negativeRange, mul(cstNegativeOne, v6), v6); |
| |
| rewriter.replaceOp(op, approximation); |
| |
| return success(); |
| } |
| |
| //----------------------------------------------------------------------------// |
| // Rsqrt approximation. |
| //----------------------------------------------------------------------------// |
| |
| namespace { |
| struct RsqrtApproximation : public OpRewritePattern<math::RsqrtOp> { |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(math::RsqrtOp op, |
| PatternRewriter &rewriter) const final; |
| }; |
| } // namespace |
| |
| LogicalResult |
| RsqrtApproximation::matchAndRewrite(math::RsqrtOp op, |
| PatternRewriter &rewriter) const { |
| if (!getElementTypeOrSelf(op.operand()).isF32()) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ArrayRef<int64_t> shape = vectorShape(op.operand()); |
| |
| // Only support already-vectorized rsqrt's. |
| if (shape.empty() || shape.back() % 8 != 0) |
| return rewriter.notifyMatchFailure(op, "unsupported operand type"); |
| |
| ImplicitLocOpBuilder builder(op->getLoc(), rewriter); |
| auto bcast = [&](Value value) -> Value { |
| return broadcast(builder, value, shape); |
| }; |
| |
| Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u)); |
| Value cstOnePointFive = bcast(f32Cst(builder, 1.5f)); |
| Value cstNegHalf = bcast(f32Cst(builder, -0.5f)); |
| Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u)); |
| |
| Value negHalf = builder.create<arith::MulFOp>(op.operand(), cstNegHalf); |
| |
| // Select only the inverse sqrt of positive normals (denormals are |
| // flushed to zero). |
| Value ltMinMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, |
| op.operand(), cstMinNormPos); |
| Value infMask = builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, |
| op.operand(), cstPosInf); |
| Value notNormalFiniteMask = builder.create<arith::OrIOp>(ltMinMask, infMask); |
| |
| // Compute an approximate result. |
| Value yApprox = handleMultidimensionalVectors( |
| builder, op->getOperands(), 8, [&builder](ValueRange operands) -> Value { |
| return builder.create<x86vector::RsqrtOp>(operands); |
| }); |
| |
| // Do a single step of Newton-Raphson iteration to improve the approximation. |
| // This uses the formula y_{n+1} = y_n * (1.5 - y_n * (0.5 * x) * y_n). |
| // It is essential to evaluate the inner term like this because forming |
| // y_n^2 may over- or underflow. |
| Value inner = builder.create<arith::MulFOp>(negHalf, yApprox); |
| Value fma = builder.create<math::FmaOp>(yApprox, inner, cstOnePointFive); |
| Value yNewton = builder.create<arith::MulFOp>(yApprox, fma); |
| |
| // Select the result of the Newton-Raphson step for positive normal arguments. |
| // For other arguments, choose the output of the intrinsic. This will |
| // return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(x) = +inf if |
| // x is zero or a positive denormalized float (equivalent to flushing positive |
| // denormalized inputs to zero). |
| Value res = builder.create<SelectOp>(notNormalFiniteMask, yApprox, yNewton); |
| rewriter.replaceOp(op, res); |
| |
| return success(); |
| } |
| |
| //----------------------------------------------------------------------------// |
| |
| void mlir::populateMathPolynomialApproximationPatterns( |
| RewritePatternSet &patterns, |
| const MathPolynomialApproximationOptions &options) { |
| patterns.add<TanhApproximation, LogApproximation, Log2Approximation, |
| Log1pApproximation, ErfPolynomialApproximation, ExpApproximation, |
| ExpM1Approximation, SinAndCosApproximation<true, math::SinOp>, |
| SinAndCosApproximation<false, math::CosOp>>( |
| patterns.getContext()); |
| if (options.enableAvx2) |
| patterns.add<RsqrtApproximation>(patterns.getContext()); |
| } |