UCP/API: Merge branch 'master' into topic/device-api-use-offsets

Author: ofirfarjun7

src/tools/perf/cuda/ucp_cuda_kernel.cu

@@ -226,7 +226,7 @@ ucp_perf_cuda_send_sync(ucp_perf_cuda_params &params, ucx_perf_counter_t idx,
                         ucp_device_request_t &req)
 {
     ucs_status_t status = ucp_perf_cuda_send_nbx<level, cmd>(params, idx, req);
-    if (status != UCS_OK) {
+    if (UCS_STATUS_IS_ERR(status)) {
         return status;
     }
 
@@ -260,7 +260,7 @@ ucp_perf_cuda_put_multi_bw_kernel(ucx_perf_cuda_context &ctx,
 
         ucp_device_request_t &req = request_mgr.get_request();
         status = ucp_perf_cuda_send_nbx<level, cmd>(params, idx, req);
-        if (status != UCS_OK) {
+        if (UCS_STATUS_IS_ERR(status)) {
             ucs_device_error("send failed: %d", status);
             goto out;
         }

src/tools/vfs/vfs_main.c

@@ -277,10 +277,12 @@ static int vfs_unlink_socket(int silent_notexist)
 }
 
 /* return 0 or the (negative) value of errno in case of error */
-static int vfs_listen(int silent_addinuse_err)
+static int vfs_listen(int silent_addrinuse_err)
 {
     int listen_fd, ret;
 
+    vfs_log("listening on socket %s", (const char*)g_sockaddr.sun_path);
+
     ret = umask(~S_IRWXU);
     if (ret < 0) {
         ret = -errno;
@@ -304,7 +306,7 @@ static int vfs_listen(int silent_addinuse_err)
                sizeof(g_sockaddr));
     if (ret < 0) {
         ret = -errno;
-        if ((errno != EADDRINUSE) || !silent_addinuse_err) {
+        if ((errno != EADDRINUSE) || !silent_addrinuse_err) {
             vfs_error("bind(%s) failed: %m", g_sockaddr.sun_path);
         }
         goto out_close;

src/ucp/api/device/ucp_device_impl.h

@@ -24,6 +24,7 @@
  */
 typedef struct ucp_device_request {
     uct_device_completion_t comp;
+    ucs_status_t            status;
     uct_device_ep_h         device_ep;
     unsigned                channel_id;
 } ucp_device_request_t;
@@ -52,9 +53,6 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep,
     if (req != nullptr) {
         comp           = &req->comp;
         req->device_ep = device_ep;
-        uct_device_completion_init(comp);
-        /* TODO: Handle multiple device posts with same req? */
-        ++comp->count;
     } else {
         comp = nullptr;
     }
@@ -64,16 +62,20 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep,
 /**
  * Macro for device put operations with retry logic
  */
-#define UCP_DEVICE_SEND_BLOCKING(_level, _uct_device_ep_send, _device_ep, ...) \
+#define UCP_DEVICE_SEND_BLOCKING(_level, _uct_device_ep_send, _device_ep, \
+                                 _req, ...) \
     ({ \
         ucs_status_t _status; \
         do { \
             _status = _uct_device_ep_send<_level>(_device_ep, __VA_ARGS__); \
             if (_status != UCS_ERR_NO_RESOURCE) { \
                 break; \
             } \
-            _status = uct_device_ep_progress<_level>(_device_ep); \
-        } while (!UCS_STATUS_IS_ERR(_status)); \
+            uct_device_ep_progress<_level>(_device_ep); \
+        } while (1); \
+        if (_req != nullptr) { \
+            _req->status = _status; \
+        } \
         _status; \
     })
 
@@ -154,10 +156,9 @@ UCS_F_DEVICE ucs_status_t ucp_device_put_single(
         return status;
     }
 
-    status = UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_single,
-                                      device_ep, uct_elem, address,
-                                      remote_address, length, flags, comp);
-    return status;
+    return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_single, device_ep,
+                                    req, uct_elem, address, remote_address,
+                                    length, flags, comp);
 }
 
 
@@ -209,8 +210,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc(
     }
 
     return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_atomic_add, device_ep,
-                                    uct_elem, inc_value, remote_address, flags,
-                                    comp);
+                                    req, uct_elem, inc_value, remote_address,
+                                    flags, comp);
 }
 
 
@@ -266,8 +267,9 @@ UCS_F_DEVICE ucs_status_t ucp_device_put_multi(
     }
 
     return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_multi, device_ep,
-                                    uct_mem_list, mem_list_h->mem_list_length,
-                                    addresses, remote_addresses, lengths,
+                                    req, uct_mem_list,
+                                    mem_list_h->mem_list_length, addresses,
+                                    remote_addresses, lengths,
                                     counter_inc_value, counter_remote_address,
                                     flags, comp);
 }
@@ -346,11 +348,12 @@ UCS_F_DEVICE ucs_status_t ucp_device_put_multi_partial(
     }
 
     return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_multi_partial,
-                                    device_ep, uct_mem_list, mem_list_indices,
-                                    mem_list_count, addresses, remote_addresses,
-                                    local_offsets, remote_offsets, lengths,
-                                    counter_index, counter_inc_value,
-                                    counter_remote_address, flags, comp);
+                                    device_ep, req, uct_mem_list,
+                                    mem_list_indices, mem_list_count, addresses,
+                                    remote_addresses, local_offsets,
+                                    remote_offsets, lengths, counter_index,
+                                    counter_inc_value, counter_remote_address,
+                                    flags, comp);
 }
 
 
@@ -417,19 +420,14 @@ UCS_F_DEVICE void ucp_device_counter_write(void *counter_ptr, uint64_t value)
 template<ucs_device_level_t level = UCS_DEVICE_LEVEL_THREAD>
 UCS_F_DEVICE ucs_status_t ucp_device_progress_req(ucp_device_request_t *req)
 {
-    ucs_status_t status;
-
-    if (ucs_likely(req->comp.count == 0)) {
-        return req->comp.status;
-    }
-
-    status = uct_device_ep_progress<level>(req->device_ep);
-    if (status != UCS_OK) {
-        return status;
+    if (ucs_likely(req->status != UCS_INPROGRESS)) {
+        return req->status;
     }
 
-    return (ucs_likely(req->comp.count == 0)) ? req->comp.status :
-                                                UCS_INPROGRESS;
+    uct_device_ep_progress<level>(req->device_ep);
+    req->status = uct_device_ep_check_completion<level>(req->device_ep,
+                                                        &req->comp);
+    return req->status;
 }
 
 #endif /* UCP_DEVICE_IMPL_H */

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