llvm/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

//===- ExpandPatterns.cpp - Code to expand various 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 various math operations.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"

usingnamespacemlir;

/// Create a float constant.
static Value createFloatConst(Location loc, Type type, APFloat value,
                              OpBuilder &b) { … }

static Value createFloatConst(Location loc, Type type, double value,
                              OpBuilder &b) { … }

/// Create an integer constant.
static Value createIntConst(Location loc, Type type, int64_t value,
                            OpBuilder &b) { … }

static Value createTruncatedFPValue(Value operand, ImplicitLocOpBuilder &b) { … }

// sinhf(float x) -> (exp(x) - exp(-x)) / 2
static LogicalResult convertSinhOp(math::SinhOp op, PatternRewriter &rewriter) { … }

// coshf(float x) -> (exp(x) + exp(-x)) / 2
static LogicalResult convertCoshOp(math::CoshOp op, PatternRewriter &rewriter) { … }

/// Expands tanh op into
/// 1-exp^{-2x} / 1+exp^{-2x}
/// To avoid overflow we exploit the reflection symmetry `tanh(-x) = -tanh(x)`.
/// We compute a "signs" value which is -1 if input is negative and +1 if input
/// is positive.  Then multiply the input by this value, guaranteeing that the
/// result is positive, which also guarantees `exp^{-2x * sign(x)}` is in (0,
/// 1]. Expand the computation on the input `x * sign(x)`, then multiply the
/// result by `sign(x)` to retain sign of the real result.
static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) { … }

// Converts math.tan to math.sin, math.cos, and arith.divf.
static LogicalResult convertTanOp(math::TanOp op, PatternRewriter &rewriter) { … }

// asinh(float x) -> log(x + sqrt(x**2 + 1))
static LogicalResult convertAsinhOp(math::AsinhOp op,
                                    PatternRewriter &rewriter) { … }

// acosh(float x) -> log(x + sqrt(x**2 - 1))
static LogicalResult convertAcoshOp(math::AcoshOp op,
                                    PatternRewriter &rewriter) { … }

// atanh(float x) -> log((1 + x) / (1 - x)) / 2
static LogicalResult convertAtanhOp(math::AtanhOp op,
                                    PatternRewriter &rewriter) { … }

static LogicalResult convertFmaFOp(math::FmaOp op, PatternRewriter &rewriter) { … }

// Converts a floorf() function to the following:
// floorf(float x) ->
//     y = (float)(int) x
//     if (x < 0) then incr = -1 else incr = 0
//     y = y + incr    <= replace this op with the floorf op.
static LogicalResult convertFloorOp(math::FloorOp op,
                                    PatternRewriter &rewriter) { … }

// Converts a ceilf() function to the following:
// ceilf(float x) ->
//      y = (float)(int) x
//      if (x > y) then incr = 1 else incr = 0
//      y = y + incr   <= replace this op with the ceilf op.
static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) { … }

// Convert `math.fpowi` to a series of `arith.mulf` operations.
// If the power is negative, we divide one by the result.
// If both the base and power are zero, the result is 1.
// In the case of non constant power, we convert the operation to `math.powf`.
static LogicalResult convertFPowIOp(math::FPowIOp op,
                                    PatternRewriter &rewriter) { … }

// Converts  Powf(float a, float b) (meaning a^b) to exp^(b * ln(a))
static LogicalResult convertPowfOp(math::PowFOp op, PatternRewriter &rewriter) { … }

// exp2f(float x) -> exp(x * ln(2))
//   Proof: Let's say 2^x = y
//   ln(2^x) = ln(y)
//   x * ln(2) = ln(y) => e ^(x*ln(2)) = y
static LogicalResult convertExp2fOp(math::Exp2Op op,
                                    PatternRewriter &rewriter) { … }

static LogicalResult convertRoundOp(math::RoundOp op,
                                    PatternRewriter &rewriter) { … }

// Converts math.ctlz to scf and arith operations. This is done
// by performing a binary search on the bits.
static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
                                   PatternRewriter &rewriter) { … }

// Convert `math.roundeven` into `math.round` + arith ops
static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
                                        PatternRewriter &rewriter) { … }

// Convert `math.rsqrt` into `arith.divf` + `math.sqrt`
static LogicalResult convertRsqrtOp(math::RsqrtOp op,
                                    PatternRewriter &rewriter) { … }

void mlir::populateExpandCtlzPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandSinhPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandCoshPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandTanPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandTanhPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandAsinhPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandAcoshPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandAtanhPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandFmaFPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandCeilFPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandExp2FPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandPowFPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandFPowIPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandRoundFPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandFloorFPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandRoundEvenPattern(RewritePatternSet &patterns) { … }

void mlir::populateExpandRsqrtPattern(RewritePatternSet &patterns) { … }