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];
                 }