UCP/RMA: Sample multi-rail zcopy throughput

src/ucp/Makefile.am

@@ -71,6 +71,8 @@ noinst_HEADERS = \
 	proto/proto.h \
 	rma/rma.h \
 	rma/rma.inl \
+	rma/rma_bw.h \
+	rma/rma_bw.inl \
 	rma/rma_rndv.h \
 	rndv/proto_rndv.h \
 	rndv/proto_rndv.inl \
@@ -149,6 +151,7 @@ libucp_la_SOURCES = \
 	rma/get_offload.c \
 	rma/put_am.c \
 	rma/put_offload.c \
+	rma/rma_bw.c \
 	rma/rma_rndv.c \
 	rma/rma_send.c \
 	rma/rma_sw.c \

src/ucp/core/ucp_context.c

@@ -586,6 +586,12 @@ static ucs_config_field_t ucp_context_config_table[] = {
    " 'n' : Select RMA/AMO lanes according to performance charasteristics",
    ucs_offsetof(ucp_context_config_t, prefer_offload), UCS_CONFIG_TYPE_BOOL},
 
+  {"RMA_BW_MEASURE", "n",
+   "Sample completed payload throughput of multi-lane PUT/GET zcopy operations. "
+   "This experimental option records measurements but does not change lane "
+   "weights.",
+   ucs_offsetof(ucp_context_config_t, rma_bw_measure), UCS_CONFIG_TYPE_BOOL},
+
   {"PROTO_OVERHEAD", "single:5ns,multi:10ns,rndv_offload:40ns,rndv_rtr:40ns,"
                      "rndv_rts:275ns,sw:40ns,rkey_ptr:0",
    "Protocol overhead", 0,

src/ucp/core/ucp_context.h

@@ -230,6 +230,8 @@ typedef struct ucp_context_config {
     int                                    reg_nb_fallback;
     /** Prefer native RMA transports for RMA/AMO protocols */
     int                                    prefer_offload;
+    /** Sample completed multi-rail RMA zcopy operations */
+    int                                    rma_bw_measure;
     /** RMA zcopy segment size */
     size_t                                 rma_zcopy_max_seg_size;
     /** Enable global VA MR */

src/ucp/core/ucp_request.c

@@ -55,6 +55,7 @@ static const char *ucp_request_flag_names[] = {
     [ucs_ilog2(UCP_REQUEST_FLAG_RNDV_GET_REQ)]          = "rndv_get_req",
     [ucs_ilog2(UCP_REQUEST_FLAG_RNDV_FLUSH)]            = "rndv_flush",
     [ucs_ilog2(UCP_REQUEST_FLAG_RNDV_START_FLUSH)]      = "rndv_start_flush",
+    [ucs_ilog2(UCP_REQUEST_FLAG_RMA_BW_SAMPLE)]         = "rma_bw_sample",
 };
 
 static ucs_memory_type_t ucp_request_get_mem_type(ucp_request_t *req)

src/ucp/core/ucp_request.h

@@ -69,7 +69,8 @@ enum {
     UCP_REQUEST_FLAG_RNDV_SEND_INTERNAL    = UCS_BIT(26),
     UCP_REQUEST_FLAG_RNDV_GET_REQ          = UCS_BIT(27),
     UCP_REQUEST_FLAG_RNDV_FLUSH            = UCS_BIT(28),
-    UCP_REQUEST_FLAG_RNDV_START_FLUSH      = UCS_BIT(29)
+    UCP_REQUEST_FLAG_RNDV_START_FLUSH      = UCS_BIT(29),
+    UCP_REQUEST_FLAG_RMA_BW_SAMPLE         = UCS_BIT(30)
 };
 
 
@@ -245,10 +246,14 @@ struct ucp_request {
                     uint64_t   remote_addr; /* Remote address */
                     ucp_rkey_h rkey; /* Remote memory key */
 
-                    struct {
-                        const uint64_t   *remote_addrs;
-                        ucp_rkey_h const *rkeys;
-                    } sgl;
+                    union {
+                        struct {
+                            const uint64_t   *remote_addrs;
+                            ucp_rkey_h const *rkeys;
+                        } sgl;
+                        /* Used only by sampled contiguous PUT/GET zcopy */
+                        ucp_rma_bw_sample_t *bw_sample;
+                    };
                 } rma;
 
                 struct {

src/ucp/core/ucp_types.h

@@ -64,6 +64,7 @@ typedef uint16_t                     ucp_worker_cfg_index_t;
 
 /* Forward declarations */
 typedef struct ucp_request            ucp_request_t;
+typedef struct ucp_rma_bw_sample      ucp_rma_bw_sample_t;
 typedef struct ucp_recv_desc          ucp_recv_desc_t;
 typedef struct ucp_address_iface_attr ucp_address_iface_attr_t;
 typedef struct ucp_address_entry      ucp_address_entry_t;

src/ucp/core/ucp_worker.c

@@ -17,6 +17,7 @@
 
 #include <ucp/core/ucp_context.h>
 #include <ucp/proto/proto_common.inl>
+#include <ucp/rma/rma_bw.h>
 #include <ucp/wireup/address.h>
 #include <ucp/wireup/wireup_cm.h>
 #include <ucp/wireup/wireup_ep.h>
@@ -2884,6 +2885,16 @@ ucs_status_t ucp_worker_create(ucp_context_h context,
         goto err_tag_match_cleanup;
     }
 
+    if (context->config.ext.rma_bw_measure) {
+        worker->rma_bw_samples = ucs_calloc(UCP_RMA_BW_MAX_ACTIVE,
+                                            sizeof(*worker->rma_bw_samples),
+                                            "rma_bw_samples");
+        if (worker->rma_bw_samples == NULL) {
+            status = UCS_ERR_NO_MEMORY;
+            goto err_am_cleanup;
+        }
+    }
+
     /* Select atomic resources */
     ucp_worker_init_atomic_tls(worker);
 
@@ -2910,6 +2921,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context,
     return UCS_OK;
 
 err_am_cleanup:
+    ucs_free(worker->rma_bw_samples);
     ucp_am_cleanup(worker);
 err_tag_match_cleanup:
     ucp_tag_match_cleanup(&worker->tm);
@@ -3186,6 +3198,7 @@ void ucp_worker_destroy(ucp_worker_h worker)
     ucs_conn_match_cleanup(&worker->conn_match_ctx);
     ucp_worker_wakeup_cleanup(worker);
     uct_worker_destroy(worker->uct);
+    ucs_free(worker->rma_bw_samples);
     ucs_async_context_cleanup(&worker->async);
     UCS_STATS_NODE_FREE(worker->tm_offload_stats);
     UCS_STATS_NODE_FREE(worker->stats);

src/ucp/core/ucp_worker.h

@@ -314,6 +314,8 @@ typedef struct ucp_worker {
     uct_worker_h                     uct;                 /* UCT worker handle */
     ucs_mpool_t                      req_mp;              /* Memory pool for requests */
     ucs_mpool_t                      rkey_mp;             /* Pool for small memory keys */
+    ucp_rma_bw_sample_t              *rma_bw_samples;     /* Sample pool */
+    ucs_time_t                       rma_bw_next_sample; /* Next sample admission */
     ucp_tl_bitmap_t                  atomic_tls;          /* Which resources can be used for atomics */
 
     int                              inprogress;

src/ucp/rma/get_offload.c

@@ -10,6 +10,7 @@
 #endif
 
 #include "rma.inl"
+#include "rma_bw.inl"
 
 #include <ucp/core/ucp_request.inl>
 #include <ucp/dt/datatype_iter.inl>
@@ -157,16 +158,18 @@ ucp_proto_t ucp_get_offload_bcopy_proto = {
 };
 
 static UCS_F_ALWAYS_INLINE ucs_status_t
-ucp_proto_get_offload_zcopy_send_func(ucp_request_t *req,
-                                      const ucp_proto_multi_lane_priv_t *lpriv,
-                                      ucp_datatype_iter_t *next_iter,
-                                      ucp_lane_index_t *lane_shift)
+ucp_proto_get_offload_zcopy_send_common(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        ucp_datatype_iter_t *next_iter, int measure)
 {
     uct_rkey_t tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey,
                                               lpriv->super.rkey_index);
-    size_t offset      = req->send.state.dt_iter.offset;
+    size_t offset          = req->send.state.dt_iter.offset;
+    uct_completion_t *comp = &req->send.state.uct_comp;
     const ucp_proto_multi_priv_t *mpriv;
+    ucp_rma_bw_frag_t *frag;
     uct_iov_t iov;
+    ucs_status_t status;
 
     ucp_datatype_iter_next_iov(&req->send.state.dt_iter,
                                ucp_proto_multi_max_payload(req, lpriv, 0),
@@ -176,18 +179,60 @@ ucp_proto_get_offload_zcopy_send_func(ucp_request_t *req,
     mpriv = req->send.proto_config->priv;
     ucp_proto_common_zcopy_adjust_min_frag(req, mpriv->min_frag, iov.length,
                                            &iov, 1, &offset);
-    return uct_ep_get_zcopy(ucp_ep_get_lane(req->send.ep, lpriv->super.lane),
-                            &iov, 1, req->send.rma.remote_addr + offset,
-                            tl_rkey, &req->send.state.uct_comp);
+    if (measure) {
+        comp = ucp_rma_bw_frag_start(req, req->send.multi_lane_idx, &frag);
+    }
+
+    status = uct_ep_get_zcopy(ucp_ep_get_lane(req->send.ep, lpriv->super.lane),
+                              &iov, 1, req->send.rma.remote_addr + offset,
+                              tl_rkey, comp);
+    if (measure) {
+        ucp_rma_bw_frag_posted(frag, iov.length, status);
+    }
+
+    return status;
+}
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_proto_get_offload_zcopy_send_func(ucp_request_t *req,
+                                      const ucp_proto_multi_lane_priv_t *lpriv,
+                                      ucp_datatype_iter_t *next_iter,
+                                      ucp_lane_index_t *lane_shift)
+{
+    return ucp_proto_get_offload_zcopy_send_common(req, lpriv, next_iter, 0);
+}
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_proto_get_offload_zcopy_sampled_send_func(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift)
+{
+    return ucp_proto_get_offload_zcopy_send_common(req, lpriv, next_iter, 1);
 }
 
 static ucs_status_t ucp_proto_get_offload_zcopy_progress(uct_pending_req_t *self)
 {
     ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct);
+    /* coverity[tainted_data_downcast] */
+    const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv;
+
+    if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) {
+        ucp_rma_bw_sample_try_start(req, mpriv->num_lanes);
+    }
+
+    if (ucs_unlikely(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE)) {
+        /* coverity[tainted_data_downcast] */
+        return ucp_proto_multi_zcopy_progress(
+                req, mpriv, ucp_proto_multi_rma_init_func,
+                UCT_MD_MEM_ACCESS_LOCAL_WRITE, UCP_DT_MASK_CONTIG_IOV,
+                ucp_proto_get_offload_zcopy_sampled_send_func,
+                ucp_request_invoke_uct_completion_success,
+                ucp_rma_bw_sample_complete);
+    }
 
     /* coverity[tainted_data_downcast] */
     return ucp_proto_multi_zcopy_progress(
-            req, req->send.proto_config->priv, ucp_proto_multi_rma_init_func,
+            req, mpriv, ucp_proto_multi_rma_init_func,
             UCT_MD_MEM_ACCESS_LOCAL_WRITE, UCP_DT_MASK_CONTIG_IOV,
             ucp_proto_get_offload_zcopy_send_func,
             ucp_request_invoke_uct_completion_success,
@@ -245,6 +290,7 @@ ucp_proto_get_offload_zcopy_probe(const ucp_proto_init_params_t *init_params)
 
 static ucs_status_t ucp_proto_get_offload_zcopy_reset(ucp_request_t *req)
 {
+    ucp_rma_bw_sample_detach(req);
     ucp_proto_get_offload_reset(req);
     return ucp_proto_offload_zcopy_reset(req);
 }

src/ucp/rma/put_offload.c

@@ -10,6 +10,7 @@
 
 #include "rma.h"
 #include "rma.inl"
+#include "rma_bw.inl"
 
 #include <ucp/core/ucp_request.inl>
 #include <ucp/dt/datatype_iter.inl>
@@ -245,26 +246,34 @@ ucp_proto_t ucp_put_offload_bcopy_proto = {
 };
 
 static UCS_F_ALWAYS_INLINE ucs_status_t
-ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req,
-                                      const ucp_proto_multi_lane_priv_t *lpriv,
-                                      ucp_datatype_iter_t *next_iter,
-                                      ucp_lane_index_t *lane_shift)
+ucp_proto_put_offload_zcopy_send_common(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        ucp_datatype_iter_t *next_iter, int measure)
 {
     ucp_ep_h ep        = req->send.ep;
     uct_ep_h uct_ep    = ucp_ep_get_lane(ep, lpriv->super.lane);
     uint64_t address   = req->send.rma.remote_addr +
                          req->send.state.dt_iter.offset;
     uct_rkey_t tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey,
                                               lpriv->super.rkey_index);
+    uct_completion_t *comp = &req->send.state.uct_comp;
+    ucp_rma_bw_frag_t *frag;
     uct_iov_t iov;
     ucs_status_t status;
 
     ucp_datatype_iter_next_iov(&req->send.state.dt_iter,
                                ucp_proto_multi_max_payload(req, lpriv, 0),
                                lpriv->super.md_index, UCP_DT_MASK_CONTIG_IOV,
                                next_iter, &iov, 1);
-    status = uct_ep_put_zcopy(uct_ep, &iov, 1, address, tl_rkey,
-                              &req->send.state.uct_comp);
+    if (measure) {
+        comp = ucp_rma_bw_frag_start(req, req->send.multi_lane_idx, &frag);
+    }
+
+    status = uct_ep_put_zcopy(uct_ep, &iov, 1, address, tl_rkey, comp);
+    if (measure) {
+        ucp_rma_bw_frag_posted(frag, iov.length, status);
+    }
+
     if (!UCS_STATUS_IS_ERR(status)) {
         ucp_proto_put_offload_update_remote_flush(ep, lpriv->flush_sys_dev_mask,
                                                   tl_rkey, uct_ep, address);
@@ -273,14 +282,47 @@ ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req,
     return status;
 }
 
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req,
+                                      const ucp_proto_multi_lane_priv_t *lpriv,
+                                      ucp_datatype_iter_t *next_iter,
+                                      ucp_lane_index_t *lane_shift)
+{
+    return ucp_proto_put_offload_zcopy_send_common(req, lpriv, next_iter, 0);
+}
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_proto_put_offload_zcopy_sampled_send_func(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift)
+{
+    return ucp_proto_put_offload_zcopy_send_common(req, lpriv, next_iter, 1);
+}
+
 static ucs_status_t
 ucp_proto_put_offload_zcopy_progress(uct_pending_req_t *self)
 {
     ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct);
+    /* coverity[tainted_data_downcast] */
+    const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv;
+
+    if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) {
+        ucp_rma_bw_sample_try_start(req, mpriv->num_lanes);
+    }
+
+    if (ucs_unlikely(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE)) {
+        /* coverity[tainted_data_downcast] */
+        return ucp_proto_multi_zcopy_progress(
+                req, mpriv, ucp_proto_multi_rma_init_func,
+                UCT_MD_MEM_ACCESS_LOCAL_READ, UCP_DT_MASK_CONTIG_IOV,
+                ucp_proto_put_offload_zcopy_sampled_send_func,
+                ucp_request_invoke_uct_completion_success,
+                ucp_rma_bw_sample_complete);
+    }
 
     /* coverity[tainted_data_downcast] */
     return ucp_proto_multi_zcopy_progress(
-            req, req->send.proto_config->priv, ucp_proto_multi_rma_init_func,
+            req, mpriv, ucp_proto_multi_rma_init_func,
             UCT_MD_MEM_ACCESS_LOCAL_READ, UCP_DT_MASK_CONTIG_IOV,
             ucp_proto_put_offload_zcopy_send_func,
             ucp_request_invoke_uct_completion_success,
@@ -345,6 +387,12 @@ ucp_proto_put_offload_zcopy_probe(const ucp_proto_init_params_t *init_params)
             init_params, 30, UCP_PROTO_COMMON_INIT_FLAG_FAILOVER, 0);
 }
 
+static ucs_status_t ucp_proto_put_offload_zcopy_reset(ucp_request_t *req)
+{
+    ucp_rma_bw_sample_detach(req);
+    return ucp_proto_offload_zcopy_reset(req);
+}
+
 ucp_proto_t ucp_put_offload_zcopy_proto = {
     .name           = "put/offload/zcopy",
     .desc           = UCP_PROTO_ZCOPY_DESC,
@@ -355,7 +403,7 @@ ucp_proto_t ucp_put_offload_zcopy_proto = {
     .query          = ucp_proto_multi_query,
     .progress       = {ucp_proto_put_offload_zcopy_progress},
     .abort          = ucp_proto_request_zcopy_abort,
-    .reset          = ucp_proto_offload_zcopy_reset
+    .reset          = ucp_proto_put_offload_zcopy_reset
 };
 
 static void