[SYCL][ESIMD] Use SPIR-V rsqrt for double

Signed-off-by: Sarnie, Nick <[email protected]>
intel · Oct 7, 2024 · 04fb5fa · 04fb5fa
1 parent 1dfc3d7
commit 04fb5fa
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 4 deletions.
diff --git a/sycl/include/sycl/ext/intel/esimd/detail/math_intrin.hpp b/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)

diff --git a/sycl/include/sycl/ext/intel/esimd/math.hpp b/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

diff --git a/sycl/test/check_device_code/esimd/rsqrt_double.cpp b/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
+}