234 lines · cpp
1//===-- MathToLibm.cpp - conversion from Math to libm calls ---------------===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Conversion/MathToLibm/MathToLibm.h"10 11#include "mlir/Dialect/Arith/IR/Arith.h"12#include "mlir/Dialect/Func/IR/FuncOps.h"13#include "mlir/Dialect/LLVMIR/LLVMDialect.h"14#include "mlir/Dialect/Math/IR/Math.h"15#include "mlir/Dialect/Utils/IndexingUtils.h"16#include "mlir/Dialect/Vector/IR/VectorOps.h"17#include "mlir/IR/BuiltinDialect.h"18#include "mlir/IR/PatternMatch.h"19#include "mlir/Transforms/DialectConversion.h"20 21namespace mlir {22#define GEN_PASS_DEF_CONVERTMATHTOLIBMPASS23#include "mlir/Conversion/Passes.h.inc"24} // namespace mlir25 26using namespace mlir;27 28namespace {29// Pattern to convert vector operations to scalar operations. This is needed as30// libm calls require scalars.31template <typename Op>32struct VecOpToScalarOp : public OpRewritePattern<Op> {33public:34 using OpRewritePattern<Op>::OpRewritePattern;35 36 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;37};38// Pattern to promote an op of a smaller floating point type to F32.39template <typename Op>40struct PromoteOpToF32 : public OpRewritePattern<Op> {41public:42 using OpRewritePattern<Op>::OpRewritePattern;43 44 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;45};46// Pattern to convert scalar math operations to calls to libm functions.47// Additionally the libm function signatures are declared.48template <typename Op>49struct ScalarOpToLibmCall : public OpRewritePattern<Op> {50public:51 using OpRewritePattern<Op>::OpRewritePattern;52 ScalarOpToLibmCall(MLIRContext *context, PatternBenefit benefit,53 StringRef floatFunc, StringRef doubleFunc)54 : OpRewritePattern<Op>(context, benefit), floatFunc(floatFunc),55 doubleFunc(doubleFunc) {};56 57 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;58 59private:60 std::string floatFunc, doubleFunc;61};62 63template <typename OpTy>64void populatePatternsForOp(RewritePatternSet &patterns, PatternBenefit benefit,65 MLIRContext *ctx, StringRef floatFunc,66 StringRef doubleFunc) {67 patterns.add<VecOpToScalarOp<OpTy>, PromoteOpToF32<OpTy>>(ctx, benefit);68 patterns.add<ScalarOpToLibmCall<OpTy>>(ctx, benefit, floatFunc, doubleFunc);69}70 71} // namespace72 73template <typename Op>74LogicalResult75VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {76 auto opType = op.getType();77 auto loc = op.getLoc();78 auto vecType = dyn_cast<VectorType>(opType);79 80 if (!vecType)81 return failure();82 if (!vecType.hasRank())83 return failure();84 auto shape = vecType.getShape();85 int64_t numElements = vecType.getNumElements();86 87 Value result = arith::ConstantOp::create(88 rewriter, loc,89 DenseElementsAttr::get(vecType,90 FloatAttr::get(vecType.getElementType(), 0.0)));91 SmallVector<int64_t> strides = computeStrides(shape);92 for (auto linearIndex = 0; linearIndex < numElements; ++linearIndex) {93 SmallVector<int64_t> positions = delinearize(linearIndex, strides);94 SmallVector<Value> operands;95 for (auto input : op->getOperands())96 operands.push_back(97 vector::ExtractOp::create(rewriter, loc, input, positions));98 Value scalarOp =99 Op::create(rewriter, loc, vecType.getElementType(), operands);100 result =101 vector::InsertOp::create(rewriter, loc, scalarOp, result, positions);102 }103 rewriter.replaceOp(op, {result});104 return success();105}106 107template <typename Op>108LogicalResult109PromoteOpToF32<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {110 auto opType = op.getType();111 if (!isa<Float16Type, BFloat16Type>(opType))112 return failure();113 114 auto loc = op.getLoc();115 auto f32 = rewriter.getF32Type();116 auto extendedOperands = llvm::to_vector(117 llvm::map_range(op->getOperands(), [&](Value operand) -> Value {118 return arith::ExtFOp::create(rewriter, loc, f32, operand);119 }));120 auto newOp = Op::create(rewriter, loc, f32, extendedOperands);121 rewriter.replaceOpWithNewOp<arith::TruncFOp>(op, opType, newOp);122 return success();123}124 125template <typename Op>126LogicalResult127ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,128 PatternRewriter &rewriter) const {129 auto module = SymbolTable::getNearestSymbolTable(op);130 auto type = op.getType();131 if (!isa<Float32Type, Float64Type>(type))132 return failure();133 134 auto name = type.getIntOrFloatBitWidth() == 64 ? doubleFunc : floatFunc;135 auto opFunc = dyn_cast_or_null<SymbolOpInterface>(136 SymbolTable::lookupSymbolIn(module, name));137 // Forward declare function if it hasn't already been138 if (!opFunc) {139 OpBuilder::InsertionGuard guard(rewriter);140 rewriter.setInsertionPointToStart(&module->getRegion(0).front());141 auto opFunctionTy = FunctionType::get(142 rewriter.getContext(), op->getOperandTypes(), op->getResultTypes());143 opFunc = func::FuncOp::create(rewriter, rewriter.getUnknownLoc(), name,144 opFunctionTy);145 opFunc.setPrivate();146 147 // By definition Math dialect operations imply LLVM's "readnone"148 // function attribute, so we can set it here to provide more149 // optimization opportunities (e.g. LICM) for backends targeting LLVM IR.150 // This will have to be changed, when strict FP behavior is supported151 // by Math dialect.152 opFunc->setAttr(LLVM::LLVMDialect::getReadnoneAttrName(),153 UnitAttr::get(rewriter.getContext()));154 }155 assert(isa<FunctionOpInterface>(SymbolTable::lookupSymbolIn(module, name)));156 157 rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(),158 op->getOperands());159 160 return success();161}162 163void mlir::populateMathToLibmConversionPatterns(RewritePatternSet &patterns,164 PatternBenefit benefit) {165 MLIRContext *ctx = patterns.getContext();166 167 populatePatternsForOp<math::AbsFOp>(patterns, benefit, ctx, "fabsf", "fabs");168 populatePatternsForOp<math::AcosOp>(patterns, benefit, ctx, "acosf", "acos");169 populatePatternsForOp<math::AcoshOp>(patterns, benefit, ctx, "acoshf",170 "acosh");171 populatePatternsForOp<math::AsinOp>(patterns, benefit, ctx, "asinf", "asin");172 populatePatternsForOp<math::AsinhOp>(patterns, benefit, ctx, "asinhf",173 "asinh");174 populatePatternsForOp<math::Atan2Op>(patterns, benefit, ctx, "atan2f",175 "atan2");176 populatePatternsForOp<math::AtanOp>(patterns, benefit, ctx, "atanf", "atan");177 populatePatternsForOp<math::AtanhOp>(patterns, benefit, ctx, "atanhf",178 "atanh");179 populatePatternsForOp<math::CbrtOp>(patterns, benefit, ctx, "cbrtf", "cbrt");180 populatePatternsForOp<math::CeilOp>(patterns, benefit, ctx, "ceilf", "ceil");181 populatePatternsForOp<math::CosOp>(patterns, benefit, ctx, "cosf", "cos");182 populatePatternsForOp<math::CoshOp>(patterns, benefit, ctx, "coshf", "cosh");183 populatePatternsForOp<math::ErfOp>(patterns, benefit, ctx, "erff", "erf");184 populatePatternsForOp<math::ErfcOp>(patterns, benefit, ctx, "erfcf", "erfc");185 populatePatternsForOp<math::ExpOp>(patterns, benefit, ctx, "expf", "exp");186 populatePatternsForOp<math::Exp2Op>(patterns, benefit, ctx, "exp2f", "exp2");187 populatePatternsForOp<math::ExpM1Op>(patterns, benefit, ctx, "expm1f",188 "expm1");189 populatePatternsForOp<math::FloorOp>(patterns, benefit, ctx, "floorf",190 "floor");191 populatePatternsForOp<math::FmaOp>(patterns, benefit, ctx, "fmaf", "fma");192 populatePatternsForOp<math::LogOp>(patterns, benefit, ctx, "logf", "log");193 populatePatternsForOp<math::Log2Op>(patterns, benefit, ctx, "log2f", "log2");194 populatePatternsForOp<math::Log10Op>(patterns, benefit, ctx, "log10f",195 "log10");196 populatePatternsForOp<math::Log1pOp>(patterns, benefit, ctx, "log1pf",197 "log1p");198 populatePatternsForOp<math::PowFOp>(patterns, benefit, ctx, "powf", "pow");199 populatePatternsForOp<math::RoundEvenOp>(patterns, benefit, ctx, "roundevenf",200 "roundeven");201 populatePatternsForOp<math::RoundOp>(patterns, benefit, ctx, "roundf",202 "round");203 populatePatternsForOp<math::SinOp>(patterns, benefit, ctx, "sinf", "sin");204 populatePatternsForOp<math::SinhOp>(patterns, benefit, ctx, "sinhf", "sinh");205 populatePatternsForOp<math::SqrtOp>(patterns, benefit, ctx, "sqrtf", "sqrt");206 populatePatternsForOp<math::RsqrtOp>(patterns, benefit, ctx, "rsqrtf",207 "rsqrt");208 populatePatternsForOp<math::TanOp>(patterns, benefit, ctx, "tanf", "tan");209 populatePatternsForOp<math::TanhOp>(patterns, benefit, ctx, "tanhf", "tanh");210 populatePatternsForOp<math::TruncOp>(patterns, benefit, ctx, "truncf",211 "trunc");212}213 214namespace {215struct ConvertMathToLibmPass216 : public impl::ConvertMathToLibmPassBase<ConvertMathToLibmPass> {217 void runOnOperation() override;218};219} // namespace220 221void ConvertMathToLibmPass::runOnOperation() {222 auto module = getOperation();223 224 RewritePatternSet patterns(&getContext());225 populateMathToLibmConversionPatterns(patterns);226 227 ConversionTarget target(getContext());228 target.addLegalDialect<arith::ArithDialect, BuiltinDialect, func::FuncDialect,229 vector::VectorDialect>();230 target.addIllegalDialect<math::MathDialect>();231 if (failed(applyPartialConversion(module, target, std::move(patterns))))232 signalPassFailure();233}234