Merge pull request #10676 from iyastreb/ucp/proto-find-lanes-callback

UCP/PROTO: Find lanes callback to minimize overhead

Author: brminich

src/ucp/proto/proto_common.c

@@ -468,13 +468,100 @@ ucp_proto_common_get_lane_perf(const ucp_proto_common_init_params_t *params,
     return status;
 }
 
+static int ucp_proto_common_find_lanes_check_mem_type(
+        const ucp_proto_common_init_params_t *params, ucp_rsc_index_t rsc_index)
+{
+    uct_perf_attr_t perf_attr  = {0};
+    ucp_worker_iface_t *wiface = ucp_worker_iface(params->super.worker, rsc_index);
+
+    ucp_proto_common_perf_attr_set_mem_type(params, &perf_attr);
+    /* TODO: Use memory reachability UCT API, when available, to check memory
+       type support */
+    return uct_iface_estimate_perf(wiface->iface, &perf_attr) == UCS_OK;
+}
+
+int
+ucp_proto_common_filter_min_frag(const ucp_proto_init_params_t *params,
+                                 ucp_lane_index_t lane, const char *lane_desc)
+{
+    const ucp_proto_common_init_params_t *common_params =
+                        ucs_derived_of(params, ucp_proto_common_init_params_t);
+    ucs_memory_type_t reg_mem_type  = common_params->reg_mem_info.type;
+    unsigned flags                  = common_params->flags;
+    ucp_context_h context           = params->worker->context;
+    ucp_rsc_index_t rsc_index       = params->ep_config_key->lanes[lane].rsc_index;
+    ucp_md_index_t md_index         = context->tl_rscs[rsc_index].md_index;
+    const uct_md_attr_v2_t *md_attr = &context->tl_mds[md_index].attr;
+    const uct_iface_attr_t *iface_attr;
+    size_t max_iov, tl_min_frag, tl_max_frag;
+
+    /* Check memory registration capabilities for zero-copy case */
+    if (reg_mem_type != UCS_MEMORY_TYPE_UNKNOWN) {
+        ucs_assertv((reg_mem_type == params->select_param->mem_type) ||
+                    !(flags & UCP_PROTO_COMMON_INIT_FLAG_SEND_ZCOPY),
+                    "flags=0x%x reg_mem_type=%s select_param->mem_type=%s",
+                    flags, ucs_memory_type_names[reg_mem_type],
+                    ucs_memory_type_names[params->select_param->mem_type]);
+
+        if (md_attr->flags & UCT_MD_FLAG_NEED_MEMH) {
+            /* Memory domain must support registration on the relevant memory
+             * type */
+            if (!(context->reg_md_map[reg_mem_type] & UCS_BIT(md_index))) {
+                ucs_trace("%s: md %s cannot register %s memory", lane_desc,
+                          context->tl_mds[md_index].rsc.md_name,
+                          ucs_memory_type_names[reg_mem_type]);
+                return 0;
+            }
+        } else if (!(md_attr->access_mem_types & UCS_BIT(reg_mem_type))) {
+            /* Memory domain which does not require a registration for zero
+             * copy operation must be able to access the relevant memory type */
+            ucs_trace("%s: no access to mem type %s", lane_desc,
+                      ucs_memory_type_names[reg_mem_type]);
+            return 0;
+        }
+    }
+
+    iface_attr = ucp_proto_common_get_iface_attr(params, lane);
+    max_iov    = ucp_proto_common_get_iface_attr_field(
+                    iface_attr, common_params->max_iov_offs, SIZE_MAX);
+    if (max_iov < common_params->min_iov) {
+        ucs_trace("%s: max iov %zu is less than min iov %zu", lane_desc,
+                  max_iov, common_params->min_iov);
+        return 0;
+    }
+
+    ucp_proto_common_get_frag_size(common_params, iface_attr, lane,
+                                   &tl_min_frag, &tl_max_frag);
+
+    /* Minimal fragment size must be 0, unless 'MIN_FRAG' flag is set */
+    if (!(flags & UCP_PROTO_COMMON_INIT_FLAG_MIN_FRAG) && (tl_min_frag > 0)) {
+        ucs_trace("%s: minimal fragment %zu is not 0", lane_desc, tl_min_frag);
+        return 0;
+    }
+
+    /* Maximal fragment size should be larger than header size */
+    if (tl_max_frag <= common_params->hdr_size) {
+        ucs_trace("%s: max fragment is too small %zu, need > %zu", lane_desc,
+                  tl_max_frag, common_params->hdr_size);
+        return 0;
+    }
+
+    if (!ucp_proto_common_find_lanes_check_mem_type(common_params, rsc_index)) {
+        ucs_trace("%s: mem type %s is not supported", lane_desc,
+                  ucs_memory_type_names[reg_mem_type]);
+        return 0;
+    }
+
+    return 1;
+}
+
 ucp_lane_index_t
 ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
