LCOV - code coverage report
Current view: top level - src/dbm - dbm_multiply_comm.c (source / functions) Hit Total Coverage
Test: CP2K Regtests (git:a2cdc02) Lines: 166 166 100.0 %
Date: 2025-04-17 08:15:26 Functions: 13 13 100.0 %

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

Generated by: LCOV version 1.15