/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2018-2026. ALL RIGHTS RESERVED.
 * See file LICENSE for terms.
 */

#ifdef HAVE_CONFIG_H
#  include "config.h"
#endif

#include "cuda_ipc.inl"
#include "cuda_ipc_cache.h"
#include "cuda_ipc_md.h"

#include <ucs/debug/log.h>
#include <ucs/debug/memtrack_int.h>
#include <ucs/profile/profile.h>
#include <ucs/sys/string.h>
#include <ucs/sys/sys.h>
#include <ucs/sys/uid.h>
#include <ucs/type/class.h>
#include <uct/api/v2/uct_v2.h>
#include <uct/cuda/base/cuda_nvml.h>
#include <uct/cuda/base/cuda_util.h>
#include <uct/api/device/uct_device_types.h>

#include <limits.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>

/* Indicates whether PID NS is contained in rkey.
 * Wire compatibility isn't broken because PID is used only for caching
 * purposes. */
#define UCT_CUDA_IPC_RKEY_FLAG_PID_NS UCS_BIT(31)

static ucs_config_field_t uct_cuda_ipc_md_config_table[] = {
    {"", "", NULL,
     ucs_offsetof(uct_cuda_ipc_md_config_t, super), UCS_CONFIG_TYPE_TABLE(uct_md_config_table)},

    {"ENABLE_MNNVL", "try",
     "Enable multi-node NVLINK capabilities.",
     ucs_offsetof(uct_cuda_ipc_md_config_t, enable_mnnvl), UCS_CONFIG_TYPE_TERNARY},

    {"CACHE_MAX_REGIONS", "inf",
     "Maximum number of regions in each per-peer CUDA IPC remote handle cache.\n"
     "Each remote peer has a separate cache; this limit applies independently\n"
     "to each one. Regions are evicted in LRU order when this limit is exceeded.",
     ucs_offsetof(uct_cuda_ipc_md_config_t, cache_max_regions),
     UCS_CONFIG_TYPE_ULUNITS},

    {"CACHE_MAX_SIZE", "inf",
     "Maximum total size of regions in each per-peer CUDA IPC remote handle cache.\n"
     "Each remote peer has a separate cache; this limit applies independently\n"
     "to each one. Regions are evicted in LRU order when this limit is exceeded.",
     ucs_offsetof(uct_cuda_ipc_md_config_t, cache_max_size),
     UCS_CONFIG_TYPE_MEMUNITS},

    {NULL}
};

static uct_cuda_ipc_dev_cache_t *uct_cuda_ipc_create_dev_cache(int dev_num)
{
    uct_cuda_ipc_dev_cache_t *cache;
    ucs_status_t status;
    int i, num_devices;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGetCount(&num_devices));
    if (UCS_OK != status) {
        return NULL;
    }

    cache = ucs_malloc(sizeof(*cache) + (num_devices * sizeof(uint8_t)),
                       "uct_cuda_ipc_dev_cache_t");
    if (cache == NULL) {
        ucs_error("failed to allocate memory for uct_cuda_ipc_dev_cache_t");
        return NULL;
    }

    cache->dev_num = dev_num;
    for (i = 0; i < num_devices; ++i) {
        cache->accessible[i] = UCS_TRY;
    }

    return cache;
}

static uct_cuda_ipc_dev_cache_t *
uct_cuda_ipc_get_dev_cache(uct_cuda_ipc_component_t *component,
                           const uct_cuda_ipc_extended_rkey_t *ext_rkey)
{
    const uct_cuda_ipc_rkey_t *rkey   = &ext_rkey->super;
    khash_t(cuda_ipc_uuid_hash) *hash = &component->uuid_hash;
    uct_cuda_ipc_uuid_hash_key_t key;
    uct_cuda_ipc_dev_cache_t *cache;
    khiter_t iter;
    int ret;

    key.uuid     = rkey->uuid;
    key.type     = rkey->ph.handle_type;
    key.is_local = uct_cuda_ipc_is_rkey_local(rkey->pid, ext_rkey->pid_ns);

    iter = kh_put(cuda_ipc_uuid_hash, hash, key, &ret);
    if (ret == UCS_KH_PUT_KEY_PRESENT) {
        return kh_val(hash, iter);
    } else if ((ret == UCS_KH_PUT_BUCKET_EMPTY) ||
               (ret == UCS_KH_PUT_BUCKET_CLEAR)) {
        cache = uct_cuda_ipc_create_dev_cache(kh_size(hash) - 1);
        if (NULL != cache) {
            kh_val(hash, iter) = cache;
        }
        return cache;
    } else {
        ucs_error("kh_put(cuda_ipc_uuid_hash) failed with %d", ret);
        return NULL;
    }
}

