UCS/RCACHE: Adds capability for merging adjacent rcache regions

src/ucp/core/ucp_mm.c

@@ -29,6 +29,9 @@
 typedef struct {
     ucs_memory_type_t mem_type;    /* Memory type */
     uint8_t           mem_flags;   /* Memory flags (ucp_mem_flags_t) */
+    ucs_sys_device_t  sys_dev;     /* Device index */
+    void              *alloc_base; /* Start of the underlying allocation */
+    size_t            alloc_len;   /* Length of the underlying allocation */
     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 */
@@ -787,13 +790,16 @@ 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,
-                    uint8_t mem_flags, ucp_md_map_t reg_md_map,
+                    const ucs_memory_info_t *mem_info, 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,
+        .mem_flags  = mem_info->mem_flags,
+        .sys_dev    = mem_info->sys_dev,
+        .alloc_base = mem_info->base_address,
+        .alloc_len  = mem_info->alloc_length,
         .reg_md_map = reg_md_map,
         .uct_flags  = uct_flags,
         .alloc_name = alloc_name
@@ -988,9 +994,10 @@ 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, uint8_t mem_flags,
-                   ucp_md_map_t reg_md_map, unsigned uct_flags,
-                   const char *alloc_name, ucp_mem_h *memh_p)
+                   size_t align, ucs_memory_type_t mem_type,
+                   const ucs_memory_info_t *mem_info, ucp_md_map_t reg_md_map,
+                   unsigned uct_flags, const char *alloc_name,
+                   ucp_mem_h *memh_p)
 {
     unsigned access_flags = UCP_MM_UCT_ACCESS_FLAGS(uct_flags);
     ucs_status_t status;
@@ -1003,7 +1010,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, mem_flags, reg_md_map, uct_flags,
+                                     mem_type, mem_info, reg_md_map, uct_flags,
                                      alloc_name, &memh);
         if (status != UCS_OK) {
             return status;
@@ -1059,8 +1066,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, mem_info.mem_flags, reg_md_map,
-                                uct_flags, alloc_name, &memh);
+                                mem_type, &mem_info, reg_md_map, uct_flags,
+                                alloc_name, &memh);
     if (status != UCS_OK) {
         goto out;
     }
@@ -1724,8 +1731,7 @@ 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_flags);
+                  reg_ctx->mem_type, reg_ctx->sys_dev, reg_ctx->mem_flags);
     memh->reg_id = context->next_memh_reg_id++;
 
     if (rcache_mem_reg_flags & UCS_RCACHE_MEM_REG_HIDE_ERRORS) {
@@ -1793,11 +1799,44 @@ static void ucp_mem_rcache_dump_region_cb(void *ctx, ucs_rcache_t *rcache,
     ucs_string_buffer_rtrim(&strb, NULL);
 }
 
+static int ucp_mem_rcache_can_merge_cb(void *arg, ucs_rcache_region_t *rregion)
+{
+    ucp_mem_rcache_reg_ctx_t *reg_ctx = arg;
+    ucp_mem_h memh                    = ucs_derived_of(rregion, ucp_mem_t);
+    ucs_pgt_addr_t start              = rregion->super.start;
+    ucs_pgt_addr_t end                = rregion->super.end;
+    ucs_pgt_addr_t alloc_start, alloc_end;
+
+    if ((memh->mem_flags != reg_ctx->mem_flags) ||
+        (memh->mem_type != reg_ctx->mem_type)) {
+        return 0;
+    }
+
+    if (memh->mem_type == UCS_MEMORY_TYPE_HOST) {
+        return 1;
+    }
+
+    if ((memh->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) ||
+        (reg_ctx->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) ||
+        (memh->sys_dev != reg_ctx->sys_dev)) {
+        return 0;
+    }
+
+    if (reg_ctx->alloc_base == NULL) {
+        return 0;
+    }
+
+    alloc_start = (uintptr_t)reg_ctx->alloc_base;
+    alloc_end   = (uintptr_t)reg_ctx->alloc_base + reg_ctx->alloc_len;
+    return (start >= alloc_start) && (end <= alloc_end);
+}
+
 static ucs_rcache_ops_t ucp_mem_rcache_ops = {
     .mem_reg     = ucp_mem_rcache_mem_reg_cb,
     .mem_dereg   = ucp_mem_rcache_mem_dereg_cb,
     .merge       = ucp_mem_rcache_merge_cb,
-    .dump_region = ucp_mem_rcache_dump_region_cb
+    .dump_region = ucp_mem_rcache_dump_region_cb,
+    .can_merge   = ucp_mem_rcache_can_merge_cb
 };
 
 static ucs_status_t
@@ -1984,6 +2023,13 @@ ucp_memh_import_slow(ucp_context_h context, ucs_rcache_t *existing_rcache,
                      ucp_mem_h user_memh,
                      ucp_unpacked_exported_memh_t *unpacked)
 {
+    ucs_memory_info_t mem_info = {
+        .type         = UCS_MEMORY_TYPE_HOST,
+        .sys_dev      = UCS_SYS_DEVICE_ID_UNKNOWN,
+        .base_address = NULL,
+        .alloc_length = -1,
+        .mem_flags    = 0
+    };
     ucs_rcache_t *rcache;
     ucs_rcache_params_t rcache_params;
     ucs_status_t status;
@@ -2023,7 +2069,8 @@ 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, 0, "", &memh);
+                                     unpacked->mem_type, &mem_info, 0, 0, "",
+                                     &memh);
         if (status != UCS_OK) {
             goto err_rcache_destroy;
         }

