Reorder ranks to improve locality in communication

Description

Right now, all ranks are distributed in a 4D hypercube where the dimensions are domain, states, k points, and other (in this order). The call to MPI_Cart_create returns a cartesian communicator where usually all ranks are assigned in a row-major way. This means that the last index runs fastest (other) and the first index runs slowest (domain).

Usually, tasks are assigned to nodes in a block fashion, i.e. the first ranks (0 to x-1) go the first node, the ranks x to 2x-1 go to the second node etc. This means that the last parallelization strategy that is used is as compact as possible, whereas the first parallelization strategy is always spread out over several nodes. This means that currently the domain parallelization is not compact, i.e. the ranks sharing the domain for the same state and k point are usually spread out over several nodes.

The problem here is that communication inside a node is faster than across nodes because it can be handled in the shared memory as opposed to going through the interconnect. Thus, the communication that happens most often should be between ranks that are as compact as possible (to avoid going through the interconnect).

In octopus, the communication that happens most often is that using the domain parallelization strategy. Thus, it should be beneficial for most applications to have a compact arrangement of the tasks in the domain parallelization. This can be achieved by reordering the ranks in the base communicator in a way that reflects the column-major order of the 4D hypercube.

Tests have shown up to 20% speed-up depending on the input data. For best performance, set ParDomains to the number of cores per node (or a fraction/multiple of this number, e.g. 1/2 or 1/4 or 2).

News snippet

Reorder ranks to improve locality in communication

Checklist

• I have checked that my code follows the Octopus coding standards
• I have added tests for all the new features added in this request.

src/basic/multicomm.F90

@@ -133,6 +133,8 @@ module multicomm_oct_m
     integer, public  :: master_comm      !< The communicator without slaves.
     integer          :: master_comm_rank !< The rank in the communicator without slaves.
     integer, public  :: slave_intercomm  !< the intercomm to communicate with slaves
+
+    logical          :: reorder_ranks    !< do we reorder ranks in a more compact way?
   end type multicomm_t
 
   !> An all-pairs communication schedule for a given group.
@@ -183,6 +185,18 @@ contains
 
     call messages_print_stress(stdout, "Parallelization")
 
+    !%Variable ReorderRanks
+    !%Default no
+    !%Type logical
+    !%Section Execution::Parallelization
+    !%Description
+    !% This variable controls whether the ranks are reorganized to have a more
+    !% compact distribution with respect to domain parallelization which needs
+    !% to communicate most often. Depending on the system, this can improve
+    !% communication speeds.
+    !%End
+    call parse_variable(namespace, 'ReorderRanks', .false., mc%reorder_ranks)
+
     call messages_obsolete_variable(namespace, 'ParallelizationStrategy')
     call messages_obsolete_variable(namespace, 'ParallelizationGroupRanks')
     
@@ -596,40 +610,46 @@ contains
       mc%full_comm = MPI_COMM_NULL
       mc%slave_intercomm = MPI_COMM_NULL
       if(mc%par_strategy /= P_STRATEGY_SERIAL) then
-        ! first, reorder the ranks
-        ! this is done to get a column-major ordering of the ranks in the
-        ! Cartesian communicator, since they a ordered row-major otherwise
-        call MPI_Comm_group(base_grp%comm, base_group, mpi_err)
-        if(mpi_err /= MPI_SUCCESS) then
-          message(1) = "Error in getting MPI group!"
-          call messages_fatal(1)
-        end if
-        ! now transpose the hypercube => get rank numbers in column-major order
-        nn = 1
-        do ii = 1, mc%group_sizes(1)
-          do jj = 1, mc%group_sizes(2)
-            do kk = 1, mc%group_sizes(3)
-              do ll = 1, mc%group_sizes(4)
-                ranks(nn) = (ll-1)*mc%group_sizes(3)*mc%group_sizes(2)*mc%group_sizes(1) &
-                          + (kk-1)*mc%group_sizes(2)*mc%group_sizes(1) &
-                          + (jj-1)*mc%group_sizes(1) + ii - 1
-                nn = nn + 1
+        if(mc%reorder_ranks) then
+          call messages_experimental('Reordering ranks')
+
+          ! first, reorder the ranks
+          ! this is done to get a column-major ordering of the ranks in the
+          ! Cartesian communicator, since they a ordered row-major otherwise
+          call MPI_Comm_group(base_grp%comm, base_group, mpi_err)
+          if(mpi_err /= MPI_SUCCESS) then
+            message(1) = "Error in getting MPI group!"
+            call messages_fatal(1)
+          end if
+          ! now transpose the hypercube => get rank numbers in column-major order
+          nn = 1
+          do ii = 1, mc%group_sizes(1)
+            do jj = 1, mc%group_sizes(2)
+              do kk = 1, mc%group_sizes(3)
+                do ll = 1, mc%group_sizes(4)
+                  ranks(nn) = (ll-1)*mc%group_sizes(3)*mc%group_sizes(2)*mc%group_sizes(1) &
+                            + (kk-1)*mc%group_sizes(2)*mc%group_sizes(1) &
+                            + (jj-1)*mc%group_sizes(1) + ii - 1
+                  nn = nn + 1
+                end do
               end do
             end do
           end do
-        end do
-        call MPI_Group_incl(base_group, base_grp%size, ranks, reorder_group, mpi_err)
-        if(mpi_err /= MPI_SUCCESS) then
-          message(1) = "Error in creating MPI group!"
-          call messages_fatal(1)
-        end if
-        ! now get the reordered communicator
-        call MPI_Comm_create(base_grp%comm, reorder_group, reorder_comm, mpi_err)
-        if(mpi_err /= MPI_SUCCESS) then
-          message(1) = "Error in creating reordered communicator!"
-          call messages_fatal(1)
+          call MPI_Group_incl(base_group, base_grp%size, ranks, reorder_group, mpi_err)
+          if(mpi_err /= MPI_SUCCESS) then
+            message(1) = "Error in creating MPI group!"
+            call messages_fatal(1)
+          end if
+          ! now get the reordered communicator
+          call MPI_Comm_create(base_grp%comm, reorder_group, reorder_comm, mpi_err)
+          if(mpi_err /= MPI_SUCCESS) then
+            message(1) = "Error in creating reordered communicator!"
+            call messages_fatal(1)
+          end if
+          call mpi_grp_init(reorder_grp, reorder_comm)
+        else
+          call mpi_grp_copy(reorder_grp, base_grp)
         end if
-        call mpi_grp_init(reorder_grp, reorder_comm)
 
         ! Multilevel parallelization is organized in a hypercube. We
         ! use an MPI Cartesian topology to generate the communicators