static ucs_status_t
uct_cuda_ipc_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr)
{
    uct_md_base_md_query(md_attr);
    md_attr->flags            = UCT_MD_FLAG_REG |
                                UCT_MD_FLAG_NEED_RKEY |
                                UCT_MD_FLAG_INVALIDATE |
                                UCT_MD_FLAG_INVALIDATE_RMA |
                                UCT_MD_FLAG_INVALIDATE_AMO |
                                UCT_MD_FLAG_MEMTYPE_COPY |
                                UCT_MD_FLAG_RKEY_PTR;
    md_attr->reg_mem_types    = UCS_BIT(UCS_MEMORY_TYPE_CUDA);
    md_attr->cache_mem_types  = UCS_BIT(UCS_MEMORY_TYPE_CUDA);
    md_attr->access_mem_types = UCS_BIT(UCS_MEMORY_TYPE_CUDA);
    md_attr->rkey_packed_size = ucs_sys_ns_is_default(UCS_SYS_NS_TYPE_PID) ?
                                        sizeof(uct_cuda_ipc_rkey_t) :
                                        sizeof(uct_cuda_ipc_extended_rkey_t);
    return UCS_OK;
}

#if HAVE_DECL_SYS_PIDFD_GETFD
static ucs_status_t
uct_cuda_ipc_mem_export_posix_fd(void *addr, uct_cuda_ipc_lkey_t *key)
{
    CUmemGenericAllocationHandle alloc_handle;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC(cuMemRetainAllocationHandle(&alloc_handle, addr),
                              UCS_LOG_LEVEL_DIAG);
    if (status != UCS_OK) {
        ucs_debug("unable to export POSIX FD handle for ptr: %p", addr);
        return status;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuMemExportToShareableHandle(
            &key->ph.handle.posix_fd.fd, alloc_handle,
            CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0));
    UCT_CUDADRV_FUNC_LOG_WARN(cuMemRelease(alloc_handle));
    if (status != UCS_OK) {
        ucs_debug("unable to export POSIX FD handle for ptr: %p", addr);
        return status;
    }

    key->ph.handle_type               = UCT_CUDA_IPC_KEY_HANDLE_TYPE_POSIX_FD;
    key->ph.handle.posix_fd.system_id = ucs_get_system_id();
    ucs_trace("posix_fd export: addr=%p fd=%d pid=%d system_id=0x%" PRIx64,
              addr, key->ph.handle.posix_fd.fd, getpid(),
              key->ph.handle.posix_fd.system_id);
    return UCS_OK;
}
#endif

#if HAVE_CUDA_FABRIC
static ucs_status_t
uct_cuda_ipc_mem_export_fabric_mempool(const void *addr,
                                       uct_cuda_ipc_lkey_t *key,
                                       CUmemoryPool mempool)
{
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuMemPoolExportToShareableHandle(
            (void*)&key->ph.handle.fabric_handle, mempool,
            CU_MEM_HANDLE_TYPE_FABRIC, 0));
    if (status != UCS_OK) {
        ucs_debug("unable to export handle for mempool ptr: %p", addr);
        return UCS_ERR_UNSUPPORTED;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(
            cuMemPoolExportPointer(&key->ph.ptr, (CUdeviceptr)key->d_bptr));
    if (status != UCS_OK) {
        return status;
    }

    key->ph.handle_type = UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL;
    ucs_trace("packed mempool handle and export pointer for %p", addr);
    return UCS_OK;
}

