UCT/GDR_COPY: Use fast-path cache lookup

Author: Artemy-Mellanox

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_rw_spinlock_read_lock(&rcache->pgt_lock);
+    region = ucs_rcache_lookup_unsafe(rcache, address, length, alignment, prot);
+    ucs_rw_spinlock_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)

src/uct/cuda/gdr_copy/gdr_copy_md.c

@@ -17,6 +17,7 @@
 #include <ucs/sys/ptr_arith.h>
 #include <ucs/debug/memtrack_int.h>
 #include <ucs/time/time.h>
+#include <ucs/memory/rcache.inl>
 #include <ucs/type/class.h>
 #include <ucs/profile/profile.h>
 #include <ucm/api/ucm.h>
@@ -357,12 +358,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;