-                            unsigned flags, ptrdiff_t max_iov_offs,
-                            size_t min_iov, ucp_lane_type_t lane_type,
-                            ucs_memory_type_t reg_mem_type,
+                            unsigned flags, ucp_lane_type_t lane_type,
                             uint64_t tl_cap_flags, ucp_lane_index_t max_lanes,
-                            ucp_lane_map_t exclude_map, ucp_lane_index_t *lanes)
+                            ucp_lane_map_t exclude_map,
+                            ucp_proto_common_filter_lane_cb_t filter,
+                            ucp_lane_index_t *lanes)
 {
     UCS_STRING_BUFFER_ONSTACK(sel_param_strb, UCP_PROTO_SELECT_PARAM_STR_MAX);
     ucp_context_h context                        = params->worker->context;
@@ -489,7 +576,6 @@ ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
     ucp_md_index_t md_index;
     ucp_lane_map_t lane_map;
     char lane_desc[64];
-    size_t max_iov;
 
     if (max_lanes == 0) {
         return 0;
@@ -551,44 +637,20 @@ ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
 
         if ((flags & UCP_PROTO_COMMON_INIT_FLAG_RKEY_PTR) &&
             !(cmpt_attr->flags & UCT_COMPONENT_FLAG_RKEY_PTR)) {
-            ucs_trace("protocol requires rkey ptr but it is not "
-                      "supported by the component");
+            ucs_trace("%s: protocol requires rkey ptr but it is not "
+                      "supported by the component", lane_desc);
             continue;
         }
 
-        /* Check memory registration capabilities for zero-copy case */
-        if (reg_mem_type != UCS_MEMORY_TYPE_UNKNOWN) {
-            ucs_assertv((reg_mem_type == select_param->mem_type) ||
-                        !(flags & UCP_PROTO_COMMON_INIT_FLAG_SEND_ZCOPY),
-                        "flags=0x%x reg_mem_type=%s select_param->mem_type=%s",
-                        flags, ucs_memory_type_names[reg_mem_type],
-                        ucs_memory_type_names[select_param->mem_type]);
-
-            if (md_attr->flags & UCT_MD_FLAG_NEED_MEMH) {
-                /* Memory domain must support registration on the relevant
-                 * memory type */
-                if (!(context->reg_md_map[reg_mem_type] & UCS_BIT(md_index))) {
-                    ucs_trace("%s: md %s cannot register %s memory", lane_desc,
-                              context->tl_mds[md_index].rsc.md_name,
-                              ucs_memory_type_names[reg_mem_type]);
-                    continue;
-                }
-            } else if (!(md_attr->access_mem_types & UCS_BIT(reg_mem_type))) {
-                /*
-                 * Memory domain which does not require a registration for zero
-                 * copy operation must be able to access the relevant memory type
-                 */
-                ucs_trace("%s: no access to mem type %s", lane_desc,
-                          ucs_memory_type_names[reg_mem_type]);
-                continue;
-            }
+        if (filter != NULL && !filter(params, lane, lane_desc)) {
+            continue;
         }
 
         /* Check remote access capabilities */
         if (flags & UCP_PROTO_COMMON_INIT_FLAG_REMOTE_ACCESS) {
             if (rkey_config_key == NULL) {
-                ucs_trace("protocol requires remote access but remote key is "
-                          "not present");
+                ucs_trace("%s: protocol requires remote access but remote key "
+                          "is not present", lane_desc);
                 goto out;
             }
 
@@ -614,12 +676,6 @@ ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
             }
         }
 
-        max_iov = ucp_proto_common_get_iface_attr_field(iface_attr,
-                                                        max_iov_offs, SIZE_MAX);
-        if (max_iov < min_iov) {
-            continue;
-        }
-
         ucs_trace("%s: added as lane %d", lane_desc, lane);
         lanes[num_lanes++] = lane;
     }
@@ -663,81 +719,6 @@ ucp_proto_common_reg_md_map(const ucp_proto_common_init_params_t *params,
     return reg_md_map;
 }
 
-static int ucp_proto_common_find_lanes_check_mem_type(
-        const ucp_proto_common_init_params_t *params, ucp_lane_index_t lane)
-{
-    uct_perf_attr_t perf_attr = {0};
-    ucp_rsc_index_t rsc_index;
-    ucp_worker_iface_t *wiface;
-
-    ucp_proto_common_perf_attr_set_mem_type(params, &perf_attr);
-
-    rsc_index = ucp_proto_common_get_rsc_index(&params->super, lane);
-    wiface    = ucp_worker_iface(params->super.worker, rsc_index);
-    /* TODO: Use memory reachability UCT API, when available, to check memory
-       type support */
-    return uct_iface_estimate_perf(wiface->iface, &perf_attr) == UCS_OK;
-}
-
-ucp_lane_index_t ucp_proto_common_find_lanes_with_min_frag(
-        const ucp_proto_common_init_params_t *params, ucp_lane_type_t lane_type,
-        uint64_t tl_cap_flags, ucp_lane_index_t max_lanes,
-        ucp_lane_map_t exclude_map, ucp_lane_index_t *lanes)
-{
-    ucp_lane_index_t lane_index, lane, num_lanes, num_valid_lanes;
-    const uct_iface_attr_t *iface_attr;
-    size_t tl_min_frag, tl_max_frag;
-    ucp_lane_index_t tmp_lanes[UCP_PROTO_MAX_LANES];
-
-    /* TODO: Request more lanes than needed in order to avoid skipping protocol
-     * if the first found candidate is filtered out. Refactor this code to pass
-     * filter callback to ucp_proto_common_find_lanes() */
-    num_lanes = ucp_proto_common_find_lanes(
-                   &params->super, params->flags, params->max_iov_offs,
-                   params->min_iov, lane_type, params->reg_mem_info.type,
-                   tl_cap_flags, ucs_max(max_lanes, 4), exclude_map, tmp_lanes);
-
-    num_valid_lanes = 0;
-    for (lane_index = 0; lane_index < num_lanes; ++lane_index) {
-        lane       = tmp_lanes[lane_index];
-        iface_attr = ucp_proto_common_get_iface_attr(&params->super, lane);
-
-        ucp_proto_common_get_frag_size(params, iface_attr, lane, &tl_min_frag,
-                                       &tl_max_frag);
-
-        /* Minimal fragment size must be 0, unless 'MIN_FRAG' flag is set */
-        if (!(params->flags & UCP_PROTO_COMMON_INIT_FLAG_MIN_FRAG) &&
-            (tl_min_frag > 0)) {
-            ucs_trace("lane[%d]: minimal fragment %zu is not 0", lane,
-                      tl_min_frag);
-            continue;
-        }
-
-        /* Maximal fragment size should be larger than header size */
-        if (tl_max_frag <= params->hdr_size) {
-            ucs_trace("lane[%d]: max fragment is too small %zu, need > %zu",
-                      lane, tl_max_frag, params->hdr_size);
-            continue;
-        }
-
-        if (!ucp_proto_common_find_lanes_check_mem_type(params, lane)) {
-            continue;
-        }
-
-        lanes[num_valid_lanes++] = lane;
-        if (num_valid_lanes >= max_lanes) {
-            break;
-        }
-    }
-
-    if (num_valid_lanes != num_lanes) {
-        ucs_assert(num_valid_lanes < num_lanes);
-        ucs_trace("selected %d/%d valid lanes", num_valid_lanes, num_lanes);
-    }
-
-    return num_valid_lanes;
-}
-
 void ucp_proto_request_zcopy_completion(uct_completion_t *self)
 {
     ucp_request_t *req = ucs_container_of(self, ucp_request_t,

src/ucp/proto/proto_common.h

@@ -283,20 +283,22 @@ ucp_proto_common_get_lane_perf(const ucp_proto_common_init_params_t *params,
                                ucp_proto_perf_node_t **perf_node_p);
 
 
-/* @return number of lanes found */
-ucp_lane_index_t ucp_proto_common_find_lanes_with_min_frag(
-        const ucp_proto_common_init_params_t *params, ucp_lane_type_t lane_type,
-        uint64_t tl_cap_flags, ucp_lane_index_t max_lanes,
-        ucp_lane_map_t exclude_map, ucp_lane_index_t *lanes);
+typedef int (*ucp_proto_common_filter_lane_cb_t)(
+                                const ucp_proto_init_params_t *params,
+                                ucp_lane_index_t lane, const char *lane_desc);
+
+
+int
+ucp_proto_common_filter_min_frag(const ucp_proto_init_params_t *params,
+                                 ucp_lane_index_t lane, const char *lane_desc);
 
 
 ucp_lane_index_t
 ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
-                            unsigned flags, ptrdiff_t max_iov_offs,
-                            size_t min_iov, ucp_lane_type_t lane_type,
-                            ucs_memory_type_t reg_mem_type,
+                            unsigned flags, ucp_lane_type_t lane_type,
                             uint64_t tl_cap_flags, ucp_lane_index_t max_lanes,
                             ucp_lane_map_t exclude_map,
+                            ucp_proto_common_filter_lane_cb_t filter,
                             ucp_lane_index_t *lanes);
 
 

src/ucp/proto/proto_multi.c

@@ -181,20 +181,21 @@ ucs_status_t ucp_proto_multi_init(const ucp_proto_multi_init_params_t *params,
     }
 
     /* Find first lane */
-    num_lanes = ucp_proto_common_find_lanes_with_min_frag(
-            &params->super, params->first.lane_type, params->first.tl_cap_flags,
-            1, 0, lanes);
+    num_lanes = ucp_proto_common_find_lanes(
+            &params->super.super, params->super.flags, params->first.lane_type,
+            params->first.tl_cap_flags, 1, 0, ucp_proto_common_filter_min_frag,
+            lanes);
     if (num_lanes == 0) {
         ucs_trace("no lanes for %s",
                   ucp_proto_id_field(params->super.super.proto_id, name));
         return UCS_ERR_NO_ELEM;
     }
 
     /* Find rest of the lanes */
-    num_lanes += ucp_proto_common_find_lanes_with_min_frag(
-            &params->super, params->middle.lane_type,
+    num_lanes += ucp_proto_common_find_lanes(
+            &params->super.super, params->super.flags, params->middle.lane_type,
             params->middle.tl_cap_flags, UCP_PROTO_MAX_LANES - 1,
-            UCS_BIT(lanes[0]), lanes + 1);
+            UCS_BIT(lanes[0]), ucp_proto_common_filter_min_frag, lanes + 1);
 
     /* Get bandwidth of all lanes and max_bandwidth */
     max_bandwidth = 0;

src/ucp/proto/proto_single.c

@@ -35,9 +35,10 @@ ucs_status_t ucp_proto_single_init(const ucp_proto_single_init_params_t *params,
         return UCS_ERR_UNSUPPORTED;
     }
 
-    num_lanes = ucp_proto_common_find_lanes_with_min_frag(
-            &params->super, params->lane_type, params->tl_cap_flags, 1,
-            params->super.exclude_map, &lane);
+    num_lanes = ucp_proto_common_find_lanes(
+            &params->super.super, params->super.flags, params->lane_type,
+            params->tl_cap_flags, 1, params->super.exclude_map,
+            ucp_proto_common_filter_min_frag, &lane);
     if (num_lanes == 0) {
         ucs_trace("no lanes for %s",
                   ucp_proto_id_field(params->super.super.proto_id, name));

src/ucp/rndv/proto_rndv.c

@@ -504,10 +504,8 @@ ucp_proto_rndv_find_ctrl_lane(const ucp_proto_init_params_t *params)
 
     num_lanes = ucp_proto_common_find_lanes(params,
                                             UCP_PROTO_COMMON_INIT_FLAG_HDR_ONLY,
-                                            UCP_PROTO_COMMON_OFFSET_INVALID, 1,
                                             UCP_LANE_TYPE_AM,
-                                            UCS_MEMORY_TYPE_UNKNOWN,
-                                            UCT_IFACE_FLAG_AM_BCOPY, 1, 0,
+                                            UCT_IFACE_FLAG_AM_BCOPY, 1, 0, NULL,
                                             &lane);
     if (num_lanes == 0) {
         ucs_debug("no active message lane for %s",