Provide API to detect C2C platform and add NVML shim (#721)

Closes #720
Partially addresses rapidsai/cudf#18272

This PR moves the work-in-progress C2C detection feature from cuDF (rapidsai/cudf#18447) to KvikIO. A few additional changes are made:
- Handle the potential problem in a multi-GPU system where CUDA and NVML enumerate devices differently.
- Update the conda yaml files to specify NVML dependency only for CUDA >=12, i.e. let the NVML shim only work for CUDA >=12.
- Use of CUDA runtime API is replaced by KvikIO's CUDA driver shim in the new PR.
- Use of `std::functional` as shim function wrappers is replaced by plain function pointers following the pattern of KvikIO's CUDA and cuFile shims.

Authors:
  - Tianyu Liu (https://github.com/kingcrimsontianyu)

Approvers:
  - Vukasin Milovanovic (https://github.com/vuule)
  - Ray Douglass (https://github.com/raydouglass)
  - Robert Maynard (https://github.com/robertmaynard)
  - Bradley Dice (https://github.com/bdice)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #721

Author: kingcrimsontianyu

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -10,6 +10,7 @@ dependencies:
 - c-compiler
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvml-dev
 - cuda-python>=12.6.2,<13.0a0
 - cuda-version=12.8
 - cupy>=12.0.0

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -10,6 +10,7 @@ dependencies:
 - c-compiler
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvml-dev
 - cuda-python>=12.6.2,<13.0a0
 - cuda-version=12.8
 - cupy>=12.0.0

conda/recipes/kvikio/recipe.yaml

@@ -98,6 +98,7 @@ requirements:
           then: libcufile ${{ cuda11_libcufile_run_version }}
       else:
         - cuda-cudart
+        - cuda-nvml-dev
         - if: linux
           then: libcufile
   ignore_run_exports:

conda/recipes/libkvikio/recipe.yaml

@@ -66,7 +66,9 @@ cache:
             then:
               - libcufile =${{ cuda11_libcufile_host_version }}
               - libcufile-dev =${{ cuda11_libcufile_host_version }}
-        else: libcufile-dev
+        else:
+          - libcufile-dev
+          - cuda-nvml-dev
 
 outputs:
   - package:
@@ -99,6 +101,7 @@ outputs:
                 - libcufile ${{ cuda11_libcufile_run_version }}
                 - libcufile-dev ${{ cuda11_libcufile_run_version }}
           else:
+            - cuda-nvml-dev
             - if: linux
               then: libcufile-dev
       ignore_run_exports:

cpp/CMakeLists.txt

@@ -69,6 +69,9 @@ if(KvikIO_REMOTE_SUPPORT)
 endif()
 
 set(cuFile_FOUND 0)
+# todo: Let nvml_FOUND be 0 if CUDA version is less than 12. This is to circumvent a potential CUDA
+# 11 build issue. Remove nvml_FOUND and KVIKIO_NVML_FOUND once CUDA 11 support is dropped.
+set(nvml_FOUND 0)
 if(KvikIO_CUDA_SUPPORT)
   rapids_find_package(
     CUDAToolkit REQUIRED
@@ -77,6 +80,14 @@ if(KvikIO_CUDA_SUPPORT)
   )
   include(cmake/thirdparty/get_nvtx.cmake)
 
+  if(CUDAToolkit_VERSION_MAJOR VERSION_GREATER_EQUAL 12)
+    if(NOT TARGET CUDA::nvml)
+      message(FATAL_ERROR "Cannot find NVML")
+    else()
+      set(nvml_FOUND 1)
+    endif()
+  endif()
+
   if(NOT TARGET CUDA::cuFile)
     message(
       WARNING "Cannot find cuFile - KvikIO will still work but won't use GPUDirect Storage (GDS)"
@@ -151,6 +162,7 @@ set(SOURCES
     "src/posix_io.cpp"
     "src/shim/cuda.cpp"
     "src/shim/cufile.cpp"
+    "src/shim/nvml.cpp"
     "src/shim/utils.cpp"
     "src/stream.cpp"
     "src/utils.cpp"
@@ -177,17 +189,21 @@ target_include_directories(
   INTERFACE "$<INSTALL_INTERFACE:include>"
 )
 
-# Notice, we do not link to cuda or cufile since KvikIO opens them manually using `dlopen()`.
+# Notice, we do not link to cuda or cufile since KvikIO opens them manually using `dlopen()`. Also
+# note that KvikIO uses an NVML shim that dynamically loads the NVML shared library at runtime. At
+# build time only the header location is needed. KvikIO is not linked against the shared library.
+# Adding the target `CUDA::toolkit` here serves this purpose.
 target_link_libraries(
   kvikio
   PUBLIC Threads::Threads BS::thread_pool ${CMAKE_DL_LIBS} $<TARGET_NAME_IF_EXISTS:nvtx3::nvtx3-cpp>
-  PRIVATE $<TARGET_NAME_IF_EXISTS:CURL::libcurl>
+  PRIVATE $<TARGET_NAME_IF_EXISTS:CURL::libcurl> CUDA::toolkit
 )
 
 target_compile_definitions(
   kvikio
   PUBLIC $<$<BOOL:${KvikIO_REMOTE_SUPPORT}>:KVIKIO_LIBCURL_FOUND>
          $<$<BOOL:${KvikIO_CUDA_SUPPORT}>:KVIKIO_CUDA_FOUND>
+         $<$<BOOL:${nvml_FOUND}>:KVIKIO_NVML_FOUND>
          $<$<BOOL:${cuFile_FOUND}>:KVIKIO_CUFILE_FOUND>
          $<$<BOOL:${cuFile_BATCH_API_FOUND}>:KVIKIO_CUFILE_BATCH_API_FOUND>
          $<$<BOOL:${cuFile_STREAM_API_FOUND}>:KVIKIO_CUFILE_STREAM_API_FOUND>

cpp/include/kvikio/shim/cuda.hpp

@@ -39,6 +39,9 @@ class cudaAPI {
   decltype(cuCtxPushCurrent)* CtxPushCurrent{nullptr};
   decltype(cuCtxPopCurrent)* CtxPopCurrent{nullptr};
   decltype(cuCtxGetCurrent)* CtxGetCurrent{nullptr};
+  decltype(cuCtxGetDevice)* CtxGetDevice{nullptr};
+  decltype(cuDeviceGetUuid_v2)* DeviceGetUuid{nullptr};
+  decltype(cuDriverGetVersion)* DriverGetVersion{nullptr};
   decltype(cuMemGetAddressRange)* MemGetAddressRange{nullptr};
   decltype(cuGetErrorName)* GetErrorName{nullptr};
   decltype(cuGetErrorString)* GetErrorString{nullptr};

cpp/include/kvikio/shim/cuda_h_wrapper.hpp

@@ -65,6 +65,9 @@ CUresult cuPointerGetAttributes(...);
 CUresult cuCtxPushCurrent(...);
 CUresult cuCtxPopCurrent(...);
 CUresult cuCtxGetCurrent(...);
+CUresult cuCtxGetDevice(...);
+CUresult cuDeviceGetUuid_v2(...);
+CUresult cuDriverGetVersion(...);
 CUresult cuMemGetAddressRange(...);
 CUresult cuGetErrorName(...);
 CUresult cuGetErrorString(...);

cpp/include/kvikio/utils.hpp

@@ -191,4 +191,46 @@ bool is_future_done(T const& future)
   return future.wait_for(std::chrono::seconds(0)) != std::future_status::timeout;
 }
 
+/**
+ * @brief Check whether the NVML shared library exists.
+ *
+ * @return Boolean answer.
+ */
+#ifdef KVIKIO_NVML_FOUND
+bool is_nvml_available();
+#else
+constexpr bool is_nvml_available() { return false; }
+#endif
+
+/**
+ * @brief Type of the device ID. On a multi-GPU system, CUDA and NVML enumerate devices in different
+ * ways.
+ */
+enum class DeviceIdType : uint8_t {
+  CUDA,  ///< CUDA device ID.
+  NVML,  ///< NVML device ID.
+};
+
+/**
+ * @brief Check if the current device has at least one active NVLink-C2C interconnect.
+ *
+ * @return Boolean answer.
+ */
+#ifdef KVIKIO_NVML_FOUND
+bool is_c2c_available(int device_idx, DeviceIdType device_id_type = DeviceIdType::CUDA);
+#else
+constexpr bool is_c2c_available(int device_idx, DeviceIdType device_id_type = DeviceIdType::CUDA)
+{
+  return false;
+}
+#endif
+
+/**
+ * @brief Shut down NVML
+ *
+ * The NVML shim singleton does not perform shutdown in the destructor. If a cleanup is desired,
+ * call this function before the return of `main()`.
+ */
+void nvml_shutdown();
+
 }  // namespace kvikio

cpp/src/shim/cuda.cpp

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include "utils.hpp"
+
 #include <stdexcept>
 
 #include <kvikio/error.hpp>
@@ -38,6 +40,9 @@ cudaAPI::cudaAPI()
   get_symbol(CtxPushCurrent, lib, KVIKIO_STRINGIFY(cuCtxPushCurrent));
   get_symbol(CtxPopCurrent, lib, KVIKIO_STRINGIFY(cuCtxPopCurrent));
   get_symbol(CtxGetCurrent, lib, KVIKIO_STRINGIFY(cuCtxGetCurrent));
+  get_symbol(CtxGetDevice, lib, KVIKIO_STRINGIFY(cuCtxGetDevice));
+  get_symbol(DeviceGetUuid, lib, KVIKIO_STRINGIFY(cuDeviceGetUuid_v2));
+  get_symbol(DriverGetVersion, lib, KVIKIO_STRINGIFY(cuDriverGetVersion));
   get_symbol(MemGetAddressRange, lib, KVIKIO_STRINGIFY(cuMemGetAddressRange));
   get_symbol(GetErrorName, lib, KVIKIO_STRINGIFY(cuGetErrorName));
   get_symbol(GetErrorString, lib, KVIKIO_STRINGIFY(cuGetErrorString));
@@ -59,15 +64,7 @@ cudaAPI& cudaAPI::instance()
 }
 
 #ifdef KVIKIO_CUDA_FOUND
-bool is_cuda_available()
-{
-  try {
-    cudaAPI::instance();
-  } catch (std::runtime_error const&) {
-    return false;
-  }
-  return true;
-}
+bool is_cuda_available() { return detail::is_available<cudaAPI>(); }
 #endif
 
 }  // namespace kvikio

cpp/src/shim/cufile.cpp

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include "utils.hpp"
+
 #include <stdexcept>
 #include <string>
 
@@ -122,15 +124,7 @@ void cuFileAPI::driver_close()
 }
 
 #ifdef KVIKIO_CUFILE_FOUND
-bool is_cufile_library_available() noexcept
-{
-  try {
-    cuFileAPI::instance();
-  } catch (...) {
-    return false;
-  }
-  return true;
-}
+bool is_cufile_library_available() noexcept { return detail::is_available<cuFileAPI>(); }
 #endif
 
 bool is_cufile_available() noexcept