UCP/UCT/UCS: Add memory registration flags

src/ucp/core/ucp_context.c

@@ -2146,6 +2146,7 @@ static ucs_status_t ucp_fill_resources(ucp_context_h context,
         context->dmabuf_mds[mem_type]           = UCP_NULL_RESOURCE;
         context->alloc_md[mem_type].md_index    = UCP_NULL_RESOURCE;
         context->alloc_md[mem_type].sys_dev     = UCS_SYS_DEVICE_ID_UNKNOWN;
+        context->alloc_md[mem_type].mem_flags   = UCS_MEM_FLAG_REGISTRABLE;
         context->alloc_md[mem_type].initialized = 0;
     }
 
@@ -2938,19 +2939,20 @@ uct_md_h ucp_context_find_tl_md(ucp_context_h context, const char *md_name)
 void ucp_memory_detect_slowpath(ucp_context_h context, const void *address,
                                 size_t length, ucs_memory_info_t *mem_info)
 {
-    uct_md_mem_attr_t mem_attr;
+    uct_md_mem_attr_v2_t mem_attr;
     ucs_status_t status;
     ucp_tl_md_t *tl_md;
     ucp_md_index_t i;
 
-    mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_MEM_TYPE |
-                          UCT_MD_MEM_ATTR_FIELD_BASE_ADDRESS |
-                          UCT_MD_MEM_ATTR_FIELD_ALLOC_LENGTH |
-                          UCT_MD_MEM_ATTR_FIELD_SYS_DEV;
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
+                          UCT_MD_MEM_ATTR_V2_FIELD_BASE_ADDRESS |
+                          UCT_MD_MEM_ATTR_V2_FIELD_ALLOC_LENGTH |
+                          UCT_MD_MEM_ATTR_V2_FIELD_SYS_DEV |
+                          UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
 
     for (i = 0; i < context->num_mem_type_detect_mds; ++i) {
         tl_md  = &context->tl_mds[context->mem_type_detect_mds[i]];
-        status = uct_md_mem_query(tl_md->md, address, length, &mem_attr);
+        status = uct_md_mem_query_v2(tl_md->md, address, length, &mem_attr);
         if (status != UCS_OK) {
             continue;
         }
@@ -2963,6 +2965,7 @@ void ucp_memory_detect_slowpath(ucp_context_h context, const void *address,
         mem_info->sys_dev      = mem_attr.sys_dev;
         mem_info->base_address = mem_attr.base_address;
         mem_info->alloc_length = mem_attr.alloc_length;
+        mem_info->mem_flags    = mem_attr.mem_flags;
         return;
     }
 

src/ucp/core/ucp_context.h

@@ -389,6 +389,7 @@ typedef struct ucp_context_alloc_md_index {
      * using ucp_memh_alloc(). */
     ucp_md_index_t   md_index;
     ucs_sys_device_t sys_dev;
+    uint8_t          mem_flags;
 } ucp_context_alloc_md_index_t;
 
 
@@ -712,6 +713,7 @@ ucp_memory_info_set_host(ucp_memory_info_t *mem_info)
 {
     mem_info->type    = UCS_MEMORY_TYPE_HOST;
     mem_info->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
+    mem_info->flags   = UCS_MEM_FLAG_REGISTRABLE;
 }
 
 static UCS_F_ALWAYS_INLINE void
@@ -771,6 +773,7 @@ ucp_memory_detect(ucp_context_h context, const void *address, size_t length,
 
     mem_info->type    = mem_info_internal.type;
     mem_info->sys_dev = mem_info_internal.sys_dev;
+    mem_info->flags   = mem_info_internal.mem_flags;
 }
 
 static UCS_F_ALWAYS_INLINE int

src/ucp/core/ucp_mm.c

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2015. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */
@@ -28,6 +28,7 @@
 /* Context for rcache memory registration callback */
 typedef struct {
     ucs_memory_type_t mem_type;    /* Memory type */
+    uint8_t           mem_flags;   /* Memory flags (ucp_mem_flags_t) */
     ucp_md_map_t      reg_md_map;  /* Map of memory domains to be registered */
     unsigned          uct_flags;   /* UCT memory registration flags */
     const char        *alloc_name; /* Memory allocation name */
@@ -40,6 +41,7 @@ ucp_mem_dummy_handle_t ucp_mem_dummy_handle = {
         .parent         = &ucp_mem_dummy_handle.memh,
         .mem_type       = UCS_MEMORY_TYPE_HOST,
         .sys_dev        = UCS_SYS_DEVICE_ID_UNKNOWN,
+        .mem_flags      = UCS_MEM_FLAG_REGISTRABLE,
         .md_map         = 0,
         .inv_md_map     = 0,
         .reg_id         = 0,
@@ -49,7 +51,8 @@ ucp_mem_dummy_handle_t ucp_mem_dummy_handle = {
 
 const ucp_memory_info_t ucp_mem_info_unknown = {
     .type    = UCS_MEMORY_TYPE_UNKNOWN,
-    .sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN
+    .sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN,
+    .flags   = UCS_MEM_FLAG_REGISTRABLE
 };
 
 
@@ -561,6 +564,7 @@ ucp_memh_register_internal(ucp_context_h context, ucp_mem_h memh,
     uct_md_mem_reg_params_t reg_params;
     uct_md_mem_attr_t mem_attr;
     ucp_md_index_t md_index;
+    uint8_t required_mem_flags;
     ucs_status_t status;
     void *reg_address;
     size_t reg_length;
@@ -575,6 +579,18 @@ ucp_memh_register_internal(ucp_context_h context, ucp_mem_h memh,
     }
 
     reg_md_map = ~memh->md_map & md_map;
+    ucs_for_each_bit(md_index, reg_md_map) {
+        required_mem_flags = context->tl_mds[md_index].attr.required_mem_flags;
+
+        if (!ucs_test_all_flags(memh->mem_flags, required_mem_flags)) {
+            ucs_trace("md[%d]=%s skipped: mem_flags=0x%x do not include "
+                      "required_mem_flags=0x%x",
+                      md_index, context->tl_mds[md_index].rsc.md_name,
+                      memh->mem_flags, required_mem_flags);
+            reg_md_map &= ~UCS_BIT(md_index);
+        }
+    }
+
     if (reg_md_map == 0) {
         return UCS_OK;
     }
@@ -715,7 +731,7 @@ void ucp_memh_disable_gva(ucp_mem_h memh, ucp_md_map_t md_map)
 static void ucp_memh_init(ucp_mem_h memh, ucp_context_h context,
                           uint8_t memh_flags, unsigned uct_flags,
                           uct_alloc_method_t method, ucs_memory_type_t mem_type,
-                          ucs_sys_device_t sys_dev)
+                          ucs_sys_device_t sys_dev, uint8_t mem_flags)
 {
     memh->md_map         = 0;
     memh->inv_md_map     = 0;
@@ -726,6 +742,7 @@ static void ucp_memh_init(ucp_mem_h memh, ucp_context_h context,
     memh->alloc_method   = method;
     memh->mem_type       = mem_type;
     memh->sys_dev        = sys_dev;
+    memh->mem_flags      = mem_flags;
 }
 
 static ucs_status_t
@@ -745,7 +762,7 @@ ucp_memh_create(ucp_context_h context, void *address, size_t length,
     memh->super.super.end   = (uintptr_t)address + length;
     ucp_memory_detect(context, address, length, &info);
     ucp_memh_init(memh, context, memh_flags, uct_flags, method, mem_type,
-                  info.sys_dev);
+                  info.sys_dev, info.flags);
 
     *memh_p = memh;
     return UCS_OK;
@@ -754,11 +771,13 @@ ucp_memh_create(ucp_context_h context, void *address, size_t length,
 static ucs_status_t
 ucp_memh_rcache_get(ucs_rcache_t *rcache, void *address, size_t length,
                     size_t alignment, ucs_memory_type_t mem_type,
-                    ucp_md_map_t reg_md_map, unsigned uct_flags,
-                    const char *alloc_name, ucp_mem_h *memh_p)
+                    uint8_t mem_flags, ucp_md_map_t reg_md_map,
+                    unsigned uct_flags, const char *alloc_name,
+                    ucp_mem_h *memh_p)
 {
     ucp_mem_rcache_reg_ctx_t reg_ctx = {
         .mem_type   = mem_type,
+        .mem_flags  = mem_flags,
         .reg_md_map = reg_md_map,
         .uct_flags  = uct_flags,
         .alloc_name = alloc_name
@@ -953,7 +972,7 @@ static size_t ucp_memh_reg_align(ucp_context_h context, ucp_md_map_t reg_md_map)
 
 static ucs_status_t
 ucp_memh_find_slow(ucp_context_h context, void *address, size_t length,
-                   size_t align, ucs_memory_type_t mem_type,
+                   size_t align, ucs_memory_type_t mem_type, uint8_t mem_flags,
                    ucp_md_map_t reg_md_map, unsigned uct_flags,
                    const char *alloc_name, ucp_mem_h *memh_p)
 {
@@ -968,7 +987,7 @@ ucp_memh_find_slow(ucp_context_h context, void *address, size_t length,
 
     for (;;) {
         status = ucp_memh_rcache_get(context->rcache, address, length, align,
-                                     mem_type, reg_md_map, uct_flags,
+                                     mem_type, mem_flags, reg_md_map, uct_flags,
                                      alloc_name, &memh);
         if (status != UCS_OK) {
             return status;
@@ -1012,8 +1031,9 @@ ucs_status_t ucp_memh_get_slow(ucp_context_h context, void *address,
     size_t reg_length;
     ucp_mem_h memh;
 
+    ucp_memory_detect_internal(context, address, length, &mem_info);
+
     if (context->config.ext.reg_whole_alloc_bitmap & UCS_BIT(mem_type)) {
-        ucp_memory_detect_internal(context, address, length, &mem_info);
         reg_address = mem_info.base_address;
         reg_length  = mem_info.alloc_length;
     } else {
@@ -1023,8 +1043,8 @@ ucs_status_t ucp_memh_get_slow(ucp_context_h context, void *address,
 
     UCP_THREAD_CS_ENTER(&context->mt_lock);
     status = ucp_memh_find_slow(context, reg_address, reg_length, reg_align,
-                                mem_type, reg_md_map, uct_flags, alloc_name,
-                                &memh);
+                                mem_type, mem_info.mem_flags, reg_md_map,
+                                uct_flags, alloc_name, &memh);
     if (status != UCS_OK) {
         goto out;
     }
@@ -1680,7 +1700,8 @@ ucp_mem_rcache_mem_reg_cb(void *ctx, ucs_rcache_t *rcache, void *arg,
     ucp_mem_h memh                    = ucs_derived_of(rregion, ucp_mem_t);
 
     ucp_memh_init(memh, context, 0, reg_ctx->uct_flags, UCT_ALLOC_METHOD_LAST,
-                  reg_ctx->mem_type, UCS_SYS_DEVICE_ID_UNKNOWN);
+                  reg_ctx->mem_type, UCS_SYS_DEVICE_ID_UNKNOWN,
+                  reg_ctx->mem_flags);
     memh->reg_id = context->next_memh_reg_id++;
 
     if (rcache_mem_reg_flags & UCS_RCACHE_MEM_REG_HIDE_ERRORS) {
@@ -1713,8 +1734,11 @@ static void ucp_mem_rcache_merge_cb(void *ctx, ucs_rcache_t *rcache, void *arg,
     ucp_mem_h memh                    = ucs_derived_of(rregion, ucp_mem_t);
 
     ucs_log_indent(+1);
-    ucs_trace("merge with memh %p uct_flags 0x%x", memh, memh->uct_flags);
+    ucs_trace("merge with memh %p uct_flags 0x%x mem_flags 0x%x/0x%x", memh,
+              memh->uct_flags, reg_ctx->mem_flags, memh->mem_flags);
     reg_ctx->uct_flags |= memh->uct_flags;
+    /* TODO: The regions shouldn't be merged if they have different mem_flags */
+    reg_ctx->mem_flags &= memh->mem_flags;
     memh->flags        &= ~UCP_MEMH_FLAG_MLOCKED;
     ucs_log_indent(-1);
 }
@@ -1849,14 +1873,16 @@ ucs_status_t ucp_mm_get_alloc_md_index(ucp_context_h context,
 
         context->alloc_md[alloc_mem_type].initialized = 1;
         context->alloc_md[alloc_mem_type].sys_dev     = mem_info.sys_dev;
+        context->alloc_md[alloc_mem_type].mem_flags   = mem_info.flags;
         context->alloc_md[alloc_mem_type].md_index    =
                 ucp_mem_get_md_index(context, mem.md, mem.method);
 
         uct_mem_free(&mem);
     }
 
-    *md_idx_p        = context->alloc_md[alloc_mem_type].md_index;
-    mem_info_p->type = alloc_mem_type;
+    *md_idx_p             = context->alloc_md[alloc_mem_type].md_index;
+    mem_info_p->type      = alloc_mem_type;
+    mem_info_p->flags     = context->alloc_md[alloc_mem_type].mem_flags;
 
     /* TODO: Extend the cache by alloc sys dev */
     if (UCP_MEM_IS_HOST(alloc_mem_type) ||
@@ -1973,7 +1999,7 @@ ucp_memh_import_slow(ucp_context_h context, ucs_rcache_t *existing_rcache,
 
         status = ucp_memh_rcache_get(rcache, unpacked->address,
                                      unpacked->length, UCS_RCACHE_MIN_ALIGNMENT,
-                                     unpacked->mem_type, 0, 0, "", &memh);
+                                     unpacked->mem_type, 0, 0, 0, "", &memh);
         if (status != UCS_OK) {
             goto err_rcache_destroy;
         }

src/ucp/core/ucp_mm.h

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2015. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
  * Copyright (C) Intel Corporation, 2023. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
@@ -76,6 +76,7 @@ typedef struct ucp_mem {
     uct_alloc_method_t  alloc_method;   /* Method used to allocate the memory */
     ucs_sys_device_t    sys_dev;        /* System device index */
     ucs_memory_type_t   mem_type;       /* Type of allocated or registered memory */
+    uint8_t             mem_flags;      /* UCS memory flags */
     ucp_md_index_t      alloc_md_index; /* Index of MD used to allocate the memory */
     uint64_t            remote_uuid;    /* Remote UUID */
     ucp_md_map_t        md_map;         /* Which MDs have valid memory handles */

src/ucp/core/ucp_mm.inl

@@ -120,6 +120,7 @@ ucp_memory_info_from_memh(ucp_mem_h memh)
 
     mem_info.type    = memh->mem_type;
     mem_info.sys_dev = memh->sys_dev;
+    mem_info.flags   = memh->mem_flags;
     return mem_info;
 }
 
@@ -179,8 +180,11 @@ ucp_memh_get_or_update(ucp_context_h context, void *address, size_t length,
     }
 
     ucs_assert((*memh_p)->parent == NULL);
-    ucs_assert((context->rcache == NULL) ||
-               ucs_test_all_flags(context->cache_md_map[mem_type], md_map));
+    ucs_assertv((context->rcache == NULL) ||
+                ucs_test_all_flags(context->cache_md_map[mem_type], md_map),
+                "mem_type=%s cache_md_map=0x%" PRIx64 " md_map=0x%" PRIx64,
+                ucs_memory_type_names[mem_type],
+                context->cache_md_map[mem_type], md_map);
 
     status = ucp_memh_register(context, *memh_p, md_map, uct_flags, alloc_name);
 out:

src/ucp/dt/dt.c

@@ -49,12 +49,13 @@ ucs_status_t ucp_dt_mem_info_verify(const char *dt_name, size_t index,
         return UCS_OK;
     }
 
-    ucs_error("inconsistent %s mem_info: [%zu]=%s-%s [0]=%s-%s count=%zu",
-              dt_name, index,
-              ucs_memory_type_names[cur->type],
-              ucs_topo_sys_device_get_name(cur->sys_dev),
+    ucs_error("inconsistent %s mem_info: [%zu]=%s-%s flags=0x%x "
+              "[0]=%s-%s flags=0x%x count=%zu",
+              dt_name, index, ucs_memory_type_names[cur->type],
+              ucs_topo_sys_device_get_name(cur->sys_dev), cur->flags,
               ucs_memory_type_names[ref->type],
-              ucs_topo_sys_device_get_name(ref->sys_dev), count);
+              ucs_topo_sys_device_get_name(ref->sys_dev), ref->flags,
+              count);
     return UCS_ERR_INVALID_PARAM;
 }
 

src/ucp/dt/dt.h

@@ -50,6 +50,7 @@ typedef struct ucp_dt_state {
 typedef struct {
     uint8_t          type;    /**< Memory type, use uint8 for compact size */
     ucs_sys_device_t sys_dev; /**< System device index */
+    uint8_t          flags;   /**< UCS memory flags */
 } ucp_memory_info_t;
 
 
@@ -73,7 +74,8 @@ ucs_status_t ucp_dt_mem_info_verify(const char *dt_name, size_t index,
 static UCS_F_ALWAYS_INLINE int
 ucp_memory_info_equal(const ucp_memory_info_t *a, const ucp_memory_info_t *b)
 {
-    return (a->type == b->type) && (a->sys_dev == b->sys_dev);
+    return (a->type == b->type) && (a->sys_dev == b->sys_dev) &&
+           (a->flags == b->flags);
 }
 
 void ucp_mem_type_pack(ucp_worker_h worker, void *dest, const void *src,

src/ucp/proto/proto_common.c

@@ -67,8 +67,9 @@ ucp_memory_info_t ucp_proto_common_select_param_mem_info(
                                    const ucp_proto_select_param_t *select_param)
 {
     ucp_memory_info_t mem_info = {
-        .type = select_param->mem_type,
-        .sys_dev = select_param->sys_dev
+        .type    = select_param->mem_type,
+        .sys_dev = select_param->sys_dev,
+        .flags   = select_param->op.mem_flags
     };
 
     return mem_info;
@@ -506,6 +507,16 @@ ucp_proto_common_filter_min_frag(const ucp_proto_init_params_t *params,
                     ucs_memory_type_names[params->select_param->mem_type]);
 
         if (md_attr->flags & UCT_MD_FLAG_NEED_MEMH) {
+            if (!ucs_test_all_flags(common_params->reg_mem_info.flags,
+                                    md_attr->required_mem_flags)) {
+                ucs_trace("%s: md %s missing required_mem_flags=0x%x for "
+                          "mem_flags=0x%x",
+                          lane_desc, context->tl_mds[md_index].rsc.md_name,
+                          md_attr->required_mem_flags,
+                          common_params->reg_mem_info.flags);
+                return 0;
+            }
+
             /* Memory domain must support registration on the relevant memory
              * type */
             if (!(context->reg_md_map[reg_mem_type] & UCS_BIT(md_index))) {
@@ -749,7 +760,9 @@ ucp_proto_common_reg_md_map(const ucp_proto_common_init_params_t *params,
            memory type, and needs a local memory handle for zero-copy
            communication */
         if ((md_attr->flags & UCT_MD_FLAG_NEED_MEMH) &&
-            (context->reg_md_map[select_param->mem_type] & UCS_BIT(md_index))) {
+            (context->reg_md_map[select_param->mem_type] & UCS_BIT(md_index)) &&
+            ucs_test_all_flags(select_param->op.mem_flags,
+                               md_attr->required_mem_flags)) {
             reg_md_map |= UCS_BIT(md_index);
         }
     }

src/ucp/proto/proto_select.h

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020-2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -90,8 +90,11 @@ struct ucp_proto_select_param {
             uint8_t         sys_dev;    /* Reply buffer system device */
         } UCS_S_PACKED reply;
 
-        /* Align struct size to uint64_t */
-        uint8_t             padding[2];
+        /* UCS memory flags of the operation buffer. Used for all operations
+         * except UCP_OP_ID_AMO_FETCH and UCP_OP_ID_AMO_CSWAP, which use the
+         * 'reply' member above instead. */
+        uint8_t             mem_flags;
+        uint8_t             padding[2]; /* Pad select key to sizeof(uint64_t) */
 
     } UCS_S_PACKED op;
 } UCS_S_PACKED;