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_CONVERTMATHTOLIBMPASS
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, PatternBenefit benefit,
54 StringRef floatFunc, StringRef doubleFunc)
55 : OpRewritePattern<Op>(context, benefit), 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, PatternBenefit benefit,
66 MLIRContext *ctx, StringRef floatFunc,
67 StringRef doubleFunc) {
68 patterns.add<VecOpToScalarOp<OpTy>, PromoteOpToF32<OpTy>>(ctx, benefit);
69 patterns.add<ScalarOpToLibmCall<OpTy>>(ctx, benefit, floatFunc, doubleFunc);
70}
71
72} // namespace
73
74template <typename Op>
75LogicalResult
76VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {
77 auto opType = op.getType();
78 auto loc = op.getLoc();
79 auto vecType = dyn_cast<VectorType>(opType);
80
81 if (!vecType)
82 return failure();
83 if (!vecType.hasRank())
84 return failure();
85 auto shape = vecType.getShape();
86 int64_t numElements = vecType.getNumElements();
87
88 Value result = rewriter.create<arith::ConstantOp>(
89 loc, DenseElementsAttr::get(
90 vecType, 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 rewriter.create<vector::ExtractOp>(loc, input, positions));
98 Value scalarOp =
99 rewriter.create<Op>(loc, vecType.getElementType(), operands);
100 result =
101 rewriter.create<vector::InsertOp>(loc, scalarOp, result, positions);
102 }
103 rewriter.replaceOp(op, {result});
104 return success();
105}
106
107template <typename Op>
108LogicalResult
109PromoteOpToF32<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 rewriter.create<arith::ExtFOp>(loc, f32, operand);
119 }));
120 auto newOp = rewriter.create<Op>(loc, f32, extendedOperands);
121 rewriter.replaceOpWithNewOp<arith::TruncFOp>(op, opType, newOp);
122 return success();
123}
124
125template <typename Op>
126LogicalResult
127ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,
128 PatternRewriter &rewriter) const {
129 auto module = SymbolTable::getNearestSymbolTable(from: 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 been
138 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 = rewriter.create<func::FuncOp>(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 more
149 // optimization opportunities (e.g. LICM) for backends targeting LLVM IR.
150 // This will have to be changed, when strict FP behavior is supported
151 // 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 ConvertMathToLibmPass
216 : public impl::ConvertMathToLibmPassBase<ConvertMathToLibmPass> {
217 void runOnOperation() override;
218};
219} // namespace
220
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

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