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