[flang-commits] [flang] [flang] Use libm routine for compile-time folding on AIX (PR #114106)

[Kelvin Li via flang-commits]

https://github.com/kkwli updated https://github.com/llvm/llvm-project/pull/114106

>From 1742f2aa7bc57123c4a7df53dc2976b0cd589bb8 Mon Sep 17 00:00:00 2001
From: Kelvin Li <kli at ca.ibm.com>
Date: Tue, 29 Oct 2024 13:36:24 -0400
Subject: [PATCH 1/2] [flang] Use libm routine for compile-time folding on AIX

---
 flang/lib/Evaluate/CMakeLists.txt         |  1 +
 flang/lib/Evaluate/intrinsics-library.cpp | 24 +++++++++++++++++++++++
 flang/lib/Evaluate/wrappers.c             | 17 ++++++++++++++++
 3 files changed, 42 insertions(+)
 create mode 100644 flang/lib/Evaluate/wrappers.c

diff --git a/flang/lib/Evaluate/CMakeLists.txt b/flang/lib/Evaluate/CMakeLists.txt
index b38f450d746ea7..8ffbfc8848da74 100644
--- a/flang/lib/Evaluate/CMakeLists.txt
+++ b/flang/lib/Evaluate/CMakeLists.txt
@@ -58,6 +58,7 @@ add_flang_library(FortranEvaluate
   tools.cpp
   type.cpp
   variable.cpp
+  wrappers.c
 
   LINK_LIBS
   FortranCommon
diff --git a/flang/lib/Evaluate/intrinsics-library.cpp b/flang/lib/Evaluate/intrinsics-library.cpp
index bb439a6bb3a746..8b16d187513690 100644
--- a/flang/lib/Evaluate/intrinsics-library.cpp
+++ b/flang/lib/Evaluate/intrinsics-library.cpp
@@ -277,6 +277,26 @@ static std::complex<HostT> StdPowF2B(
   return std::pow(x, y);
 }
 
+#ifdef _AIX
+extern "C" {
+void csqrtf_wrapper(const float[], float[]);
+void csqrt_wrapper(const double[], double[]);
+} // extern "C"
+
+template <typename HostT>
+static std::complex<HostT> CSQRT(const std::complex<HostT> &x) {
+  HostT y[2]{x.real(), x.imag()};
+  HostT r[2];
+  if constexpr (std::is_same_v<HostT, float>) {
+    csqrtf_wrapper(y, r);
+  } else if constexpr (std::is_same_v<HostT, double>) {
+    csqrt_wrapper(y, r);
+  }
+  std::complex<HostT> res(r[0], r[1]);
+  return res;
+}
+#endif
+
 template <typename HostT>
 struct HostRuntimeLibrary<std::complex<HostT>, LibraryVersion::Libm> {
   using F = FuncPointer<std::complex<HostT>, const std::complex<HostT> &>;
@@ -302,7 +322,11 @@ struct HostRuntimeLibrary<std::complex<HostT>, LibraryVersion::Libm> {
       FolderFactory<F2B, F2B{StdPowF2B}>::Create("pow"),
       FolderFactory<F, F{std::sin}>::Create("sin"),
       FolderFactory<F, F{std::sinh}>::Create("sinh"),
+#ifdef _AIX
+      FolderFactory<F, F{CSQRT}>::Create("sqrt"),
+#else
       FolderFactory<F, F{std::sqrt}>::Create("sqrt"),
+#endif
       FolderFactory<F, F{std::tan}>::Create("tan"),
       FolderFactory<F, F{std::tanh}>::Create("tanh"),
   };
diff --git a/flang/lib/Evaluate/wrappers.c b/flang/lib/Evaluate/wrappers.c
new file mode 100644
index 00000000000000..b0823cac45ae4b
--- /dev/null
+++ b/flang/lib/Evaluate/wrappers.c
@@ -0,0 +1,17 @@
+#include <complex.h>
+
+void csqrtf_wrapper(const float x[2], float res[2])
+{
+  float complex c = x[0] + I * x[1];
+  float complex r = csqrtf(c);
+  res[0] = crealf(r);
+  res[1] = cimagf(r);
+}
+
+void csqrt_wrapper(const double x[2], double res[2])
+{
+  double complex c = x[0] + I * x[1];
+  double complex r = csqrt(c);
+  res[0] = creal(r);
+  res[1] = cimag(r);
+}

>From d5219a9792ca177f8b4bca79a9ab718c6076ec64 Mon Sep 17 00:00:00 2001
From: Kelvin Li <kli at ca.ibm.com>
Date: Tue, 29 Oct 2024 14:23:26 -0400
Subject: [PATCH 2/2] Add copyright note and fix format

---
 flang/lib/Evaluate/wrappers.c | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/flang/lib/Evaluate/wrappers.c b/flang/lib/Evaluate/wrappers.c
index b0823cac45ae4b..360d638366b466 100644
--- a/flang/lib/Evaluate/wrappers.c
+++ b/flang/lib/Evaluate/wrappers.c
@@ -1,15 +1,21 @@
+//===-- lib/Evaluate/wrappers.c -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include <complex.h>
 
-void csqrtf_wrapper(const float x[2], float res[2])
-{
+void csqrtf_wrapper(const float x[2], float res[2]) {
   float complex c = x[0] + I * x[1];
   float complex r = csqrtf(c);
   res[0] = crealf(r);
   res[1] = cimagf(r);
 }
 
-void csqrt_wrapper(const double x[2], double res[2])
-{
+void csqrt_wrapper(const double x[2], double res[2]) {
   double complex c = x[0] + I * x[1];
   double complex r = csqrt(c);
   res[0] = creal(r);