|
| 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-exception |
| 6 | +// |
| 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 | + |
| 19 | +namespace fir { |
| 20 | +#define GEN_PASS_DEF_CONVERTCOMPLEXPOW |
| 21 | +#include "flang/Optimizer/Transforms/Passes.h.inc" |
| 22 | +} // namespace fir |
| 23 | + |
| 24 | +using namespace mlir; |
| 25 | + |
| 26 | +namespace { |
| 27 | +class ConvertComplexPowPass |
| 28 | + : public fir::impl::ConvertComplexPowBase<ConvertComplexPowPass> { |
| 29 | +public: |
| 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 | +} // namespace |
| 37 | + |
| 38 | +// Helper to declare or get a math library function. |
| 39 | +static 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 | + |
| 50 | +static bool isZero(Value v) { |
| 51 | + if ( auto cst = v. getDefiningOp<arith::ConstantOp>()) |
| 52 | + if (auto attr = dyn_cast<FloatAttr>(cst.getValue())) |
| 53 | + return attr.getValue().isZero(); |
| 54 | + return false; |
| 55 | +} |
| 56 | + |
| 57 | +void ConvertComplexPowPass::runOnOperation() { |
| 58 | + ModuleOp mod = getOperation(); |
| 59 | + fir::FirOpBuilder builder(mod, fir::getKindMapping(mod)); |
| 60 | + |
| 61 | + mod.walk([&](complex::PowOp op) { |
| 62 | + builder.setInsertionPoint(op); |
| 63 | + Location loc = op.getLoc(); |
| 64 | + auto complexTy = cast<ComplexType>(op.getType()); |
| 65 | + auto elemTy = complexTy.getElementType(); |
| 66 | + |
| 67 | + Value base = op.getLhs(); |
| 68 | + Value rhs = op.getRhs(); |
| 69 | + |
| 70 | + Value intExp; |
| 71 | + if (auto create = rhs.getDefiningOp<complex::CreateOp>()) { |
| 72 | + if (isZero(create.getImaginary())) { |
| 73 | + if (auto conv = create.getReal().getDefiningOp<fir::ConvertOp>()) { |
| 74 | + if (auto intTy = dyn_cast<IntegerType>(conv.getValue().getType())) |
| 75 | + intExp = conv.getValue(); |
| 76 | + } |
| 77 | + } |
| 78 | + } |
| 79 | + |
| 80 | + func::FuncOp callee; |
| 81 | + SmallVector<Value> args; |
| 82 | + if (intExp) { |
| 83 | + unsigned realBits = cast<FloatType>(elemTy).getWidth(); |
| 84 | + unsigned intBits = cast<IntegerType>(intExp.getType()).getWidth(); |
| 85 | + auto funcTy = builder.getFunctionType( |
| 86 | + {complexTy, builder.getIntegerType(intBits)}, {complexTy}); |
| 87 | + if (realBits == 32 && intBits == 32) |
| 88 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(cpowi), funcTy); |
| 89 | + else if (realBits == 32 && intBits == 64) |
| 90 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(cpowk), funcTy); |
| 91 | + else if (realBits == 64 && intBits == 32) |
| 92 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(zpowi), funcTy); |
| 93 | + else if (realBits == 64 && intBits == 64) |
| 94 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(zpowk), funcTy); |
| 95 | + else if (realBits == 128 && intBits == 32) |
| 96 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(cqpowi), funcTy); |
| 97 | + else if (realBits == 128 && intBits == 64) |
| 98 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(cqpowk), funcTy); |
| 99 | + else |
| 100 | + return; |
| 101 | + args = {base, intExp}; |
| 102 | + } else { |
| 103 | + unsigned realBits = cast<FloatType>(elemTy).getWidth(); |
| 104 | + auto funcTy = |
| 105 | + builder.getFunctionType({complexTy, complexTy}, {complexTy}); |
| 106 | + if (realBits == 32) |
| 107 | + callee = getOrDeclare(builder, loc, "cpowf", funcTy); |
| 108 | + else if (realBits == 64) |
| 109 | + callee = getOrDeclare(builder, loc, "cpow", funcTy); |
| 110 | + else if (realBits == 128) |
| 111 | + callee = getOrDeclare(builder, loc, RTNAME_STRING(CPowF128), funcTy); |
| 112 | + else |
| 113 | + return; |
| 114 | + args = {base, rhs}; |
| 115 | + } |
| 116 | + |
| 117 | + auto call = fir::CallOp::create(builder, loc, callee, args); |
| 118 | + if (auto fmf = op.getFastmathAttr()) |
| 119 | + call.setFastmathAttr(fmf); |
| 120 | + op.replaceAllUsesWith(call.getResult(0)); |
| 121 | + op.erase(); |
| 122 | + }); |
| 123 | +} |
0 commit comments