Line data Source code
1 : /*----------------------------------------------------------------------------*/
2 : /* CP2K: A general program to perform molecular dynamics simulations */
3 : /* Copyright 2000-2024 CP2K developers group <https://cp2k.org> */
4 : /* */
5 : /* SPDX-License-Identifier: BSD-3-Clause */
6 : /*----------------------------------------------------------------------------*/
7 :
8 : #include "dbm_multiply_comm.h"
9 :
10 : #include <assert.h>
11 : #include <stdlib.h>
12 : #include <string.h>
13 :
14 : #include "dbm_hyperparams.h"
15 : #include "dbm_mempool.h"
16 : #include "dbm_mpi.h"
17 :
18 : /*******************************************************************************
19 : * \brief Returns the larger of two given integer (missing from the C standard)
20 : * \author Ole Schuett
21 : ******************************************************************************/
22 641196 : static inline int imax(int x, int y) { return (x > y ? x : y); }
23 :
24 : /*******************************************************************************
25 : * \brief Private routine for computing greatest common divisor of two numbers.
26 : * \author Ole Schuett
27 : ******************************************************************************/
28 320246 : static int gcd(const int a, const int b) {
29 320246 : if (a == 0)
30 : return b;
31 163277 : return gcd(b % a, a); // Euclid's algorithm.
32 : }
33 :
34 : /*******************************************************************************
35 : * \brief Private routine for computing least common multiple of two numbers.
36 : * \author Ole Schuett
37 : ******************************************************************************/
38 156969 : static int lcm(const int a, const int b) { return (a * b) / gcd(a, b); }
39 :
40 : /*******************************************************************************
41 : * \brief Private routine for computing the sum of the given integers.
42 : * \author Ole Schuett
43 : ******************************************************************************/
44 641196 : static inline int isum(const int n, const int input[n]) {
45 641196 : int output = 0;
46 1322352 : for (int i = 0; i < n; i++) {
47 681156 : output += input[i];
48 : }
49 641196 : return output;
50 : }
51 :
52 : /*******************************************************************************
53 : * \brief Private routine for computing the cumulative sums of given numbers.
54 : * \author Ole Schuett
55 : ******************************************************************************/
56 1602990 : static inline void icumsum(const int n, const int input[n], int output[n]) {
57 1602990 : output[0] = 0;
58 1682910 : for (int i = 1; i < n; i++) {
59 79920 : output[i] = output[i - 1] + input[i - 1];
60 : }
61 1602990 : }
62 :
63 : /*******************************************************************************
64 : * \brief Private struct used for planing during pack_matrix.
65 : * \author Ole Schuett
66 : ******************************************************************************/
67 : typedef struct {
68 : const dbm_block_t *blk; // source block
69 : int rank; // target mpi rank
70 : int row_size;
71 : int col_size;
72 : } plan_t;
73 :
74 : /*******************************************************************************
75 : * \brief Private routine for planing packs.
76 : * \author Ole Schuett
77 : ******************************************************************************/
78 313938 : static void create_pack_plans(const bool trans_matrix, const bool trans_dist,
79 : const dbm_matrix_t *matrix,
80 : const dbm_mpi_comm_t comm,
81 : const dbm_dist_1d_t *dist_indices,
82 : const dbm_dist_1d_t *dist_ticks, const int nticks,
83 : const int npacks, plan_t *plans_per_pack[npacks],
84 : int nblks_per_pack[npacks],
85 : int ndata_per_pack[npacks]) {
86 :
87 313938 : memset(nblks_per_pack, 0, npacks * sizeof(int));
88 313938 : memset(ndata_per_pack, 0, npacks * sizeof(int));
89 :
90 313938 : #pragma omp parallel
91 : {
92 : // 1st pass: Compute number of blocks that will be send in each pack.
93 : int nblks_mythread[npacks];
94 : memset(nblks_mythread, 0, npacks * sizeof(int));
95 : #pragma omp for schedule(static)
96 : for (int ishard = 0; ishard < dbm_get_num_shards(matrix); ishard++) {
97 : dbm_shard_t *shard = &matrix->shards[ishard];
98 : for (int iblock = 0; iblock < shard->nblocks; iblock++) {
99 : const dbm_block_t *blk = &shard->blocks[iblock];
100 : const int sum_index = (trans_matrix) ? blk->row : blk->col;
101 : const int itick = (1021 * sum_index) % nticks; // 1021 = a random prime
102 : const int ipack = itick / dist_ticks->nranks;
103 : nblks_mythread[ipack]++;
104 : }
105 : }
106 :
107 : // Sum nblocks across threads and allocate arrays for plans.
108 : #pragma omp critical
109 : for (int ipack = 0; ipack < npacks; ipack++) {
110 : nblks_per_pack[ipack] += nblks_mythread[ipack];
111 : nblks_mythread[ipack] = nblks_per_pack[ipack];
112 : }
113 : #pragma omp barrier
114 : #pragma omp for
115 : for (int ipack = 0; ipack < npacks; ipack++) {
116 : plans_per_pack[ipack] = malloc(nblks_per_pack[ipack] * sizeof(plan_t));
117 : }
118 :
119 : // 2nd pass: Plan where to send each block.
120 : int ndata_mythread[npacks];
121 : memset(ndata_mythread, 0, npacks * sizeof(int));
122 : #pragma omp for schedule(static) // Need static to match previous loop.
123 : for (int ishard = 0; ishard < dbm_get_num_shards(matrix); ishard++) {
124 : dbm_shard_t *shard = &matrix->shards[ishard];
125 : for (int iblock = 0; iblock < shard->nblocks; iblock++) {
126 : const dbm_block_t *blk = &shard->blocks[iblock];
127 : const int free_index = (trans_matrix) ? blk->col : blk->row;
128 : const int sum_index = (trans_matrix) ? blk->row : blk->col;
129 : const int itick = (1021 * sum_index) % nticks; // Same mapping as above.
130 : const int ipack = itick / dist_ticks->nranks;
131 : // Compute rank to which this block should be sent.
132 : const int coord_free_idx = dist_indices->index2coord[free_index];
133 : const int coord_sum_idx = itick % dist_ticks->nranks;
134 : const int coords[2] = {(trans_dist) ? coord_sum_idx : coord_free_idx,
135 : (trans_dist) ? coord_free_idx : coord_sum_idx};
136 : const int rank = dbm_mpi_cart_rank(comm, coords);
137 : const int row_size = matrix->row_sizes[blk->row];
138 : const int col_size = matrix->col_sizes[blk->col];
139 : ndata_mythread[ipack] += row_size * col_size;
140 : // Create plan.
141 : const int iplan = --nblks_mythread[ipack];
142 : plans_per_pack[ipack][iplan].blk = blk;
143 : plans_per_pack[ipack][iplan].rank = rank;
144 : plans_per_pack[ipack][iplan].row_size = row_size;
145 : plans_per_pack[ipack][iplan].col_size = col_size;
146 : }
147 : }
148 : #pragma omp critical
149 : for (int ipack = 0; ipack < npacks; ipack++) {
150 : ndata_per_pack[ipack] += ndata_mythread[ipack];
151 : }
152 : } // end of omp parallel region
153 313938 : }
154 :
155 : /*******************************************************************************
156 : * \brief Private routine for filling send buffers.
157 : * \author Ole Schuett
158 : ******************************************************************************/
159 320598 : static void fill_send_buffers(
160 : const dbm_matrix_t *matrix, const bool trans_matrix, const int nblks_send,
161 : const int ndata_send, plan_t plans[nblks_send], const int nranks,
162 : int blks_send_count[nranks], int data_send_count[nranks],
163 : int blks_send_displ[nranks], int data_send_displ[nranks],
164 : dbm_pack_block_t blks_send[nblks_send], double data_send[ndata_send]) {
165 :
166 320598 : memset(blks_send_count, 0, nranks * sizeof(int));
167 320598 : memset(data_send_count, 0, nranks * sizeof(int));
168 :
169 320598 : #pragma omp parallel
170 : {
171 : // 3th pass: Compute per rank nblks and ndata.
172 : int nblks_mythread[nranks], ndata_mythread[nranks];
173 : memset(nblks_mythread, 0, nranks * sizeof(int));
174 : memset(ndata_mythread, 0, nranks * sizeof(int));
175 : #pragma omp for schedule(static)
176 : for (int iblock = 0; iblock < nblks_send; iblock++) {
177 : const plan_t *plan = &plans[iblock];
178 : nblks_mythread[plan->rank] += 1;
179 : ndata_mythread[plan->rank] += plan->row_size * plan->col_size;
180 : }
181 :
182 : // Sum nblks and ndata across threads.
183 : #pragma omp critical
184 : for (int irank = 0; irank < nranks; irank++) {
185 : blks_send_count[irank] += nblks_mythread[irank];
186 : data_send_count[irank] += ndata_mythread[irank];
187 : nblks_mythread[irank] = blks_send_count[irank];
188 : ndata_mythread[irank] = data_send_count[irank];
189 : }
190 : #pragma omp barrier
191 :
192 : // Compute send displacements.
193 : #pragma omp master
194 : {
195 : icumsum(nranks, blks_send_count, blks_send_displ);
196 : icumsum(nranks, data_send_count, data_send_displ);
197 : const int m = nranks - 1;
198 : assert(nblks_send == blks_send_displ[m] + blks_send_count[m]);
199 : assert(ndata_send == data_send_displ[m] + data_send_count[m]);
200 : }
201 : #pragma omp barrier
202 :
203 : // 4th pass: Fill blks_send and data_send arrays.
204 : #pragma omp for schedule(static) // Need static to match previous loop.
205 : for (int iblock = 0; iblock < nblks_send; iblock++) {
206 : const plan_t *plan = &plans[iblock];
207 : const dbm_block_t *blk = plan->blk;
208 : const int ishard = dbm_get_shard_index(matrix, blk->row, blk->col);
209 : const dbm_shard_t *shard = &matrix->shards[ishard];
210 : const double *blk_data = &shard->data[blk->offset];
211 : const int row_size = plan->row_size, col_size = plan->col_size;
212 : const int irank = plan->rank;
213 :
214 : // The blk_send_data is ordered by rank, thread, and block.
215 : // data_send_displ[irank]: Start of data for irank within blk_send_data.
216 : // ndata_mythread[irank]: Current threads offset within data for irank.
217 : nblks_mythread[irank] -= 1;
218 : ndata_mythread[irank] -= row_size * col_size;
219 : const int offset = data_send_displ[irank] + ndata_mythread[irank];
220 : const int jblock = blks_send_displ[irank] + nblks_mythread[irank];
221 :
222 : double norm = 0.0; // Compute norm as double...
223 : if (trans_matrix) {
224 : // Transpose block to allow for outer-product style multiplication.
225 : for (int i = 0; i < row_size; i++) {
226 : for (int j = 0; j < col_size; j++) {
227 : const double element = blk_data[j * row_size + i];
228 : norm += element * element;
229 : data_send[offset + i * col_size + j] = element;
230 : }
231 : }
232 : blks_send[jblock].free_index = plan->blk->col;
233 : blks_send[jblock].sum_index = plan->blk->row;
234 : } else {
235 : for (int i = 0; i < row_size * col_size; i++) {
236 : const double element = blk_data[i];
237 : norm += element * element;
238 : data_send[offset + i] = element;
239 : }
240 : blks_send[jblock].free_index = plan->blk->row;
241 : blks_send[jblock].sum_index = plan->blk->col;
242 : }
243 : blks_send[jblock].norm = (float)norm; // ...store norm as float.
244 :
245 : // After the block exchange data_recv_displ will be added to the offsets.
246 : blks_send[jblock].offset = offset - data_send_displ[irank];
247 : }
248 : } // end of omp parallel region
249 320598 : }
250 :
251 : /*******************************************************************************
252 : * \brief Private comperator passed to qsort to compare two blocks by sum_index.
253 : * \author Ole Schuett
254 : ******************************************************************************/
255 80294198 : static int compare_pack_blocks_by_sum_index(const void *a, const void *b) {
256 80294198 : const dbm_pack_block_t *blk_a = (const dbm_pack_block_t *)a;
257 80294198 : const dbm_pack_block_t *blk_b = (const dbm_pack_block_t *)b;
258 80294198 : return blk_a->sum_index - blk_b->sum_index;
259 : }
260 :
261 : /*******************************************************************************
262 : * \brief Private routine for post-processing received blocks.
263 : * \author Ole Schuett
264 : ******************************************************************************/
265 320598 : static void postprocess_received_blocks(
266 : const int nranks, const int nshards, const int nblocks_recv,
267 : const int blks_recv_count[nranks], const int blks_recv_displ[nranks],
268 : const int data_recv_displ[nranks],
269 320598 : dbm_pack_block_t blks_recv[nblocks_recv]) {
270 :
271 320598 : int nblocks_per_shard[nshards], shard_start[nshards];
272 320598 : memset(nblocks_per_shard, 0, nshards * sizeof(int));
273 320598 : dbm_pack_block_t *blocks_tmp =
274 320598 : malloc(nblocks_recv * sizeof(dbm_pack_block_t));
275 :
276 320598 : #pragma omp parallel
277 : {
278 : // Add data_recv_displ to recveived block offsets.
279 : for (int irank = 0; irank < nranks; irank++) {
280 : #pragma omp for
281 : for (int i = 0; i < blks_recv_count[irank]; i++) {
282 : blks_recv[blks_recv_displ[irank] + i].offset += data_recv_displ[irank];
283 : }
284 : }
285 :
286 : // First use counting sort to group blocks by their free_index shard.
287 : int nblocks_mythread[nshards];
288 : memset(nblocks_mythread, 0, nshards * sizeof(int));
289 : #pragma omp for schedule(static)
290 : for (int iblock = 0; iblock < nblocks_recv; iblock++) {
291 : blocks_tmp[iblock] = blks_recv[iblock];
292 : const int ishard = blks_recv[iblock].free_index % nshards;
293 : nblocks_mythread[ishard]++;
294 : }
295 : #pragma omp critical
296 : for (int ishard = 0; ishard < nshards; ishard++) {
297 : nblocks_per_shard[ishard] += nblocks_mythread[ishard];
298 : nblocks_mythread[ishard] = nblocks_per_shard[ishard];
299 : }
300 : #pragma omp barrier
301 : #pragma omp master
302 : icumsum(nshards, nblocks_per_shard, shard_start);
303 : #pragma omp barrier
304 : #pragma omp for schedule(static) // Need static to match previous loop.
305 : for (int iblock = 0; iblock < nblocks_recv; iblock++) {
306 : const int ishard = blocks_tmp[iblock].free_index % nshards;
307 : const int jblock = --nblocks_mythread[ishard] + shard_start[ishard];
308 : blks_recv[jblock] = blocks_tmp[iblock];
309 : }
310 :
311 : // Then sort blocks within each shard by their sum_index.
312 : #pragma omp for
313 : for (int ishard = 0; ishard < nshards; ishard++) {
314 : if (nblocks_per_shard[ishard] > 1) {
315 : qsort(&blks_recv[shard_start[ishard]], nblocks_per_shard[ishard],
316 : sizeof(dbm_pack_block_t), &compare_pack_blocks_by_sum_index);
317 : }
318 : }
319 : } // end of omp parallel region
320 :
321 320598 : free(blocks_tmp);
322 320598 : }
323 :
324 : /*******************************************************************************
325 : * \brief Private routine for redistributing a matrix along selected dimensions.
326 : * \author Ole Schuett
327 : ******************************************************************************/
328 313938 : static dbm_packed_matrix_t pack_matrix(const bool trans_matrix,
329 : const bool trans_dist,
330 : const dbm_matrix_t *matrix,
331 : const dbm_distribution_t *dist,
332 313938 : const int nticks) {
333 :
334 313938 : assert(dbm_mpi_comms_are_similar(matrix->dist->comm, dist->comm));
335 :
336 : // The row/col indicies are distributed along one cart dimension and the
337 : // ticks are distributed along the other cart dimension.
338 313938 : const dbm_dist_1d_t *dist_indices = (trans_dist) ? &dist->cols : &dist->rows;
339 313938 : const dbm_dist_1d_t *dist_ticks = (trans_dist) ? &dist->rows : &dist->cols;
340 :
341 : // Allocate packed matrix.
342 313938 : const int nsend_packs = nticks / dist_ticks->nranks;
343 313938 : assert(nsend_packs * dist_ticks->nranks == nticks);
344 313938 : dbm_packed_matrix_t packed;
345 313938 : packed.dist_indices = dist_indices;
346 313938 : packed.dist_ticks = dist_ticks;
347 313938 : packed.nsend_packs = nsend_packs;
348 313938 : packed.send_packs = malloc(nsend_packs * sizeof(dbm_pack_t));
349 :
350 : // Plan all packs.
351 313938 : plan_t *plans_per_pack[nsend_packs];
352 313938 : int nblks_send_per_pack[nsend_packs], ndata_send_per_pack[nsend_packs];
353 313938 : create_pack_plans(trans_matrix, trans_dist, matrix, dist->comm, dist_indices,
354 : dist_ticks, nticks, nsend_packs, plans_per_pack,
355 : nblks_send_per_pack, ndata_send_per_pack);
356 :
357 : // Allocate send buffers for maximum number of blocks/data over all packs.
358 313938 : int nblks_send_max = 0, ndata_send_max = 0;
359 634536 : for (int ipack = 0; ipack < nsend_packs; ++ipack) {
360 320598 : nblks_send_max = imax(nblks_send_max, nblks_send_per_pack[ipack]);
361 320598 : ndata_send_max = imax(ndata_send_max, ndata_send_per_pack[ipack]);
362 : }
363 313938 : dbm_pack_block_t *blks_send =
364 313938 : dbm_mpi_alloc_mem(nblks_send_max * sizeof(dbm_pack_block_t));
365 313938 : double *data_send = dbm_mempool_host_malloc(ndata_send_max * sizeof(double));
366 :
367 : // Cannot parallelize over packs (there might be too few of them).
368 634536 : for (int ipack = 0; ipack < nsend_packs; ipack++) {
369 : // Fill send buffers according to plans.
370 320598 : const int nranks = dist->nranks;
371 320598 : int blks_send_count[nranks], data_send_count[nranks];
372 320598 : int blks_send_displ[nranks], data_send_displ[nranks];
373 320598 : fill_send_buffers(matrix, trans_matrix, nblks_send_per_pack[ipack],
374 : ndata_send_per_pack[ipack], plans_per_pack[ipack], nranks,
375 : blks_send_count, data_send_count, blks_send_displ,
376 : data_send_displ, blks_send, data_send);
377 320598 : free(plans_per_pack[ipack]);
378 :
379 : // 1st communication: Exchange block counts.
380 320598 : int blks_recv_count[nranks], blks_recv_displ[nranks];
381 320598 : dbm_mpi_alltoall_int(blks_send_count, 1, blks_recv_count, 1, dist->comm);
382 320598 : icumsum(nranks, blks_recv_count, blks_recv_displ);
383 320598 : const int nblocks_recv = isum(nranks, blks_recv_count);
384 :
385 : // 2nd communication: Exchange blocks.
386 320598 : dbm_pack_block_t *blks_recv =
387 320598 : dbm_mpi_alloc_mem(nblocks_recv * sizeof(dbm_pack_block_t));
388 320598 : int blks_send_count_byte[nranks], blks_send_displ_byte[nranks];
389 320598 : int blks_recv_count_byte[nranks], blks_recv_displ_byte[nranks];
390 661176 : for (int i = 0; i < nranks; i++) { // TODO: this is ugly!
391 340578 : blks_send_count_byte[i] = blks_send_count[i] * sizeof(dbm_pack_block_t);
392 340578 : blks_send_displ_byte[i] = blks_send_displ[i] * sizeof(dbm_pack_block_t);
393 340578 : blks_recv_count_byte[i] = blks_recv_count[i] * sizeof(dbm_pack_block_t);
394 340578 : blks_recv_displ_byte[i] = blks_recv_displ[i] * sizeof(dbm_pack_block_t);
395 : }
396 320598 : dbm_mpi_alltoallv_byte(
397 : blks_send, blks_send_count_byte, blks_send_displ_byte, blks_recv,
398 320598 : blks_recv_count_byte, blks_recv_displ_byte, dist->comm);
399 :
400 : // 3rd communication: Exchange data counts.
401 : // TODO: could be computed from blks_recv.
402 320598 : int data_recv_count[nranks], data_recv_displ[nranks];
403 320598 : dbm_mpi_alltoall_int(data_send_count, 1, data_recv_count, 1, dist->comm);
404 320598 : icumsum(nranks, data_recv_count, data_recv_displ);
405 320598 : const int ndata_recv = isum(nranks, data_recv_count);
406 :
407 : // 4th communication: Exchange data.
408 320598 : double *data_recv = dbm_mempool_host_malloc(ndata_recv * sizeof(double));
409 320598 : dbm_mpi_alltoallv_double(data_send, data_send_count, data_send_displ,
410 : data_recv, data_recv_count, data_recv_displ,
411 320598 : dist->comm);
412 :
413 : // Post-process received blocks and assemble them into a pack.
414 320598 : postprocess_received_blocks(nranks, dist_indices->nshards, nblocks_recv,
415 : blks_recv_count, blks_recv_displ,
416 : data_recv_displ, blks_recv);
417 320598 : packed.send_packs[ipack].nblocks = nblocks_recv;
418 320598 : packed.send_packs[ipack].data_size = ndata_recv;
419 320598 : packed.send_packs[ipack].blocks = blks_recv;
420 320598 : packed.send_packs[ipack].data = data_recv;
421 : }
422 :
423 : // Deallocate send buffers.
424 313938 : dbm_mpi_free_mem(blks_send);
425 313938 : dbm_mempool_free(data_send);
426 :
427 : // Allocate pack_recv.
428 313938 : int max_nblocks = 0, max_data_size = 0;
429 634536 : for (int ipack = 0; ipack < packed.nsend_packs; ipack++) {
430 320598 : max_nblocks = imax(max_nblocks, packed.send_packs[ipack].nblocks);
431 320598 : max_data_size = imax(max_data_size, packed.send_packs[ipack].data_size);
432 : }
433 313938 : dbm_mpi_max_int(&max_nblocks, 1, packed.dist_ticks->comm);
434 313938 : dbm_mpi_max_int(&max_data_size, 1, packed.dist_ticks->comm);
435 313938 : packed.max_nblocks = max_nblocks;
436 313938 : packed.max_data_size = max_data_size;
437 627876 : packed.recv_pack.blocks =
438 313938 : dbm_mpi_alloc_mem(packed.max_nblocks * sizeof(dbm_pack_block_t));
439 627876 : packed.recv_pack.data =
440 313938 : dbm_mempool_host_malloc(packed.max_data_size * sizeof(double));
441 :
442 313938 : return packed; // Ownership of packed transfers to caller.
443 : }
444 :
445 : /*******************************************************************************
446 : * \brief Private routine for sending and receiving the pack for the given tick.
447 : * \author Ole Schuett
448 : ******************************************************************************/
449 327258 : static dbm_pack_t *sendrecv_pack(const int itick, const int nticks,
450 : dbm_packed_matrix_t *packed) {
451 327258 : const int nranks = packed->dist_ticks->nranks;
452 327258 : const int my_rank = packed->dist_ticks->my_rank;
453 :
454 : // Compute send rank and pack.
455 327258 : const int itick_of_rank0 = (itick + nticks - my_rank) % nticks;
456 327258 : const int send_rank = (my_rank + nticks - itick_of_rank0) % nranks;
457 327258 : const int send_itick = (itick_of_rank0 + send_rank) % nticks;
458 327258 : const int send_ipack = send_itick / nranks;
459 327258 : assert(send_itick % nranks == my_rank);
460 :
461 : // Compute receive rank and pack.
462 327258 : const int recv_rank = itick % nranks;
463 327258 : const int recv_ipack = itick / nranks;
464 :
465 327258 : if (send_rank == my_rank) {
466 320598 : assert(send_rank == recv_rank && send_ipack == recv_ipack);
467 320598 : return &packed->send_packs[send_ipack]; // Local pack, no mpi needed.
468 : } else {
469 6660 : const dbm_pack_t *send_pack = &packed->send_packs[send_ipack];
470 :
471 : // Exchange blocks.
472 13320 : const int nblocks_in_bytes = dbm_mpi_sendrecv_byte(
473 6660 : /*sendbuf=*/send_pack->blocks,
474 6660 : /*sendcound=*/send_pack->nblocks * sizeof(dbm_pack_block_t),
475 : /*dest=*/send_rank,
476 : /*sendtag=*/send_ipack,
477 6660 : /*recvbuf=*/packed->recv_pack.blocks,
478 6660 : /*recvcount=*/packed->max_nblocks * sizeof(dbm_pack_block_t),
479 : /*source=*/recv_rank,
480 : /*recvtag=*/recv_ipack,
481 6660 : /*comm=*/packed->dist_ticks->comm);
482 :
483 6660 : assert(nblocks_in_bytes % sizeof(dbm_pack_block_t) == 0);
484 6660 : packed->recv_pack.nblocks = nblocks_in_bytes / sizeof(dbm_pack_block_t);
485 :
486 : // Exchange data.
487 13320 : packed->recv_pack.data_size = dbm_mpi_sendrecv_double(
488 6660 : /*sendbuf=*/send_pack->data,
489 6660 : /*sendcound=*/send_pack->data_size,
490 : /*dest=*/send_rank,
491 : /*sendtag=*/send_ipack,
492 : /*recvbuf=*/packed->recv_pack.data,
493 : /*recvcount=*/packed->max_data_size,
494 : /*source=*/recv_rank,
495 : /*recvtag=*/recv_ipack,
496 6660 : /*comm=*/packed->dist_ticks->comm);
497 :
498 6660 : return &packed->recv_pack;
499 : }
500 : }
501 :
502 : /*******************************************************************************
503 : * \brief Private routine for releasing a packed matrix.
504 : * \author Ole Schuett
505 : ******************************************************************************/
506 313938 : static void free_packed_matrix(dbm_packed_matrix_t *packed) {
507 313938 : dbm_mpi_free_mem(packed->recv_pack.blocks);
508 313938 : dbm_mempool_free(packed->recv_pack.data);
509 634536 : for (int ipack = 0; ipack < packed->nsend_packs; ipack++) {
510 320598 : dbm_mpi_free_mem(packed->send_packs[ipack].blocks);
511 320598 : dbm_mempool_free(packed->send_packs[ipack].data);
512 : }
513 313938 : free(packed->send_packs);
514 313938 : }
515 :
516 : /*******************************************************************************
517 : * \brief Internal routine for creating a communication iterator.
518 : * \author Ole Schuett
519 : ******************************************************************************/
520 156969 : dbm_comm_iterator_t *dbm_comm_iterator_start(const bool transa,
521 : const bool transb,
522 : const dbm_matrix_t *matrix_a,
523 : const dbm_matrix_t *matrix_b,
524 : const dbm_matrix_t *matrix_c) {
525 :
526 156969 : dbm_comm_iterator_t *iter = malloc(sizeof(dbm_comm_iterator_t));
527 156969 : iter->dist = matrix_c->dist;
528 :
529 : // During each communication tick we'll fetch a pack_a and pack_b.
530 : // Since the cart might be non-squared, the number of communication ticks is
531 : // chosen as the least common multiple of the cart's dimensions.
532 156969 : iter->nticks = lcm(iter->dist->rows.nranks, iter->dist->cols.nranks);
533 156969 : iter->itick = 0;
534 :
535 : // 1.arg=source dimension, 2.arg=target dimension, false=rows, true=columns.
536 156969 : iter->packed_a =
537 156969 : pack_matrix(transa, false, matrix_a, iter->dist, iter->nticks);
538 156969 : iter->packed_b =
539 156969 : pack_matrix(!transb, true, matrix_b, iter->dist, iter->nticks);
540 :
541 156969 : return iter;
542 : }
543 :
544 : /*******************************************************************************
545 : * \brief Internal routine for retriving next pair of packs from given iterator.
546 : * \author Ole Schuett
547 : ******************************************************************************/
548 320598 : bool dbm_comm_iterator_next(dbm_comm_iterator_t *iter, dbm_pack_t **pack_a,
549 : dbm_pack_t **pack_b) {
550 320598 : if (iter->itick >= iter->nticks) {
551 : return false; // end of iterator reached
552 : }
553 :
554 : // Start each rank at a different tick to spread the load on the sources.
555 163629 : const int shift = iter->dist->rows.my_rank + iter->dist->cols.my_rank;
556 163629 : const int shifted_itick = (iter->itick + shift) % iter->nticks;
557 163629 : *pack_a = sendrecv_pack(shifted_itick, iter->nticks, &iter->packed_a);
558 163629 : *pack_b = sendrecv_pack(shifted_itick, iter->nticks, &iter->packed_b);
559 :
560 163629 : iter->itick++;
561 163629 : return true;
562 : }
563 :
564 : /*******************************************************************************
565 : * \brief Internal routine for releasing the given communication iterator.
566 : * \author Ole Schuett
567 : ******************************************************************************/
568 156969 : void dbm_comm_iterator_stop(dbm_comm_iterator_t *iter) {
569 156969 : free_packed_matrix(&iter->packed_a);
570 156969 : free_packed_matrix(&iter->packed_b);
571 156969 : free(iter);
572 156969 : }
573 :
574 : // EOF
|
