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

source code of mlir/lib/Conversion/MathToLibm/MathToLibm.cpp