Add bfloat16-container support to clamp

i.e. sycl::marray<bfloat16, N>, sycl::vec<bfloat16, N>

Author: joeatodd

sycl/include/syclcompat/math.hpp

@@ -81,6 +81,28 @@ inline sycl::ext::oneapi::bfloat16 clamp(sycl::ext::oneapi::bfloat16 val,
     return max_val;
   return val;
 }
+
+template <typename T, int Size>
+inline std::enable_if_t<std::is_same_v<T, sycl::ext::oneapi::bfloat16>,
+                        sycl::vec<T, Size>>
+clamp(sycl::vec<T, Size> val, sycl::vec<T, Size> min_val,
+      sycl::vec<T, Size> max_val) {
+  return [&val, &min_val, &max_val]<int... I>(std::integer_sequence<int, I...>) {
+    return sycl::vec<T, Size>{
+        clamp<sycl::ext::oneapi::bfloat16>(val[I], min_val[I], max_val[I])...};
+  }(std::make_integer_sequence<int, Size>{});
+}
+
+template <typename T, std::size_t Size>
+inline std::enable_if_t<std::is_same_v<T, sycl::ext::oneapi::bfloat16>,
+                        sycl::marray<T, Size>>
+clamp(sycl::marray<T, Size> val, sycl::marray<T, Size> min_val,
+      sycl::marray<T, Size> max_val) {
+  return [&val, &min_val, &max_val]<std::size_t... I>(std::index_sequence<I...>) {
+    return sycl::marray<T, Size>{
+        clamp<sycl::ext::oneapi::bfloat16>(val[I], min_val[I], max_val[I])...};
+  }(std::make_index_sequence<Size>{});
+}
 #endif
 
 template <typename VecT, class BinaryOperation, class = void>

sycl/test-e2e/syclcompat/math/math_ops.cpp

@@ -317,6 +317,24 @@ template <typename ValueT> void test_clamp() {
       .template launch_test<clamp_kernel<ValueT>>(op3, expect3);
 }
 
+template <template <typename T, int Dim> typename ContainerT, typename ValueT> void test_container_clamp() {
+  std::cout << __PRETTY_FUNCTION__ << std::endl;
+
+  constexpr syclcompat::dim3 grid{1};
+  constexpr syclcompat::dim3 threads{1};
+  ValueT op1 = static_cast<ValueT>(7);
+  ValueT expect1 = static_cast<ValueT>(7);
+
+  ValueT op2 = static_cast<ValueT>(MAX_CLAMP + 1);
+  ValueT expect2 = static_cast<ValueT>(MAX_CLAMP);
+
+  using ContT = ContainerT<ValueT, 2>;
+  const ContT op4{op1, op2};
+  const ContT expect4{expect1, expect2};
+  UnaryOpTestLauncher<ContT>(grid, threads)
+      .template launch_test<clamp_kernel<ContT>>(op4, expect4);
+}
+
 int main() {
   INSTANTIATE_ALL_TYPES(value_type_list, test_syclcompat_max);
   INSTANTIATE_ALL_TYPES(value_type_list, test_syclcompat_min);
@@ -356,6 +374,8 @@ int main() {
   INSTANTIATE_ALL_CONTAINER_TYPES(fp_type_list, sycl::marray, test_isnan);
 
   INSTANTIATE_ALL_TYPES(value_type_list, test_clamp);
+  INSTANTIATE_ALL_CONTAINER_TYPES(vec_type_list, sycl::vec, test_container_clamp);
+  INSTANTIATE_ALL_CONTAINER_TYPES(marray_type_list, sycl::marray, test_container_clamp);
 
   return 0;
 }