Provide a raft::copy overload for mdspan-to-mdspan copies

cpp/CMakeLists.txt

@@ -22,7 +22,8 @@ include(rapids-find)
 
 option(BUILD_CPU_ONLY "Build CPU only components. Applies to RAFT ANN benchmarks currently" OFF)
 
-# workaround for rapids_cuda_init_architectures not working for arch detection with enable_language(CUDA)
+# workaround for rapids_cuda_init_architectures not working for arch detection with
+# enable_language(CUDA)
 set(lang_list "CXX")
 
 if(NOT BUILD_CPU_ONLY)
@@ -286,7 +287,8 @@ endif()
 set_target_properties(raft_compiled PROPERTIES EXPORT_NAME compiled)
 
 if(RAFT_COMPILE_LIBRARY)
-  add_library(raft_objs OBJECT
+  add_library(
+    raft_objs OBJECT
     src/core/logger.cpp
     src/distance/detail/pairwise_matrix/dispatch_canberra_double_double_double_int.cu
     src/distance/detail/pairwise_matrix/dispatch_canberra_float_float_float_int.cu
@@ -452,18 +454,21 @@ if(RAFT_COMPILE_LIBRARY)
     src/spatial/knn/detail/fused_l2_knn_int64_t_float.cu
     src/spatial/knn/detail/fused_l2_knn_uint32_t_float.cu
     src/util/memory_pool.cpp
-    )
+  )
   set_target_properties(
     raft_objs
     PROPERTIES CXX_STANDARD 17
                CXX_STANDARD_REQUIRED ON
                CUDA_STANDARD 17
                CUDA_STANDARD_REQUIRED ON
-               POSITION_INDEPENDENT_CODE ON)
+               POSITION_INDEPENDENT_CODE ON
+  )
 
   target_compile_definitions(raft_objs PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY")
-  target_compile_options(raft_objs PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:${RAFT_CXX_FLAGS}>"
-                                           "$<$<COMPILE_LANGUAGE:CUDA>:${RAFT_CUDA_FLAGS}>")
+  target_compile_options(
+    raft_objs PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:${RAFT_CXX_FLAGS}>"
+                      "$<$<COMPILE_LANGUAGE:CUDA>:${RAFT_CUDA_FLAGS}>"
+  )
 
   add_library(raft_lib SHARED $<TARGET_OBJECTS:raft_objs>)
   add_library(raft_lib_static STATIC $<TARGET_OBJECTS:raft_objs>)
@@ -477,13 +482,15 @@ if(RAFT_COMPILE_LIBRARY)
   )
 
   foreach(target raft_lib raft_lib_static raft_objs)
-    target_link_libraries(${target} PUBLIC
-      raft::raft
-      ${RAFT_CTK_MATH_DEPENDENCIES} # TODO: Once `raft::resources` is used everywhere, this
-                                    # will just be cublas
-      $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>)
+    target_link_libraries(
+      ${target}
+      PUBLIC raft::raft
+             ${RAFT_CTK_MATH_DEPENDENCIES} # TODO: Once `raft::resources` is used everywhere, this
+                                           # will just be cublas
+             $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
+    )
 