static ucs_status_t uct_cuda_ipc_mem_export_fabric(void *addr,
                                                   uct_cuda_ipc_lkey_t *key,
                                                   CUmemoryPool mempool)
{
    CUmemGenericAllocationHandle handle;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC(cuMemRetainAllocationHandle(&handle, addr),
                              UCS_LOG_LEVEL_DIAG);
    if (status == UCS_OK) {
        status = UCT_CUDADRV_FUNC_LOG_ERR(
                cuMemExportToShareableHandle(&key->ph.handle.fabric_handle,
                                             handle, CU_MEM_HANDLE_TYPE_FABRIC,
                                             0));
        UCT_CUDADRV_FUNC_LOG_WARN(cuMemRelease(handle));
        if (status != UCS_OK) {
            ucs_debug("unable to export handle for VMM ptr: %p", addr);
            return UCS_ERR_UNSUPPORTED;
        }

        key->ph.handle_type = UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM;
        ucs_trace("packed vmm fabric handle for %p", addr);
        return UCS_OK;
    }

    if (mempool == NULL) {
        return UCS_ERR_INVALID_ADDR;
    }

    return uct_cuda_ipc_mem_export_fabric_mempool(addr, key, mempool);
}
#endif /* HAVE_CUDA_FABRIC */

static ucs_status_t
uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
                         uct_cuda_ipc_lkey_t **key_p)
{
    uct_cuda_ipc_lkey_t *key;
    ucs_status_t status;
    int is_ctx_pushed;
    CUdevice cuda_device;
#if HAVE_CUDA_FABRIC || HAVE_DECL_SYS_PIDFD_GETFD
#define UCT_CUDA_IPC_QUERY_NUM_ATTRS 3
    CUmemoryPool mempool;
    CUpointer_attribute attr_type[UCT_CUDA_IPC_QUERY_NUM_ATTRS];
    void *attr_data[UCT_CUDA_IPC_QUERY_NUM_ATTRS];
    int legacy_capable;
    uint64_t allowed_handle_types;
#endif

    key = ucs_calloc(1, sizeof(*key), "uct_cuda_ipc_lkey_t");
    if (key == NULL) {
        return UCS_ERR_NO_MEMORY;
    }

    status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)addr, &cuda_device,
                                             &is_ctx_pushed);
    if (status != UCS_OK) {
        goto out;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(
            cuMemGetAddressRange(&key->d_bptr, &key->b_len, (CUdeviceptr)addr));
    if (status != UCS_OK) {
        goto out_pop_ctx;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuPointerGetAttribute(&key->ph.buffer_id,
                                      CU_POINTER_ATTRIBUTE_BUFFER_ID,
                                      (CUdeviceptr)addr));
    if (status != UCS_OK) {
        goto out_pop_ctx;
    }

#if HAVE_CUDA_FABRIC || HAVE_DECL_SYS_PIDFD_GETFD
    /* cuda_ipc can handle VMM, mallocasync, and legacy pinned device so need to
     * pack appropriate handle */

    attr_type[0] = CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE;
    attr_data[0] = &legacy_capable;
    attr_type[1] = CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES;
    attr_data[1] = &allowed_handle_types;
    attr_type[2] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE;
    attr_data[2] = &mempool;

    status = UCT_CUDADRV_FUNC_LOG_ERR(
            cuPointerGetAttributes(ucs_static_array_size(attr_data), attr_type,
                attr_data, (CUdeviceptr)addr));
    if (status != UCS_OK) {
        goto out_pop_ctx;
    }

    if (legacy_capable) {
        goto legacy_path;
    }

#if HAVE_CUDA_FABRIC
    if (allowed_handle_types & CU_MEM_HANDLE_TYPE_FABRIC) {
        status = uct_cuda_ipc_mem_export_fabric(addr, key, mempool);
        if (status == UCS_OK) {
            goto common_path;
        }
        if (status != UCS_ERR_UNSUPPORTED) {
            goto out_pop_ctx;
        }
    }
#endif /* HAVE_CUDA_FABRIC */

#if HAVE_DECL_SYS_PIDFD_GETFD
    if (allowed_handle_types & CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR) {
        status = uct_cuda_ipc_mem_export_posix_fd(addr, key);
        if (status == UCS_OK) {
            goto common_path;
        }
    }
#endif /* HAVE_DECL_SYS_PIDFD_GETFD */

    key->ph.handle_type = UCT_CUDA_IPC_KEY_HANDLE_TYPE_NO_IPC;
    goto common_path;
#endif /* HAVE_CUDA_FABRIC || HAVE_DECL_SYS_PIDFD_GETFD */
legacy_path:
    key->ph.handle_type = UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY;
    status              = UCT_CUDADRV_FUNC_LOG_ERR(
            cuIpcGetMemHandle(&key->ph.handle.legacy, (CUdeviceptr)addr));
    if (status != UCS_OK) {
        goto out_pop_ctx;
    }

common_path:
    ucs_list_add_tail(&memh->list, &key->link);
    ucs_trace("registered addr:%p/%p length:%zd type:%u dev_num:%d "
              "buffer_id:%llu",
              addr, (void*)key->d_bptr, key->b_len, key->ph.handle_type,
              cuda_device, key->ph.buffer_id);

    memh->dev_num = cuda_device;
    *key_p        = key;
    status        = UCS_OK;

out_pop_ctx:
    uct_cuda_ipc_check_and_pop_ctx(is_ctx_pushed);
out:
    if (status != UCS_OK) {
        ucs_free(key);
    }

    return status;
}

