Search reductions use correct branch for float16

constexpr branch logic accounted for floating point types but not sycl::half, which meant NaNs were not propagating for float16 data
IntelPython · Nov 3, 2023 · 119d43d · 119d43d
1 parent 5709f99
commit 119d43d
Showing 1 changed file with 15 additions and 5 deletions.
diff --git a/dpctl/tensor/libtensor/include/kernels/reductions.hpp b/dpctl/tensor/libtensor/include/kernels/reductions.hpp
@@ -3476,7 +3476,9 @@ struct SequentialSearchReduction
                             idx_val = static_cast<outT>(m);
                         }
                     }
-                    else if constexpr (std::is_floating_point_v<argT>) {
+                    else if constexpr (std::is_floating_point_v<argT> ||
+                                       std::is_same_v<argT, sycl::half>)
+                    {
                         if (val < red_val || std::isnan(val)) {
                             red_val = val;
                             idx_val = static_cast<outT>(m);
@@ -3501,7 +3503,9 @@ struct SequentialSearchReduction
                             idx_val = static_cast<outT>(m);
                         }
                     }
-                    else if constexpr (std::is_floating_point_v<argT>) {
+                    else if constexpr (std::is_floating_point_v<argT> ||
+                                       std::is_same_v<argT, sycl::half>)
+                    {
                         if (val > red_val || std::isnan(val)) {
                             red_val = val;
                             idx_val = static_cast<outT>(m);
@@ -3789,7 +3793,9 @@ struct CustomSearchReduction
                                 }
                             }
                         }
-                        else if constexpr (std::is_floating_point_v<argT>) {
+                        else if constexpr (std::is_floating_point_v<argT> ||
+                                           std::is_same_v<argT, sycl::half>)
+                        {
                             if (val < local_red_val || std::isnan(val)) {
                                 local_red_val = val;
                                 if constexpr (!First) {
@@ -3833,7 +3839,9 @@ struct CustomSearchReduction
                                 }
                             }
                         }
-                        else if constexpr (std::is_floating_point_v<argT>) {
+                        else if constexpr (std::is_floating_point_v<argT> ||
+                                           std::is_same_v<argT, sycl::half>)
+                        {
                             if (val > local_red_val || std::isnan(val)) {
                                 local_red_val = val;
                                 if constexpr (!First) {
@@ -3876,7 +3884,9 @@ struct CustomSearchReduction
                             ? local_idx
                             : idx_identity_;
         }
-        else if constexpr (std::is_floating_point_v<argT>) {
+        else if constexpr (std::is_floating_point_v<argT> ||
+                           std::is_same_v<argT, sycl::half>)
+        {
             // equality does not hold for NaNs, so check here
             local_idx =
                 (red_val_over_wg == local_red_val || std::isnan(local_red_val))