Implemented vec_cast in dpctl::tensor::type_utils

Where kernel for contiguous data now uses vec_cast

Author: ndgrigorian

dpctl/tensor/libtensor/include/kernels/where.hpp

@@ -80,8 +80,6 @@ class WhereContigFunctor
         T *dst_data = reinterpret_cast<T *>(dst_cp);
         const condT *cond_data = reinterpret_cast<const condT *>(cond_cp);
 
-        using dpctl::tensor::type_utils::convert_impl;
-
         using dpctl::tensor::type_utils::is_complex;
         if constexpr (is_complex<condT>::value || is_complex<T>::value) {
             std::uint8_t sgSize = ndit.get_sub_group().get_local_range()[0];
@@ -92,6 +90,7 @@ class WhereContigFunctor
                  offset < std::min(nelems, base + sgSize * (n_vecs * vec_sz));
                  offset += sgSize)
             {
+                using dpctl::tensor::type_utils::convert_impl;
                 bool check = convert_impl<bool, condT>(cond_data[offset]);
                 dst_data[offset] = check ? x1_data[offset] : x2_data[offset];
             }
@@ -115,7 +114,6 @@ class WhereContigFunctor
                 using cond_ptrT =
                     sycl::multi_ptr<const condT,
                                     sycl::access::address_space::global_space>;
-
                 sycl::vec<T, vec_sz> dst_vec;
                 sycl::vec<T, vec_sz> x1_vec;
                 sycl::vec<T, vec_sz> x2_vec;
@@ -127,18 +125,32 @@ class WhereContigFunctor
                     x1_vec = sg.load<vec_sz>(x_ptrT(&x1_data[idx]));
                     x2_vec = sg.load<vec_sz>(x_ptrT(&x2_data[idx]));
                     cond_vec = sg.load<vec_sz>(cond_ptrT(&cond_data[idx]));
-
+                    if constexpr (std::is_same_v<bool, condT>) {
+#pragma unroll
+                        for (std::uint8_t k = 0; k < vec_sz; ++k) {
+                            bool check = cond_vec[k];
+                            dst_vec[k] = check ? x1_vec[k] : x2_vec[k];
+                        }
+                    }
+                    else {
+                        using dpctl::tensor::type_utils::vec_cast;
+                        sycl::vec<bool, vec_sz> tmp =
+                            vec_cast<bool,
+                                     typename decltype(cond_vec)::element_type,
+                                     vec_sz>(cond_vec);
 #pragma unroll
-                    for (std::uint8_t k = 0; k < vec_sz; ++k) {
-                        bool check = convert_impl<bool, condT>(cond_vec[k]);
-                        dst_vec[k] = check ? x1_vec[k] : x2_vec[k];
+                        for (std::uint8_t k = 0; k < vec_sz; ++k) {
+                            bool check = tmp[k];
+                            dst_vec[k] = check ? x1_vec[k] : x2_vec[k];
+                        }
                     }
                     sg.store<vec_sz>(dst_ptrT(&dst_data[idx]), dst_vec);
                 }
             }
             else {
                 for (size_t k = base + sg.get_local_id()[0]; k < nelems;
                      k += sgSize) {
+                    using dpctl::tensor::type_utils::convert_impl;
                     bool check = convert_impl<bool, condT>(cond_data[k]);
                     dst_data[k] = check ? x1_data[k] : x2_data[k];
                 }

dpctl/tensor/libtensor/include/utils/type_utils.hpp

@@ -41,6 +41,21 @@ template <class T> struct is_complex<std::complex<T>> : std::true_type
 {
 };
 
+template <typename Op, typename Vec, std::size_t... I>
+auto vec_cast_impl(const Vec &v, std::index_sequence<I...>)
+{
+    return Op{v[I]...};
+}
+
+template <typename dstT,
+          typename srcT,
+          std::size_t N,
+          typename Indices = std::make_index_sequence<N>>
+auto vec_cast(const sycl::vec<srcT, N> &s)
+{
+    return vec_cast_impl<sycl::vec<dstT, N>, sycl::vec<srcT, N>>(s, Indices{});
+}
+
 template <typename dstTy, typename srcTy> dstTy convert_impl(const srcTy &v)
 {
     if constexpr (std::is_same<dstTy, srcTy>::value) {