Line data Source code
1 : /*----------------------------------------------------------------------------*/
2 : /* CP2K: A general program to perform molecular dynamics simulations */
3 : /* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
4 : /* */
5 : /* SPDX-License-Identifier: BSD-3-Clause */
6 : /*----------------------------------------------------------------------------*/
7 : #include "dbm_multiply_comm.h"
8 : #include "../mpiwrap/cp_mpi.h"
9 : #include "../offload/offload_mempool.h"
10 :
11 : #include <assert.h>
12 : #include <limits.h>
13 : #include <omp.h>
14 : #include <stdlib.h>
15 : #include <string.h>
16 :
17 : #if 1
18 : #define DBM_MULTIPLY_COMM_MEMPOOL
19 : #endif
20 :
21 : /*******************************************************************************
22 : * \brief Private routine for allocating a communication buffer, which is taken
23 : * from the memory pool (DBM_MULTIPLY_COMM_MEMPOOL) rather than mapping
24 : * fresh pages for every multiplication.
25 : * \author Hans Pabst
26 : ******************************************************************************/
27 4269224 : static void *allocate_comm_buffer(const size_t size) {
28 : #if defined(DBM_MULTIPLY_COMM_MEMPOOL)
29 4269224 : return offload_mempool_host_malloc(size);
30 : #else
31 : return cp_mpi_alloc_mem(size);
32 : #endif
33 : }
34 :
35 : /*******************************************************************************
36 : * \brief Private routine for releasing a buffer from allocate_comm_buffer.
37 : * \author Hans Pabst
38 : ******************************************************************************/
39 4269224 : static void free_comm_buffer(void *mem) {
40 : #if defined(DBM_MULTIPLY_COMM_MEMPOOL)
41 4269224 : offload_mempool_host_free(mem);
42 : #else
43 : cp_mpi_free_mem(mem);
44 : #endif
45 4269224 : }
46 :
47 : /*******************************************************************************
48 : * \brief Private routine for computing greatest common divisor of two numbers.
49 : * \author Ole Schuett
50 : ******************************************************************************/
51 734778 : static int gcd(const int a, const int b) {
52 734778 : if (a == 0) {
53 : return b;
54 : }
55 384840 : return gcd(b % a, a); // Euclid's algorithm.
56 : }
57 :
58 : /*******************************************************************************
59 : * \brief Private routine for computing least common multiple of two numbers.
60 : * \author Ole Schuett
61 : ******************************************************************************/
62 349938 : static int lcm(const int a, const int b) { return (a * b) / gcd(a, b); }
63 :
64 : /*******************************************************************************
65 : * \brief Private routine for converting element counts to byte counts.
66 : * \author Hans Pabst
67 : ******************************************************************************/
68 3429216 : static int checked_byte_count(const int nelements, const size_t element_size) {
69 : // MPI displacements/counts are int; unchecked multiplication could overflow
70 : // silently and corrupt the alltoallv layout.
71 3429216 : assert(0 <= nelements);
72 3429216 : assert(element_size <= INT_MAX);
73 3429216 : assert(nelements <= INT_MAX / (int)element_size);
74 3429216 : return nelements * (int)element_size;
75 : }
76 :
77 : /*******************************************************************************
78 : * \brief Private routine for computing the sum of the given integers.
79 : * \author Ole Schuett
80 : ******************************************************************************/
81 : static inline int isum(const int n, const int input[n]) {
82 : int output = 0;
83 4828968 : for (int i = 0; i < n; i++) {
84 1679624 : output += input[i];
85 : }
86 1469720 : return output;
87 : }
88 :
89 : /*******************************************************************************
90 : * \brief Private routine for computing the cumulative sums of given numbers.
91 : * \author Ole Schuett and Hans Pabst
92 : ******************************************************************************/
93 1469720 : static inline void icumsum(const int n, const int input[n], int output[n]) {
94 : // Exclusive prefix sum using carried accumulators to avoid a load from
95 : // output[i-1] each iteration (removes loop-carried memory dependency).
96 1469720 : int oval = output[0] = 0, ival = input[0];
97 1679624 : for (int i = 1; i < n; i++) {
98 209904 : output[i] = (oval += ival);
99 209904 : ival = input[i];
100 : }
101 : }
102 :
103 : /*******************************************************************************
104 : * \brief Private routine computing received data counts from block metadata.
105 : * \author Hans Pabst
106 : ******************************************************************************/
107 734860 : static void compute_data_recv_count(const int nranks,
108 : const int blks_recv_count[nranks],
109 : const int blks_recv_displ[nranks],
110 : const int free_index_sizes[],
111 : const int sum_index_sizes[],
112 : const dbm_pack_block_t blks_recv[],
113 : int data_recv_count[nranks]) {
114 : // Derives per-rank data counts locally from the already-received block
115 : // metadata, eliminating a collective alltoall for exchanging data amounts.
116 734860 : memset(data_recv_count, 0, nranks * sizeof(int));
117 1574672 : for (int irank = 0; irank < nranks; irank++) {
118 15199429 : for (int i = 0; i < blks_recv_count[irank]; i++) {
119 14359617 : const dbm_pack_block_t *const blk =
120 14359617 : &blks_recv[blks_recv_displ[irank] + i];
121 14359617 : const int block_size =
122 14359617 : free_index_sizes[blk->free_index] * sum_index_sizes[blk->sum_index];
123 14359617 : assert(block_size >= 0);
124 14359617 : assert(data_recv_count[irank] <= INT_MAX - block_size);
125 14359617 : data_recv_count[irank] += block_size;
126 : }
127 : }
128 734860 : }
129 :
130 : /*******************************************************************************
131 : * \brief Private struct used for planing during pack_matrix.
132 : * \author Ole Schuett
133 : ******************************************************************************/
134 : typedef struct {
135 : const dbm_block_t *blk; // source block
136 : int rank; // target mpi rank
137 : int row_size;
138 : int col_size;
139 : } plan_t;
140 :
141 : /*******************************************************************************
142 : * \brief Private routine for calculating tick indices in pack plans.
143 : * \author Maximilian Graml
144 : ******************************************************************************/
145 : static inline unsigned long long calculate_tick_index(int sum_index,
146 : int nticks) {
147 : // 1021 is used as a random prime to scramble the index
148 : return ((unsigned long long)sum_index * 1021ULL) % (unsigned long long)nticks;
149 : }
150 :
151 : /*******************************************************************************
152 : * \brief Private routine for planing packs.
153 : * \author Ole Schuett
154 : ******************************************************************************/
155 699876 : static void create_pack_plans(const bool trans_matrix, const bool trans_dist,
156 : const dbm_matrix_t *matrix,
157 : const cp_mpi_comm_t comm,
158 : const dbm_dist_1d_t *dist_indices,
159 : const dbm_dist_1d_t *dist_ticks, const int nticks,
160 : const int npacks, plan_t *plans_per_pack[npacks],
161 : int nblks_per_pack[npacks],
162 : int ndata_per_pack[npacks]) {
163 699876 : memset(nblks_per_pack, 0, npacks * sizeof(int));
164 699876 : memset(ndata_per_pack, 0, npacks * sizeof(int));
165 :
166 699876 : #pragma omp parallel
167 : {
168 : // 1st pass: Compute number of blocks that will be send in each pack.
169 : int nblks_mythread[npacks];
170 : memset(nblks_mythread, 0, npacks * sizeof(int));
171 : #pragma omp for schedule(static)
172 : for (int ishard = 0; ishard < dbm_get_num_shards(matrix); ishard++) {
173 : dbm_shard_t *shard = &matrix->shards[ishard];
174 : for (int iblock = 0; iblock < shard->nblocks; iblock++) {
175 : const dbm_block_t *blk = &shard->blocks[iblock];
176 : const int sum_index = (trans_matrix) ? blk->row : blk->col;
177 : unsigned long long itick64 = calculate_tick_index(sum_index, nticks);
178 : const int ipack = itick64 / dist_ticks->nranks;
179 : nblks_mythread[ipack]++;
180 : }
181 : }
182 :
183 : // Sum nblocks across threads and allocate arrays for plans.
184 : #pragma omp critical
185 : for (int ipack = 0; ipack < npacks; ipack++) {
186 : nblks_per_pack[ipack] += nblks_mythread[ipack];
187 : nblks_mythread[ipack] = nblks_per_pack[ipack];
188 : }
189 : #pragma omp barrier
190 : #pragma omp for
191 : for (int ipack = 0; ipack < npacks; ipack++) {
192 : const int nblks = nblks_per_pack[ipack];
193 : plans_per_pack[ipack] = malloc(nblks * sizeof(plan_t));
194 : assert(plans_per_pack[ipack] != NULL || nblks == 0);
195 : }
196 :
197 : // 2nd pass: Plan where to send each block.
198 : int ndata_mythread[npacks];
199 : memset(ndata_mythread, 0, npacks * sizeof(int));
200 : #pragma omp for schedule(static) // Need static to match previous loop.
201 : for (int ishard = 0; ishard < dbm_get_num_shards(matrix); ishard++) {
202 : dbm_shard_t *shard = &matrix->shards[ishard];
203 : for (int iblock = 0; iblock < shard->nblocks; iblock++) {
204 : const dbm_block_t *blk = &shard->blocks[iblock];
205 : const int free_index = (trans_matrix) ? blk->col : blk->row;
206 : const int sum_index = (trans_matrix) ? blk->row : blk->col;
207 : unsigned long long itick64 = calculate_tick_index(sum_index, nticks);
208 : const int ipack = itick64 / dist_ticks->nranks;
209 : // Compute rank to which this block should be sent.
210 : const int coord_free_idx = dist_indices->index2coord[free_index];
211 : const int coord_sum_idx = itick64 % dist_ticks->nranks;
212 : const int coords[2] = {(trans_dist) ? coord_sum_idx : coord_free_idx,
213 : (trans_dist) ? coord_free_idx : coord_sum_idx};
214 : const int rank = cp_mpi_cart_rank(comm, coords);
215 : const int row_size = matrix->row_sizes[blk->row];
216 : const int col_size = matrix->col_sizes[blk->col];
217 : ndata_mythread[ipack] += row_size * col_size;
218 : // Create plan.
219 : const int iplan = --nblks_mythread[ipack];
220 : plans_per_pack[ipack][iplan].blk = blk;
221 : plans_per_pack[ipack][iplan].rank = rank;
222 : plans_per_pack[ipack][iplan].row_size = row_size;
223 : plans_per_pack[ipack][iplan].col_size = col_size;
224 : }
225 : }
226 : #pragma omp critical
227 : for (int ipack = 0; ipack < npacks; ipack++) {
228 : ndata_per_pack[ipack] += ndata_mythread[ipack];
229 : }
230 : } // end of omp parallel region
231 699876 : }
232 :
233 : /*******************************************************************************
234 : * \brief Private routine for filling send buffers.
235 : * \author Ole Schuett
236 : ******************************************************************************/
237 734860 : static void fill_send_buffers(
238 : const dbm_matrix_t *matrix, const bool trans_matrix, const int nblks_send,
239 : const int ndata_send, plan_t plans[nblks_send], const int nranks,
240 : int blks_send_count[nranks], int data_send_count[nranks],
241 : int blks_send_displ[nranks], int data_send_displ[nranks],
242 : dbm_pack_block_t blks_send[nblks_send], double data_send[ndata_send]) {
243 734860 : memset(blks_send_count, 0, nranks * sizeof(int));
244 734860 : memset(data_send_count, 0, nranks * sizeof(int));
245 :
246 734860 : #pragma omp parallel
247 : {
248 : // 3th pass: Compute per rank nblks and ndata.
249 : int nblks_mythread[nranks], ndata_mythread[nranks];
250 : memset(nblks_mythread, 0, nranks * sizeof(int));
251 : memset(ndata_mythread, 0, nranks * sizeof(int));
252 : #pragma omp for schedule(static)
253 : for (int iblock = 0; iblock < nblks_send; iblock++) {
254 : const plan_t *plan = &plans[iblock];
255 : nblks_mythread[plan->rank] += 1;
256 : ndata_mythread[plan->rank] += plan->row_size * plan->col_size;
257 : }
258 :
259 : // Sum nblks and ndata across threads.
260 : #pragma omp critical
261 : for (int irank = 0; irank < nranks; irank++) {
262 : blks_send_count[irank] += nblks_mythread[irank];
263 : data_send_count[irank] += ndata_mythread[irank];
264 : nblks_mythread[irank] = blks_send_count[irank];
265 : ndata_mythread[irank] = data_send_count[irank];
266 : }
267 : #pragma omp barrier
268 :
269 : // Compute send displacements.
270 : #pragma omp single
271 : {
272 : icumsum(nranks, blks_send_count, blks_send_displ);
273 : icumsum(nranks, data_send_count, data_send_displ);
274 : const int m = nranks - 1;
275 : assert(nblks_send == blks_send_displ[m] + blks_send_count[m]);
276 : assert(ndata_send == data_send_displ[m] + data_send_count[m]);
277 : }
278 : #pragma omp barrier
279 :
280 : // 4th pass: Fill blks_send and data_send arrays.
281 : #pragma omp for schedule(static) // Need static to match previous loop.
282 : for (int iblock = 0; iblock < nblks_send; iblock++) {
283 : const plan_t *const plan = &plans[iblock];
284 : const dbm_block_t *const blk = plan->blk;
285 : const int ishard = dbm_get_shard_index(matrix, blk->row, blk->col);
286 : const dbm_shard_t *const shard = &matrix->shards[ishard];
287 : const double *blk_data = &shard->data[blk->offset];
288 : const int row_size = plan->row_size, col_size = plan->col_size;
289 : const int plan_size = row_size * col_size;
290 : const int irank = plan->rank;
291 :
292 : // The blk_send_data is ordered by rank, thread, and block.
293 : // data_send_displ[irank]: Start of data for irank within blk_send_data.
294 : // ndata_mythread[irank]: Current threads offset within data for irank.
295 : nblks_mythread[irank] -= 1;
296 : ndata_mythread[irank] -= plan_size;
297 : const int offset = data_send_displ[irank] + ndata_mythread[irank];
298 : const int jblock = blks_send_displ[irank] + nblks_mythread[irank];
299 :
300 : double norm = 0.0; // Compute norm as double...
301 : if (trans_matrix) {
302 : // Transpose block to allow for outer-product style multiplication.
303 : for (int i = 0; i < row_size; i++) {
304 : for (int j = 0; j < col_size; j++) {
305 : const double element = blk_data[j * row_size + i];
306 : data_send[offset + i * col_size + j] = element;
307 : norm += element * element;
308 : }
309 : }
310 : blks_send[jblock].free_index = plan->blk->col;
311 : blks_send[jblock].sum_index = plan->blk->row;
312 : } else {
313 : for (int i = 0; i < plan_size; i++) {
314 : const double element = blk_data[i];
315 : data_send[offset + i] = element;
316 : norm += element * element;
317 : }
318 : blks_send[jblock].free_index = plan->blk->row;
319 : blks_send[jblock].sum_index = plan->blk->col;
320 : }
321 : blks_send[jblock].norm = (float)norm; // ...store norm as float.
322 :
323 : // After the block exchange data_recv_displ will be added to the offsets.
324 : blks_send[jblock].offset = offset - data_send_displ[irank];
325 : }
326 : } // end of omp parallel region
327 734860 : }
328 :
329 : /*******************************************************************************
330 : * \brief Private comperator passed to qsort to compare two blocks by sum_index.
331 : * \author Ole Schuett
332 : ******************************************************************************/
333 73911881 : static int compare_pack_blocks_by_sum_index(const void *a, const void *b) {
334 73911881 : const dbm_pack_block_t *blk_a = (const dbm_pack_block_t *)a;
335 73911881 : const dbm_pack_block_t *blk_b = (const dbm_pack_block_t *)b;
336 73911881 : return blk_a->sum_index - blk_b->sum_index;
337 : }
338 :
339 : /*******************************************************************************
340 : * \brief Private routine for post-processing received blocks.
341 : * \author Ole Schuett
342 : ******************************************************************************/
343 734860 : static void postprocess_received_blocks(
344 : const int nranks, const int nshards, const int nblocks_recv,
345 : const int blks_recv_count[nranks], const int blks_recv_displ[nranks],
346 : const int data_recv_displ[nranks],
347 734860 : dbm_pack_block_t blks_recv[nblocks_recv]) {
348 734860 : int nblocks_per_shard[nshards], shard_start[nshards];
349 734860 : memset(nblocks_per_shard, 0, nshards * sizeof(int));
350 734860 : dbm_pack_block_t *blocks_tmp =
351 734860 : malloc(nblocks_recv * sizeof(dbm_pack_block_t));
352 734860 : assert(blocks_tmp != NULL || nblocks_recv == 0);
353 :
354 734860 : #pragma omp parallel
355 : {
356 : // Add data_recv_displ to recveived block offsets.
357 : for (int irank = 0; irank < nranks; irank++) {
358 : #pragma omp for
359 : for (int i = 0; i < blks_recv_count[irank]; i++) {
360 : blks_recv[blks_recv_displ[irank] + i].offset += data_recv_displ[irank];
361 : }
362 : }
363 :
364 : // First use counting sort to group blocks by their free_index shard.
365 : int nblocks_mythread[nshards];
366 : memset(nblocks_mythread, 0, nshards * sizeof(int));
367 : #pragma omp for schedule(static)
368 : for (int iblock = 0; iblock < nblocks_recv; iblock++) {
369 : blocks_tmp[iblock] = blks_recv[iblock];
370 : const int ishard = blks_recv[iblock].free_index % nshards;
371 : nblocks_mythread[ishard]++;
372 : }
373 : #pragma omp critical
374 : for (int ishard = 0; ishard < nshards; ishard++) {
375 : nblocks_per_shard[ishard] += nblocks_mythread[ishard];
376 : nblocks_mythread[ishard] = nblocks_per_shard[ishard];
377 : }
378 : #pragma omp barrier
379 : #pragma omp single
380 : icumsum(nshards, nblocks_per_shard, shard_start);
381 : #pragma omp barrier
382 : #pragma omp for schedule(static) // Need static to match previous loop.
383 : for (int iblock = 0; iblock < nblocks_recv; iblock++) {
384 : const int ishard = blocks_tmp[iblock].free_index % nshards;
385 : const int jblock = --nblocks_mythread[ishard] + shard_start[ishard];
386 : blks_recv[jblock] = blocks_tmp[iblock];
387 : }
388 :
389 : // Then sort blocks within each shard by their sum_index.
390 : #pragma omp for
391 : for (int ishard = 0; ishard < nshards; ishard++) {
392 : if (nblocks_per_shard[ishard] > 1) {
393 : qsort(&blks_recv[shard_start[ishard]], nblocks_per_shard[ishard],
394 : sizeof(dbm_pack_block_t), &compare_pack_blocks_by_sum_index);
395 : }
396 : }
397 : } // end of omp parallel region
398 :
399 734860 : free(blocks_tmp);
400 734860 : }
401 :
402 : /*******************************************************************************
403 : * \brief Private routine for redistributing a matrix along selected dimensions.
404 : * \author Ole Schuett
405 : ******************************************************************************/
406 699876 : static dbm_packed_matrix_t pack_matrix(const bool trans_matrix,
407 : const bool trans_dist,
408 : const dbm_matrix_t *restrict matrix,
409 : const dbm_distribution_t *restrict dist,
410 : const int nticks,
411 699876 : dbm_comm_iterator_t *iter) {
412 699876 : assert(cp_mpi_comms_are_similar(matrix->dist->comm, dist->comm));
413 :
414 : // The row/col indicies are distributed along one cart dimension and the
415 : // ticks are distributed along the other cart dimension.
416 699876 : const dbm_dist_1d_t *dist_indices = (trans_dist) ? &dist->cols : &dist->rows;
417 699876 : const dbm_dist_1d_t *dist_ticks = (trans_dist) ? &dist->rows : &dist->cols;
418 699876 : const int *free_index_sizes =
419 : (trans_matrix) ? matrix->col_sizes : matrix->row_sizes;
420 699876 : const int *sum_index_sizes =
421 : (trans_matrix) ? matrix->row_sizes : matrix->col_sizes;
422 :
423 : // Allocate packed matrix.
424 699876 : const int nsend_packs = nticks / dist_ticks->nranks;
425 699876 : assert(nsend_packs * dist_ticks->nranks == nticks);
426 699876 : dbm_packed_matrix_t packed;
427 699876 : packed.dist_indices = dist_indices;
428 699876 : packed.dist_ticks = dist_ticks;
429 699876 : packed.nsend_packs = nsend_packs;
430 699876 : packed.send_packs = malloc(nsend_packs * sizeof(dbm_pack_t));
431 699876 : assert(packed.send_packs != NULL || nsend_packs == 0);
432 :
433 : // Plan all packs.
434 699876 : plan_t *plans_per_pack[nsend_packs];
435 699876 : int nblks_send_per_pack[nsend_packs], ndata_send_per_pack[nsend_packs];
436 699876 : create_pack_plans(trans_matrix, trans_dist, matrix, dist->comm, dist_indices,
437 : dist_ticks, nticks, nsend_packs, plans_per_pack,
438 : nblks_send_per_pack, ndata_send_per_pack);
439 :
440 : // Allocate send buffers for maximum number of blocks/data over all packs.
441 699876 : int nblks_send_max = 0, ndata_send_max = 0;
442 1434736 : for (int ipack = 0; ipack < nsend_packs; ++ipack) {
443 734860 : nblks_send_max = imax(nblks_send_max, nblks_send_per_pack[ipack]);
444 734860 : ndata_send_max = imax(ndata_send_max, ndata_send_per_pack[ipack]);
445 : }
446 699876 : dbm_pack_block_t *blks_send =
447 699876 : allocate_comm_buffer(nblks_send_max * sizeof(dbm_pack_block_t));
448 699876 : double *data_send = allocate_comm_buffer(ndata_send_max * sizeof(double));
449 :
450 : // Cannot parallelize over packs (there might be too few of them).
451 1434736 : for (int ipack = 0; ipack < nsend_packs; ipack++) {
452 : // Fill send buffers according to plans.
453 734860 : const int nranks = dist->nranks;
454 734860 : int blks_send_count[nranks], data_send_count[nranks];
455 734860 : int blks_send_displ[nranks], data_send_displ[nranks];
456 734860 : fill_send_buffers(matrix, trans_matrix, nblks_send_per_pack[ipack],
457 734860 : ndata_send_per_pack[ipack], plans_per_pack[ipack], nranks,
458 : blks_send_count, data_send_count, blks_send_displ,
459 : data_send_displ, blks_send, data_send);
460 734860 : free(plans_per_pack[ipack]);
461 :
462 : // 1st communication: Exchange block counts.
463 734860 : double tick = omp_get_wtime();
464 734860 : int blks_recv_count[nranks], blks_recv_displ[nranks];
465 734860 : cp_mpi_alltoall_int(blks_send_count, 1, blks_recv_count, 1, dist->comm);
466 2309532 : icumsum(nranks, blks_recv_count, blks_recv_displ);
467 734860 : const int nblocks_recv = isum(nranks, blks_recv_count);
468 :
469 : // 2nd communication: Exchange blocks.
470 734860 : dbm_pack_block_t *blks_recv =
471 734860 : allocate_comm_buffer(nblocks_recv * sizeof(dbm_pack_block_t));
472 734860 : int blks_send_count_byte[nranks], blks_send_displ_byte[nranks];
473 734860 : int blks_recv_count_byte[nranks], blks_recv_displ_byte[nranks];
474 1574672 : for (int i = 0; i < nranks; i++) { // TODO: this is ugly!
475 1679624 : blks_send_count_byte[i] =
476 839812 : checked_byte_count(blks_send_count[i], sizeof(dbm_pack_block_t));
477 1679624 : blks_send_displ_byte[i] =
478 839812 : checked_byte_count(blks_send_displ[i], sizeof(dbm_pack_block_t));
479 1679624 : blks_recv_count_byte[i] =
480 839812 : checked_byte_count(blks_recv_count[i], sizeof(dbm_pack_block_t));
481 839812 : blks_recv_displ_byte[i] =
482 839812 : checked_byte_count(blks_recv_displ[i], sizeof(dbm_pack_block_t));
483 : }
484 1574672 : for (int i = 0; i < nranks; i++) {
485 839812 : iter->bytes_alltoall += blks_send_count_byte[i] +
486 839812 : (int64_t)data_send_count[i] * sizeof(double);
487 : }
488 734860 : cp_mpi_alltoallv_byte(blks_send, blks_send_count_byte, blks_send_displ_byte,
489 : blks_recv, blks_recv_count_byte, blks_recv_displ_byte,
490 : dist->comm);
491 :
492 : // 3rd compute data counts from the received block metadata.
493 734860 : int data_recv_count[nranks], data_recv_displ[nranks];
494 734860 : compute_data_recv_count(nranks, blks_recv_count, blks_recv_displ,
495 : free_index_sizes, sum_index_sizes, blks_recv,
496 : data_recv_count);
497 2309532 : icumsum(nranks, data_recv_count, data_recv_displ);
498 734860 : const int ndata_recv = isum(nranks, data_recv_count);
499 :
500 : // 4th communication: Exchange data.
501 734860 : double *data_recv = allocate_comm_buffer(ndata_recv * sizeof(double));
502 734860 : cp_mpi_alltoallv_double(data_send, data_send_count, data_send_displ,
503 : data_recv, data_recv_count, data_recv_displ,
504 : dist->comm);
505 :
506 : // 5th post-process received blocks and assemble them into a pack.
507 734860 : const double tock = omp_get_wtime();
508 734860 : iter->seconds_alltoall += tock - tick;
509 734860 : postprocess_received_blocks(nranks, dist_indices->nshards, nblocks_recv,
510 : blks_recv_count, blks_recv_displ,
511 : data_recv_displ, blks_recv);
512 734860 : iter->seconds_sort += omp_get_wtime() - tock;
513 734860 : packed.send_packs[ipack].nblocks = nblocks_recv;
514 734860 : packed.send_packs[ipack].data_size = ndata_recv;
515 734860 : packed.send_packs[ipack].blocks = blks_recv;
516 734860 : packed.send_packs[ipack].data = data_recv;
517 : }
518 :
519 : // Deallocate send buffers.
520 699876 : free_comm_buffer(blks_send);
521 699876 : free_comm_buffer(data_send);
522 :
523 : // Allocate pack_recv.
524 699876 : int max_nblocks = 0, max_data_size = 0;
525 1434736 : for (int ipack = 0; ipack < packed.nsend_packs; ipack++) {
526 734860 : max_nblocks = imax(max_nblocks, packed.send_packs[ipack].nblocks);
527 734860 : max_data_size = imax(max_data_size, packed.send_packs[ipack].data_size);
528 : }
529 699876 : const double tick = omp_get_wtime();
530 699876 : cp_mpi_max_int(&max_nblocks, 1, packed.dist_ticks->comm);
531 699876 : cp_mpi_max_int(&max_data_size, 1, packed.dist_ticks->comm);
532 699876 : iter->seconds_alltoall += omp_get_wtime() - tick;
533 699876 : packed.max_nblocks = max_nblocks;
534 699876 : packed.max_data_size = max_data_size;
535 1399752 : packed.recv_pack.blocks =
536 699876 : allocate_comm_buffer(packed.max_nblocks * sizeof(dbm_pack_block_t));
537 1399752 : packed.recv_pack.data =
538 699876 : allocate_comm_buffer(packed.max_data_size * sizeof(double));
539 :
540 699876 : return packed; // Ownership of packed transfers to caller.
541 : }
542 :
543 : /*******************************************************************************
544 : * \brief Private routine for sending and receiving the pack for the given tick.
545 : * \author Ole Schuett
546 : ******************************************************************************/
547 769844 : static dbm_pack_t *sendrecv_pack(const int itick, const int nticks,
548 : dbm_packed_matrix_t *packed, int64_t *bytes) {
549 769844 : const int nranks = packed->dist_ticks->nranks;
550 769844 : const int my_rank = packed->dist_ticks->my_rank;
551 :
552 : // Compute send rank and pack.
553 769844 : const int itick_of_rank0 = (itick + nticks - my_rank) % nticks;
554 769844 : const int send_rank = (my_rank + nticks - itick_of_rank0) % nranks;
555 769844 : const int send_itick = (itick_of_rank0 + send_rank) % nticks;
556 769844 : const int send_ipack = send_itick / nranks;
557 769844 : assert(send_itick % nranks == my_rank);
558 :
559 : // Compute receive rank and pack.
560 769844 : const int recv_rank = itick % nranks;
561 769844 : const int recv_ipack = itick / nranks;
562 :
563 769844 : dbm_pack_t *send_pack = &packed->send_packs[send_ipack];
564 769844 : if (send_rank == my_rank) {
565 734860 : assert(send_rank == recv_rank && send_ipack == recv_ipack);
566 : return send_pack; // Local pack, no mpi needed.
567 : } else {
568 34984 : *bytes += (int64_t)send_pack->nblocks * sizeof(dbm_pack_block_t) +
569 34984 : (int64_t)send_pack->data_size * sizeof(double);
570 : // Exchange blocks.
571 104952 : const int nblocks_in_bytes = cp_mpi_sendrecv_byte(
572 34984 : /*sendbuf=*/send_pack->blocks,
573 : /*sendcound=*/
574 34984 : checked_byte_count(send_pack->nblocks, sizeof(dbm_pack_block_t)),
575 : /*dest=*/send_rank,
576 : /*sendtag=*/send_ipack,
577 34984 : /*recvbuf=*/packed->recv_pack.blocks,
578 : /*recvcount=*/
579 34984 : checked_byte_count(packed->max_nblocks, sizeof(dbm_pack_block_t)),
580 : /*source=*/recv_rank,
581 : /*recvtag=*/recv_ipack,
582 34984 : /*comm=*/packed->dist_ticks->comm);
583 :
584 34984 : assert(nblocks_in_bytes % sizeof(dbm_pack_block_t) == 0);
585 34984 : packed->recv_pack.nblocks = nblocks_in_bytes / sizeof(dbm_pack_block_t);
586 :
587 : // Exchange data.
588 69968 : packed->recv_pack.data_size = cp_mpi_sendrecv_double(
589 34984 : /*sendbuf=*/send_pack->data,
590 34984 : /*sendcound=*/send_pack->data_size,
591 : /*dest=*/send_rank,
592 : /*sendtag=*/send_ipack,
593 : /*recvbuf=*/packed->recv_pack.data,
594 34984 : /*recvcount=*/packed->max_data_size,
595 : /*source=*/recv_rank,
596 : /*recvtag=*/recv_ipack,
597 34984 : /*comm=*/packed->dist_ticks->comm);
598 :
599 34984 : return &packed->recv_pack;
600 : }
601 : }
602 :
603 : /*******************************************************************************
604 : * \brief Private routine for releasing a packed matrix.
605 : * \author Ole Schuett
606 : ******************************************************************************/
607 699876 : static void free_packed_matrix(dbm_packed_matrix_t *packed) {
608 699876 : free_comm_buffer(packed->recv_pack.blocks);
609 699876 : free_comm_buffer(packed->recv_pack.data);
610 1434736 : for (int ipack = 0; ipack < packed->nsend_packs; ipack++) {
611 734860 : free_comm_buffer(packed->send_packs[ipack].blocks);
612 734860 : free_comm_buffer(packed->send_packs[ipack].data);
613 : }
614 699876 : free(packed->send_packs);
615 699876 : }
616 :
617 : /*******************************************************************************
618 : * \brief Internal routine for creating a communication iterator.
619 : * \author Ole Schuett
620 : ******************************************************************************/
621 349938 : dbm_comm_iterator_t *dbm_comm_iterator_start(const bool transa,
622 : const bool transb,
623 : const dbm_matrix_t *matrix_a,
624 : const dbm_matrix_t *matrix_b,
625 : const dbm_matrix_t *matrix_c) {
626 349938 : dbm_comm_iterator_t *iter = malloc(sizeof(dbm_comm_iterator_t));
627 349938 : assert(iter != NULL);
628 349938 : iter->dist = matrix_c->dist;
629 :
630 : // During each communication tick we'll fetch a pack_a and pack_b.
631 : // Since the cart might be non-squared, the number of communication ticks is
632 : // chosen as the least common multiple of the cart's dimensions.
633 349938 : iter->nticks = lcm(iter->dist->rows.nranks, iter->dist->cols.nranks);
634 349938 : iter->itick = 0;
635 349938 : iter->seconds_alltoall = iter->seconds_sort = 0.0;
636 349938 : iter->bytes_alltoall = iter->bytes_shift = 0;
637 :
638 : // 1.arg=source dimension, 2.arg=target dimension, false=rows, true=columns.
639 349938 : iter->packed_a =
640 349938 : pack_matrix(transa, false, matrix_a, iter->dist, iter->nticks, iter);
641 349938 : iter->packed_b =
642 349938 : pack_matrix(!transb, true, matrix_b, iter->dist, iter->nticks, iter);
643 :
644 349938 : return iter;
645 : }
646 :
647 : /*******************************************************************************
648 : * \brief Internal routine for retrieving next pair of packs of given iterator.
649 : * \author Ole Schuett
650 : ******************************************************************************/
651 734860 : bool dbm_comm_iterator_next(dbm_comm_iterator_t *iter, dbm_pack_t **pack_a,
652 : dbm_pack_t **pack_b) {
653 734860 : if (iter->itick >= iter->nticks) {
654 : return false; // end of iterator reached
655 : }
656 :
657 : // Start each rank at a different tick to spread the load on the sources.
658 384922 : const int shift = iter->dist->rows.my_rank + iter->dist->cols.my_rank;
659 384922 : const int itick = (iter->itick + shift) % iter->nticks;
660 769844 : *pack_a =
661 384922 : sendrecv_pack(itick, iter->nticks, &iter->packed_a, &iter->bytes_shift);
662 769844 : *pack_b =
663 384922 : sendrecv_pack(itick, iter->nticks, &iter->packed_b, &iter->bytes_shift);
664 :
665 384922 : ++iter->itick;
666 384922 : return true;
667 : }
668 :
669 : /*******************************************************************************
670 : * \brief Internal routine for releasing the given communication iterator.
671 : * \author Ole Schuett
672 : ******************************************************************************/
673 349938 : void dbm_comm_iterator_stop(dbm_comm_iterator_t *iter) {
674 349938 : free_packed_matrix(&iter->packed_a);
675 349938 : free_packed_matrix(&iter->packed_b);
676 349938 : free(iter);
677 349938 : }
678 :
679 : // EOF
|
