UCT/ZE/COPY: Allocate DMA-BUF-exportable memory for RDMA

src/uct/ib/base/ib_md.c

@@ -2,6 +2,7 @@
  * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
  * Copyright (C) The University of Tennessee and The University
  *               of Tennessee Research Foundation. 2016. ALL RIGHTS RESERVED.
+ * Copyright (C) Intel Corporation, 2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -1449,10 +1450,6 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
         /* check if ROCM KFD driver is loaded */
         uct_ib_check_gpudirect_driver(md, "/dev/kfd", UCS_MEMORY_TYPE_ROCM);
 
-        /* Check if Intel Xe driver is loaded */
-        uct_ib_check_gpudirect_driver(md, "/sys/module/xe/srcversion",
-                                      UCS_MEMORY_TYPE_ZE_DEVICE);
-
         /* Check for HabanaLabs Gaudi DMABuf support */
         uct_ib_check_gpudirect_driver(md, "/dev/accel/accel0",
                                       UCS_MEMORY_TYPE_GAUDI);
@@ -1466,15 +1463,20 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
         !(md->cap_flags & UCT_MD_FLAG_REG_DMABUF) &&
         (md_config->enable_gpudirect_rdma == UCS_YES)) {
         ucs_error("%s: Couldn't enable GPUDirect RDMA. Please make sure "
-                  "nv_peer_mem, amdgpu plugin, or Intel Xe driver is "
-                  "installed correctly, or dmabuf is supported.",
+                  "nv_peer_mem or amdgpu plugin is installed correctly, or "
+                  "dmabuf is supported.",
                   uct_ib_device_name(&md->dev));
         status = UCS_ERR_UNSUPPORTED;
         goto err_cleanup_device;
     }
 
+    /* Limit the number of SGEs dumped by the packet tracer, so that it does not
+     * dereference non-host memory from the CPU. DMA-BUF registration makes such
+     * memory registrable without it being listed in 'reg_mem_types', so the
+     * limit also applies to host-only runs on a DMA-BUF capable device. */
     md->dev.max_zcopy_log_sge = INT_MAX;
-    if (md->reg_mem_types & ~UCS_BIT(UCS_MEMORY_TYPE_HOST)) {
+    if ((md->reg_mem_types & ~UCS_BIT(UCS_MEMORY_TYPE_HOST)) ||
+        (md->cap_flags & UCT_MD_FLAG_REG_DMABUF)) {
         md->dev.max_zcopy_log_sge = 1;
     }
 

src/uct/ze/copy/ze_copy_md.c

@@ -22,6 +22,11 @@
 #include <ucs/memory/memtype_cache.h>
 
 
+/* Memory types which can be exported as a DMA-BUF file descriptor */
+#define UCT_ZE_COPY_MD_DMABUF_MEM_TYPES \
+    (UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE))
+
+
 static ucs_config_field_t uct_ze_copy_md_config_table[] = {
     {"", "", NULL, ucs_offsetof(uct_ze_copy_md_config_t, super),
      UCS_CONFIG_TYPE_TABLE(uct_md_config_table)},
@@ -31,11 +36,44 @@ static ucs_config_field_t uct_ze_copy_md_config_table[] = {
      ucs_offsetof(uct_ze_copy_md_config_t, device_ordinal),
      UCS_CONFIG_TYPE_INT},
 
+    {"DMABUF", "try",
+     "Enable using cross-device dmabuf file descriptor.\n"
+     "Level Zero memory may have to be allocated with external DMA-BUF\n"
+     "export enabled. Without it, registration can succeed but an RDMA\n"
+     "device may access incorrect data, and no error is reported. Set to\n"
+     "'n' if the application allocates Level Zero memory without\n"
+     "requesting external export; device memory is then staged through\n"
+     "host memory for RDMA.",
+     ucs_offsetof(uct_ze_copy_md_config_t, enable_dmabuf),
+     UCS_CONFIG_TYPE_TERNARY},
+
     {NULL}
 };
 
+static int uct_ze_copy_md_is_dmabuf_supported(ze_device_handle_t ze_device)
+{
+    ze_device_external_memory_properties_t props = {
+        .stype = ZE_STRUCTURE_TYPE_DEVICE_EXTERNAL_MEMORY_PROPERTIES
+    };
+    int supported;
+
+    if (UCT_ZE_FUNC_LOG_DEBUG(
+                zeDeviceGetExternalMemoryProperties(ze_device, &props)) !=
+        UCS_OK) {
+        return 0;
+    }
+
+    supported = !!(props.memoryAllocationExportTypes &
+                   ZE_EXTERNAL_MEMORY_TYPE_FLAG_DMA_BUF);
+    ucs_debug("ze device %p: DMA-BUF export is%s supported", ze_device,
+              supported ? "" : " not");
+    return supported;
+}
+
 static ucs_status_t uct_ze_copy_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr)
 {
+    uct_ze_copy_md_t *ze_md = ucs_derived_of(md, uct_ze_copy_md_t);
+
     uct_md_base_md_query(md_attr);
     md_attr->flags            = UCT_MD_FLAG_REG | UCT_MD_FLAG_ALLOC;
     md_attr->reg_mem_types    = UCS_BIT(UCS_MEMORY_TYPE_HOST) |
@@ -55,8 +93,8 @@ static ucs_status_t uct_ze_copy_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr)
     md_attr->detect_mem_types = UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) |
                                 UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE) |
                                 UCS_BIT(UCS_MEMORY_TYPE_ZE_MANAGED);
