Adds sycl::vec overloads to abs, cos, expm1, log, log1p, and sqrt

ndgrigorian · ndgrigorian · commit 323ed6db9ddf · 2023-06-22T13:59:11.000-07:00
diff --git a/dpctl/tensor/libtensor/include/kernels/elementwise_functions/abs.hpp b/dpctl/tensor/libtensor/include/kernels/elementwise_functions/abs.hpp
@@ -51,13 +51,15 @@ namespace py = pybind11;
 namespace td_ns = dpctl::tensor::type_dispatch;
 
 using dpctl::tensor::type_utils::is_complex;
+using dpctl::tensor::type_utils::vec_cast;
 
 template <typename argT, typename resT> struct AbsFunctor
 {
 
     using is_constant = typename std::false_type;
     // constexpr resT constant_value = resT{};
-    using supports_vec = typename std::false_type;
+    using supports_vec = typename std::negation<
+        std::disjunction<is_complex<resT>, is_complex<argT>>>;
     using supports_sg_loadstore = typename std::negation<
         std::disjunction<is_complex<resT>, is_complex<argT>>>;
 
@@ -75,6 +77,40 @@ template <typename argT, typename resT> struct AbsFunctor
             return std::abs(x);
         }
     }
+
+    template <int vec_sz>
+    sycl::vec<resT, vec_sz> operator()(const sycl::vec<argT, vec_sz> &in)
+    {
+        if constexpr (std::is_integral<argT>::value) {
+            if constexpr (std::is_same_v<argT, bool> ||
+                          std::is_unsigned<argT>::value) {
+                return in;
+            }
+            else {
+                auto const &res_vec = sycl::abs(in);
+                using deducedT = typename std::remove_cv_t<
+                    std::remove_reference_t<decltype(res_vec)>>::element_type;
+                if constexpr (std::is_same_v<resT, deducedT>) {
+                    return res_vec;
+                }
+                else {
+
+                    return vec_cast<resT, deducedT, vec_sz>(res_vec);
+                }
+            }
+        }
+        else {
+            auto const &res_vec = sycl::fabs(in);
+            using deducedT = typename std::remove_cv_t<
+                std::remove_reference_t<decltype(res_vec)>>::element_type;
+            if constexpr (std::is_same_v<resT, deducedT>) {
+                return res_vec;
+            }
+            else {
+                return vec_cast<resT, deducedT, vec_sz>(res_vec);
+            }
+        }
+    }
 };
 
 template <typename argT,
diff --git a/dpctl/tensor/libtensor/include/kernels/elementwise_functions/cos.hpp b/dpctl/tensor/libtensor/include/kernels/elementwise_functions/cos.hpp
@@ -49,6 +49,7 @@ namespace py = pybind11;
 namespace td_ns = dpctl::tensor::type_dispatch;
 
 using dpctl::tensor::type_utils::is_complex;
+using dpctl::tensor::type_utils::vec_cast;
 
 template <typename argT, typename resT> struct CosFunctor
 {
@@ -58,7 +59,8 @@ template <typename argT, typename resT> struct CosFunctor
     // constant value, if constant
     // constexpr resT constant_value = resT{};
     // is function defined for sycl::vec
-    using supports_vec = typename std::false_type;
+    using supports_vec = typename std::negation<
+        std::disjunction<is_complex<resT>, is_complex<argT>>>;
     // do both argTy and resTy support sugroup store/load operation
     using supports_sg_loadstore = typename std::negation<
         std::disjunction<is_complex<resT>, is_complex<argT>>>;
@@ -67,6 +69,20 @@ template <typename argT, typename resT> struct CosFunctor
     {
         return std::cos(in);
     }
+
+    template <int vec_sz>
+    sycl::vec<resT, vec_sz> operator()(const sycl::vec<argT, vec_sz> &in)
+    {
+        auto const &res_vec = sycl::cos(in);
+        using deducedT = typename std::remove_cv_t<
+            std::remove_reference_t<decltype(res_vec)>>::element_type;
+        if constexpr (std::is_same_v<resT, deducedT>) {
+            return res_vec;
+        }
+        else {
+            return vec_cast<resT, deducedT, vec_sz>(res_vec);
+        }
+    }
 };
 
 template <typename argTy,
diff --git a/dpctl/tensor/libtensor/include/kernels/elementwise_functions/expm1.hpp b/dpctl/tensor/libtensor/include/kernels/elementwise_functions/expm1.hpp
@@ -50,6 +50,7 @@ namespace py = pybind11;
 namespace td_ns = dpctl::tensor::type_dispatch;
 
 using dpctl::tensor::type_utils::is_complex;
+using dpctl::tensor::type_utils::vec_cast;
 
 template <typename argT, typename resT> struct Expm1Functor
 {
@@ -59,7 +60,8 @@ template <typename argT, typename resT> struct Expm1Functor
     // constant value, if constant
     // constexpr resT constant_value = resT{};
     // is function defined for sycl::vec
-    using supports_vec = typename std::false_type;
+    using supports_vec = typename std::negation<
+        std::disjunction<is_complex<resT>, is_complex<argT>>>;
     // do both argTy and resTy support sugroup store/load operation
     using supports_sg_loadstore = typename std::negation<
         std::disjunction<is_complex<resT>, is_complex<argT>>>;
@@ -88,6 +90,20 @@ template <typename argT, typename resT> struct Expm1Functor
             return std::expm1(in);
         }
     }
