UCP/RKEY: Pack sender flush flag where applicable. (#10909)

Author: tvegas1

src/ucp/core/ucp_mm.c

@@ -40,7 +40,6 @@ 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,
-        .packed_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN,
         .md_map         = 0,
         .inv_md_map     = 0,
         .reg_id         = 0,
@@ -727,13 +726,6 @@ 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;
-
-    /* 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

src/ucp/core/ucp_mm.h

@@ -36,14 +36,21 @@ enum {
     /*
      * Memory handle was imported and points to some peer's memory buffer.
      */
-    UCP_MEMH_FLAG_IMPORTED     = UCS_BIT(0),
-    UCP_MEMH_FLAG_MLOCKED      = UCS_BIT(1),
-    UCP_MEMH_FLAG_HAS_AUTO_GVA = UCS_BIT(2),
+    UCP_MEMH_FLAG_IMPORTED           = UCS_BIT(0),
+    UCP_MEMH_FLAG_MLOCKED            = UCS_BIT(1),
+    UCP_MEMH_FLAG_HAS_AUTO_GVA       = UCS_BIT(2),
 
     /**
      * Avoid using registration cache for the particular memory region.
      */
-    UCP_MEMH_FLAG_NO_RCACHE    = UCS_BIT(3)
+    UCP_MEMH_FLAG_NO_RCACHE          = UCS_BIT(3),
+
+    /**
+     * Track if sender-side flush is needed, check is only done when needed
+     * and cached.
+     */
+    UCP_MEMH_FLAG_SEND_FLUSH_CHECKED = UCS_BIT(4),
+    UCP_MEMH_FLAG_SEND_FLUSH_NEEDED  = UCS_BIT(5)
 };
 
 
@@ -68,7 +75,6 @@ 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_rkey.c

@@ -150,29 +150,38 @@ ucp_sys_dev_map_t ucp_memh_sys_dev_map(ucp_mem_h memh)
     return 0;
 }
 
-ucs_sys_device_t ucp_rkey_pack_sys_dev(ucp_mem_h memh)
+static int ucp_memh_send_flush_is_needed(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);
+    if (memh->flags & UCP_MEMH_FLAG_SEND_FLUSH_CHECKED) {
+        return !!(memh->flags & UCP_MEMH_FLAG_SEND_FLUSH_NEEDED);
+    }
+
+    memh->flags |= UCP_MEMH_FLAG_SEND_FLUSH_CHECKED;
 
-    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;
+    if (memh->sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) {
+        ucs_assert(memh->sys_dev <= 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, memh->sys_dev)) {
+                    /*
+                     * PUT operation on such device will require remote flush
+                     * when using network devices.
+                     * Set a flag for the peer to recognize it.
+                     */
+                    memh->flags |= UCP_MEMH_FLAG_SEND_FLUSH_NEEDED;
+                    return 1;
+                }
             }
         }
     }
 