src/ucs/memory/rcache.c

@@ -114,6 +114,9 @@ ucs_config_field_t ucs_config_rcache_table[] = {
      "Purge registration cache upon fork",
      ucs_offsetof(ucs_rcache_config_t, purge_on_fork), UCS_CONFIG_TYPE_BOOL},
 
+    {"RCACHE_MERGE_ADJACENT", "n", "Merge adjacent registration cache regions",
+     ucs_offsetof(ucs_rcache_config_t, merge_adjacent), UCS_CONFIG_TYPE_BOOL},
+
     {NULL}
 };
 
@@ -191,8 +194,14 @@ void ucs_rcache_set_params(ucs_rcache_params_t *rcache_params,
     rcache_params->max_regions        = rcache_config->max_regions;
     rcache_params->max_size           = rcache_config->max_size;
     rcache_params->max_unreleased     = rcache_config->max_unreleased;
-    rcache_params->flags              = !rcache_config->purge_on_fork ? 0 :
-                                        UCS_RCACHE_FLAG_PURGE_ON_FORK;
+
+    if (rcache_config->purge_on_fork) {
+        rcache_params->flags |= UCS_RCACHE_FLAG_PURGE_ON_FORK;
+    }
+
+    if (rcache_config->merge_adjacent) {
+        rcache_params->flags |= UCS_RCACHE_FLAG_MERGE_ADJACENT;
+    }
 }
 
 static size_t ucs_rcache_stat_max_pow2()
@@ -817,11 +826,76 @@ ucs_rcache_check_overlap_one(ucs_rcache_t *rcache, ucs_pgt_addr_t *start,
     return UCS_OK;
 }
 
+static int ucs_rcache_check_adj_size(ucs_pgt_addr_t start, ucs_pgt_addr_t end,
+                                     const ucs_rcache_region_t *region_adj)
+{
+    return (end - start) >= (region_adj->super.end - region_adj->super.start);
+}
+
+static int ucs_rcache_can_merge_adj_region(ucs_rcache_t *rcache, void *arg,
+                                           int prot,
+                                           ucs_rcache_region_t *region)
+{
+    return (region->flags & UCS_RCACHE_REGION_FLAG_REGISTERED) &&
+           (region->prot == prot) && rcache->params.ops->can_merge(arg, region);
+}
+
+static void ucs_rcache_check_adj_regions(ucs_rcache_t *rcache, void *arg,
+                                         ucs_pgt_addr_t start,
+                                         ucs_pgt_addr_t end, int prot,
+                                         ucs_list_link_t *list)
+{
+    ucs_rcache_region_t *region_left  = NULL;
+    ucs_rcache_region_t *region_right = NULL;
+    int can_merge_left                = 0;
+    int can_merge_right               = 0;
+    ucs_pgt_region_t *pgt_left, *pgt_right;
+
+    if (rcache->params.ops->can_merge == NULL) {
+        return;
+    }
+
+    pgt_left  = ucs_pgtable_lookup(&rcache->pgtable, start - 1);
+    pgt_right = ucs_pgtable_lookup(&rcache->pgtable, end);
+
+    if (pgt_left != NULL) {
+        region_left    = ucs_derived_of(pgt_left, ucs_rcache_region_t);
+        can_merge_left = ucs_rcache_can_merge_adj_region(rcache, arg, prot,
+                                                         region_left);
+    }
+
+    if (pgt_right != NULL) {
+        region_right    = ucs_derived_of(pgt_right, ucs_rcache_region_t);
+        can_merge_right = ucs_rcache_can_merge_adj_region(rcache, arg, prot,
+                                                          region_right);
+    }
+
+    /* Fill in a gap between two existing registrations */
+    if (can_merge_left && can_merge_right) {
+        ucs_list_add_tail(list, &region_left->tmp_list);
+        ucs_list_add_tail(list, &region_right->tmp_list);
+        return;
+    }
+
+    if ((region_left != NULL) && can_merge_left) {
+        if (ucs_rcache_check_adj_size(start, end, region_left)) {
+            ucs_list_add_tail(list, &region_left->tmp_list);
+        }
+    }
+
+    if ((region_right != NULL) && can_merge_right) {
+        if (ucs_rcache_check_adj_size(start, end, region_right)) {
+            ucs_list_add_tail(list, &region_right->tmp_list);
+        }
+    }
+}
+
 /* Lock must be held */
 static ucs_status_t
