[Flang] Lower the sqrt intrinsics

[kiranchandramohan]

Summary

The intrinsic computes the square root for real and complex numbers. By
default they are lowered to runtime calls to libpgmath. With the llvm
option, it can be lowered to llvm intrinsics (not all types .eg. complex
are supported for llvm lowering).

This is part of the upstreaming effort from the fir-dev branch in [1].
[1] https://github.com/flang-compiler/f18-llvm-project

Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: Jean Perier <[e-mail]>

[schweitz]

flang/lib/Lower/IntrinsicCall.cpp

1018

What happens with KIND=10? Does it fall back to a TODO message?

[kiranchandramohan]

flang/lib/Lower/IntrinsicCall.cpp

1018

Yes, it hits a TODO, but code is generated as well.

But fir-dev has the following behaviour. A release build adds a convert to 32 bit and calls llvm.sqrt.f32. A debug build will trigger an assertion.
Assertion `!bestMatchDistance.isLosingPrecision() && "runtime selection loses precision"' failed. I have posted the following issue to capture this.
https://github.com/flang-compiler/f18-llvm-project/issues/1497

flang/lib/Lower/IntrinsicCall.cpp

 //===-- IntrinsicCall.cpp -------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
static constexpr RuntimeFunction llvmIntrinsics[] = {
     {"nint", "llvm.lround.i32.f64", genIntF64FuncType<32>},
     {"nint", "llvm.lround.i32.f32", genIntF32FuncType<32>},
     {"pow", "llvm.pow.f32", genF32F32F32FuncType},
     {"pow", "llvm.pow.f64", genF64F64F64FuncType},
     {"sign", "llvm.copysign.f32", genF32F32F32FuncType},
     {"sign", "llvm.copysign.f64", genF64F64F64FuncType},
     {"sign", "llvm.copysign.f80", genF80F80F80FuncType},
     {"sign", "llvm.copysign.f128", genF128F128F128FuncType},
+    {"sqrt", "llvm.sqrt.f32", genF32F32FuncType},
+    {"sqrt", "llvm.sqrt.f64", genF64F64FuncType},
 };
 
 // This helper class computes a "distance" between two function types.
 // The distance measures how many narrowing conversions of actual arguments
 // and result of "from" must be made in order to use "to" instead of "from".
 // For instance, the distance between ACOS(REAL(10)) and ACOS(REAL(8)) is
 // greater than the one between ACOS(REAL(10)) and ACOS(REAL(16)). This means
 // if no implementation of ACOS(REAL(10)) is available, it is better to use

flang/test/Lower/llvm-math.f90

! CHECK-NEXT:   %2 = fir.call @llvm.pow.f64(%0, %1) : (f64, f64) -> f64
       RETURN
       END
 
 ! CHECK-LABEL: func @_QPpowf_wrapper(
 ! CHECK-SAME: %{{.*}}: !fir.ref<f32>{{.*}}, %{{.*}}: !fir.ref<f32>{{.*}}, %{{.*}}: !fir.ref<f32>{{.*}}) {
 ! CHECK-NEXT:   %0 = fir.load %arg0 : !fir.ref<f32>
 ! CHECK-NEXT:   %1 = fir.load %arg1 : !fir.ref<f32>
 ! CHECK-NEXT:   %2 = fir.call @llvm.pow.f32(%0, %1) : (f32, f32) -> f32
+
+      SUBROUTINE SQRT_WRAPPER(IN, OUT)
+      REAL :: IN, OUT
+      OUT = SQRT(IN)
+      END SUBROUTINE
+
+! CHECK-LABEL: func @_QPsqrt_wrapper(
+! CHECK-SAME: %[[IN_REF:.*]]: !fir.ref<f32> {{.*}}, %[[OUT_REF:.*]]: !fir.ref<f32> {{.*}}) {
+! CHECK-NEXT:  %[[IN:.*]] = fir.load %[[IN_REF]] : !fir.ref<f32>
+! CHECK-NEXT:  %[[RESULT:.*]] = fir.call @llvm.sqrt.f32(%[[IN]]) : (f32) -> f32
+! CHECK-NEXT:  fir.store %[[RESULT]] to %[[OUT_REF]] : !fir.ref<f32>
+
+      SUBROUTINE SQRTD_WRAPPER(IN, OUT)
+      REAL(KIND=8) :: IN, OUT
+      OUT = SQRT(IN)
+      END SUBROUTINE
+
+! CHECK-LABEL: func @_QPsqrtd_wrapper(
+! CHECK-SAME: %[[IN_REF:.*]]: !fir.ref<f64> {{.*}}, %[[OUT_REF:.*]]: !fir.ref<f64> {{.*}}) {
+! CHECK-NEXT:  %[[IN:.*]] = fir.load %[[IN_REF]] : !fir.ref<f64>
+! CHECK-NEXT:  %[[RESULT:.*]] = fir.call @llvm.sqrt.f64(%[[IN]]) : (f64) -> f64
+! CHECK-NEXT:  fir.store %[[RESULT]] to %[[OUT_REF]] : !fir.ref<f64>

flang/test/Lower/sqrt.f90

+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-fir %s -o - | FileCheck %s
+
+! CHECK-LABEL: sqrt_testr
+subroutine sqrt_testr(a, b)
+  real :: a, b
+  ! CHECK: fir.call {{.*}}mth_i_sqrt
+  b = sqrt(a)
+end subroutine
+
+! CHECK-LABEL: sqrt_testd
+subroutine sqrt_testd(a, b)
+  real(kind=8) :: a, b
+  ! CHECK: fir.call {{.*}}mth_i_dsqrt
+  b = sqrt(a)
+end subroutine
+
+! CHECK-LABEL: sqrt_testc
+subroutine sqrt_testc(z)
+  complex :: z
+  ! CHECK: fir.call {{.*}}fc_sqrt
+  z = sqrt(z)
+end subroutine
+
+! CHECK-LABEL: sqrt_testcd
+subroutine sqrt_testcd(z)
+  complex(kind=8) :: z
+  ! CHECK: fir.call {{.*}}fz_sqrt
+  z = sqrt(z)
+end subroutine
+