UCT/CUDA: Prepare GDRCopy support for cudaMallocAsync

config/m4/gdrcopy.m4

@@ -39,7 +39,8 @@ AS_IF([test "x$with_gdrcopy" != "xno"],
             ], [gdrcopy_happy="no"])
 
         AS_IF([test "x$gdrcopy_happy" = "xyes"], [
-            AC_CHECK_DECLS([gdr_pin_buffer_v2, gdr_copy_to_mapping], [], [],
+            AC_CHECK_DECLS([gdr_pin_buffer_v2, gdr_copy_to_mapping,
+                            gdr_get_attribute, GDR_ATTR_USING_DMA_BUF_MMAP], [], [],
                             [#include "gdrapi.h"])
         ])
 
@@ -49,6 +50,7 @@ AS_IF([test "x$with_gdrcopy" != "xno"],
 
         AS_IF([test "x$gdrcopy_happy" = "xyes"],
             [
+                AC_DEFINE([HAVE_GDR_COPY], 1, [GDRCopy support])
                 AC_SUBST(GDR_COPY_CPPFLAGS, "-I$ucx_check_gdrcopy_dir/include/ ")
                 AC_SUBST(GDR_COPY_LIBS, "-lgdrapi")
 

src/uct/cuda/base/cuda_util.c

@@ -32,6 +32,11 @@ const char *uct_cuda_cu_get_error_string(CUresult result)
     return error_str;
 }
 
+ucs_status_t uct_cuda_driver_get_version(int *version_p)
+{
+    return UCT_CUDADRV_FUNC_LOG_ERR(cuDriverGetVersion(version_p));
+}
+
 ucs_status_t uct_cuda_find_device_by_bus_id(const ucs_sys_bus_id_t *bus_id,
                                             ucs_sys_device_t *sys_dev)
 {

src/uct/cuda/base/cuda_util.h

@@ -19,6 +19,10 @@
 const char *uct_cuda_cu_get_error_string(CUresult result);
 
 
+/** Query the CUDA driver API version, for example 13050 for CUDA 13.5. */
+ucs_status_t uct_cuda_driver_get_version(int *version_p);
+
+
 #define UCT_CUDADRV_LOG(_func, _log_level, _result) \
     ucs_log((_log_level), "%s failed: %s", UCS_PP_MAKE_STRING(_func), \
             uct_cuda_cu_get_error_string(_result))

src/uct/cuda/gdr_copy/Makefile.am

@@ -6,7 +6,8 @@
 if HAVE_GDR_COPY
 
 module_LTLIBRARIES              = libuct_cuda_gdrcopy.la
-libuct_cuda_gdrcopy_la_CPPFLAGS = $(BASE_CPPFLAGS) $(GDR_COPY_CPPFLAGS)
+libuct_cuda_gdrcopy_la_CPPFLAGS = $(BASE_CPPFLAGS) $(CUDA_CPPFLAGS) \
+                                  $(GDR_COPY_CPPFLAGS)
 libuct_cuda_gdrcopy_la_CFLAGS   = $(BASE_CFLAGS) $(LT_CFLAGS)
 libuct_cuda_gdrcopy_la_LDFLAGS  = $(GDR_COPY_LDFLAGS) \
                                   -version-info $(SOVERSION) \

src/uct/cuda/gdr_copy/gdr_copy_md.c

@@ -9,6 +9,7 @@
 
 #include "gdr_copy_md.h"
 
+#include <dlfcn.h>
 #include <string.h>
 #include <limits.h>
 #include <ucs/debug/log.h>
@@ -23,6 +24,7 @@
 #include <ucm/api/ucm.h>
 #include <uct/api/v2/uct_v2.h>
 #include <uct/cuda/base/cuda_md.h>
+#include <uct/cuda/base/cuda_util.h>
 
 
 #define UCT_GDR_COPY_RCACHE_OVERHEAD_AUTO 50.0e-9
@@ -65,6 +67,83 @@ uct_gdr_copy_use_pcie_params_get(ucs_ternary_auto_value_t use_pcie,
 #endif
 }
 
+static int uct_gdr_copy_using_dmabuf(gdr_t gdr_ctx)
+{
+#if HAVE_DECL_GDR_GET_ATTRIBUTE && HAVE_DECL_GDR_ATTR_USING_DMA_BUF_MMAP
+    typedef int (*uct_gdr_get_attribute_func_t)(gdr_t, gdr_attr_t, int*);
+    uct_gdr_get_attribute_func_t get_attribute;
+    Dl_info lib_info, attr_info;
+    void *dl_handle;
+    const char *error;
+    int using_dmabuf;
+    int ret;
+
+    /* Resolve the optional query in the runtime which created this context. */
+    if (dladdr(gdr_open, &lib_info) == 0) {
+        ucs_debug("could not identify GDRCopy runtime providing gdr_open");
+        return 0;
+    }
+
+    dl_handle = dlopen(lib_info.dli_fname, RTLD_LAZY | RTLD_NOLOAD);
+    if (dl_handle == NULL) {
+        error = dlerror();
+        ucs_debug("could not open loaded GDRCopy runtime: %s",
+                  (error != NULL) ? error : "unknown error");
+        return 0;
+    }
+
+    get_attribute = (uct_gdr_get_attribute_func_t)dlsym(dl_handle,
+                                                        "gdr_get_attribute");
+    if (get_attribute == NULL) {
+        ucs_debug("gdr_get_attribute is not available in loaded GDRCopy "
+                  "runtime");
+        goto err_dlclose;
+    }
+
+    if (dladdr(get_attribute, &attr_info) == 0) {
+        ucs_debug("could not identify GDRCopy runtime providing "
+                  "gdr_get_attribute");
+        goto err_dlclose;
+    }
+
+    if (attr_info.dli_fbase != lib_info.dli_fbase) {
+        ucs_debug("gdr_get_attribute library %p differs from gdr_open "
+                  "library %p",
+                  attr_info.dli_fbase, lib_info.dli_fbase);
+        goto err_dlclose;
+    }
+
+    ret = get_attribute(gdr_ctx, GDR_ATTR_USING_DMA_BUF_MMAP, &using_dmabuf);
+    dlclose(dl_handle);
+    if (ret != 0) {
+        ucs_debug("failed to query GDRCopy DMA-BUF backend: %d", ret);
+        return 0;
+    }
+
+    return using_dmabuf;
+
+err_dlclose:
+    dlclose(dl_handle);
+    return 0;
+#else
+    return 0;
+#endif
+}
+
+int uct_gdr_copy_cuda_async_detect(gdr_t gdr_ctx)
+{
+    int major, minor, driver_version;
+
+    gdr_runtime_get_version(&major, &minor);
+    if ((major < 2) || ((major == 2) && (minor < 6)) ||
+        !uct_gdr_copy_using_dmabuf(gdr_ctx) ||
+        (uct_cuda_driver_get_version(&driver_version) != UCS_OK)) {
+        return 0;
+    }
+
+    return driver_version >= UCT_GDR_COPY_CUDA_ASYNC_MIN_DRIVER_VERSION;
+}
+
 typedef struct {
     pthread_mutex_t   lock;
     unsigned          refcount;
@@ -227,7 +306,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_gdr_copy_mem_reg_internal,
 
     ret = gdr_get_info(md->gdrcpy_ctx, mem_hndl->mh, &mem_hndl->info);
     if (ret) {
-        ucs_error("gdr_get_info failed. ret:%d", ret);
+        ucs_log(log_level, "gdr_get_info failed. ret:%d", ret);
         goto unmap_buffer;
     }
 
@@ -280,6 +359,7 @@ static ucs_status_t
 uct_gdr_copy_mem_reg(uct_md_h uct_md, void *address, size_t length,
                      const uct_md_mem_reg_params_t *params, uct_mem_h *memh_p)
 {
+    unsigned flags = UCT_MD_MEM_REG_FIELD_VALUE(params, flags, FIELD_FLAGS, 0);
     uct_gdr_copy_mem_t *mem_hndl = NULL;
     ucs_status_t status;
 
@@ -292,7 +372,7 @@ uct_gdr_copy_mem_reg(uct_md_h uct_md, void *address, size_t length,
     }
 
     ucs_ptr_check_align(address, length, GPU_PAGE_SIZE);
-    status = uct_gdr_copy_mem_reg_internal(uct_md, address, length, 0,
+    status = uct_gdr_copy_mem_reg_internal(uct_md, address, length, flags,
                                            mem_hndl);
     if (status != UCS_OK) {
         ucs_free(mem_hndl);

src/uct/cuda/gdr_copy/gdr_copy_md.h

@@ -13,6 +13,14 @@
 extern uct_component_t uct_gdr_copy_component;
 
 
+/* CUDA 13.5 fixes DMA-BUF CPU mapping of stream-ordered allocations. */
+#define UCT_GDR_COPY_CUDA_ASYNC_MIN_DRIVER_VERSION 13050
+
+
+/** Check the loaded runtime, selected backend and CUDA driver prerequisites. */
+int uct_gdr_copy_cuda_async_detect(gdr_t gdr_ctx);
+
+
 /**
  * @brief gdr_copy MD descriptor
  */

test/gtest/Makefile.am

@@ -301,6 +301,11 @@ endif
 
 gtest_CPPFLAGS += \
 	$(CUDA_CPPFLAGS)
+
+if HAVE_GDR_COPY
+gtest_CPPFLAGS += $(GDR_COPY_CPPFLAGS)
+endif
+
 gtest_LDFLAGS += \
 	$(CUDA_LDFLAGS)
 gtest_LDADD += \

test/gtest/uct/test_md.cc

@@ -25,8 +25,13 @@ extern "C" {
 #if HAVE_IB
 #include <uct/ib/base/ib_md.h>
 #endif
+#if HAVE_CUDA && HAVE_GDR_COPY
+#include <uct/cuda/gdr_copy/gdr_copy_md.h>
+#include <cuda.h>
+#endif
 }
 #include <sys/resource.h>
+#include <dlfcn.h>
 #include <net/if_arp.h>
 #include <ifaddrs.h>
 #include <netdb.h>
@@ -1372,7 +1377,7 @@ UCS_TEST_P(test_cuda, sparse_regions)
 
 UCT_MD_INSTANTIATE_TEST_CASE(test_cuda)
 
-#if HAVE_DECL_GDR_PIN_BUFFER_V2
+#if HAVE_CUDA && HAVE_GDR_COPY
 
 class test_gdr_copy : public test_md {
 protected:
@@ -1410,6 +1415,8 @@ class test_gdr_copy : public test_md {
     }
 };
 
+#if HAVE_DECL_GDR_PIN_BUFFER_V2
+
 UCS_TEST_SKIP_COND_P(test_gdr_copy, gdr_copy_reg_cuda_default_pin,
                      !check_caps(UCT_MD_FLAG_REG), "GDR_COPY_USE_PCIE?=auto")
 {
@@ -1439,6 +1446,87 @@ UCS_TEST_SKIP_COND_P(test_gdr_copy, gdr_copy_reg_cuda_try_pcie_pin,
     ASSERT_UCS_OK(register_mem());
 }
 
+#endif /* HAVE_DECL_GDR_PIN_BUFFER_V2 */
+
+UCS_TEST_SKIP_COND_P(test_gdr_copy, hide_registration_error,
+                     !check_caps(UCT_MD_FLAG_REG), "GDR_COPY_RCACHE?=no") {
+    void *address;
+    ucs::handle<void*> buffer;
+
+    const int ret = posix_memalign(&address, REG_SIZE, REG_SIZE);
+    if (ret == 0) {
+        buffer.reset(address, free);
+    }
+
+    ASSERT_EQ(0, ret);
+
+    for (unsigned flags : {0u, unsigned(UCT_MD_MEM_FLAG_HIDE_ERRORS)}) {
+        ucs_status_t status;
+
+        m_errors.clear();
+        {
+            scoped_log_handler slh(hide_errors_logger);
+            uct_mem_h memh;
+
+            status = reg_mem(UCT_MD_MEM_ACCESS_ALL | flags, address, REG_SIZE,
+                             &memh);
+        }
+        ASSERT_UCS_STATUS_EQ(UCS_ERR_IO_ERROR, status);
+        if (flags == 0) {
+            EXPECT_GT(m_errors.size(), 0u);
+        } else {
+            EXPECT_EQ(0u, m_errors.size());
+        }
+    }
+}
+
+UCS_TEST_SKIP_COND_P(test_gdr_copy, cuda_async_qualification,
+                     !check_caps(UCT_MD_FLAG_REG)) {
+    using detect_func_t = int (*)(gdr_t);
+    Dl_info plugin_info;
+    int expected = 0;
+
+    ASSERT_NE(0, dladdr(reinterpret_cast<void*>(md()->ops->mem_reg),
+                        &plugin_info));
+    ucs::handle<void*> plugin(dlopen(plugin_info.dli_fname,
+                                     RTLD_LAZY | RTLD_NOLOAD),
+                              [](void *handle) {
+                                  if (handle != nullptr) {
+                                      dlclose(handle);
+                                  }
+                              });
+    ASSERT_NE(nullptr, plugin.get());
+    auto detect = reinterpret_cast<detect_func_t>(
+            dlsym(plugin.get(), "uct_gdr_copy_cuda_async_detect"));
+    ASSERT_NE(nullptr, detect);
+
+#if HAVE_DECL_GDR_GET_ATTRIBUTE && HAVE_DECL_GDR_ATTR_USING_DMA_BUF_MMAP
+    using version_func_t   = void (*)(int*, int*);
+    using attribute_func_t = int (*)(gdr_t, gdr_attr_t, int*);
+    auto version           = reinterpret_cast<version_func_t>(
+            dlsym(plugin.get(), "gdr_runtime_get_version"));
+    auto attribute         = reinterpret_cast<attribute_func_t>(
+            dlsym(plugin.get(), "gdr_get_attribute"));
+    int major, minor, backend, driver;
+
+    ASSERT_NE(nullptr, version);
+    version(&major, &minor);
+    if (((major > 2) || ((major == 2) && (minor >= 6))) &&
+        (attribute != nullptr) &&
+        (cuDriverGetVersion(&driver) == CUDA_SUCCESS) &&
+        (driver >= UCT_GDR_COPY_CUDA_ASYNC_MIN_DRIVER_VERSION) &&
+        (attribute(ucs_derived_of(md(), uct_gdr_copy_md_t)->gdrcpy_ctx,
+                   GDR_ATTR_USING_DMA_BUF_MMAP, &backend) == 0)) {
+        expected = !!backend;
+    }
+#endif
+
+    EXPECT_EQ(expected,
+              detect(ucs_derived_of(md(), uct_gdr_copy_md_t)->gdrcpy_ctx));
+}
+
+#if HAVE_DECL_GDR_PIN_BUFFER_V2
+
 /* Pinning goes through nvidia_p2p_get_pages, so memory which is not GPUDirect
  * RDMA capable must never be offered to this memory domain */
 UCS_TEST_P(test_gdr_copy, gdr_copy_requires_gdr_capable)
@@ -1457,6 +1545,8 @@ UCS_TEST_SKIP_COND_P(test_gdr_copy, gdr_copy_reg_async_cuda_fail,
     ASSERT_UCS_STATUS_EQ(UCS_ERR_IO_ERROR, register_async_mem());
 }
 
+#endif /* HAVE_DECL_GDR_PIN_BUFFER_V2 */
+
 _UCT_MD_INSTANTIATE_TEST_CASE(test_gdr_copy, gdr_copy)
 
-#endif /* HAVE_DECL_GDR_PIN_BUFFER_V2 */
+#endif /* HAVE_CUDA && HAVE_GDR_COPY */