UCX/RNDV/CUDA: RNDV protocol improvements for CUDA

Author: bureddy

src/ucp/core/ucp_context.c

@@ -256,6 +256,10 @@ static ucs_config_field_t ucp_config_table[] = {
    "RNDV fragment size \n",
    ucs_offsetof(ucp_config_t, ctx.rndv_frag_size), UCS_CONFIG_TYPE_MEMUNITS},
 
+  {"RNDV_PIPELINE_SEND_THRESH", "inf",
+   "RNDV size threshold to enable sender side pipeline for mem type\n",
+   ucs_offsetof(ucp_config_t, ctx.rndv_pipeline_send_thresh), UCS_CONFIG_TYPE_MEMUNITS},
+
   {"MEMTYPE_CACHE", "y",
    "Enable memory type (cuda/rocm) cache \n",
    ucs_offsetof(ucp_config_t, ctx.enable_memtype_cache), UCS_CONFIG_TYPE_BOOL},

src/ucp/core/ucp_context.h

@@ -61,6 +61,8 @@ typedef struct ucp_context_config {
     size_t                                 seg_size;
     /** RNDV pipeline fragment size */
     size_t                                 rndv_frag_size;
+    /** RNDV pipline send threshold */
+    size_t                                 rndv_pipeline_send_thresh;
     /** Threshold for using tag matching offload capabilities. Smaller buffers
      *  will not be posted to the transport. */
     size_t                                 tm_thresh;

src/ucp/core/ucp_request.h

@@ -154,12 +154,15 @@ struct ucp_request {
                 } proxy;
 
                 struct {
-                    uint64_t             remote_address; /* address of the sender's data buffer */
-                    uintptr_t            remote_request; /* pointer to the sender's request */
-                    ucp_request_t        *rreq;          /* receive request on the recv side */
-                    ucp_rkey_h           rkey;           /* key for remote send buffer */
-                    ucp_lane_map_t       lanes_map;      /* used lanes map */
-                    ucp_lane_index_t     lane_count;     /* number of lanes used in transaction */
+                    uint64_t             remote_address;  /* address of the sender's data buffer */
+                    uintptr_t            remote_request;  /* pointer to the sender's request */
+                    ucp_request_t        *rreq;           /* receive request on the recv side */
+                    ucp_rkey_h           rkey;            /* key for remote send buffer */
+                    ucp_lane_map_t       lanes_map_avail; /* used lanes map */
+                    ucp_lane_map_t       lanes_map_all;   /* actual lanes map */
+                    uint8_t              lanes_count;     /* actual lanes count */
+                    uint8_t              rkey_index[UCP_MAX_LANES];
+
                 } rndv_get;
 
                 struct {

src/ucp/tag/offload.c

@@ -561,8 +561,7 @@ ucs_status_t ucp_tag_offload_sw_rndv(uct_pending_req_t *self)
     rndv_rts_hdr = ucs_alloca(rndv_hdr_len);
     packed_len   = ucp_tag_rndv_rts_pack(rndv_rts_hdr, req);
     ucs_assert((rndv_rts_hdr->address != 0) || !UCP_DT_IS_CONTIG(req->send.datatype) ||
-               !ucp_rndv_is_get_zcopy(req->send.mem_type,
-                                      ep->worker->context->config.ext.rndv_mode));
+               !ucp_rndv_is_get_zcopy(req, ep->worker->context));
     return uct_ep_tag_rndv_request(ep->uct_eps[req->send.lane],
                                    req->send.msg_proto.tag.tag,
                                    rndv_rts_hdr, packed_len, 0);