+
+    template <int vec_sz>
+    sycl::vec<resT, vec_sz> operator()(const sycl::vec<argT, vec_sz> &in)
+    {
+        auto const &res_vec = sycl::expm1(in);
+        using deducedT = typename std::remove_cv_t<
+            std::remove_reference_t<decltype(res_vec)>>::element_type;
+        if constexpr (std::is_same_v<resT, deducedT>) {
+            return res_vec;
+        }
+        else {
+            return vec_cast<resT, deducedT, vec_sz>(res_vec);
+        }
+    }
 };
 
 template <typename argTy,
diff --git a/dpctl/tensor/libtensor/include/kernels/elementwise_functions/log.hpp b/dpctl/tensor/libtensor/include/kernels/elementwise_functions/log.hpp
@@ -49,6 +49,7 @@ namespace py = pybind11;
 namespace td_ns = dpctl::tensor::type_dispatch;
 
 using dpctl::tensor::type_utils::is_complex;
+using dpctl::tensor::type_utils::vec_cast;
 
 template <typename argT, typename resT> struct LogFunctor
 {
@@ -58,7 +59,8 @@ template <typename argT, typename resT> struct LogFunctor
     // constant value, if constant
     // constexpr resT constant_value = resT{};
     // is function defined for sycl::vec
-    using supports_vec = typename std::false_type;
+    using supports_vec = typename std::negation<
+        std::disjunction<is_complex<resT>, is_complex<argT>>>;
     // do both argTy and resTy support sugroup store/load operation
     using supports_sg_loadstore = typename std::negation<
         std::disjunction<is_complex<resT>, is_complex<argT>>>;
@@ -67,6 +69,20 @@ template <typename argT, typename resT> struct LogFunctor
     {
         return std::log(in);
     }
+
+    template <int vec_sz>
+    sycl::vec<resT, vec_sz> operator()(const sycl::vec<argT, vec_sz> &in)
+    {
+        auto const &res_vec = sycl::log(in);
+        using deducedT = typename std::remove_cv_t<
+            std::remove_reference_t<decltype(res_vec)>>::element_type;
+        if constexpr (std::is_same_v<resT, deducedT>) {
+            return res_vec;
+        }
+        else {
+            return vec_cast<resT, deducedT, vec_sz>(res_vec);
+        }
+    }
 };
 
 template <typename argTy,
diff --git a/dpctl/tensor/libtensor/include/kernels/elementwise_functions/log1p.hpp b/dpctl/tensor/libtensor/include/kernels/elementwise_functions/log1p.hpp
@@ -50,6 +50,7 @@ namespace py = pybind11;
 namespace td_ns = dpctl::tensor::type_dispatch;
 
 using dpctl::tensor::type_utils::is_complex;
+using dpctl::tensor::type_utils::vec_cast;
 
 // TODO: evaluate precision against alternatives
 template <typename argT, typename resT> struct Log1pFunctor
@@ -60,7 +61,8 @@ template <typename argT, typename resT> struct Log1pFunctor
     // constant value, if constant
     // constexpr resT constant_value = resT{};
     // is function defined for sycl::vec
-    using supports_vec = typename std::false_type;
+    using supports_vec = typename std::negation<
+        std::disjunction<is_complex<resT>, is_complex<argT>>>;
     // do both argTy and resTy support sugroup store/load operation
     using supports_sg_loadstore = typename std::negation<
         std::disjunction<is_complex<resT>, is_complex<argT>>>;
@@ -99,6 +101,20 @@ template <typename argT, typename resT> struct Log1pFunctor
             return std::log1p(in);
         }
     }
+
+    template <int vec_sz>
+    sycl::vec<resT, vec_sz> operator()(const sycl::vec<argT, vec_sz> &in)
+    {
+        auto const &res_vec = sycl::log1p(in);
+        using deducedT = typename std::remove_cv_t<
+            std::remove_reference_t<decltype(res_vec)>>::element_type;
+        if constexpr (std::is_same_v<resT, deducedT>) {
+            return res_vec;
+        }
+        else {
+            return vec_cast<resT, deducedT, vec_sz>(res_vec);
+        }
+    }
 };
 
 template <typename argTy,
diff --git a/dpctl/tensor/libtensor/include/kernels/elementwise_functions/sqrt.hpp b/dpctl/tensor/libtensor/include/kernels/elementwise_functions/sqrt.hpp
@@ -50,6 +50,7 @@ namespace py = pybind11;
 namespace td_ns = dpctl::tensor::type_dispatch;
 
 using dpctl::tensor::type_utils::is_complex;
+using dpctl::tensor::type_utils::vec_cast;
 
 template <typename argT, typename resT> struct SqrtFunctor
 {
@@ -59,7 +60,8 @@ template <typename argT, typename resT> struct SqrtFunctor
     // constant value, if constant
     // constexpr resT constant_value = resT{};
     // is function defined for sycl::vec
-    using supports_vec = typename std::false_type;
+    using supports_vec = typename std::negation<
+        std::disjunction<is_complex<resT>, is_complex<argT>>>;
     // do both argTy and resTy support sugroup store/load operation
     using supports_sg_loadstore = typename std::negation<
         std::disjunction<is_complex<resT>, is_complex<argT>>>;
@@ -68,6 +70,20 @@ template <typename argT, typename resT> struct SqrtFunctor
     {
         return std::sqrt(in);
     }
+
+    template <int vec_sz>
+    sycl::vec<resT, vec_sz> operator()(const sycl::vec<argT, vec_sz> &in)
+    {
+        auto const &res_vec = sycl::sqrt(in);
+        using deducedT = typename std::remove_cv_t<
+            std::remove_reference_t<decltype(res_vec)>>::element_type;
+        if constexpr (std::is_same_v<resT, deducedT>) {
+            return res_vec;
+        }
+        else {
+            return vec_cast<resT, deducedT, vec_sz>(res_vec);
+        }
+    }
 };
 
 template <typename argTy,
