145 lines · cpp
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-exception6//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_CONVERTCOMPLEXTOLIBM18#include "mlir/Conversion/Passes.h.inc"19} // namespace mlir20 21using namespace mlir;22 23namespace {24// Functor to resolve the function name corresponding to the given complex25// result type.26struct ComplexTypeResolver {27 std::optional<bool> operator()(Type type) const {28 auto complexType = cast<ComplexType>(type);29 auto elementType = complexType.getElementType();30 if (!isa<Float32Type, Float64Type>(elementType))31 return {};32 33 return elementType.getIntOrFloatBitWidth() == 64;34 }35};36 37// Functor to resolve the function name corresponding to the given float result38// type.39struct FloatTypeResolver {40 std::optional<bool> operator()(Type type) const {41 auto elementType = cast<FloatType>(type);42 if (!isa<Float32Type, Float64Type>(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 the52// 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} // namespace68 69template <typename Op, typename TypeResolver>70LogicalResult ScalarOpToLibmCall<Op, TypeResolver>::matchAndRewrite(71 Op op, PatternRewriter &rewriter) const {72 auto module = SymbolTable::getNearestSymbolTable(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 been82 if (!opFunc) {83 OpBuilder::InsertionGuard guard(rewriter);84 rewriter.setInsertionPointToStart(&module->getRegion(0).front());85 auto opFunctionTy = FunctionType::get(86 rewriter.getContext(), op->getOperandTypes(), op->getResultTypes());87 opFunc = func::FuncOp::create(rewriter, 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>>(patterns.getContext(),102 "cpowf", "cpow", benefit);103 patterns.add<ScalarOpToLibmCall<complex::SqrtOp>>(patterns.getContext(),104 "csqrtf", "csqrt", benefit);105 patterns.add<ScalarOpToLibmCall<complex::TanhOp>>(patterns.getContext(),106 "ctanhf", "ctanh", benefit);107 patterns.add<ScalarOpToLibmCall<complex::CosOp>>(patterns.getContext(),108 "ccosf", "ccos", benefit);109 patterns.add<ScalarOpToLibmCall<complex::SinOp>>(patterns.getContext(),110 "csinf", "csin", benefit);111 patterns.add<ScalarOpToLibmCall<complex::ConjOp>>(patterns.getContext(),112 "conjf", "conj", benefit);113 patterns.add<ScalarOpToLibmCall<complex::LogOp>>(patterns.getContext(),114 "clogf", "clog", benefit);115 patterns.add<ScalarOpToLibmCall<complex::AbsOp, FloatTypeResolver>>(116 patterns.getContext(), "cabsf", "cabs", benefit);117 patterns.add<ScalarOpToLibmCall<complex::AngleOp, FloatTypeResolver>>(118 patterns.getContext(), "cargf", "carg", benefit);119 patterns.add<ScalarOpToLibmCall<complex::TanOp>>(patterns.getContext(),120 "ctanf", "ctan", benefit);121}122 123namespace {124struct ConvertComplexToLibmPass125 : public impl::ConvertComplexToLibmBase<ConvertComplexToLibmPass> {126 void runOnOperation() override;127};128} // namespace129 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(applyPartialConversion(module, target, std::move(patterns))))143 signalPassFailure();144}145