Move cpow constant optimisation to Fortran lowering.

Author: TIFitis

flang/lib/Optimizer/Builder/IntrinsicCall.cpp

@@ -1280,6 +1280,14 @@ mlir::Value genComplexPow(fir::FirOpBuilder &builder, mlir::Location loc,
                           const MathOperation &mathOp,
                           mlir::FunctionType mathLibFuncType,
                           llvm::ArrayRef<mlir::Value> args) {
+  if (auto expInt = fir::getIntIfConstant(args[1]))
+    if (*expInt >= 2 && *expInt <= 8) {
+      mlir::Value result = args[0];
+      for (int i = 1; i < *expInt; ++i)
+        result = builder.create<mlir::complex::MulOp>(loc, result, args[0]);
+      return builder.createConvert(loc, mathLibFuncType.getResult(0), result);
+    }
+
   bool canUseApprox = mlir::arith::bitEnumContainsAny(
       builder.getFastMathFlags(), mlir::arith::FastMathFlags::afn);
   bool isAMDGPU = fir::getTargetTriple(builder.getModule()).isAMDGCN();

flang/test/Lower/amdgcn-complex.f90

@@ -18,10 +18,19 @@ subroutine cexpf_test(a, b)
    b = exp(a)
 end subroutine
 
-! CHECK-LABEL: func @_QPpow_test(
-! CHECK: complex.pow
+! CHECK-LABEL: func @_QPpow_test1(
+! CHECK: complex.mul
+! CHECK-NOT: complex.pow
 ! CHECK-NOT: fir.call @_FortranAcpowi
-subroutine pow_test(a, b)
+subroutine pow_test1(a, b)
    complex :: a, b
    a = b**2
-end subroutine pow_test
+end subroutine pow_test1
+
+! CHECK-LABEL: func @_QPpow_test2(
+! CHECK: complex.pow
+! CHECK-NOT: fir.call @_FortranAcpowi
+subroutine pow_test2(a, b, c)
+   complex :: a, b, c
+   a = b**c
+end subroutine pow_test2

flang/test/Lower/power-operator.f90

@@ -143,3 +143,11 @@ subroutine pow_c8_c8(x, y, z)
   ! PRECISE: call @cpow
 end subroutine
 
+! CHECK-LABEL: pow_const
+subroutine pow_const(a, b)
+  complex :: a, b
+  ! CHECK-NOT: complex.pow
+  ! CHECK-NOT: @_FortranAcpowi
+  ! CHECK-COUNT-3: complex.mul
+  a = b**4
+end subroutine

mlir/lib/Conversion/ComplexToROCDLLibraryCalls/ComplexToROCDLLibraryCalls.cpp

@@ -7,7 +7,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/ComplexToROCDLLibraryCalls/ComplexToROCDLLibraryCalls.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
@@ -59,52 +58,12 @@ struct ComplexOpToROCDLLibraryCalls : public OpRewritePattern<Op> {
 };
 
 // Rewrite complex.pow(z, w) -> complex.exp(w * complex.log(z))
-// Rewrite complex.pow(z, i) -> z * z ... * z for 2 >= i <=8
 struct PowOpToROCDLLibraryCalls : public OpRewritePattern<complex::PowOp> {
   using OpRewritePattern<complex::PowOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(complex::PowOp op,
                                 PatternRewriter &rewriter) const final {
     auto loc = op.getLoc();
-
-    auto peelConst = [&](Value val) -> std::optional<TypedAttr> {
-      while (val) {
-        Operation *defOp = val.getDefiningOp();
-        if (!defOp)
-          return std::nullopt;
-
-        if (auto constVal = dyn_cast<arith::ConstantOp>(defOp))
-          return dyn_cast<TypedAttr>(constVal.getValue());
-
-        if (defOp->getName().getStringRef() == "fir.convert" &&
-            defOp->getNumOperands() == 1) {
-          val = defOp->getOperand(0);
-          continue;
-        }
-        return std::nullopt;
-      }
-      return std::nullopt;
-    };
-
-    if (auto createOp = op.getRhs().getDefiningOp<complex::CreateOp>()) {
-      auto image = peelConst(createOp.getImaginary());
-      auto real = peelConst(createOp.getReal());
-      if (image && real) {
-        auto imagFloat = dyn_cast<FloatAttr>(*image);
-        if (imagFloat && imagFloat.getValue().isZero()) {
-          auto realInt = dyn_cast<IntegerAttr>(*real);
-          if (realInt && realInt.getInt() >= 2 && realInt.getInt() <= 8) {
-            Value base = op.getLhs();
-            Value result = base;
-            for (int i = 1; i < realInt.getInt(); ++i)
-              result = rewriter.create<complex::MulOp>(loc, result, base);
-            rewriter.replaceOp(op, result);
-            return success();
-          }
-        }
-      }
-    }
-
     Value logBase = rewriter.create<complex::LogOp>(loc, op.getLhs());
     Value mul = rewriter.create<complex::MulOp>(loc, op.getRhs(), logBase);
     Value exp = rewriter.create<complex::ExpOp>(loc, mul);

mlir/test/Conversion/ComplexToROCDLLibraryCalls/complex-to-rocdl-library-calls.mlir

@@ -68,28 +68,6 @@ func.func @pow_caller(%z: complex<f32>, %w: complex<f32>) -> complex<f32> {
   return %r : complex<f32>
 }
 
-// CHECK-LABEL: @pow_int_caller
-func.func @pow_int_caller(%f : complex<f32>, %d : complex<f64>)
-    ->(complex<f32>, complex<f64>) {
-  // CHECK-NOT: call @__ocml
-  // CHECK: %[[M2:.*]] = complex.mul %{{.*}}, %{{.*}} : complex<f32>
-  %c2_i32 = arith.constant 2 : i32
-  %c2r = "fir.convert"(%c2_i32) : (i32) -> f32
-  %c2i = arith.constant 0.0 : f32
-  %c2 = complex.create %c2r, %c2i : complex<f32>
-  %p2 = complex.pow %f, %c2 : complex<f32>
-  // CHECK-NOT: call @__ocml
-  // CHECK: %[[M3A:.*]] = complex.mul %{{.*}}, %{{.*}} : complex<f64>
-  // CHECK: %[[M3B:.*]] = complex.mul %[[M3A]], %{{.*}} : complex<f64>
-  %c3_i32 = arith.constant 3 : i32
-  %c3r = "fir.convert"(%c3_i32) : (i32) -> f64
-  %c3i = arith.constant 0.0 : f64
-  %c3 = complex.create %c3r, %c3i : complex<f64>
-  %p3 = complex.pow %d, %c3 : complex<f64>
-  // CHECK: return %[[M2]], %[[M3B]]
-  return %p2, %p3 : complex<f32>, complex<f64>
-}
-
 //CHECK-LABEL: @sin_caller
 func.func @sin_caller(%f: complex<f32>, %d: complex<f64>) -> (complex<f32>, complex<f64>) {
   // CHECK: %[[SF:.*]] = call @__ocml_csin_f32(%{{.*}})