UCP/CORE: Rkey pack featuring Direct NIC.

Author: rakhmets

src/ucp/core/ucp_ep.c

@@ -229,6 +229,7 @@ static ucp_ep_h ucp_ep_allocate(ucp_worker_h worker, const char *peer_name)
     ep->ext->unflushed_lanes              = 0;
     ep->ext->fence_seq                    = 0;
     ep->ext->uct_eps                      = NULL;
+    ep->ext->flush_sys_dev_map            = 0;
 
     UCS_STATIC_ASSERT(sizeof(ep->ext->ep_match) >=
                       sizeof(ep->ext->flush_state));
@@ -537,8 +538,9 @@ void ucp_ep_flush_state_reset(ucp_ep_h ep)
                 (flush_state->cmpl_sn == 0) &&
                 ucs_hlist_is_empty(&flush_state->reqs)));
 
-    flush_state->send_sn = 0;
-    flush_state->cmpl_sn = 0;
+    flush_state->send_sn         = 0;
+    flush_state->cmpl_sn         = 0;
+    flush_state->mem_in_progress = 0;
     ucs_hlist_head_init(&flush_state->reqs);
     ucp_ep_update_flags(ep, UCP_EP_FLAG_FLUSH_STATE_VALID, 0);
 }

src/ucp/core/ucp_ep.h

@@ -500,6 +500,7 @@ typedef struct {
                               are waiting for remote completion */
     uint32_t         send_sn; /* Sequence number of sent operations */
     uint32_t         cmpl_sn; /* Sequence number of completions */
+    uint32_t         mem_in_progress; /* Track ongoing memory flushes for this endpoint */
 } ucp_ep_flush_state_t;
 
 
@@ -554,6 +555,12 @@ typedef struct ucp_ep_ext {
      * structure. TODO allocate this array dynamically.
      */
     uct_ep_h                     *uct_eps;
+
+
+    /**
+     * Map of system devices that require a flush operation
+     */
+    ucp_sys_dev_map_t             flush_sys_dev_map;
 } ucp_ep_ext_t;
 
 
@@ -935,4 +942,16 @@ ucs_status_t ucp_ep_realloc_lanes(ucp_ep_h ep, unsigned new_num_lanes);
  */
 void ucp_ep_set_cfg_index(ucp_ep_h ep, ucp_worker_cfg_index_t cfg_index);
 
+
+/**
+ * @brief Progress function for memory specific remote flushing.
+ *
+ * This call starts and progresses all memory specific remote flushes.
+ *
+ * @param[in] self        Pending request tracking the flush.
+ *
+ * @return Error code as defined by @ref ucs_status_t
+ */
+ucs_status_t ucp_ep_flush_mem_progress(uct_pending_req_t *self);
+
 #endif

src/ucp/core/ucp_mm.c

