[REVIEW] Add Fused L2 Expanded KNN kernel

ci/prtest.config

@@ -1,6 +1,6 @@
 RUN_CUGRAPH_LIBCUGRAPH_TESTS=OFF
 RUN_CUGRAPH_PYTHON_TESTS=OFF
 
-RUN_CUML_LIBCUML_TESTS=OFF
-RUN_CUML_PRIMS_TESTS=OFF
-RUN_CUML_PYTHON_TESTS=OFF
+RUN_CUML_LIBCUML_TESTS=ON
+RUN_CUML_PRIMS_TESTS=ON
+RUN_CUML_PYTHON_TESTS=ON

cpp/include/raft/device_atomics.cuh

@@ -179,10 +179,16 @@ struct genericAtomicOperationImpl<T, Op, 4> {
   __forceinline__ __device__ T operator()(T* addr, T const& update_value,
                                           Op op) {
     using T_int = unsigned int;
-
     T old_value = *addr;
     T assumed{old_value};
 
+    if (std::is_same<T, float>{} &&
+        (std::is_same<Op, DeviceMax>{} || std::is_same<Op, DeviceMin>{})) {
+      if (isnan(update_value)) {
+        return update_value;
+      }
+    }
+
     do {
       assumed = old_value;
       const T new_value = op(old_value, update_value);
@@ -191,7 +197,6 @@ struct genericAtomicOperationImpl<T, Op, 4> {
                             type_reinterpret<T_int, T>(assumed),
                             type_reinterpret<T_int, T>(new_value));
       old_value = type_reinterpret<T, T_int>(ret);
-
     } while (assumed != old_value);
 
     return old_value;
@@ -423,7 +428,6 @@ struct typesAtomicCASImpl<T, 4> {
     T_int ret = atomicCAS(reinterpret_cast<T_int*>(addr),
                           type_reinterpret<T_int, T>(compare),
                           type_reinterpret<T_int, T>(update_value));
-
     return type_reinterpret<T, T_int>(ret);
   }
 };
@@ -549,6 +553,17 @@ __forceinline__ __device__ T atomicMax(T* address, T val) {
     address, val, raft::device_atomics::detail::DeviceMax{});
 }
 
+// fp32 only atomicMax.
+__forceinline__ __device__ float customAtomicMax(float* address, float val) {
+  float old;
+  old = (val >= 0)
+          ? __int_as_float(atomicMax((int*)address, __float_as_int(val)))
+          : __uint_as_float(
+              atomicMin((unsigned int*)address, __float_as_uint(val)));
+
+  return old;
+}
+
 /**
  * @brief Overloads for `atomicCAS`
  *

cpp/include/raft/sparse/op/reduce.cuh

@@ -67,7 +67,7 @@ __global__ void max_duplicates_kernel(const value_idx *src_rows,
 
   if (tid < nnz) {
     value_idx idx = index[tid];
-    atomicMax(&out_vals[idx], src_vals[tid]);
+    customAtomicMax(&out_vals[idx], src_vals[tid]);
     out_rows[idx] = src_rows[tid];
     out_cols[idx] = src_cols[tid];
   }