-    #So consumers know when using libraft.so/libraft.a
+    # So consumers know when using libraft.so/libraft.a
     target_compile_definitions(${target} PUBLIC "RAFT_COMPILED")
     # ensure CUDA symbols aren't relocated to the middle of the debug build binaries
     target_link_options(${target} PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")

cpp/bench/prims/CMakeLists.txt

@@ -32,6 +32,7 @@ function(ConfigureBench)
     PRIVATE raft::raft
             raft_internal
             $<$<BOOL:${ConfigureBench_LIB}>:raft::compiled>
+            ${RAFT_CTK_MATH_DEPENDENCIES}
             benchmark::benchmark
             Threads::Threads
             $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
@@ -73,6 +74,8 @@ function(ConfigureBench)
 endfunction()
 
 if(BUILD_PRIMS_BENCH)
+  ConfigureBench(NAME CORE_BENCH PATH bench/prims/core/copy.cu bench/prims/main.cpp)
+
   ConfigureBench(
     NAME CLUSTER_BENCH PATH bench/prims/cluster/kmeans_balanced.cu bench/prims/cluster/kmeans.cu
     bench/prims/main.cpp OPTIONAL LIB EXPLICIT_INSTANTIATE_ONLY

cpp/include/raft/core/copy.cuh

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <raft/core/detail/copy.hpp>
+namespace raft {
+/**
+ * @brief Copy data from one mdspan to another with the same extents
+ *
+ * This function copies data from one mdspan to another, regardless of whether
+ * or not the mdspans have the same layout, memory type (host/device/managed)
+ * or data type. So long as it is possible to convert the data type from source
+ * to destination, and the extents are equal, this function should be able to
+ * perform the copy. Any necessary device operations will be stream-ordered via the CUDA stream
+ * provided by the `raft::resources` argument.
+ *
+ * This header includes a custom kernel used for copying data between
+ * completely arbitrary mdspans on device. To compile this function in a
+ * non-CUDA translation unit, `raft/core/copy.hpp` may be used instead. The
+ * pure C++ header will correctly compile even without a CUDA compiler.
+ * Depending on the specialization, this CUDA header may invoke the kernel and
+ * therefore require a CUDA compiler.
+ *
+ * Limitations: Currently this function does not support copying directly
+ * between two arbitrary mdspans on different CUDA devices. It is assumed that the caller sets the
+ * correct CUDA device. Furthermore, host-to-host copies that require a transformation of the
+ * underlying memory layout are currently not performant, although they are supported.
+ *
+ * Note that when copying to an mdspan with a non-unique layout (i.e. the same
+ * underlying memory is addressed by different element indexes), the source
+ * data must contain non-unique values for every non-unique destination
+ * element. If this is not the case, the behavior is undefined. Some copies
+ * to non-unique layouts which are well-defined will nevertheless fail with an
+ * exception to avoid race conditions in the underlying copy.
+ *
+ * @tparam DstType An mdspan type for the destination container.
+ * @tparam SrcType An mdspan type for the source container
+ * @param res raft::resources used to provide a stream for copies involving the
+ * device.
+ * @param dst The destination mdspan.
+ * @param src The source mdspan.
+ */
+template <typename DstType, typename SrcType>
+detail::mdspan_copyable_with_kernel_t<DstType, SrcType> copy(resources const& res,
+                                                             DstType&& dst,
+                                                             SrcType&& src)
+{
+  detail::copy(res, std::forward<DstType>(dst), std::forward<SrcType>(src));
+}
+
+#ifndef RAFT_NON_CUDA_COPY_IMPLEMENTED
+#define RAFT_NON_CUDA_COPY_IMPLEMENTED
+template <typename DstType, typename SrcType>
+detail::mdspan_uncopyable_with_kernel_t<DstType, SrcType> copy(resources const& res,
+                                                               DstType&& dst,
+                                                               SrcType&& src)
+{
+  detail::copy(res, std::forward<DstType>(dst), std::forward<SrcType>(src));
+}
+#endif
+}  // namespace raft

cpp/include/raft/core/copy.hpp

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include <raft/core/detail/copy.hpp>
+namespace raft {
+
+#ifndef RAFT_NON_CUDA_COPY_IMPLEMENTED
+#define RAFT_NON_CUDA_COPY_IMPLEMENTED
+/**
+ * @brief Copy data from one mdspan to another with the same extents
+ *
+ * This function copies data from one mdspan to another, regardless of whether
+ * or not the mdspans have the same layout, memory type (host/device/managed)
+ * or data type. So long as it is possible to convert the data type from source
+ * to destination, and the extents are equal, this function should be able to
+ * perform the copy.
+ *
+ * This header does _not_ include the custom kernel used for copying data
+ * between completely arbitrary mdspans on device. For arbitrary copies of this
+ * kind, `#include <raft/core/copy.cuh>` instead. Specializations of this
+ * function that require the custom kernel will be SFINAE-omitted when this
+ * header is used instead of `copy.cuh`. This header _does_ support
+ * device-to-device copies that can be performed with cuBLAS or a
+ * straightforward cudaMemcpy. Any necessary device operations will be stream-ordered via the CUDA
+ * stream provided by the `raft::resources` argument.
+ *
+ * Limitations: Currently this function does not support copying directly
+ * between two arbitrary mdspans on different CUDA devices. It is assumed that the caller sets the
+ * correct CUDA device. Furthermore, host-to-host copies that require a transformation of the
+ * underlying memory layout are currently not performant, although they are supported.
+ *
+ * Note that when copying to an mdspan with a non-unique layout (i.e. the same
+ * underlying memory is addressed by different element indexes), the source
+ * data must contain non-unique values for every non-unique destination
+ * element. If this is not the case, the behavior is undefined. Some copies
+ * to non-unique layouts which are well-defined will nevertheless fail with an
+ * exception to avoid race conditions in the underlying copy.
+ *
+ * @tparam DstType An mdspan type for the destination container.
+ * @tparam SrcType An mdspan type for the source container
+ * @param res raft::resources used to provide a stream for copies involving the
+ * device.
+ * @param dst The destination mdspan.
+ * @param src The source mdspan.
+ */
+template <typename DstType, typename SrcType>
+detail::mdspan_uncopyable_with_kernel_t<DstType, SrcType> copy(resources const& res,
+                                                               DstType&& dst,
+                                                               SrcType&& src)
+{
+  detail::copy(res, std::forward<DstType>(dst), std::forward<SrcType>(src));
+}
+#endif
+
+}  // namespace raft

cpp/include/raft/core/cuda_support.hpp

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+namespace raft {
+#ifndef RAFT_DISABLE_CUDA
+auto constexpr static const CUDA_ENABLED = true;
+#else
+auto constexpr static const CUDA_ENABLED = false;
+#endif
+}  // namespace raft