-ucs_rcache_check_overlap(ucs_rcache_t *rcache, void *arg, ucs_pgt_addr_t *start,
-                         ucs_pgt_addr_t *end, size_t *alignment, int *prot,
-                         int *merged, ucs_rcache_region_t **region_p)
+ucs_rcache_check_neighbors(ucs_rcache_t *rcache, void *arg,
+                           ucs_pgt_addr_t *start, ucs_pgt_addr_t *end,
+                           size_t *alignment, int *prot, int *merged,
+                           ucs_rcache_region_t **region_p)
 {
     ucs_rcache_region_t *region, *tmp;
     ucs_pgt_addr_t old_start, old_end;
@@ -874,6 +948,12 @@ ucs_rcache_check_overlap(ucs_rcache_t *rcache, void *arg, ucs_pgt_addr_t *start,
         ucs_list_head_init(&region_list);
         ucs_rcache_find_regions(rcache, *start, old_start - 1, &region_list);
         ucs_rcache_find_regions(rcache, old_end, *end - 1, &region_list);
+
+        if ((rcache->params.flags & UCS_RCACHE_FLAG_MERGE_ADJACENT) &&
+            (ucs_list_is_empty(&region_list))) {
+            ucs_rcache_check_adj_regions(rcache, arg, *start, *end, *prot,
+                                         &region_list);
+        }
     } while (!ucs_list_is_empty(&region_list));
 
     return UCS_OK;
@@ -937,7 +1017,7 @@ ucs_status_t ucs_rcache_create_region(ucs_rcache_t *rcache, void *address,
     /* Check overlap with existing regions */
     /* coverity[double_unlock] */
     /* coverity[double_lock] */
-    status = UCS_PROFILE_CALL(ucs_rcache_check_overlap, rcache, arg, &start,
+    status = UCS_PROFILE_CALL(ucs_rcache_check_neighbors, rcache, arg, &start,
                               &end, &alignment, &prot, &merged, &region);
     if (status == UCS_ERR_ALREADY_EXISTS) {
         /* Found a matching region (it could have been added after we released

src/ucs/memory/rcache.h

@@ -58,10 +58,11 @@ enum {
  * Rcache flags.
  */
 enum {
-    UCS_RCACHE_FLAG_NO_PFN_CHECK  = UCS_BIT(0), /**< PFN check not supported for this rcache */
-    UCS_RCACHE_FLAG_PURGE_ON_FORK = UCS_BIT(1), /**< purge rcache on fork */
-    UCS_RCACHE_FLAG_SYNC_EVENTS   = UCS_BIT(2), /**< Synchronize memory events handling */
-    UCS_RCACHE_FLAG_NEED_LRU_LOCK = UCS_BIT(3), /**< rcache not protected by other lock */
+    UCS_RCACHE_FLAG_NO_PFN_CHECK   = UCS_BIT(0), /**< PFN check not supported for this rcache */
+    UCS_RCACHE_FLAG_PURGE_ON_FORK  = UCS_BIT(1), /**< purge rcache on fork */
+    UCS_RCACHE_FLAG_SYNC_EVENTS    = UCS_BIT(2), /**< Synchronize memory events handling */
+    UCS_RCACHE_FLAG_NEED_LRU_LOCK  = UCS_BIT(3), /**< rcache not protected by other lock */
+    UCS_RCACHE_FLAG_MERGE_ADJACENT = UCS_BIT(4)  /**< Merge adjacent rcache regions */
 };
 
 /*
@@ -137,6 +138,20 @@ struct ucs_rcache_ops {
     void                   (*dump_region)(void *context, ucs_rcache_t *rcache,
                                           ucs_rcache_region_t *region,
                                           char *buf, size_t max);
+
+    /**
+     * Check to see if a new memory region can be safely merged with an
+     * existing memory registration by checking if they share memory type,
+     * memory flags, and memory allocation.
+     *
+     * @param [in]  arg      Custom argument passed to @ref ucs_rcache_get().
+     * @param [in]  region   Existing region to be checked for merge
+     *                       compatibility.
+     *
+     * @return 1 if new region is compatible with existing region. 0 if not.
+     */
+    int                     (*can_merge)(void *arg,
+                                         ucs_rcache_region_t *region);
 };
 
 
@@ -168,6 +183,7 @@ struct ucs_rcache_config {
     size_t        max_size;       /**< Maximal size of mapped memory */
     size_t        max_unreleased; /**< Threshold for triggering a cleanup */
     int           purge_on_fork;  /**< Enable/disable rcache purge on fork */
+    int           merge_adjacent; /**< Enable/disable rcache adjacent region merge */
 };
 
 

src/uct/rocm/copy/rocm_copy_md.c

@@ -453,7 +453,7 @@ uct_rocm_copy_md_open(uct_component_h component, const char *md_name,
         rcache_params.ucm_event_priority = md_config->rcache.event_prio;
         rcache_params.context            = md;
         rcache_params.ops                = &uct_rocm_copy_rcache_ops;
-        rcache_params.flags              = UCS_RCACHE_FLAG_PURGE_ON_FORK;
+        rcache_params.flags             |= UCS_RCACHE_FLAG_PURGE_ON_FORK;
 
         status = ucs_rcache_create(&rcache_params, "rocm_copy", NULL, &md->rcache);
         if (status == UCS_OK) {