112 lines · cpp
1//===- ConvertComplexPow.cpp - Convert complex.pow to library 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 "flang/Common/static-multimap-view.h"10#include "flang/Optimizer/Builder/FIRBuilder.h"11#include "flang/Optimizer/Dialect/FIRDialect.h"12#include "flang/Optimizer/Transforms/Passes.h"13#include "flang/Runtime/entry-names.h"14#include "mlir/Dialect/Arith/IR/Arith.h"15#include "mlir/Dialect/Complex/IR/Complex.h"16#include "mlir/Dialect/Func/IR/FuncOps.h"17#include "mlir/Pass/Pass.h"18 19namespace fir {20#define GEN_PASS_DEF_CONVERTCOMPLEXPOW21#include "flang/Optimizer/Transforms/Passes.h.inc"22} // namespace fir23 24using namespace mlir;25 26namespace {27class ConvertComplexPowPass28 : public fir::impl::ConvertComplexPowBase<ConvertComplexPowPass> {29public:30 void getDependentDialects(DialectRegistry ®istry) const override {31 registry.insert<fir::FIROpsDialect, complex::ComplexDialect,32 arith::ArithDialect, func::FuncDialect>();33 }34 void runOnOperation() override;35};36} // namespace37 38// Helper to declare or get a math library function.39static func::FuncOp getOrDeclare(fir::FirOpBuilder &builder, Location loc,40 StringRef name, FunctionType type) {41 if (auto func = builder.getNamedFunction(name))42 return func;43 auto func = builder.createFunction(loc, name, type);44 func->setAttr(fir::getSymbolAttrName(), builder.getStringAttr(name));45 func->setAttr(fir::FIROpsDialect::getFirRuntimeAttrName(),46 builder.getUnitAttr());47 return func;48}49 50void ConvertComplexPowPass::runOnOperation() {51 ModuleOp mod = getOperation();52 fir::FirOpBuilder builder(mod, fir::getKindMapping(mod));53 54 mod.walk([&](Operation *op) {55 if (auto powIop = dyn_cast<complex::PowiOp>(op)) {56 builder.setInsertionPoint(powIop);57 Location loc = powIop.getLoc();58 auto complexTy = cast<ComplexType>(powIop.getType());59 auto elemTy = complexTy.getElementType();60 Value base = powIop.getLhs();61 Value intExp = powIop.getRhs();62 func::FuncOp callee;63 unsigned realBits = cast<FloatType>(elemTy).getWidth();64 unsigned intBits = cast<IntegerType>(intExp.getType()).getWidth();65 auto funcTy = builder.getFunctionType(66 {complexTy, builder.getIntegerType(intBits)}, {complexTy});67 if (realBits == 32 && intBits == 32)68 callee = getOrDeclare(builder, loc, RTNAME_STRING(cpowi), funcTy);69 else if (realBits == 32 && intBits == 64)70 callee = getOrDeclare(builder, loc, RTNAME_STRING(cpowk), funcTy);71 else if (realBits == 64 && intBits == 32)72 callee = getOrDeclare(builder, loc, RTNAME_STRING(zpowi), funcTy);73 else if (realBits == 64 && intBits == 64)74 callee = getOrDeclare(builder, loc, RTNAME_STRING(zpowk), funcTy);75 else if (realBits == 128 && intBits == 32)76 callee = getOrDeclare(builder, loc, RTNAME_STRING(cqpowi), funcTy);77 else if (realBits == 128 && intBits == 64)78 callee = getOrDeclare(builder, loc, RTNAME_STRING(cqpowk), funcTy);79 else80 return;81 auto call = fir::CallOp::create(builder, loc, callee, {base, intExp});82 if (auto fmf = powIop.getFastmathAttr())83 call.setFastmathAttr(fmf);84 powIop.replaceAllUsesWith(call.getResult(0));85 powIop.erase();86 } else if (auto powOp = dyn_cast<complex::PowOp>(op)) {87 builder.setInsertionPoint(powOp);88 Location loc = powOp.getLoc();89 auto complexTy = cast<ComplexType>(powOp.getType());90 auto elemTy = complexTy.getElementType();91 unsigned realBits = cast<FloatType>(elemTy).getWidth();92 func::FuncOp callee;93 auto funcTy =94 builder.getFunctionType({complexTy, complexTy}, {complexTy});95 if (realBits == 32)96 callee = getOrDeclare(builder, loc, "cpowf", funcTy);97 else if (realBits == 64)98 callee = getOrDeclare(builder, loc, "cpow", funcTy);99 else if (realBits == 128)100 callee = getOrDeclare(builder, loc, RTNAME_STRING(CPowF128), funcTy);101 else102 return;103 auto call = fir::CallOp::create(builder, loc, callee,104 {powOp.getLhs(), powOp.getRhs()});105 if (auto fmf = powOp.getFastmathAttr())106 call.setFastmathAttr(fmf);107 powOp.replaceAllUsesWith(call.getResult(0));108 powOp.erase();109 }110 });111}112