1//===-- ComplexToLibm.cpp - conversion from Complex 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/ComplexToLibm/ComplexToLibm.h"
10
11#include "mlir/Dialect/Complex/IR/Complex.h"
12#include "mlir/Dialect/Func/IR/FuncOps.h"
13#include "mlir/IR/PatternMatch.h"
14#include <optional>
15
16namespace mlir {
17#define GEN_PASS_DEF_CONVERTCOMPLEXTOLIBM
18#include "mlir/Conversion/Passes.h.inc"
19} // namespace mlir
20
21using namespace mlir;
22
23namespace {
24// Functor to resolve the function name corresponding to the given complex
25// result type.
26struct ComplexTypeResolver {
27 std::optional<bool> operator()(Type type) const {
28 auto complexType = cast<ComplexType>(Val&: type);
29 auto elementType = complexType.getElementType();
30 if (!isa<Float32Type, Float64Type>(Val: elementType))
31 return {};
32
33 return elementType.getIntOrFloatBitWidth() == 64;
34 }
35};
36
37// Functor to resolve the function name corresponding to the given float result
38// type.
39struct FloatTypeResolver {
40 std::optional<bool> operator()(Type type) const {
41 auto elementType = cast<FloatType>(Val&: type);
42 if (!isa<Float32Type, Float64Type>(Val: elementType))
43 return {};
44
45 return elementType.getIntOrFloatBitWidth() == 64;
46 }
47};
48
49// Pattern to convert scalar complex operations to calls to libm functions.
50// Additionally the libm function signatures are declared.
51// TypeResolver is a functor returning the libm function name according to the
52// expected type double or float.
53template <typename Op, typename TypeResolver = ComplexTypeResolver>
54struct ScalarOpToLibmCall : public OpRewritePattern<Op> {
55public:
56 using OpRewritePattern<Op>::OpRewritePattern;
57 ScalarOpToLibmCall(MLIRContext *context, StringRef floatFunc,
58 StringRef doubleFunc, PatternBenefit benefit)
59 : OpRewritePattern<Op>(context, benefit), floatFunc(floatFunc),
60 doubleFunc(doubleFunc){};
61
62 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;
63
64private:
65 std::string floatFunc, doubleFunc;
66};
67} // namespace
68
69template <typename Op, typename TypeResolver>
70LogicalResult ScalarOpToLibmCall<Op, TypeResolver>::matchAndRewrite(
71 Op op, PatternRewriter &rewriter) const {
72 auto module = SymbolTable::getNearestSymbolTable(from: op);
73 auto isDouble = TypeResolver()(op.getType());
74 if (!isDouble.has_value())
75 return failure();
76
77 auto name = *isDouble ? doubleFunc : floatFunc;
78
79 auto opFunc = dyn_cast_or_null<SymbolOpInterface>(
80 SymbolTable::lookupSymbolIn(module, name));
81 // Forward declare function if it hasn't already been
82 if (!opFunc) {
83 OpBuilder::InsertionGuard guard(rewriter);
84 rewriter.setInsertionPointToStart(&module->getRegion(0).front());
85 auto opFunctionTy = FunctionType::get(
86 context: rewriter.getContext(), inputs: op->getOperandTypes(), results: op->getResultTypes());
87 opFunc = rewriter.create<func::FuncOp>(rewriter.getUnknownLoc(), name,
88 opFunctionTy);
89 opFunc.setPrivate();
90 }
91 assert(isa<FunctionOpInterface>(SymbolTable::lookupSymbolIn(module, name)));
92
93 rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(),
94 op->getOperands());
95
96 return success();
97}
98
99void mlir::populateComplexToLibmConversionPatterns(RewritePatternSet &patterns,
100 PatternBenefit benefit) {
101 patterns.add<ScalarOpToLibmCall<complex::PowOp>>(arg: patterns.getContext(),
102 args: "cpowf", args: "cpow", args&: benefit);
103 patterns.add<ScalarOpToLibmCall<complex::SqrtOp>>(arg: patterns.getContext(),
104 args: "csqrtf", args: "csqrt", args&: benefit);
105 patterns.add<ScalarOpToLibmCall<complex::TanhOp>>(arg: patterns.getContext(),
106 args: "ctanhf", args: "ctanh", args&: benefit);
107 patterns.add<ScalarOpToLibmCall<complex::CosOp>>(arg: patterns.getContext(),
108 args: "ccosf", args: "ccos", args&: benefit);
109 patterns.add<ScalarOpToLibmCall<complex::SinOp>>(arg: patterns.getContext(),
110 args: "csinf", args: "csin", args&: benefit);
111 patterns.add<ScalarOpToLibmCall<complex::ConjOp>>(arg: patterns.getContext(),
112 args: "conjf", args: "conj", args&: benefit);
113 patterns.add<ScalarOpToLibmCall<complex::LogOp>>(arg: patterns.getContext(),
114 args: "clogf", args: "clog", args&: benefit);
115 patterns.add<ScalarOpToLibmCall<complex::AbsOp, FloatTypeResolver>>(
116 arg: patterns.getContext(), args: "cabsf", args: "cabs", args&: benefit);
117 patterns.add<ScalarOpToLibmCall<complex::AngleOp, FloatTypeResolver>>(
118 arg: patterns.getContext(), args: "cargf", args: "carg", args&: benefit);
119 patterns.add<ScalarOpToLibmCall<complex::TanOp>>(arg: patterns.getContext(),
120 args: "ctanf", args: "ctan", args&: benefit);
121}
122
123namespace {
124struct ConvertComplexToLibmPass
125 : public impl::ConvertComplexToLibmBase<ConvertComplexToLibmPass> {
126 void runOnOperation() override;
127};
128} // namespace
129
130void ConvertComplexToLibmPass::runOnOperation() {
131 auto module = getOperation();
132
133 RewritePatternSet patterns(&getContext());
134 populateComplexToLibmConversionPatterns(patterns, /*benefit=*/1);
135
136 ConversionTarget target(getContext());
137 target.addLegalDialect<func::FuncDialect>();
138 target.addIllegalOp<complex::PowOp, complex::SqrtOp, complex::TanhOp,
139 complex::CosOp, complex::SinOp, complex::ConjOp,
140 complex::LogOp, complex::AbsOp, complex::AngleOp,
141 complex::TanOp>();
142 if (failed(Result: applyPartialConversion(op: module, target, patterns: std::move(patterns))))
143 signalPassFailure();
144}
145

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