UCT/GDR_COPY: Optimize memory registration in fast-path

openucx · Dec 2, 2024 · b504c8e · b504c8e
1 parent 0129be1
commit b504c8e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/src/ucs/memory/rcache.inl b/src/ucs/memory/rcache.inl
@@ -8,6 +8,7 @@
 #define UCS_RCACHE_INL_
 
 #include "rcache_int.h"
+#include <ucs/profile/profile.h>
 
 static UCS_F_ALWAYS_INLINE int
 ucs_rcache_region_test(ucs_rcache_region_t *region, int prot, size_t alignment)
@@ -80,6 +81,17 @@ ucs_rcache_lookup_unsafe(ucs_rcache_t *rcache, void *address, size_t length,
     return region;
 }
 
+static UCS_F_ALWAYS_INLINE ucs_rcache_region_t *
+ucs_rcache_lookup(ucs_rcache_t *rcache, void *address, size_t length,
+                  size_t alignment, int prot)
+{
+    ucs_rcache_region_t *region;
+
+    ucs_rwlock_read_lock(&rcache->pgt_lock);
+    region = ucs_rcache_lookup_unsafe(rcache, address, length, alignment, prot);
+    ucs_rwlock_read_unlock(&rcache->pgt_lock);
+    return region;
+}
 
 static UCS_F_ALWAYS_INLINE void
 ucs_rcache_region_put_unsafe(ucs_rcache_t *rcache, ucs_rcache_region_t *region)

diff --git a/src/uct/cuda/gdr_copy/gdr_copy_md.c b/src/uct/cuda/gdr_copy/gdr_copy_md.c
@@ -16,6 +16,7 @@
 #include <ucs/sys/sys.h>
 #include <ucs/sys/ptr_arith.h>
 #include <ucs/debug/memtrack_int.h>
+#include <ucs/memory/rcache.inl>
 #include <ucs/type/class.h>
 #include <ucs/profile/profile.h>
 #include <ucm/api/ucm.h>
@@ -342,12 +343,19 @@ uct_gdr_copy_mem_rcache_reg(uct_md_h uct_md, void *address, size_t length,
     ucs_status_t status;
     uct_gdr_copy_mem_t *memh;
 
+    rregion = ucs_rcache_lookup(md->rcache, address, length, GPU_PAGE_SIZE,
+                                 PROT_READ | PROT_WRITE);
+    if (rregion != NULL) {
+        goto out;
+    }
+
     status = ucs_rcache_get(md->rcache, address, length, GPU_PAGE_SIZE,
                             PROT_READ | PROT_WRITE, &flags, &rregion);
     if (status != UCS_OK) {
         return status;
     }
 
+out:
     ucs_assert(rregion->refcount > 0);
     memh    = &ucs_derived_of(rregion, uct_gdr_copy_rcache_region_t)->memh;
     *memh_p = memh;