@@ -35,9 +35,15 @@ typedef struct {
 
 ucp_mem_dummy_handle_t ucp_mem_dummy_handle = {
     .memh = {
-        .alloc_method = UCT_ALLOC_METHOD_LAST,
+        .alloc_method   = UCT_ALLOC_METHOD_LAST,
         .alloc_md_index = UCP_NULL_RESOURCE,
-        .parent = &ucp_mem_dummy_handle.memh,
+        .parent         = &ucp_mem_dummy_handle.memh,
+        .mem_type       = UCS_MEMORY_TYPE_HOST,
+        .sys_dev        = UCS_SYS_DEVICE_ID_UNKNOWN,
+        .packed_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN,
+        .md_map         = 0,
+        .inv_md_map     = 0,
+        .reg_id         = 0,
     },
     .uct = { UCT_MEM_HANDLE_NULL }
 };
@@ -709,7 +715,8 @@ 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)
+                          uct_alloc_method_t method, ucs_memory_type_t mem_type,
+                          ucs_sys_device_t sys_dev)
 {
     memh->md_map         = 0;
     memh->inv_md_map     = 0;
@@ -719,6 +726,14 @@ static void ucp_memh_init(ucp_mem_h memh, ucp_context_h context,
     memh->alloc_md_index = UCP_NULL_RESOURCE;
     memh->alloc_method   = method;
     memh->mem_type       = mem_type;
+    memh->sys_dev        = sys_dev;
+
+    /* Cache sys_dev in a format packed to rkey to minimize overhead during
+     * rndv protocols. TODO remove if using another method to mark rkey with
+     * remote flush requirement. */
+    memh->packed_sys_dev = (sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) ?
+                                   UCS_SYS_DEVICE_ID_UNKNOWN :
+                                   ucp_rkey_pack_sys_dev(memh);
 }
 
 static ucs_status_t
@@ -736,11 +751,9 @@ ucp_memh_create(ucp_context_h context, void *address, size_t length,
 
     memh->super.super.start = (uintptr_t)address;
     memh->super.super.end   = (uintptr_t)address + length;
-    ucp_memh_init(memh, context, memh_flags, uct_flags, method, mem_type);
-
-    ucp_memory_detect(context, ucp_memh_address(memh), ucp_memh_length(memh),
-                      &info);
-    memh->sys_dev = info.sys_dev;
+    ucp_memory_detect(context, address, length, &info);
+    ucp_memh_init(memh, context, memh_flags, uct_flags, method, mem_type,
+                  info.sys_dev);
 
     *memh_p = memh;
     return UCS_OK;
@@ -1632,7 +1645,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);
+                  reg_ctx->mem_type, UCS_SYS_DEVICE_ID_UNKNOWN);
     memh->reg_id = context->next_memh_reg_id++;
 
     if (rcache_mem_reg_flags & UCS_RCACHE_MEM_REG_HIDE_ERRORS) {

src/ucp/core/ucp_mm.h

@@ -68,6 +68,7 @@ typedef struct ucp_mem {
     ucp_context_h       context;        /* UCP context that owns a memory handle */
     uct_alloc_method_t  alloc_method;   /* Method used to allocate the memory */
     ucs_sys_device_t    sys_dev;        /* System device index */
+    ucs_sys_device_t    packed_sys_dev; /* System device index */
     ucs_memory_type_t   mem_type;       /* Type of allocated or registered memory */
     ucp_md_index_t      alloc_md_index; /* Index of MD used to allocate the memory */
     uint64_t            remote_uuid;    /* Remote UUID */

src/ucp/core/ucp_request.h

@@ -121,6 +121,24 @@ enum {
 };
 
 
+/**
+ * Memory flush related structure
+ */
+typedef struct {
+    /* Number of 0-read flush to complete overall */
+    int              count;
+
+    /* Number of 0-read flush started */
+    int              started;
+
+    /* Shared completion to track remaining */
+    uct_completion_t uct_comp;
+
+    /* List of memory areas to track 0-read flush operations */
+    ucp_mem_area_t   *entries;
+} ucp_mem_flush_t;
+
+
 /**
  * Request in progress.
  */
@@ -335,6 +353,7 @@ struct ucp_request {
                     uint8_t            sw_done;
                     uint8_t            num_lanes; /* How many lanes are being flushed */
                     ucp_lane_map_t     started_lanes; /* Which lanes need were flushed */
+                    ucp_mem_flush_t    mem; /* Memory specific flushes */
                 } flush;
 
                 struct {

src/ucp/core/ucp_rkey.c

@@ -125,6 +125,31 @@ ucp_rkey_unpack_distance(const ucp_rkey_packed_distance_t *packed_distance,
     distance->bandwidth = UCS_FP8_UNPACK(BANDWIDTH, packed_distance->bandwidth);
 }
 
+ucs_sys_device_t ucp_rkey_pack_sys_dev(ucp_mem_h memh)
+{
+    ucs_sys_device_t sys_dev_packed = memh->sys_dev;
+    ucp_md_index_t md_index;
+    ucp_sys_dev_map_t sys_dev_map;
+    ucs_sys_device_t sys_dev;
+
+    ucs_assert(sys_dev_packed <= UCP_SYS_DEVICE_MAX_PACKED);
+
+    ucs_for_each_bit(md_index, memh->md_map) {
+        sys_dev_map = memh->context->tl_mds[md_index].sys_dev_map;
+        ucs_for_each_bit(sys_dev, sys_dev_map) {
+            if (ucs_topo_is_sibling(sys_dev, sys_dev_packed)) {
+                /* PUT operation on such rkey requires remote flush.
+                 * Set a flag for the peer to recognize it. */
+                sys_dev_packed |= UCP_SYS_DEVICE_FLUSH_BIT;
+                goto out;
+            }
+        }
+    }
+
+out:
+    return sys_dev_packed;
+}
+
 UCS_PROFILE_FUNC(ssize_t, ucp_rkey_pack_memh,
                  (context, md_map, memh, address, length, mem_info, sys_dev_map,
                   sys_distance, uct_flags, buffer),
@@ -184,7 +209,7 @@ UCS_PROFILE_FUNC(ssize_t, ucp_rkey_pack_memh,
     }
 
     /* Pack system device id */
-    *ucs_serialize_next(&p, uint8_t) = mem_info->sys_dev;
+    *ucs_serialize_next(&p, uint8_t) = memh->packed_sys_dev;
 
     /* Pack distance from sys_dev to each device in distance_dev_map */
     ucs_for_each_bit(sys_dev, sys_dev_map) {
@@ -577,12 +602,11 @@ ucp_memh_packed_size(ucp_mem_h memh, uint64_t flags, int rkey_compat)
                                              ucp_memh_export_md_map(memh));
     }
 
-    if (rkey_compat) {
-        return ucp_rkey_packed_size(context, memh->md_map,
-                                    UCS_SYS_DEVICE_ID_UNKNOWN, 0);
+    if (!rkey_compat) {
+        ucs_fatal("packing rkey using ucp_memh_pack() is unsupported");
     }
 
-    ucs_fatal("packing rkey using ucp_memh_pack() is unsupported");
+    return ucp_rkey_packed_size(context, memh->md_map, memh->sys_dev, 0);
 }
 
 static ssize_t ucp_memh_do_pack(ucp_mem_h memh, uint64_t flags,
@@ -594,15 +618,15 @@ static ssize_t ucp_memh_do_pack(ucp_mem_h memh, uint64_t flags,
         return ucp_memh_exported_pack(memh, memh_buffer);
     }
 
-    if (rkey_compat) {
-        mem_info.type    = memh->mem_type;
-        mem_info.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
-        return ucp_rkey_pack_memh(memh->context, memh->md_map, memh,
-                                  ucp_memh_address(memh), ucp_memh_length(memh),
-                                  &mem_info, 0, NULL, 0, memh_buffer);
+    if (!rkey_compat) {
+        ucs_fatal("packing rkey using ucp_memh_pack() is unsupported");
     }
 
-    ucs_fatal("packing rkey using ucp_memh_pack() is unsupported");
+    mem_info.type    = memh->mem_type;
+    mem_info.sys_dev = memh->sys_dev;
+    return ucp_rkey_pack_memh(memh->context, memh->md_map, memh,
+                              ucp_memh_address(memh), ucp_memh_length(memh),
+                              &mem_info, 0, NULL, 0, memh_buffer);
 }
 
 int ucp_memh_buffer_is_dummy(const void *exported_memh_buffer)

src/ucp/core/ucp_rkey.h

@@ -65,6 +65,10 @@ struct ucp_rkey_config_key {
 };
 
 
+#define UCP_SYS_DEVICE_FLUSH_BIT UCS_BIT(7)
+#define UCP_SYS_DEVICE_MAX_PACKED UCP_SYS_DEVICE_FLUSH_BIT - 1
+
+
 /**
  * Rkey configuration
  */
@@ -238,4 +242,7 @@ void ucp_rkey_proto_select_dump(ucp_worker_h worker,
                                 ucp_worker_cfg_index_t rkey_cfg_index,
                                 ucs_string_buffer_t *strb);
 
+
+ucs_sys_device_t ucp_rkey_pack_sys_dev(ucp_mem_h memh);
+
 #endif

src/ucp/core/ucp_rkey.inl

@@ -64,4 +64,12 @@ ucp_ep_rkey_unpack_reachable(ucp_ep_h ep, const void *buffer, size_t length,
                                        UCS_SYS_DEVICE_ID_UNKNOWN, rkey_p);
 }
 
+static UCS_F_ALWAYS_INLINE int
+ucp_rkey_need_remote_flush(const ucp_rkey_config_key_t *key)
+{
+    return (key->sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) &&
+           (key->sys_dev & UCP_SYS_DEVICE_FLUSH_BIT);
+
+}
+
 #endif

src/ucp/core/ucp_worker.c

@@ -2504,6 +2504,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context,
     ucs_list_head_init(&worker->internal_eps);
     kh_init_inplace(ucp_worker_rkey_config, &worker->rkey_config_hash);
     kh_init_inplace(ucp_worker_discard_uct_ep_hash, &worker->discard_uct_ep_hash);
+    kh_init_inplace(ucp_worker_remote_flush, &worker->remote_flush_hash);
     worker->counters.ep_creations         = 0;
     worker->counters.ep_creation_failures = 0;
     worker->counters.ep_closures          = 0;
@@ -2714,6 +2715,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context,
     kh_destroy_inplace(ucp_worker_discard_uct_ep_hash,
                        &worker->discard_uct_ep_hash);
     kh_destroy_inplace(ucp_worker_rkey_config, &worker->rkey_config_hash);
+    kh_destroy_inplace(ucp_worker_remote_flush, &worker->remote_flush_hash);
     ucp_worker_destroy_configs(worker);
     ucs_free(worker);
     return status;
@@ -2963,6 +2965,7 @@ void ucp_worker_destroy(ucp_worker_h worker)
     ucs_strided_alloc_cleanup(&worker->ep_alloc);
     kh_destroy_inplace(ucp_worker_discard_uct_ep_hash,
                        &worker->discard_uct_ep_hash);
+    kh_destroy_inplace(ucp_worker_remote_flush, &worker->remote_flush_hash);
     kh_destroy_inplace(ucp_worker_rkey_config, &worker->rkey_config_hash);
     ucp_worker_destroy_configs(worker);
     ucs_free(worker);

src/ucp/core/ucp_worker.h

@@ -230,6 +230,27 @@ KHASH_TYPE(ucp_worker_mpool_hash, ucp_worker_mpool_key_t, ucs_mpool_t);
 typedef khash_t(ucp_worker_mpool_hash) ucp_worker_mpool_hash_t;
 
 
+typedef struct {
+    ucp_ep_h         ep;
+    ucs_sys_device_t sys_dev;
+} ucp_worker_remote_flush_key_t;
+
+
+/**
+ * Remote memory to be used for device targeted remote flushing.
+ */
+typedef struct {
+    uct_rkey_t uct_rkey;
+    uct_ep_t   *uct_ep;
+    uint64_t   address;
+} ucp_mem_area_t;
+
+
+/* Hash map to find what remote devices ep needs to flush */
+KHASH_TYPE(ucp_worker_remote_flush, ucp_worker_remote_flush_key_t,
+           ucp_mem_area_t);
+
+
 /* EP configurations storage */
 UCS_ARRAY_DECLARE_TYPE(ucp_ep_config_arr_t, unsigned, ucp_ep_config_t);
 
@@ -342,6 +363,7 @@ typedef struct ucp_worker {
                                                              mapping */
     UCS_PTR_MAP_T(request)           request_map;         /* UCP requests key to
                                                              ptr mapping */
+    kh_ucp_worker_remote_flush_t     remote_flush_hash;
 
     ucp_ep_config_arr_t              ep_config; /* EP configurations storage */