TL/NCCL: add user buffer registration via memmap

Author: wfaderhold21

src/components/tl/nccl/tl_nccl.c

@@ -45,8 +45,7 @@ static ucs_config_field_t ucc_tl_nccl_context_config_table[] = {
     {"SYNC", "auto",
      "Determines how UCC tests completion of NCCL collective",
      ucs_offsetof(ucc_tl_nccl_context_config_t, sync_type),
-     UCS_CONFIG_TYPE_ENUM(ucc_tl_nccl_completion_sync_names)
-    },
+     UCS_CONFIG_TYPE_ENUM(ucc_tl_nccl_completion_sync_names)},
 
     {"BLOCKING", "yes",
      "If set to no will use non-blocking mode communicator behavior, "
@@ -59,6 +58,12 @@ static ucs_config_field_t ucc_tl_nccl_context_config_table[] = {
      ucc_offsetof(ucc_tl_nccl_context_config_t, nccl_lazy_init),
      UCC_CONFIG_TYPE_BOOL},
 
+    {"ENABLE_UBR", "try",
+     "Enable NCCL User Buffer Registration for zero-copy operations. "
+     "Requires NCCL v2.19+.",
+     ucc_offsetof(ucc_tl_nccl_context_config_t, enable_ubr),
+     UCC_CONFIG_TYPE_TERNARY},
+
     {NULL}};
 
 UCC_CLASS_DEFINE_NEW_FUNC(ucc_tl_nccl_lib_t, ucc_base_lib_t,

src/components/tl/nccl/tl_nccl.h

@@ -12,6 +12,8 @@
 #include "components/tl/ucc_tl.h"
 #include "components/tl/ucc_tl_log.h"
 #include "utils/ucc_mpool.h"
+#include "utils/ucc_list.h"
+#include "utils/ucc_spinlock.h"
 
 #include <cuda_runtime.h>
 #if CUDART_VERSION >= 11000
@@ -44,6 +46,8 @@
 #define UCC_TL_NCCL_PROFILE_REQUEST_FREE UCC_PROFILE_REQUEST_FREE
 #define NCCL_VERSION_COMM_INIT_NB NCCL_VERSION(2,14,3)
 #define NCCL_USE_NON_BLOCKING NCCL_VERSION_CODE >= NCCL_VERSION_COMM_INIT_NB
+#define NCCL_VERSION_UBR NCCL_VERSION(2,19,0)
+#define NCCL_HAS_UBR (NCCL_VERSION_CODE >= NCCL_VERSION_UBR)
 
 enum {
     TL_NCCL_COMM_STATE_ERROR,
@@ -76,6 +80,7 @@ typedef struct ucc_tl_nccl_context_config {
     ucc_tl_nccl_completion_sync_type_t sync_type;
     int                                nccl_cfg_blocking;
     int                                nccl_lazy_init;
+    int                                enable_ubr;
 } ucc_tl_nccl_context_config_t;
 
 typedef struct ucc_tl_nccl_lib {
@@ -84,11 +89,33 @@ typedef struct ucc_tl_nccl_lib {
 UCC_CLASS_DECLARE(ucc_tl_nccl_lib_t, const ucc_base_lib_params_t *,
                   const ucc_base_config_t *);
 
+/* Data structure for NCCL memory handle with lazy registration support */
+typedef struct ucc_tl_nccl_memh_data {
+    void             *address;
+    size_t            length;
+    /* Array of NCCL comms this buffer is registered with (NULL = invalidated) */
+    ncclComm_t       *registered_comms;
+    /* Array of NCCL handles from ncclCommRegister */
+    void            **nccl_handles;
+    /* Number of comms in the array */
+    int               num_comms;
+    /* Allocated size of the array */
+    int               max_comms;
+    /* Protects registered_comms/nccl_handles/num_comms/max_comms */
+    ucc_spinlock_t    lock;
+    /* Linked into ucc_tl_nccl_context_t::memh_list */
+    ucc_list_link_t   list_elem;
+} ucc_tl_nccl_memh_data_t;
+
 typedef struct ucc_tl_nccl_context {
     ucc_tl_context_t             super;
     ucc_tl_nccl_context_config_t cfg;
     ucc_mpool_t                  req_mp;
     void                        *scratch_buf;
+    int                          ubr_available;
+    /* List of all live ucc_tl_nccl_memh_data_t objects; protected by memh_lock */
+    ucc_list_link_t              memh_list;
+    ucc_spinlock_t               memh_lock;
 } ucc_tl_nccl_context_t;
 UCC_CLASS_DECLARE(ucc_tl_nccl_context_t, const ucc_base_context_params_t *,
                   const ucc_base_config_t *);

src/components/tl/nccl/tl_nccl_coll.c

@@ -9,6 +9,7 @@
 #include "components/mc/ucc_mc.h"
 #include "components/ec/ucc_ec.h"
 #include "core/ucc_ee.h"
+#include "core/ucc_context.h"
 #include "utils/ucc_compiler_def.h"
 #include "utils/ucc_math.h"
 #include "utils/ucc_coll_utils.h"
@@ -127,6 +128,164 @@ static inline ucc_status_t ucc_tl_nccl_check_and_convert_buffer_reduction(
     return UCC_OK;
 }
 
+#if NCCL_HAS_UBR
+/* Helper function to lazily register a memory region with NCCL communicator */
+static inline ucc_status_t ucc_tl_nccl_lazy_register_memh(
+    void *buffer, size_t length, ucc_tl_nccl_team_t *team,
+    ucc_mem_map_mem_h memh)
+{
+    ucc_tl_nccl_context_t   *ctx = UCC_TL_NCCL_TEAM_CTX(team);
+    ucc_tl_nccl_memh_data_t *m_data;
+    ucc_mem_map_memh_t      *mem_handle;
+    ncclResult_t             nccl_status;
+    ncclComm_t              *new_comms;
+    void                   **new_handles;
+    void                    *nccl_handle;
+    int                      i, new_max;
+    uintptr_t                buf_start, buf_end, region_start, region_end;
+
+    /* Skip if UBR is not available or memh not provided */
+    if (!ctx->ubr_available || !memh) {
+        return UCC_OK;
+    }
+
+    mem_handle = (ucc_mem_map_memh_t *)memh;
+    m_data     = NULL;
+    for (i = 0; i < mem_handle->num_tls; i++) {
+        if (strcmp(mem_handle->tl_h[i].tl_name, "nccl") == 0) {
+            m_data = (ucc_tl_nccl_memh_data_t *)mem_handle->tl_h[i].tl_data;
+            break;
+        }
+    }
+
+    if (!m_data) {
+        /* No NCCL memh data - buffer not registered with TL/NCCL */
+        return UCC_OK;
+    }
+
+    if (length > (UINTPTR_MAX - (uintptr_t)buffer)) {
+        tl_error(UCC_TL_TEAM_LIB(team), "NCCL UBR: buffer size causes overflow");
+        return UCC_ERR_INVALID_PARAM;
+    }
+
+    /* Verify that the entire buffer is within the registered memory region */
+    buf_start    = (uintptr_t)buffer;
+    buf_end      = buf_start + length;
+    region_start = (uintptr_t)m_data->address;
+    region_end   = region_start + m_data->length;
+
+    if (buf_start < region_start || buf_end > region_end) {
+        tl_error(
+            UCC_TL_TEAM_LIB(team),
+            "NCCL UBR: buffer [%p, %p) is outside registered region [%p, %p)",
+            buffer,
+            (void *)buf_end,
+            m_data->address,
+            (void *)region_end);
+        return UCC_ERR_INVALID_PARAM;
+    }
+
+    /* Verify team communicator is initialized */
+    if (!team->nccl_comm) {
+        tl_debug(UCC_TL_TEAM_LIB(team),
+                 "NCCL UBR: communicator not initialized, skipping registration");
+        return UCC_OK;
+    }
+
+    ucc_spin_lock(&m_data->lock);
+
+    /* Check if already registered with this communicator */
+    for (i = 0; i < m_data->num_comms; i++) {
+        if (m_data->registered_comms[i] == team->nccl_comm) {
+            /* Already registered */
+            ucc_spin_unlock(&m_data->lock);
+            return UCC_OK;
+        }
+    }
+
+    /* Need to register the memory region with this communicator.
+     * Release the lock while calling into NCCL (potentially slow), then
+     * re-acquire to update the bookkeeping arrays. */
+    ucc_spin_unlock(&m_data->lock);
+
+    nccl_status = ncclCommRegister(
+        team->nccl_comm, m_data->address, m_data->length, &nccl_handle);
+    if (nccl_status != ncclSuccess) {
+        tl_warn(
+            UCC_TL_TEAM_LIB(team),
+            "NCCL UBR: failed to register region %p, size %zu: %s",
+            m_data->address,
+            m_data->length,
+            ncclGetErrorString(nccl_status));
+        /* Don't fail - UBR is an optimization */
+        return UCC_OK;
+    }
+
+    ucc_spin_lock(&m_data->lock);
+
+    /* Re-check after re-acquiring the lock: another thread may have registered
+     * while we were in ncclCommRegister. */
+    for (i = 0; i < m_data->num_comms; i++) {
+        if (m_data->registered_comms[i] == team->nccl_comm) {
+            ucc_spin_unlock(&m_data->lock);
+            /* Undo the redundant registration */
+            ncclCommDeregister(team->nccl_comm, nccl_handle);
+            return UCC_OK;
+        }
+    }
+
+    /* Add this comm and handle to the registered lists */
+    if (m_data->num_comms >= m_data->max_comms) {
+        /* Need to grow the arrays.  Allocate both new buffers before touching
+         * m_data so that a failure on either leaves the struct untouched. */
+        new_max     = (m_data->max_comms == 0) ? 4 : (m_data->max_comms * 2);
+        new_comms   = (ncclComm_t *)ucc_malloc(new_max * sizeof(ncclComm_t),
+                                               "nccl_registered_comms");
+        new_handles = (void **)ucc_malloc(new_max * sizeof(void *),
+                                          "nccl_handles");
+        if (!new_comms || !new_handles) {
+            ucc_free(new_comms);
+            ucc_free(new_handles);
+            ucc_spin_unlock(&m_data->lock);
+            tl_error(UCC_TL_TEAM_LIB(team),
+                     "failed to grow NCCL UBR bookkeeping arrays");
+            ncclCommDeregister(team->nccl_comm, nccl_handle);
+            return UCC_ERR_NO_MEMORY;
+        }
+        /* Copy existing entries into the new buffers, then swap. */
+        if (m_data->num_comms > 0) {
+            memcpy(new_comms, m_data->registered_comms,
+                   m_data->num_comms * sizeof(ncclComm_t));
+            memcpy(new_handles, m_data->nccl_handles,
+                   m_data->num_comms * sizeof(void *));
+        }
+        ucc_free(m_data->registered_comms);
+        ucc_free(m_data->nccl_handles);
+        m_data->registered_comms = new_comms;
+        m_data->nccl_handles     = new_handles;
+        m_data->max_comms        = new_max;
+    }
+
+    m_data->registered_comms[m_data->num_comms] = team->nccl_comm;
+    m_data->nccl_handles[m_data->num_comms]     = nccl_handle;
+    m_data->num_comms++;
+
+    ucc_spin_unlock(&m_data->lock);
+
+    tl_debug(
+        UCC_TL_TEAM_LIB(team),
+        "NCCL UBR: lazily registered region %p, size %zu with comm %p "
+        "(for buffer [%p, %p))",
+        m_data->address,
+        m_data->length,
+        team->nccl_comm,
+        buffer,
+        (void *)buf_end);
+
+    return UCC_OK;
+}
+#endif
+
 ucc_status_t ucc_tl_nccl_init_task(ucc_base_coll_args_t *coll_args,
                                    ucc_base_team_t *team,
                                    ucc_tl_nccl_task_t **coll_task)
@@ -176,6 +335,102 @@ ucc_status_t ucc_tl_nccl_init_task(ucc_base_coll_args_t *coll_args,
         }
     }
 
+#if NCCL_HAS_UBR
+    /* Lazily register memory regions if they were pre-mapped and UBR is enabled */
+    if (nccl_ctx->ubr_available) {
+        ucc_mem_map_mem_h src_memh = NULL;
+        ucc_mem_map_mem_h dst_memh = NULL;
+        ucc_rank_t        grank    = UCC_TL_TEAM_RANK(nccl_team);
+        ucc_count_t       total_count;
+        ucc_datatype_t    dt;
+        int               is_src_v_type;
+        int               is_dst_v_type;
+
+        /* SCATTERV: src is info_v (root distributes variable chunks),
+         *           dst is info   (each rank receives one contiguous chunk).
+         * GATHERV:  src is info   (each rank contributes one contiguous chunk),
+         *           dst is info_v (root collects variable chunks).
+         * ALLGATHERV: src is info, dst is info_v.
+         * ALLTOALLV:  src is info_v, dst is info_v. */
+        is_src_v_type = (coll_args->args.coll_type == UCC_COLL_TYPE_ALLTOALLV ||
+                         coll_args->args.coll_type == UCC_COLL_TYPE_SCATTERV);
+        is_dst_v_type = (coll_args->args.coll_type == UCC_COLL_TYPE_ALLTOALLV ||
+                         coll_args->args.coll_type == UCC_COLL_TYPE_ALLGATHERV ||
+                         coll_args->args.coll_type == UCC_COLL_TYPE_GATHERV);
+
+        /* Register source buffer's memory region if memh provided */
+        if (coll_args->args.mask & UCC_COLL_ARGS_FIELD_MEM_MAP_SRC_MEMH) {
+            /* Check if global or local memh */
+            if ((coll_args->args.mask & UCC_COLL_ARGS_FIELD_FLAGS) &&
+                (coll_args->args.flags & UCC_COLL_ARGS_FLAG_SRC_MEMH_GLOBAL)) {
+                src_memh = coll_args->args.src_memh.global_memh[grank];
+            } else {
+                src_memh = coll_args->args.src_memh.local_memh;
+            }
+
+            if (is_src_v_type) {
+                total_count = ucc_coll_args_get_v_buffer_size(
+                    &coll_args->args,
+                    coll_args->args.src.info_v.counts,
+                    coll_args->args.src.info_v.displacements,
+                    UCC_TL_TEAM_SIZE(nccl_team));
+                dt = coll_args->args.src.info_v.datatype;
+            } else {
+                total_count = coll_args->args.src.info.count;
+                dt          = coll_args->args.src.info.datatype;
+            }
+            status = ucc_tl_nccl_lazy_register_memh(
+                is_src_v_type ? coll_args->args.src.info_v.buffer
+                              : coll_args->args.src.info.buffer,
+                total_count * ucc_dt_size(dt),
+                nccl_team,
+                src_memh);
+            if (ucc_unlikely(status != UCC_OK)) {
+                tl_error(
+                    team->context->lib,
+                    "NCCL UBR: lazy_register failed with status %d",
+                    status);
+                ucc_tl_nccl_free_task(task);
+                return status;
+            }
+        }
+
+        /* Register destination buffer's memory region if memh provided */
+        if (coll_args->args.mask & UCC_COLL_ARGS_FIELD_MEM_MAP_DST_MEMH) {
+            /* Check if global or local memh */
+            if ((coll_args->args.mask & UCC_COLL_ARGS_FIELD_FLAGS) &&
+                (coll_args->args.flags & UCC_COLL_ARGS_FLAG_DST_MEMH_GLOBAL)) {
+                dst_memh = coll_args->args.dst_memh.global_memh[grank];
+            } else {
+                dst_memh = coll_args->args.dst_memh.local_memh;
+            }
+
+            if (is_dst_v_type) {
+                total_count = ucc_coll_args_get_v_buffer_size(
+                    &coll_args->args,
+                    coll_args->args.dst.info_v.counts,
+                    coll_args->args.dst.info_v.displacements,
+                    UCC_TL_TEAM_SIZE(nccl_team));
+                dt = coll_args->args.dst.info_v.datatype;
+            } else {
+                total_count = coll_args->args.dst.info.count;
+                dt          = coll_args->args.dst.info.datatype;
+            }
+
+            status = ucc_tl_nccl_lazy_register_memh(
+                is_dst_v_type ? coll_args->args.dst.info_v.buffer
+                              : coll_args->args.dst.info.buffer,
+                total_count * ucc_dt_size(dt),
+                nccl_team,
+                dst_memh);
+            if (ucc_unlikely(status != UCC_OK)) {
+                ucc_tl_nccl_free_task(task);
+                return status;
+            }
+        }
+    }
+#endif
+
     *coll_task = task;
     return UCC_OK;
 }