[SYCL][ESIMD] Add support for esimd_ballot function (#4658)

* [SYCL][ESIMD] Add support for esimd_ballot function

Signed-off-by: Sergey Dmitriev <[email protected]>

Author: sndmitriev

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

@@ -1667,6 +1667,27 @@ ESIMD_NODEBUG ESIMD_INLINE
   return esimd_pack_mask(src_0);
 }
 
+/// Compare source vector elements against zero and return a bitfield combining
+/// the comparison result. The representative bit in the result is set if
+/// corresponding source vector element is non-zero, and is unset otherwise.
+/// @param mask the source operand to be compared with zero.
+/// @return an \c uint, where each bit is set if the corresponding element of
+/// the source operand is non-zero and unset otherwise.
+template <typename T, int N>
+ESIMD_NODEBUG ESIMD_INLINE typename sycl::detail::enable_if_t<
+    detail::is_type<T, ushort, uint> && (N > 0 && N <= 32), uint>
+esimd_ballot(simd<T, N> mask) {
+  simd_mask<N> cmp = (mask != 0);
+  if constexpr (N == 8 || N == 16 || N == 32) {
+    return __esimd_pack_mask<N>(cmp.data());
+  } else {
+    constexpr int N1 = (N <= 8 ? 8 : N <= 16 ? 16 : 32);
+    simd<uint16_t, N1> res = 0;
+    res.template select<N, 1>() = cmp.data();
+    return __esimd_pack_mask<N1>(res.data());
+  }
+}
+
 /// Count number of bits set in the source operand per element.
 /// @param src0 the source operand to count bits in.
 /// @return a vector of \c uint32_t, where each element is set to bit count of

sycl/test/esimd/esimd_math.cpp

@@ -72,3 +72,26 @@ bool test_esimd_trunc() __attribute__((sycl_device)) {
 
   return (vfr[0] == 1.f) && (vsr[0] == 1) && (sfr == 2.f) && (ssr == 2);
 }
+
+bool test_esimd_ballot() __attribute__((sycl_device)) {
+  simd<ushort, 4> vus4 = {1, 0, 3, 0};
+  simd<ushort, 8> vus8 = {1, 0, 3, 0, 5, 0, 7, 0};
+  simd<ushort, 20> vus20 = {1, 0, 3, 0, 5, 0, 7, 0, 9, 0,
+                            1, 0, 3, 0, 5, 0, 7, 0, 9, 0};
+
+  uint mus4 = esimd_ballot(vus4);
+  uint mus8 = esimd_ballot(vus8);
+  uint mus20 = esimd_ballot(vus20);
+
+  simd<uint, 4> vui4 = {1, 0, 3, 0};
+  simd<uint, 16> vui16 = {1, 0, 3, 0, 5, 0, 7, 0, 1, 0, 3, 0, 5, 0, 7, 0};
+  simd<uint, 20> vui20 = {1, 0, 3, 0, 5, 0, 7, 0, 9, 0,
+                          1, 0, 3, 0, 5, 0, 7, 0, 9, 0};
+
+  uint mui4 = esimd_ballot(vui4);
+  uint mui16 = esimd_ballot(vui16);
+  uint mui20 = esimd_ballot(vui20);
+
+  return (mus4 == 0x5) && (mus8 == 0x55) && (mus20 = 0x55555) &&
+         (mui4 == 0x5) && (mui16 == 0x5555) && (mui20 = 0x55555);
+}