static ucs_status_t
uct_cuda_ipc_mkey_pack(uct_md_h md, uct_mem_h tl_memh, void *address,
                       size_t length, const uct_md_mkey_pack_params_t *params,
                       void *mkey_buffer)
{
    uct_cuda_ipc_rkey_t *packed = mkey_buffer;
    uct_cuda_ipc_memh_t *memh   = tl_memh;
    uct_cuda_ipc_extended_rkey_t *ext_rkey;
    uct_cuda_ipc_lkey_t *key;
    ucs_status_t status;

    ucs_list_for_each(key, &memh->list, link) {
        if (((uintptr_t)address >= key->d_bptr) &&
            ((uintptr_t)address < (key->d_bptr + key->b_len))) {
            goto found;
        }
    }

    status = uct_cuda_ipc_mem_add_reg(address, memh, &key);
    if (status != UCS_OK) {
        return status;
    }

found:
    ucs_assertv(((uintptr_t)address + length) <= (key->d_bptr + key->b_len),
                "buffer 0x%lx..0x%lx region 0x%llx..0x%llx", (uintptr_t)address,
                (uintptr_t)address + length, key->d_bptr, key->d_bptr +
                key->b_len);

    packed->pid    = memh->pid;
    packed->ph     = key->ph;
    packed->d_bptr = key->d_bptr;
    packed->b_len  = key->b_len;

    if (!ucs_sys_ns_is_default(UCS_SYS_NS_TYPE_PID)) {
        ext_rkey         = (uct_cuda_ipc_extended_rkey_t*)packed;
        ext_rkey->pid_ns = memh->pid_ns;

        ucs_assertv_always(!(packed->pid & UCT_CUDA_IPC_RKEY_FLAG_PID_NS),
                           "pid=%d", packed->pid);
        packed->pid |= UCT_CUDA_IPC_RKEY_FLAG_PID_NS;
    }

    return UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGetUuid(&packed->uuid,
                                                    memh->dev_num));
}

static ucs_status_t
uct_cuda_ipc_is_peer_accessible(uct_cuda_ipc_component_t *component,
                                uct_cuda_ipc_unpacked_rkey_t *rkey,
                                CUdevice cu_dev)
{
    ucs_status_t status;
    void *d_mapped;
    uct_cuda_ipc_dev_cache_t *cache;
    uint8_t *accessible;

    pthread_mutex_lock(&component->lock);

    cache = uct_cuda_ipc_get_dev_cache(component, &rkey->super);
    if (ucs_unlikely(NULL == cache)) {
        status = UCS_ERR_NO_RESOURCE;
        goto err;
    }

    /* use unique id for stream id; this means that relative remote
     * device number of multiple peers do not map on the same stream and reduces
     * stream sequentialization */
    rkey->stream_id = cache->dev_num;
    accessible      = &cache->accessible[cu_dev];
    if (ucs_unlikely(*accessible == UCS_TRY)) { /* unchecked, add to cache */

        /* Check if peer is reachable by trying to open memory handle. This is
         * necessary when the device is not visible through CUDA_VISIBLE_DEVICES
         * and checking peer accessibility through CUDA driver API is not
         * possible.
         * Previously, reachability was checked by opening a memory handle
         * and immediately closing it as the handle to memory handle cache
         * was not not globally visible. Doing this with multiple threads is an
         * issue as a thread may first check reachability, and later open the
         * handle, and save mapped pointer in cache as part of a put/get
         * operation. At this point another thread can then close the same
         * memory handle as part of reachability check. This leads to a
         * cuMemcpyAsync error when accessing the mapped pointer as part of
         * put/get operation.
         * Now, we immediately insert into cache to save on calling
         * OpenMemHandle for the same handle because the cache is globally
         * accessible using rkey->pid. */
        status = uct_cuda_ipc_map_memhandle(&rkey->super, cu_dev, &d_mapped,
                                            UCS_LOG_LEVEL_DEBUG);

        *accessible = ((status == UCS_OK) || (status == UCS_ERR_ALREADY_EXISTS))
                      ? UCS_YES : UCS_NO;
    }

    status = (*accessible == UCS_YES) ? UCS_OK : UCS_ERR_UNREACHABLE;

err:
    pthread_mutex_unlock(&component->lock);
    return status;
}

UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack,
                 (component, rkey_buffer, params, rkey_p, handle_p),
                 uct_component_t *component, const void *rkey_buffer,
                 const uct_rkey_unpack_params_t *params, uct_rkey_t *rkey_p,
                 void **handle_p)
{
    uct_cuda_ipc_component_t *com     = ucs_derived_of(component,
                                                       uct_cuda_ipc_component_t);
    const uct_cuda_ipc_rkey_t *packed = rkey_buffer;
    uct_cuda_ipc_unpacked_rkey_t *unpacked;
    ucs_sys_device_t sys_dev;
    CUdevice cuda_device;
    CUdevice avail_cuda_device;
    ucs_status_t status;
    const uct_cuda_ipc_extended_rkey_t *ext_rkey;

    sys_dev = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, sys_device,
                              SYS_DEVICE, UCS_SYS_DEVICE_ID_UNKNOWN);

    unpacked = ucs_malloc(sizeof(*unpacked), "uct_cuda_ipc_unpacked_rkey_t");
    if (NULL == unpacked) {
        ucs_error("failed to allocate memory for uct_cuda_ipc_unpacked_rkey_t");
        status = UCS_ERR_NO_MEMORY;
        goto err;
    }

    unpacked->super.super      = *packed;
    unpacked->super.super.pid &= ~UCT_CUDA_IPC_RKEY_FLAG_PID_NS;

    /* Check if PID NS exists before using it (for wire compatibility) */
    if (packed->pid & UCT_CUDA_IPC_RKEY_FLAG_PID_NS) {
        ext_rkey               = (uct_cuda_ipc_extended_rkey_t*)packed;
        unpacked->super.pid_ns = ext_rkey->pid_ns;
    } else {
        unpacked->super.pid_ns = ucs_sys_get_default_ns(UCS_SYS_NS_TYPE_PID);
    }

    status = uct_cuda_ctx_primary_push_avail(0, sys_dev, &cuda_device,
                                             &avail_cuda_device,
                                             UCS_LOG_LEVEL_DEBUG);
    if (status != UCS_OK) {
        status = UCS_ERR_UNREACHABLE;
        goto err_free_key;
    }

    status = uct_cuda_ipc_is_peer_accessible(com, unpacked, avail_cuda_device);
    if (cuda_device != avail_cuda_device) {
        uct_cuda_ctx_primary_pop_and_release(avail_cuda_device);
    }
    if (status != UCS_OK) {
        goto err_free_key;
    }

    *handle_p = NULL;
    *rkey_p   = (uintptr_t) unpacked;
    return UCS_OK;

err_free_key:
    ucs_free(unpacked);
err:
    return status;
}

static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component,
                                              uct_rkey_t rkey, void *handle)
{
    ucs_assert(NULL == handle);
    ucs_free((void *)rkey);
    return UCS_OK;
}

static ucs_status_t
uct_cuda_ipc_mem_reg(uct_md_h md, void *address, size_t length,
                     const uct_md_mem_reg_params_t *params, uct_mem_h *memh_p)
{
    uct_cuda_ipc_memh_t *memh;

    memh = ucs_malloc(sizeof(*memh), "uct_cuda_ipc_memh_t");
    if (NULL == memh) {
        ucs_error("failed to allocate memory for uct_cuda_ipc_memh_t");
        return UCS_ERR_NO_MEMORY;
    }

    /* dev_num is initialized during pack in uct_cuda_ipc_mem_add_reg */
    memh->dev_num = CU_DEVICE_INVALID;
    memh->pid     = getpid();
    memh->pid_ns  = ucs_sys_get_ns(UCS_SYS_NS_TYPE_PID);
    ucs_list_head_init(&memh->list);

    *memh_p = memh;
    return UCS_OK;
}