-out:
-    return sys_dev_packed;
+    return 0;
 }
 
 UCS_PROFILE_FUNC(ssize_t, ucp_rkey_pack_memh,
@@ -231,8 +240,14 @@ UCS_PROFILE_FUNC(ssize_t, ucp_rkey_pack_memh,
 
     if (md_map != 0) {
         /* Since UCX 1.20: always pack sys_dev for non-empty rkeys. */
-         ucs_assert(memh != NULL);
-         *ucs_serialize_next(&p, uint8_t) = memh->packed_sys_dev;
+        ucs_assert(memh != NULL);
+
+        sys_dev = memh->sys_dev;
+        if (ucp_memh_send_flush_is_needed(memh)) {
+            sys_dev |= UCP_SYS_DEVICE_FLUSH_BIT;
+        }
+
+        *ucs_serialize_next(&p, uint8_t) = sys_dev;
     }
 
     if ((mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) || (md_map == 0)) {
@@ -829,15 +844,24 @@ ucp_rkey_unpack_lanes_distance(const ucp_ep_config_key_t *ep_config_key,
     }
 }
 
-static UCS_F_ALWAYS_INLINE ucs_sys_device_t
+static UCS_F_ALWAYS_INLINE void
 ucp_rkey_extract_sys_dev(const ucp_ep_config_t *ep_config, ucp_rkey_h rkey,
-                         const void **buffer_p, const void *buffer_end)
+                         const void **buffer_p, const void *buffer_end,
+                         ucp_rkey_config_key_t *rkey_config_key)
 {
     if ((*buffer_p < buffer_end) ||
         ((ep_config->key.dst_version > 19) && (rkey->md_map != 0))) {
-        return *ucs_serialize_next(buffer_p, const uint8_t);
+        rkey_config_key->sys_dev = *ucs_serialize_next(buffer_p, const uint8_t);
+    } else {
+        rkey_config_key->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
+    }
+
+    if ((rkey_config_key->sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) &&
+        (rkey_config_key->sys_dev & UCP_SYS_DEVICE_FLUSH_BIT)) {
+        rkey_config_key->flags    = UCP_RKEY_CONFIG_FLAG_FLUSH;
+        rkey_config_key->sys_dev &= ~UCP_SYS_DEVICE_FLUSH_BIT;
     } else {
-        return UCS_SYS_DEVICE_ID_UNKNOWN;
+        rkey_config_key->flags = 0;
     }
 }
 
@@ -864,8 +888,8 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rkey_proto_resolve,
     rkey_config_key.md_map             = rkey->md_map;
     rkey_config_key.mem_type           = rkey->mem_type;
     rkey_config_key.unreachable_md_map = unreachable_md_map;
-    rkey_config_key.sys_dev            = ucp_rkey_extract_sys_dev(
-                                           ep_config, rkey, &p, buffer_end);
+
+    ucp_rkey_extract_sys_dev(ep_config, rkey, &p, buffer_end, &rkey_config_key);
 
     /* Starting with UCX v1.20, lane distances are always packed if sys_dev is
      * not UNKNOWN. Even if the rkey length is not explicitly passed to the API,

src/ucp/core/ucp_rkey.h

@@ -44,6 +44,15 @@ enum {
 };
 
 
+/**
+ * Rkey config flags
+ */
+enum {
+    UCP_RKEY_CONFIG_FLAG_FLUSH    = UCS_BIT(0)  /* Put and atomic operations on this rkey
+                                                   require remote flush */
+};
+
+
 /**
  * Rkey configuration key
  */
@@ -57,6 +66,9 @@ struct ucp_rkey_config_key {
     /* Remote system device id */
     ucs_sys_device_t       sys_dev;
 
+    /* Rkey specific flags, like @a UCP_RKEY_CONFIG_FLAG_FLUSH */
+    uint8_t                flags;
+
     /* Remote memory type */
     ucs_memory_type_t      mem_type;
 
@@ -65,8 +77,8 @@ 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
+#define UCP_SYS_DEVICE_FLUSH_BIT  UCS_BIT(7)
+#define UCP_SYS_DEVICE_MAX_PACKED (UCP_SYS_DEVICE_FLUSH_BIT - 1)
 
 
 /**
@@ -243,7 +255,4 @@ 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

@@ -20,7 +20,8 @@ ucp_rkey_config_hash_func(ucp_rkey_config_key_t rkey_config_key)
             (rkey_config_key.unreachable_md_map << 32)) ^
            (rkey_config_key.ep_cfg_index << 8) ^
            (rkey_config_key.sys_dev << 16) ^
-           (rkey_config_key.mem_type << 24);
+           (rkey_config_key.mem_type << 24) ^
+           rkey_config_key.flags;
 }
 
 static UCS_F_ALWAYS_INLINE int
@@ -31,6 +32,7 @@ ucp_rkey_config_is_equal(ucp_rkey_config_key_t rkey_config_key1,
            (rkey_config_key1.ep_cfg_index == rkey_config_key2.ep_cfg_index) &&
            (rkey_config_key1.sys_dev == rkey_config_key2.sys_dev) &&
            (rkey_config_key1.mem_type == rkey_config_key2.mem_type) &&
+           (rkey_config_key1.flags == rkey_config_key2.flags) &&
            (rkey_config_key1.unreachable_md_map ==
             rkey_config_key2.unreachable_md_map);
 }
@@ -67,9 +69,7 @@ ucp_ep_rkey_unpack_reachable(ucp_ep_h ep, const void *buffer, size_t length,
 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);
-
+    return key->flags & UCP_RKEY_CONFIG_FLAG_FLUSH;
 }
 
 #endif

src/ucp/proto/proto_multi.c

@@ -141,14 +141,20 @@ ucp_proto_multi_select_bw_lanes(const ucp_proto_init_params_t *params,
      * path ratio */
 }
 
-static ucp_sys_dev_map_t
-ucp_proto_multi_init_flush_sys_dev_mask(const ucp_rkey_config_key_t *key)
+static ucp_sys_dev_map_t ucp_proto_multi_init_flush_sys_dev_mask(
+        const ucp_proto_multi_init_params_t *params, ucp_lane_index_t lane)
 {
-    if (key == NULL || !ucp_rkey_need_remote_flush(key)) {
+    const ucp_rkey_config_key_t *key   = params->super.super.rkey_config_key;
+    const uct_iface_attr_t *iface_attr =
+        ucp_proto_common_get_iface_attr(&params->super.super, lane);
+
+    if ((key == NULL) || !ucp_rkey_need_remote_flush(key) ||
+        !(iface_attr->cap.flags & UCT_IFACE_FLAG_GET_BCOPY) ||
+        !ucp_proto_common_is_net_dev(&params->super.super, lane)) {
         return 0;
     }
 
-    return UCS_BIT(key->sys_dev & ~UCP_SYS_DEVICE_FLUSH_BIT);
+    return UCS_BIT(key->sys_dev);
 }
 
 static ucp_lane_index_t ucp_proto_multi_filter_net_devices(
@@ -456,8 +462,8 @@ ucs_status_t ucp_proto_multi_init(const ucp_proto_multi_init_params_t *params,
         lpriv->max_frag_sum   = mpriv->max_frag_sum;
         lpriv->opt_align      = ucp_proto_multi_get_lane_opt_align(params, lane);
         mpriv->align_thresh   = ucs_max(mpriv->align_thresh, lpriv->opt_align);
-        lpriv->flush_sys_dev_mask = ucp_proto_multi_init_flush_sys_dev_mask(
-                params->super.super.rkey_config_key);
+        lpriv->flush_sys_dev_mask =
+                ucp_proto_multi_init_flush_sys_dev_mask(params, lane);
     }
     ucs_assert(mpriv->num_lanes == ucs_popcount(selection.lane_map));
 

src/ucp/rma/flush.c

@@ -243,6 +243,7 @@ static ucs_status_t ucp_ep_flush_mem_start(ucp_request_t *req)
 
     ep->ext->flush_sys_dev_map = 0;
 
+    req->send.flush.mem.count           = count;
     req->send.flush.mem.started         = 0;
     req->send.uct.func                  = ucp_ep_flush_mem_progress;
     req->send.flush.mem.uct_comp.func   = ucp_ep_flush_mem_completion;

src/ucp/rndv/proto_rndv.c

@@ -158,6 +158,7 @@ static ucs_status_t ucp_proto_rndv_ctrl_select_remote_proto(
     rkey_config_key.ep_cfg_index = ep_cfg_index;
     rkey_config_key.sys_dev      = params->super.reg_mem_info.sys_dev;
     rkey_config_key.mem_type     = params->super.reg_mem_info.type;
+    rkey_config_key.flags        = 0;
 
     rkey_config_key.unreachable_md_map = 0;
 

test/gtest/ucp/test_ucp_proto.cc

@@ -116,6 +116,7 @@ test_ucp_proto::create_rkey_config_key(ucp_md_map_t md_map)
     rkey_config_key.mem_type           = UCS_MEMORY_TYPE_HOST;
     rkey_config_key.sys_dev            = UCS_SYS_DEVICE_ID_UNKNOWN;
     rkey_config_key.unreachable_md_map = 0;
+    rkey_config_key.flags              = 0;
 
     return rkey_config_key;
 }