This source file includes following definitions.
- pml_ucx_generic_datatype_start_pack
- pml_ucx_generic_datatype_start_unpack
- pml_ucx_generic_datatype_packed_size
- pml_ucx_generic_datatype_pack
- pml_ucx_generic_datatype_unpack
- pml_ucx_generic_datatype_finish
- mca_pml_ucx_datatype_attr_del_fn
- mca_pml_ucx_init_datatype
- mca_pml_ucx_convertor_construct
- mca_pml_ucx_convertor_destruct
1
2
3
4
5
6
7
8
9
10 #include "pml_ucx_datatype.h"
11
12 #include "ompi/runtime/mpiruntime.h"
13 #include "ompi/attribute/attribute.h"
14
15 #include <inttypes.h>
16
17
18 static void* pml_ucx_generic_datatype_start_pack(void *context, const void *buffer,
19 size_t count)
20 {
21 ompi_datatype_t *datatype = context;
22 mca_pml_ucx_convertor_t *convertor;
23
24 convertor = (mca_pml_ucx_convertor_t *)PML_UCX_FREELIST_GET(&ompi_pml_ucx.convs);
25
26 OMPI_DATATYPE_RETAIN(datatype);
27 convertor->datatype = datatype;
28 opal_convertor_copy_and_prepare_for_send(ompi_proc_local_proc->super.proc_convertor,
29 &datatype->super, count, buffer, 0,
30 &convertor->opal_conv);
31 return convertor;
32 }
33
34 static void* pml_ucx_generic_datatype_start_unpack(void *context, void *buffer,
35 size_t count)
36 {
37 ompi_datatype_t *datatype = context;
38 mca_pml_ucx_convertor_t *convertor;
39
40 convertor = (mca_pml_ucx_convertor_t *)PML_UCX_FREELIST_GET(&ompi_pml_ucx.convs);
41
42 OMPI_DATATYPE_RETAIN(datatype);
43 convertor->datatype = datatype;
44 convertor->offset = 0;
45 opal_convertor_copy_and_prepare_for_recv(ompi_proc_local_proc->super.proc_convertor,
46 &datatype->super, count, buffer, 0,
47 &convertor->opal_conv);
48 return convertor;
49 }
50
51 static size_t pml_ucx_generic_datatype_packed_size(void *state)
52 {
53 mca_pml_ucx_convertor_t *convertor = state;
54 size_t size;
55
56 opal_convertor_get_packed_size(&convertor->opal_conv, &size);
57 return size;
58 }
59
60 static size_t pml_ucx_generic_datatype_pack(void *state, size_t offset,
61 void *dest, size_t max_length)
62 {
63 mca_pml_ucx_convertor_t *convertor = state;
64 uint32_t iov_count;
65 struct iovec iov;
66 size_t length;
67
68 iov_count = 1;
69 iov.iov_base = dest;
70 iov.iov_len = max_length;
71
72 opal_convertor_set_position(&convertor->opal_conv, &offset);
73 length = max_length;
74 opal_convertor_pack(&convertor->opal_conv, &iov, &iov_count, &length);
75 return length;
76 }
77
78 static ucs_status_t pml_ucx_generic_datatype_unpack(void *state, size_t offset,
79 const void *src, size_t length)
80 {
81 mca_pml_ucx_convertor_t *convertor = state;
82
83 uint32_t iov_count;
84 struct iovec iov;
85 opal_convertor_t conv;
86
87 iov_count = 1;
88 iov.iov_base = (void*)src;
89 iov.iov_len = length;
90
91
92
93 if (offset != convertor->offset) {
94 OBJ_CONSTRUCT(&conv, opal_convertor_t);
95 opal_convertor_copy_and_prepare_for_recv(ompi_proc_local_proc->super.proc_convertor,
96 &convertor->datatype->super,
97 convertor->opal_conv.count,
98 convertor->opal_conv.pBaseBuf, 0,
99 &conv);
100 opal_convertor_set_position(&conv, &offset);
101 opal_convertor_unpack(&conv, &iov, &iov_count, &length);
102 opal_convertor_cleanup(&conv);
103 OBJ_DESTRUCT(&conv);
104
105 convertor->offset = 0;
106 } else {
107 opal_convertor_unpack(&convertor->opal_conv, &iov, &iov_count, &length);
108 convertor->offset += length;
109 }
110 return UCS_OK;
111 }
112
113 static void pml_ucx_generic_datatype_finish(void *state)
114 {
115 mca_pml_ucx_convertor_t *convertor = state;
116
117 opal_convertor_cleanup(&convertor->opal_conv);
118 OMPI_DATATYPE_RELEASE(convertor->datatype);
119 PML_UCX_FREELIST_RETURN(&ompi_pml_ucx.convs, &convertor->super);
120 }
121
122 static ucp_generic_dt_ops_t pml_ucx_generic_datatype_ops = {
123 .start_pack = pml_ucx_generic_datatype_start_pack,
124 .start_unpack = pml_ucx_generic_datatype_start_unpack,
125 .packed_size = pml_ucx_generic_datatype_packed_size,
126 .pack = pml_ucx_generic_datatype_pack,
127 .unpack = pml_ucx_generic_datatype_unpack,
128 .finish = pml_ucx_generic_datatype_finish
129 };
130
131 int mca_pml_ucx_datatype_attr_del_fn(ompi_datatype_t* datatype, int keyval,
132 void *attr_val, void *extra)
133 {
134 ucp_datatype_t ucp_datatype = (ucp_datatype_t)attr_val;
135
136 PML_UCX_ASSERT((uint64_t)ucp_datatype == datatype->pml_data);
137
138 ucp_dt_destroy(ucp_datatype);
139 datatype->pml_data = PML_UCX_DATATYPE_INVALID;
140 return OMPI_SUCCESS;
141 }
142
143 ucp_datatype_t mca_pml_ucx_init_datatype(ompi_datatype_t *datatype)
144 {
145 ucp_datatype_t ucp_datatype;
146 ucs_status_t status;
147 ptrdiff_t lb;
148 size_t size;
149 int ret;
150
151 ompi_datatype_type_lb(datatype, &lb);
152
153 if ((datatype->super.flags & OPAL_DATATYPE_FLAG_CONTIGUOUS) &&
154 (datatype->super.flags & OPAL_DATATYPE_FLAG_NO_GAPS) &&
155 (lb == 0))
156 {
157 ompi_datatype_type_size(datatype, &size);
158 PML_UCX_ASSERT(size > 0);
159 datatype->pml_data = ucp_dt_make_contig(size);
160 return datatype->pml_data;
161 }
162
163 status = ucp_dt_create_generic(&pml_ucx_generic_datatype_ops,
164 datatype, &ucp_datatype);
165 if (status != UCS_OK) {
166 PML_UCX_ERROR("Failed to create UCX datatype for %s", datatype->name);
167 ompi_mpi_abort(&ompi_mpi_comm_world.comm, 1);
168 }
169
170 datatype->pml_data = ucp_datatype;
171
172
173
174
175 if (ompi_datatype_is_predefined(datatype)) {
176 PML_UCX_ASSERT(datatype->id < OMPI_DATATYPE_MAX_PREDEFINED);
177 ompi_pml_ucx.predefined_types[datatype->id] = ucp_datatype;
178 } else {
179 ret = ompi_attr_set_c(TYPE_ATTR, datatype, &datatype->d_keyhash,
180 ompi_pml_ucx.datatype_attr_keyval,
181 (void*)ucp_datatype, false);
182 if (ret != OMPI_SUCCESS) {
183 PML_UCX_ERROR("Failed to add UCX datatype attribute for %s: %d",
184 datatype->name, ret);
185 ompi_mpi_abort(&ompi_mpi_comm_world.comm, 1);
186 }
187 }
188
189 PML_UCX_VERBOSE(7, "created generic UCX datatype 0x%"PRIx64, ucp_datatype)
190
191 return ucp_datatype;
192 }
193
194 static void mca_pml_ucx_convertor_construct(mca_pml_ucx_convertor_t *convertor)
195 {
196 OBJ_CONSTRUCT(&convertor->opal_conv, opal_convertor_t);
197 }
198
199 static void mca_pml_ucx_convertor_destruct(mca_pml_ucx_convertor_t *convertor)
200 {
201 OBJ_DESTRUCT(&convertor->opal_conv);
202 }
203
204 OBJ_CLASS_INSTANCE(mca_pml_ucx_convertor_t,
205 opal_free_list_item_t,
206 mca_pml_ucx_convertor_construct,
207 mca_pml_ucx_convertor_destruct);