-    md_attr->dmabuf_mem_types = UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) |
-                                UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE);
+    md_attr->dmabuf_mem_types = ze_md->dmabuf_supported ?
+                                UCT_ZE_COPY_MD_DMABUF_MEM_TYPES : 0;
     md_attr->max_alloc        = SIZE_MAX;
     return UCS_OK;
 }
@@ -66,23 +104,41 @@ uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p,
                       ucs_memory_type_t mem_type, ucs_sys_device_t sys_dev,
                       unsigned flags, const char *alloc_name, uct_mem_h *memh_p)
 {
-    uct_ze_copy_md_t *md                = ucs_derived_of(tl_md, uct_ze_copy_md_t);
-    ze_host_mem_alloc_desc_t host_desc  = {
+    uct_ze_copy_md_t *md = ucs_derived_of(tl_md, uct_ze_copy_md_t);
+    /* Request export only for 'dmabuf_mem_types', so that other types cannot
+     * fail because of it. A DMA-BUF of memory not allocated for export may not
+     * describe the requested range, and an RDMA device may read incorrect data
+     * through it. */
+    ze_external_memory_export_desc_t export_desc = {
+        .stype = ZE_STRUCTURE_TYPE_EXTERNAL_MEMORY_EXPORT_DESC,
+        .flags = ZE_EXTERNAL_MEMORY_TYPE_FLAG_DMA_BUF
+    };
+    void *export_desc_p                          = NULL;
+    ze_host_mem_alloc_desc_t host_desc           = {
         .stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC
     };
-    ze_device_mem_alloc_desc_t dev_desc = {
+    ze_device_mem_alloc_desc_t dev_desc          = {
         .stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC
     };
-    size_t alignment                    = ucs_get_page_size();
+    size_t alignment                             = ucs_get_page_size();
     ucs_status_t status;
 
+    if (md->dmabuf_supported &&
+        (UCS_BIT(mem_type) & UCT_ZE_COPY_MD_DMABUF_MEM_TYPES)) {
+        export_desc_p = &export_desc;
+    }
+
     switch (mem_type) {
     case UCS_MEMORY_TYPE_ZE_HOST:
+        host_desc.pNext = export_desc_p;
+
         status = UCT_ZE_FUNC_LOG_ERR(zeMemAllocHost(md->ze_context, &host_desc,
                                                     *length_p, alignment,
                                                     address_p));
         break;
     case UCS_MEMORY_TYPE_ZE_DEVICE:
+        dev_desc.pNext = export_desc_p;
+
         status = UCT_ZE_FUNC_LOG_ERR(zeMemAllocDevice(md->ze_context, &dev_desc,
                                                       *length_p, alignment,
                                                       md->ze_device,
@@ -198,14 +254,15 @@ static ucs_status_t uct_ze_copy_md_mem_query(uct_md_h md, const void *addr,
                                              const size_t length,
                                              uct_md_mem_attr_v2_t *mem_attr_p)
 {
-    int dmabuf_fd    = UCT_DMABUF_FD_INVALID;
-    int *dmabuf_fd_p = NULL;
+    uct_ze_copy_md_t *ze_md = ucs_derived_of(md, uct_ze_copy_md_t);
+    int dmabuf_fd           = UCT_DMABUF_FD_INVALID;
+    int *dmabuf_fd_p        = NULL;
     ucs_memory_info_t mem_info;
     ucs_status_t status;
 
-    if (mem_attr_p->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) {
-        mem_attr_p->dmabuf_fd = UCT_DMABUF_FD_INVALID;
-        dmabuf_fd_p           = &dmabuf_fd;
+    if ((mem_attr_p->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) &&
+        ze_md->dmabuf_supported) {
+        dmabuf_fd_p = &dmabuf_fd;
     }
 
     status = uct_ze_copy_md_query_attributes(md, addr, length, &mem_info,
@@ -302,6 +359,8 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name,
     ze_context_desc_t context_desc = {};
     ze_result_t ret;
     const uct_ze_subdevice_t *subdevice;
+    int dmabuf_supported;
+    ucs_status_t status;
 
     ze_driver = uct_ze_base_get_driver();
     if (ze_driver == NULL) {
@@ -318,33 +377,47 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name,
     subdevice = uct_ze_base_get_subdevice_by_global_id(config->device_ordinal);
     if (subdevice == NULL) {
         ucs_error("Failed to get sub-device at ordinal %d", config->device_ordinal);
-        ucs_free(md);
-        return UCS_ERR_NO_DEVICE;
+        status = UCS_ERR_NO_DEVICE;
+        goto err_free_md;
     }
 
     /* Get the actual device handle from the sub-device */
     md->ze_device = uct_ze_base_get_device_handle_from_subdevice(subdevice);
     if (md->ze_device == NULL) {
         ucs_error("Failed to get device handle for sub-device %d", config->device_ordinal);
-        ucs_free(md);
-        return UCS_ERR_NO_DEVICE;
+        status = UCS_ERR_NO_DEVICE;
+        goto err_free_md;
+    }
+
+    dmabuf_supported = uct_ze_copy_md_is_dmabuf_supported(md->ze_device);
+    if ((config->enable_dmabuf == UCS_YES) && !dmabuf_supported) {
+        ucs_error("%s: DMA-BUF is not supported by the Level Zero device",
+                  md_name);
+        status = UCS_ERR_UNSUPPORTED;
+        goto err_free_md;
     }
 
     ret = zeContextCreate(ze_driver, &context_desc, &md->ze_context);
     if (ret != ZE_RESULT_SUCCESS) {
         ucs_error("zeContextCreate failed with error %x", ret);
-        ucs_free(md);
-        return UCS_ERR_NO_DEVICE;
+        status = UCS_ERR_NO_DEVICE;
+        goto err_free_md;
     }
 
-    md->super.ops       = &md_ops;
-    md->super.component = &uct_ze_copy_component;
+    md->dmabuf_supported = (config->enable_dmabuf != UCS_NO) &&
+                           dmabuf_supported;
+    md->super.ops        = &md_ops;
+    md->super.component  = &uct_ze_copy_component;
 
     *md_p = (uct_md_h)md;
 
     /* Install the memtype cache handler before Level Zero allocations. */
     ucs_memtype_cache_global_create();
     return UCS_OK;
+
+err_free_md:
+    ucs_free(md);
+    return status;
 }
 
 uct_component_t uct_ze_copy_component = {

src/uct/ze/copy/ze_copy_md.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) Intel Corporation, 2023-2024. ALL RIGHTS RESERVED.
+ * Copyright (C) Intel Corporation, 2023-2026. ALL RIGHTS RESERVED.
  * See file LICENSE for terms.
  */
 
@@ -21,15 +21,17 @@ typedef struct uct_ze_copy_md {
     uct_md_t            super; /**< Domain info */
     ze_context_handle_t ze_context;
     ze_device_handle_t  ze_device;
+    int                 dmabuf_supported; /**< Expose memory as DMA-BUF */
 } uct_ze_copy_md_t;
 
 
 /**
  * ze copy domain configuration.
  */
 typedef struct uct_ze_copy_md_config {
-    uct_md_config_t super;
-    int             device_ordinal;
+    uct_md_config_t          super;
+    int                      device_ordinal;
+    ucs_ternary_auto_value_t enable_dmabuf;
 } uct_ze_copy_md_config_t;
 
 

test/gtest/common/mem_buffer.cc

@@ -416,6 +416,26 @@ bool mem_buffer::cuda_mem_is_gdr_capable()
 #endif
 }
 
+ucs_status_t mem_buffer::query_dmabuf_export(uct_md_h md,
+                                             ucs_memory_type_t mem_type,
+                                             bool *exported_p)
+{
+    uct_md_mem_attr_t mem_attr = {};
+    mem_buffer buffer(ucs_get_page_size(), mem_type);
+    ucs_status_t status;
+
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_DMABUF_FD;
+    status = uct_md_mem_query(md, buffer.ptr(), buffer.size(), &mem_attr);
+    if (status == UCS_OK) {
+        *exported_p = (mem_attr.dmabuf_fd != UCT_DMABUF_FD_INVALID);
+        if (*exported_p) {
+            close(mem_attr.dmabuf_fd);
+        }
+    }
+
+    return status;
+}
+
 #if HAVE_ZE
 namespace {
 

test/gtest/common/mem_buffer.h

@@ -1,12 +1,14 @@
 /**
  * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
+ * Copyright (C) Intel Corporation, 2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
 
 #ifndef GTEST_MEM_BUFFER_H_
 #define GTEST_MEM_BUFFER_H_
 
+#include <uct/api/uct.h>
 #include <ucs/memory/memory_type.h>
 #include <ucs/sys/math.h>
 #include <stdint.h>
@@ -117,6 +119,12 @@ class mem_buffer {
      * which it does not under an MPS MLOPart device */
     static bool cuda_mem_is_gdr_capable();
 
+    /* Query 'md' for a DMA-BUF of a page of 'mem_type' memory allocated by
+     * mem_buffer, and set 'exported_p' to whether it returned one */
+    static ucs_status_t query_dmabuf_export(uct_md_h md,
+                                            ucs_memory_type_t mem_type,
+                                            bool *exported_p);
+
     /* Return free memory on the BAR1 / GPU. If GPU is not used
      * SIZE_MAX is returned */
     static size_t get_bar1_free_size()