UCP/PROTO: Protocol variants

src/ucp/core/ucp_context.c

@@ -504,6 +504,14 @@ static ucs_config_field_t ucp_context_config_table[] = {
    "directory.",
    ucs_offsetof(ucp_context_config_t, proto_info_dir), UCS_CONFIG_TYPE_STRING},
 
+  {"PROTO_VARIANTS", "n",
+   "Enable multiple variants of UCP protocols, meaning that a single protocol\n"
+   "may have multiple variants (optimized for latency or bandwidth) for the same\n"
+   "operation. The value is interpreted as follows:\n"
+   " 'y'          : Enable multiple variants\n"
+   " 'n'          : Disable multiple variants\n",
+   ucs_offsetof(ucp_context_config_t, proto_variants_enable), UCS_CONFIG_TYPE_BOOL},
+
   {"REG_NONBLOCK_MEM_TYPES", "",
    "Perform only non-blocking memory registration for these memory types.\n"
    "Non-blocking registration means that the page registration may be\n"

src/ucp/core/ucp_context.h

@@ -184,6 +184,8 @@ typedef struct ucp_context_config {
     char                                   *select_distance_md;
     /** Directory to write protocol selection information */
     char                                   *proto_info_dir;
+    /** Enable multiple variants of UCP protocols */
+    int                                    proto_variants_enable;
     /** Memory types that perform non-blocking registration by default */
     uint64_t                               reg_nb_mem_types;
     /** Prefer native RMA transports for RMA/AMO protocols */

src/ucp/core/ucp_worker.c

@@ -130,6 +130,14 @@ static ucs_mpool_ops_t ucp_rkey_mpool_ops = {
     .obj_str       = NULL
 };
 
+static ucs_mpool_ops_t ucp_proto_select_mpool_ops = {
+    .chunk_alloc   = ucs_mpool_chunk_malloc,
+    .chunk_release = ucs_mpool_chunk_free,
+    .obj_init      = NULL,
+    .obj_cleanup   = NULL,
+    .obj_str       = NULL
+};
+
 #define ucp_worker_discard_uct_ep_hash_key(_uct_ep) \
     kh_int64_hash_func((uintptr_t)(_uct_ep))
 
@@ -2006,6 +2014,16 @@ static ucs_status_t ucp_worker_init_mpools(ucp_worker_h worker)
         worker->flags |= UCP_WORKER_FLAG_AM_MPOOL_INITIALIZED;
     }
 
+    ucs_mpool_params_reset(&mp_params);
+    mp_params.elem_size       = sizeof(ucp_proto_lane_select_t);
+    mp_params.elems_per_chunk = 8;
+    mp_params.ops             = &ucp_proto_select_mpool_ops;
+    mp_params.name            = "ucp_proto_lane_select";
+    status = ucs_mpool_init(&mp_params, &worker->proto_select_mp);
+    if (status != UCS_OK) {
+        goto err_reg_mp_cleanup;
+    }
+
     return UCS_OK;
 
 err_reg_mp_cleanup:
@@ -2043,6 +2061,7 @@ static void ucp_worker_destroy_mpools(ucp_worker_h worker)
     }
     ucs_mpool_cleanup(&worker->req_mp,
                       !(worker->flags & UCP_WORKER_FLAG_IGNORE_REQUEST_LEAK));
+    ucs_mpool_cleanup(&worker->proto_select_mp, 1);
 }
 
 static unsigned ucp_worker_ep_config_free_cb(void *arg)

src/ucp/core/ucp_worker.h

@@ -339,6 +339,7 @@ typedef struct ucp_worker {
     ucp_worker_cm_t                  *cms;                /* Array of CMs, one for each component */
     ucs_mpool_set_t                  am_mps;              /* Memory pool set for AM receives */
     ucs_mpool_t                      reg_mp;              /* Registered memory pool */
+    ucs_mpool_t                      proto_select_mp;     /* Protocol selection memory pool */
     ucp_worker_mpool_hash_t          mpool_hash;          /* Hash table of memory pools */
     ucs_queue_head_t                 rkey_ptr_reqs;       /* Queue of submitted RKEY PTR requests that
                                                            * are in-progress */