[SYCL][ESIMD] Use SPIR-V rsqrt for double (#15609)

The driver added a perf optimization for this case and asked us to use
`__spirv_ocl_rsqrt` instead of `inv(sqrt)` or `__spirv_ocl_native_rsqrt`
for double `rsqrt` only.
There is already an E2E test for rsqrt.

Signed-off-by: Sarnie, Nick <[email protected]>

Author: sarnex

sycl/include/sycl/ext/intel/esimd/detail/math_intrin.hpp

@@ -102,6 +102,11 @@ extern __DPCPP_SYCL_EXTERNAL __ESIMD_raw_vec_t(T, N)
     __spirv_ocl_fmax(__ESIMD_raw_vec_t(T, N),
                      __ESIMD_raw_vec_t(T, N)) __ESIMD_INTRIN_END;
 
+template <typename T> extern __DPCPP_SYCL_EXTERNAL T __spirv_ocl_rsqrt(T);
+template <typename T, int N>
+extern __DPCPP_SYCL_EXTERNAL __ESIMD_raw_vec_t(T, N)
+    __spirv_ocl_rsqrt(__ESIMD_raw_vec_t(T, N)) __ESIMD_INTRIN_END;
+
 // saturation intrinsics
 template <typename T0, typename T1, int SZ>
 __ESIMD_INTRIN __ESIMD_raw_vec_t(T0, SZ)

sycl/include/sycl/ext/intel/esimd/math.hpp

@@ -456,20 +456,24 @@ __ESIMD_UNARY_INTRINSIC_DEF(__ESIMD_EMATH_SPIRV_COND, cos, cos)
 template <class T, int N, class Sat = saturation_off_tag>
 __ESIMD_API std::enable_if_t<std::is_same_v<T, double>, simd<double, N>>
 rsqrt(simd<T, N> src, Sat sat = {}) {
+  __ESIMD_DNS::vector_type_t<__ESIMD_DNS::__raw_t<double>, N> res =
+      __spirv_ocl_rsqrt<__ESIMD_DNS::__raw_t<double>, N>(src.data());
   if constexpr (std::is_same_v<Sat, saturation_off_tag>)
-    return inv(sqrt(src));
+    return res;
   else
-    return esimd::saturate<double>(inv(sqrt(src)));
+    return esimd::saturate<double>(simd<double, N>(res));
 }
 
 /** Scalar version.                                                       */
 template <class T, class Sat = saturation_off_tag>
 __ESIMD_API std::enable_if_t<std::is_same_v<T, double>, double>
 rsqrt(T src, Sat sat = {}) {
+  __ESIMD_DNS::__raw_t<double> res =
+      __spirv_ocl_rsqrt<__ESIMD_DNS::__raw_t<double>>(src);
   if constexpr (std::is_same_v<Sat, saturation_off_tag>)
-    return inv(sqrt(src));
+    return res;
   else
-    return esimd::saturate<double>(inv(sqrt(src)));
+    return esimd::saturate<double>(simd<double, 1>(res))[0];
 }
 
 #undef __ESIMD_UNARY_INTRINSIC_DEF

sycl/test/check_device_code/esimd/rsqrt_double.cpp

@@ -0,0 +1,13 @@
+// RUN: %clangxx -fsycl -fsycl-device-only -S -emit-llvm %s -o - | \
+// RUN: FileCheck %s --implicit-check-not=recip
+
+// This test checks that all we use SPIR-V rsqrt for doubles.
+
+#include <sycl/ext/intel/esimd.hpp>
+#include <sycl/sycl.hpp>
+
+SYCL_ESIMD_KERNEL SYCL_EXTERNAL void kernel() {
+  __ESIMD_NS::simd<double, 16> v(0, 1);
+  v = __ESIMD_NS::rsqrt(v);
+  // CHECK: __spirv_ocl_rsqrt
+}