static ucs_status_t
uct_cuda_ipc_mem_dereg(uct_md_h md, const uct_md_mem_dereg_params_t *params)
{
    uct_cuda_ipc_memh_t *memh = params->memh;
    uct_cuda_ipc_lkey_t *key, *tmp;

    UCT_MD_MEM_DEREG_CHECK_PARAMS(params, 0);

    ucs_list_for_each_safe(key, tmp, &memh->list, link) {
        if (key->ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_POSIX_FD) {
            close(key->ph.handle.posix_fd.fd);
        }
        ucs_free(key);
    }

    ucs_free(memh);
    return UCS_OK;
}

static int
uct_cuda_ipc_md_check_fabric_info(uct_cuda_ipc_md_t *md,
                                  ucs_ternary_auto_value_t mnnvl_enable)
{
#if !HAVE_NVML_FABRIC_INFO
    static int mnnvl_supported = 0;
#else
    static int mnnvl_supported = -1;
    nvmlGpuFabricInfoV_t fabric_info;
    nvmlDevice_t device;
    ucs_status_t status;
    char buf[64];
    int64_t *cluster_uuid;

    if (mnnvl_supported != -1) {
        goto out;
    }

    if (mnnvl_enable == UCS_NO) {
        mnnvl_supported = 0;
        goto out;
    }

    status = UCT_CUDA_NVML_WRAP_CALL(nvmlDeviceGetHandleByIndex, 0, &device);
    if (status != UCS_OK) {
        goto out_not_supported;
    }

    fabric_info.version = nvmlGpuFabricInfo_v2;
    status              = UCT_CUDA_NVML_WRAP_CALL(nvmlDeviceGetGpuFabricInfoV,
                                                  device, &fabric_info);
    if (status != UCS_OK) {
        goto out_not_supported;
    }

    ucs_debug("fabric_info: state=%u status=%u uuid=%s", fabric_info.state,
              fabric_info.status,
              ucs_str_dump_hex(fabric_info.clusterUuid,
                               NVML_GPU_FABRIC_UUID_LEN, buf, sizeof(buf),
                               SIZE_MAX));

    cluster_uuid = (int64_t*)fabric_info.clusterUuid;
    if ((fabric_info.state == NVML_GPU_FABRIC_STATE_COMPLETED) &&
        (fabric_info.status == NVML_SUCCESS) &&
        (cluster_uuid[0] | cluster_uuid[1]) != 0) {
        mnnvl_supported = 1;
        goto out;
    }

out_not_supported:
    mnnvl_supported = 0;
out:
#endif
    if ((mnnvl_enable == UCS_YES) && !mnnvl_supported) {
        ucs_error("multi-node NVLINK support is requested but not supported");
    } else {
        ucs_debug("multi-node NVLINK support is %s", mnnvl_supported ?
                  "enabled" : "disabled");
    }

    return mnnvl_supported;
}

static void uct_cuda_ipc_md_close(uct_md_h md)
{
    ucs_free(md);
}

static ucs_status_t
uct_cuda_ipc_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
                              uct_device_mem_elem_t *mem_elem_p)
{
    uct_cuda_ipc_unpacked_rkey_t *key = (uct_cuda_ipc_unpacked_rkey_t*)rkey;
    uct_cuda_ipc_md_device_mem_element_t *cuda_ipc_md_mem_element =
            (uct_cuda_ipc_md_device_mem_element_t*)mem_elem_p;
    ucs_status_t status;
    CUdevice cuda_device;
    void *mapped_addr;

    if (UCT_CUDADRV_FUNC_LOG_DEBUG(cuCtxGetDevice(&cuda_device)) != UCS_OK) {
        return UCS_ERR_UNREACHABLE;
    }

    status = uct_cuda_ipc_map_memhandle(&key->super, cuda_device, &mapped_addr,
                                        UCS_LOG_LEVEL_ERROR);
    if (ucs_unlikely(status != UCS_OK)) {
        return status;
    }

    cuda_ipc_md_mem_element->mapped_offset =
            UCS_PTR_BYTE_DIFF(key->super.super.d_bptr, mapped_addr);

    return UCS_OK;
}

