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 <assert.h>
       9             : #include <math.h>
      10             : #include <omp.h>
      11             : #include <stdbool.h>
      12             : #include <stddef.h>
      13             : #include <stdlib.h>
      14             : #include <string.h>
      15             : 
      16             : #include "dbm_distribution.h"
      17             : #include "dbm_hyperparams.h"
      18             : 
      19             : /*******************************************************************************
      20             :  * \brief Private routine for creating a new one dimensional distribution.
      21             :  * \author Ole Schuett
      22             :  ******************************************************************************/
      23     1614324 : static void dbm_dist_1d_new(dbm_dist_1d_t *dist, const int length,
      24             :                             const int coords[length], const dbm_mpi_comm_t comm,
      25             :                             const int nshards) {
      26     1614324 :   dist->comm = comm;
      27     1614324 :   dist->nshards = nshards;
      28     1614324 :   dist->my_rank = dbm_mpi_comm_rank(comm);
      29     1614324 :   dist->nranks = dbm_mpi_comm_size(comm);
      30     1614324 :   dist->length = length;
      31     1614324 :   dist->index2coord = malloc(length * sizeof(int));
      32     1614324 :   assert(dist->index2coord != NULL);
      33     1614324 :   memcpy(dist->index2coord, coords, length * sizeof(int));
      34             : 
      35             :   // Check that cart coordinates and ranks are equivalent.
      36     1614324 :   int cart_dims[1], cart_periods[1], cart_coords[1];
      37     1614324 :   dbm_mpi_cart_get(comm, 1, cart_dims, cart_periods, cart_coords);
      38     1614324 :   assert(dist->nranks == cart_dims[0]);
      39     1614324 :   assert(dist->my_rank == cart_coords[0]);
      40             : 
      41             :   // Count local rows/columns.
      42    20370896 :   for (int i = 0; i < length; i++) {
      43    18756572 :     assert(0 <= coords[i] && coords[i] < dist->nranks);
      44    18756572 :     if (coords[i] == dist->my_rank) {
      45    17838461 :       dist->nlocals++;
      46             :     }
      47             :   }
      48             : 
      49             :   // Store local rows/columns.
      50     1614324 :   dist->local_indicies = malloc(dist->nlocals * sizeof(int));
      51     1614324 :   assert(dist->local_indicies != NULL);
      52             :   int j = 0;
      53    20370896 :   for (int i = 0; i < length; i++) {
      54    18756572 :     if (coords[i] == dist->my_rank) {
      55    17838461 :       dist->local_indicies[j++] = i;
      56             :     }
      57             :   }
      58     1614324 :   assert(j == dist->nlocals);
      59     1614324 : }
      60             : 
      61             : /*******************************************************************************
      62             :  * \brief Private routine for releasing a one dimensional distribution.
      63             :  * \author Ole Schuett
      64             :  ******************************************************************************/
      65     1614324 : static void dbm_dist_1d_free(dbm_dist_1d_t *dist) {
      66     1614324 :   free(dist->index2coord);
      67     1614324 :   free(dist->local_indicies);
      68     1614324 :   dbm_mpi_comm_free(&dist->comm);
      69     1614324 : }
      70             : 
      71             : /*******************************************************************************
      72             :  * \brief Returns the larger of two given integer (missing from the C standard)
      73             :  * \author Ole Schuett
      74             :  ******************************************************************************/
      75      807162 : static inline int imax(int x, int y) { return (x > y ? x : y); }
      76             : 
      77             : /*******************************************************************************
      78             :  * \brief Private routine for finding the optimal number of shard rows.
      79             :  * \author Ole Schuett
      80             :  ******************************************************************************/
      81      807162 : static int find_best_nrow_shards(const int nshards, const int nrows,
      82             :                                  const int ncols) {
      83      807162 :   const double target = (double)imax(nrows, 1) / (double)imax(ncols, 1);
      84      807162 :   int best_nrow_shards = nshards;
      85      807162 :   double best_error = fabs(log(target / (double)nshards));
      86             : 
      87     1614324 :   for (int nrow_shards = 1; nrow_shards <= nshards; nrow_shards++) {
      88      807162 :     const int ncol_shards = nshards / nrow_shards;
      89      807162 :     if (nrow_shards * ncol_shards != nshards)
      90           0 :       continue; // Not a factor of nshards.
      91      807162 :     const double ratio = (double)nrow_shards / (double)ncol_shards;
      92      807162 :     const double error = fabs(log(target / ratio));
      93      807162 :     if (error < best_error) {
      94           0 :       best_error = error;
      95           0 :       best_nrow_shards = nrow_shards;
      96             :     }
      97             :   }
      98      807162 :   return best_nrow_shards;
      99             : }
     100             : 
     101             : /*******************************************************************************
     102             :  * \brief Creates a new two dimensional distribution.
     103             :  * \author Ole Schuett
     104             :  ******************************************************************************/
     105      807162 : void dbm_distribution_new(dbm_distribution_t **dist_out, const int fortran_comm,
     106             :                           const int nrows, const int ncols,
     107             :                           const int row_dist[nrows],
     108             :                           const int col_dist[ncols]) {
     109      807162 :   assert(omp_get_num_threads() == 1);
     110      807162 :   dbm_distribution_t *dist = calloc(1, sizeof(dbm_distribution_t));
     111      807162 :   dist->ref_count = 1;
     112             : 
     113      807162 :   dist->comm = dbm_mpi_comm_f2c(fortran_comm);
     114      807162 :   dist->my_rank = dbm_mpi_comm_rank(dist->comm);
     115      807162 :   dist->nranks = dbm_mpi_comm_size(dist->comm);
     116             : 
     117      807162 :   const int row_dim_remains[2] = {1, 0};
     118      807162 :   const dbm_mpi_comm_t row_comm = dbm_mpi_cart_sub(dist->comm, row_dim_remains);
     119             : 
     120      807162 :   const int col_dim_remains[2] = {0, 1};
     121      807162 :   const dbm_mpi_comm_t col_comm = dbm_mpi_cart_sub(dist->comm, col_dim_remains);
     122             : 
     123      807162 :   const int nshards = SHARDS_PER_THREAD * omp_get_max_threads();
     124      807162 :   const int nrow_shards = find_best_nrow_shards(nshards, nrows, ncols);
     125      807162 :   const int ncol_shards = nshards / nrow_shards;
     126             : 
     127      807162 :   dbm_dist_1d_new(&dist->rows, nrows, row_dist, row_comm, nrow_shards);
     128      807162 :   dbm_dist_1d_new(&dist->cols, ncols, col_dist, col_comm, ncol_shards);
     129             : 
     130      807162 :   assert(*dist_out == NULL);
     131      807162 :   *dist_out = dist;
     132      807162 : }
     133             : 
     134             : /*******************************************************************************
     135             :  * \brief Increases the reference counter of the given distribution.
     136             :  * \author Ole Schuett
     137             :  ******************************************************************************/
     138     2309297 : void dbm_distribution_hold(dbm_distribution_t *dist) {
     139     2309297 :   assert(dist->ref_count > 0);
     140     2309297 :   dist->ref_count++;
     141     2309297 : }
     142             : 
     143             : /*******************************************************************************
     144             :  * \brief Decreases the reference counter of the given distribution.
     145             :  * \author Ole Schuett
     146             :  ******************************************************************************/
     147     3116459 : void dbm_distribution_release(dbm_distribution_t *dist) {
     148     3116459 :   assert(dist->ref_count > 0);
     149     3116459 :   dist->ref_count--;
     150     3116459 :   if (dist->ref_count == 0) {
     151      807162 :     dbm_dist_1d_free(&dist->rows);
     152      807162 :     dbm_dist_1d_free(&dist->cols);
     153      807162 :     free(dist);
     154             :   }
     155     3116459 : }
     156             : 
     157             : /*******************************************************************************
     158             :  * \brief Returns the rows of the given distribution.
     159             :  * \author Ole Schuett
     160             :  ******************************************************************************/
     161      316538 : void dbm_distribution_row_dist(const dbm_distribution_t *dist, int *nrows,
     162             :                                const int **row_dist) {
     163      316538 :   assert(dist->ref_count > 0);
     164      316538 :   *nrows = dist->rows.length;
     165      316538 :   *row_dist = dist->rows.index2coord;
     166      316538 : }
     167             : 
     168             : /*******************************************************************************
     169             :  * \brief Returns the columns of the given distribution.
     170             :  * \author Ole Schuett
     171             :  ******************************************************************************/
     172      316538 : void dbm_distribution_col_dist(const dbm_distribution_t *dist, int *ncols,
     173             :                                const int **col_dist) {
     174      316538 :   assert(dist->ref_count > 0);
     175      316538 :   *ncols = dist->cols.length;
     176      316538 :   *col_dist = dist->cols.index2coord;
     177      316538 : }
     178             : 
     179             : /*******************************************************************************
     180             :  * \brief Returns the MPI rank on which the given block should be stored.
     181             :  * \author Ole Schuett
     182             :  ******************************************************************************/
     183    91689736 : int dbm_distribution_stored_coords(const dbm_distribution_t *dist,
     184             :                                    const int row, const int col) {
     185    91689736 :   assert(dist->ref_count > 0);
     186    91689736 :   assert(0 <= row && row < dist->rows.length);
     187    91689736 :   assert(0 <= col && col < dist->cols.length);
     188    91689736 :   int coords[2] = {dist->rows.index2coord[row], dist->cols.index2coord[col]};
     189    91689736 :   return dbm_mpi_cart_rank(dist->comm, coords);
     190             : }
     191             : 
     192             : // EOF