@@ -359,14 +359,21 @@ ucp_proto_common_get_lane_perf(const ucp_proto_common_init_params_t *params,
359359 context -> config .est_num_ppn ,
360360 context -> config .est_num_eps );
361361
362- perf_attr .field_mask = UCT_PERF_ATTR_FIELD_OPERATION |
363- UCT_PERF_ATTR_FIELD_SEND_PRE_OVERHEAD |
364- UCT_PERF_ATTR_FIELD_SEND_POST_OVERHEAD |
365- UCT_PERF_ATTR_FIELD_RECV_OVERHEAD |
366- UCT_PERF_ATTR_FIELD_BANDWIDTH |
367- UCT_PERF_ATTR_FIELD_PATH_BANDWIDTH |
368- UCT_PERF_ATTR_FIELD_LATENCY ;
369- perf_attr .operation = params -> send_op ;
362+ perf_attr .field_mask = UCT_PERF_ATTR_FIELD_OPERATION |
363+ UCT_PERF_ATTR_FIELD_LOCAL_MEMORY_TYPE |
364+ UCT_PERF_ATTR_FIELD_SEND_PRE_OVERHEAD |
365+ UCT_PERF_ATTR_FIELD_SEND_POST_OVERHEAD |
366+ UCT_PERF_ATTR_FIELD_RECV_OVERHEAD |
367+ UCT_PERF_ATTR_FIELD_BANDWIDTH |
368+ UCT_PERF_ATTR_FIELD_PATH_BANDWIDTH |
369+ UCT_PERF_ATTR_FIELD_LATENCY ;
370+ perf_attr .operation = params -> send_op ;
371+ perf_attr .local_memory_type = params -> reg_mem_info .type ;
372+
373+ if (params -> super .rkey_config_key != NULL ) {
374+ perf_attr .field_mask |= UCT_PERF_ATTR_FIELD_REMOTE_MEMORY_TYPE ;
375+ perf_attr .remote_memory_type = params -> super .rkey_config_key -> mem_type ;
376+ }
370377
371378 status = ucp_worker_iface_estimate_perf (wiface , & perf_attr );
372379 if (status != UCS_OK ) {
0 commit comments