static ucs_status_t
uct_cuda_ipc_md_open(uct_component_t *component, const char *md_name,
                     const uct_md_config_t *config, uct_md_h *md_p)
{
    static uct_md_ops_t md_ops = {
        .close              = uct_cuda_ipc_md_close,
        .query              = uct_cuda_ipc_md_query,
        .mem_alloc          = (uct_md_mem_alloc_func_t)ucs_empty_function_return_unsupported,
        .mem_free           = (uct_md_mem_free_func_t)ucs_empty_function_return_unsupported,
        .mem_advise         = (uct_md_mem_advise_func_t)ucs_empty_function_return_unsupported,
        .mem_reg            = uct_cuda_ipc_mem_reg,
        .mem_dereg          = uct_cuda_ipc_mem_dereg,
        .mem_query          = (uct_md_mem_query_func_t)ucs_empty_function_return_unsupported,
        .mkey_pack          = uct_cuda_ipc_mkey_pack,
        .mem_elem_pack      = uct_cuda_ipc_md_mem_elem_pack,
        .mem_attach         = (uct_md_mem_attach_func_t)ucs_empty_function_return_unsupported,
        .detect_memory_type = (uct_md_detect_memory_type_func_t)ucs_empty_function_return_unsupported
    };
    uct_cuda_ipc_md_config_t *ipc_config = ucs_derived_of(config,
                                                          uct_cuda_ipc_md_config_t);
    uct_cuda_ipc_md_t* md;

    md = ucs_calloc(1, sizeof(*md), "uct_cuda_ipc_md");
    if (md == NULL) {
        return UCS_ERR_NO_MEMORY;
    }

    md->super.ops       = &md_ops;
    md->super.component = &uct_cuda_ipc_component.super;
    md->enable_mnnvl    = uct_cuda_ipc_md_check_fabric_info(
                                                  md, ipc_config->enable_mnnvl);

    uct_cuda_ipc_cache_set_global_limits(ipc_config->cache_max_regions,
                                         ipc_config->cache_max_size);

    *md_p                 = &md->super;

    return UCS_OK;
}

ucs_status_t uct_cuda_ipc_rkey_ptr(uct_component_t *component, uct_rkey_t rkey,
                                   void *handle, uint64_t raddr, void **laddr_p)
{
    CUdevice cu_dev;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxGetDevice(&cu_dev));
    if (ucs_unlikely(status != UCS_OK)) {
        return status;
    }

    return uct_cuda_ipc_get_remote_address((uct_cuda_ipc_extended_rkey_t*)rkey,
                                           raddr, cu_dev, laddr_p, NULL);
}

uct_cuda_ipc_component_t uct_cuda_ipc_component = {
    .super = {
        .query_md_resources = uct_cuda_base_query_md_resources,
        .md_open            = uct_cuda_ipc_md_open,
        .cm_open            = (uct_component_cm_open_func_t)ucs_empty_function_return_unsupported,
        .rkey_unpack        = uct_cuda_ipc_rkey_unpack,
        .rkey_ptr           = uct_cuda_ipc_rkey_ptr,
        .rkey_release       = uct_cuda_ipc_rkey_release,
        .rkey_compare       = uct_base_rkey_compare,
        .name               = "cuda_ipc",
        .md_config          = {
            .name           = "Cuda-IPC memory domain",
            .prefix         = "CUDA_IPC_",
            .table          = uct_cuda_ipc_md_config_table,
            .size           = sizeof(uct_cuda_ipc_md_config_t),
        },
        .cm_config          = UCS_CONFIG_EMPTY_GLOBAL_LIST_ENTRY,
        .tl_list            = UCT_COMPONENT_TL_LIST_INITIALIZER(&uct_cuda_ipc_component.super),
        .flags              = 0,
        .md_vfs_init        =
                (uct_component_md_vfs_init_func_t)ucs_empty_function
    },
    .uuid_hash              = KHASH_STATIC_INITIALIZER,
    .lock                   = PTHREAD_MUTEX_INITIALIZER
};
UCT_COMPONENT_REGISTER(&uct_cuda_ipc_component.super);

UCS_STATIC_CLEANUP {
    uct_cuda_ipc_dev_cache_t *cache;

    kh_foreach_value(&uct_cuda_ipc_component.uuid_hash, cache, {
        ucs_free(cache);
    })
    kh_destroy_inplace(cuda_ipc_uuid_hash, &uct_cuda_ipc_component.uuid_hash);
    pthread_mutex_destroy(&uct_cuda_ipc_component.lock);
}
