ucc/src/components/tl/ucp/tl_ucp_ep.h at 70480e7dcc43d74a50f1d1bcf01113df3a8365b6 · openucx/ucc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
/**
 * Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See file LICENSE for terms.
 */

#include "config.h"

#ifndef UCC_TL_UCP_EP_H_
#define UCC_TL_UCP_EP_H_
#include "ucc/api/ucc.h"
#include <ucp/api/ucp.h>
#include "tl_ucp.h"
#include "core/ucc_team.h"

/* TL/UCP endpoint address layout: (ucp_addrlen may vary per proc)

   [worker->ucp_addrlen][ucp_worker_address][onesided_info]
       8 bytes    ucp_addrlen bytes

    If a special service worker is set through UCC_TL_UCP_SERVICE_TLS:
   [worker->ucp_addrlen][ucp_worker_address][service_worker->ucp_addrlen][ucp_service_worker_address][onesided_info]
       8 bytes    ucp_addrlen bytes      8 bytes        service.ucp_addrlen bytes
*/
#define TL_UCP_EP_ADDRLEN_SIZE 8
#define TL_UCP_EP_ADDR_WORKER_LEN(_addr) (*((uint64_t*)(_addr)))
#define TL_UCP_EP_ADDR_WORKER(_addr)     PTR_OFFSET((_addr), TL_UCP_EP_ADDRLEN_SIZE)
#define TL_UCP_EP_OFFSET_WORKER_INFO(_addr)                                    \
    PTR_OFFSET((_addr),                                                        \
               TL_UCP_EP_ADDRLEN_SIZE + TL_UCP_EP_ADDR_WORKER_LEN(_addr))
#define TL_UCP_EP_ADDR_WORKER_SERVICE(_addr)                                   \
    TL_UCP_EP_ADDR_WORKER(TL_UCP_EP_OFFSET_WORKER_INFO(_addr))
#define TL_UCP_EP_ADDR_ONESIDED_INFO(_addr, _ctx)                              \
    _ctx->cfg.service_worker                                                   \
        ? TL_UCP_EP_OFFSET_WORKER_INFO(TL_UCP_EP_OFFSET_WORKER_INFO(_addr))    \
        : TL_UCP_EP_OFFSET_WORKER_INFO(_addr)

typedef struct ucc_tl_ucp_context ucc_tl_ucp_context_t;
typedef struct ucc_tl_ucp_team    ucc_tl_ucp_team_t;

ucc_status_t ucc_tl_ucp_connect_team_ep(ucc_tl_ucp_team_t *team,
                                        ucc_rank_t team_rank, ucp_ep_h *ep);

void         ucc_tl_ucp_close_eps(
            ucc_tl_ucp_worker_t *worker, ucc_tl_ucp_context_t *ctx);

static inline ucc_context_addr_header_t *
ucc_tl_ucp_get_team_ep_header(ucc_tl_ucp_team_t *team, ucc_rank_t core_rank)

{
    return ucc_get_team_ep_header(UCC_TL_CORE_CTX(team), UCC_TL_CORE_TEAM(team),
                                  core_rank);
}

static inline ucc_status_t ucc_tl_ucp_get_ep(ucc_tl_ucp_team_t *team,
                                             ucc_rank_t rank, ucp_ep_h *ep)
{
    ucc_context_addr_header_t *h        = NULL;
    ucc_rank_t                 ctx_rank = 0;
    ucc_status_t               status;
    ucc_rank_t                 core_rank;
    core_rank = ucc_ep_map_eval(UCC_TL_TEAM_MAP(team), rank);
    if (team->worker->eps) {
        ucc_team_t *core_team = UCC_TL_CORE_TEAM(team);
        /* Core super.super.team ptr is NULL for service_team
           which has scope == UCC_CL_LAST + 1*/
        ucc_assert((NULL != core_team) || UCC_TL_IS_SERVICE_TEAM(team));
        ctx_rank = core_team ? ucc_get_ctx_rank(core_team, core_rank)
                       : core_rank;
        *ep      = team->worker->eps[ctx_rank];
    } else {
        h   = ucc_tl_ucp_get_team_ep_header(team, core_rank);
        *ep = tl_ucp_hash_get(team->worker->ep_hash, h->ctx_id);
    }
    if (NULL == (*ep)) {
        /* Not connected yet */
        status = ucc_tl_ucp_connect_team_ep(team, core_rank, ep);
        if (ucc_unlikely(UCC_OK != status)) {
            tl_error(UCC_TL_TEAM_LIB(team), "failed to connect team ep");
            *ep = NULL;
            return status;
        }
        if (!h) {
            team->worker->eps[ctx_rank] = *ep;
        } else {
            tl_ucp_hash_put(team->worker->ep_hash, h->ctx_id, *ep);
        }
    }
    return UCC_OK;
}

#endif