1
2
3
4
5
6
7
8
9
10 #ifndef PML_UCX_H_
11 #define PML_UCX_H_
12
13 #include "ompi_config.h"
14 #include "ompi/request/request.h"
15 #include "ompi/mca/pml/pml.h"
16 #include "ompi/mca/pml/base/base.h"
17 #include "ompi/datatype/ompi_datatype.h"
18 #include "ompi/datatype/ompi_datatype_internal.h"
19 #include "ompi/communicator/communicator.h"
20 #include "ompi/request/request.h"
21 #include "opal/mca/common/ucx/common_ucx.h"
22
23 #include <ucp/api/ucp.h>
24 #include "pml_ucx_freelist.h"
25
26 #define PML_UCX_ASSERT MCA_COMMON_UCX_ASSERT
27 #define PML_UCX_ERROR MCA_COMMON_UCX_ERROR
28 #define PML_UCX_VERBOSE MCA_COMMON_UCX_VERBOSE
29
30
31 typedef struct mca_pml_ucx_module mca_pml_ucx_module_t;
32 typedef struct pml_ucx_persistent_request mca_pml_ucx_persistent_request_t;
33 typedef struct pml_ucx_convertor mca_pml_ucx_convertor_t;
34
35
36
37
38
39 struct mca_pml_ucx_module {
40 mca_pml_base_module_t super;
41
42
43 ucp_context_h ucp_context;
44 ucp_worker_h ucp_worker;
45
46
47 int datatype_attr_keyval;
48 ucp_datatype_t predefined_types[OMPI_DATATYPE_MPI_MAX_PREDEFINED];
49
50
51 mca_pml_ucx_freelist_t persistent_reqs;
52 ompi_request_t completed_send_req;
53 size_t request_size;
54 int num_disconnect;
55
56
57 mca_pml_ucx_freelist_t convs;
58
59 int priority;
60 };
61
62 extern mca_pml_base_component_2_0_0_t mca_pml_ucx_component;
63 extern mca_pml_ucx_module_t ompi_pml_ucx;
64
65 int mca_pml_ucx_open(void);
66 int mca_pml_ucx_close(void);
67 int mca_pml_ucx_init(int enable_mpi_threads);
68 int mca_pml_ucx_cleanup(void);
69
70 int mca_pml_ucx_add_procs(struct ompi_proc_t **procs, size_t nprocs);
71 int mca_pml_ucx_del_procs(struct ompi_proc_t **procs, size_t nprocs);
72
73 int mca_pml_ucx_enable(bool enable);
74 int mca_pml_ucx_progress(void);
75
76 int mca_pml_ucx_add_comm(struct ompi_communicator_t* comm);
77 int mca_pml_ucx_del_comm(struct ompi_communicator_t* comm);
78
79 int mca_pml_ucx_irecv_init(void *buf, size_t count, ompi_datatype_t *datatype,
80 int src, int tag, struct ompi_communicator_t* comm,
81 struct ompi_request_t **request);
82
83 int mca_pml_ucx_irecv(void *buf, size_t count, ompi_datatype_t *datatype,
84 int src, int tag, struct ompi_communicator_t* comm,
85 struct ompi_request_t **request);
86
87 int mca_pml_ucx_recv(void *buf, size_t count, ompi_datatype_t *datatype, int src,
88 int tag, struct ompi_communicator_t* comm,
89 ompi_status_public_t* status);
90
91 int mca_pml_ucx_isend_init(const void *buf, size_t count, ompi_datatype_t *datatype,
92 int dst, int tag, mca_pml_base_send_mode_t mode,
93 struct ompi_communicator_t* comm,
94 struct ompi_request_t **request);
95
96 int mca_pml_ucx_isend(const void *buf, size_t count, ompi_datatype_t *datatype,
97 int dst, int tag, mca_pml_base_send_mode_t mode,
98 struct ompi_communicator_t* comm,
99 struct ompi_request_t **request);
100
101 int mca_pml_ucx_send(const void *buf, size_t count, ompi_datatype_t *datatype, int dst,
102 int tag, mca_pml_base_send_mode_t mode,
103 struct ompi_communicator_t* comm);
104
105 int mca_pml_ucx_iprobe(int src, int tag, struct ompi_communicator_t* comm,
106 int *matched, ompi_status_public_t* status);
107
108 int mca_pml_ucx_probe(int src, int tag, struct ompi_communicator_t* comm,
109 ompi_status_public_t* status);
110
111 int mca_pml_ucx_improbe(int src, int tag, struct ompi_communicator_t* comm,
112 int *matched, struct ompi_message_t **message,
113 ompi_status_public_t* status);
114
115 int mca_pml_ucx_mprobe(int src, int tag, struct ompi_communicator_t* comm,
116 struct ompi_message_t **message,
117 ompi_status_public_t* status);
118
119 int mca_pml_ucx_imrecv(void *buf, size_t count, ompi_datatype_t *datatype,
120 struct ompi_message_t **message,
121 struct ompi_request_t **request);
122
123 int mca_pml_ucx_mrecv(void *buf, size_t count, ompi_datatype_t *datatype,
124 struct ompi_message_t **message,
125 ompi_status_public_t* status);
126
127 int mca_pml_ucx_start(size_t count, ompi_request_t** requests);
128
129 int mca_pml_ucx_dump(struct ompi_communicator_t* comm, int verbose);
130
131
132 #endif