Fix group_size /= 1

cp2k · Oct 27, 2022 · 479ae5a · 479ae5a
1 parent 0078135
commit 479ae5a
Showing 1 changed file with 38 additions and 24 deletions.
diff --git a/src/rpa_grad.F b/src/rpa_grad.F
@@ -1927,17 +1927,22 @@ SUBROUTINE sos_mp2_grad_finalize(sos_mp2_work_occ, sos_mp2_work_virt, para_env,
          ! Symmetrize P_ab
          IF (para_env_sub%num_pe > 1) THEN
             BLOCK
-               INTEGER :: my_B_start, my_B_end, my_B_size, send_a_start, send_a_end, send_a_size, &
+               INTEGER :: my_B_start, my_B_end, send_a_start, send_a_end, send_a_size, &
                           recv_a_start, recv_a_end, recv_a_size, proc_shift, proc_send, proc_recv
-               REAL(KIND=dp), DIMENSION(:, :), ALLOCATABLE :: buffer_send, buffer_recv
+               REAL(KIND=dp), DIMENSION(:), ALLOCATABLE, TARGET :: buffer_send_1D
+               REAL(KIND=dp), DIMENSION(:, :), POINTER :: buffer_send
+               REAL(KIND=dp), DIMENSION(:, :), ALLOCATABLE :: buffer_recv
                TYPE(group_dist_d1_type)                           :: gd_virtual_sub
+
                CALL create_group_dist(gd_virtual_sub, para_env_sub%num_pe, virtual(ispin))
+               my_B_start = itmp(1)
+               my_B_end = itmp(2)
 
                mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end) = &
                   0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end) &
                           + TRANSPOSE(mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end)))
 
-               ALLOCATE (buffer_send(my_B_size, maxsize(gd_virtual_sub)))
+               ALLOCATE (buffer_send_1D(my_B_size*maxsize(gd_virtual_sub)))
                ALLOCATE (buffer_recv(my_B_size, maxsize(gd_virtual_sub)))
 
                DO proc_shift = 1, para_env_sub%num_pe - 1
@@ -1948,16 +1953,18 @@ SUBROUTINE sos_mp2_grad_finalize(sos_mp2_work_occ, sos_mp2_work_virt, para_env,
                   CALL get_group_dist(gd_virtual_sub, proc_send, send_a_start, send_a_end, send_a_size)
                   CALL get_group_dist(gd_virtual_sub, proc_recv, recv_a_start, recv_a_end, recv_a_size)
 
-                  buffer_send(:, :send_a_size) = mp2_env%ri_grad%P_ab(ispin)%array(:, send_a_start:send_a_end)
-                  CALL mp_sendrecv(buffer_send(:, :send_a_size), proc_send, &
+                  buffer_send(1:send_a_size, 1:my_B_size) => buffer_send_1D(1:my_B_size*send_a_size)
+
+                  buffer_send(:send_a_size, :) = TRANSPOSE(mp2_env%ri_grad%P_ab(ispin)%array(:, send_a_start:send_a_end))
+                  CALL mp_sendrecv(buffer_send(:send_a_size, :), proc_send, &
                                    buffer_recv(:, :recv_a_size), proc_recv, para_env_sub%group)
 
                   mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) = &
-                     0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) + buffer_send)
+                     0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) + buffer_recv(:, 1:recv_a_size))
 
                END DO
 
-               DEALLOCATE (buffer_send, buffer_recv)
+               DEALLOCATE (buffer_send_1D, buffer_recv)
 
                CALL release_group_dist(gd_virtual_sub)
             END BLOCK
@@ -2121,32 +2128,39 @@ SUBROUTINE rpa_grad_work_finalize(rpa_work, mp2_env, homo, virtual, para_env, pa
 
          ! Symmetrize P_ab
          IF (para_env_sub%num_pe > 1) THEN
+            BLOCK
+               REAL(KIND=dp), DIMENSION(:), ALLOCATABLE, TARGET :: buffer_send_1D
+               REAL(KIND=dp), DIMENSION(:, :), POINTER :: buffer_send
+               REAL(KIND=dp), DIMENSION(:, :), ALLOCATABLE :: buffer_recv
 
-            mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end) = &
-               0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end) &
-                       + TRANSPOSE(mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end)))
+               mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end) = &
+                  0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end) &
+                          + TRANSPOSE(mp2_env%ri_grad%P_ab(ispin)%array(:, my_B_start:my_B_end)))
 
-            ALLOCATE (buffer_send(my_B_size, maxsize(gd_virtual_sub)))
-            ALLOCATE (buffer_recv(my_B_size, maxsize(gd_virtual_sub)))
+               ALLOCATE (buffer_send_1D(my_B_size*maxsize(gd_virtual_sub)))
+               ALLOCATE (buffer_recv(my_B_size, maxsize(gd_virtual_sub)))
 
-            DO proc_shift = 1, para_env_sub%num_pe - 1
+               DO proc_shift = 1, para_env_sub%num_pe - 1
 
-               proc_send = MODULO(para_env_sub%mepos + proc_shift, para_env_sub%num_pe)
-               proc_recv = MODULO(para_env_sub%mepos - proc_shift, para_env_sub%num_pe)
+                  proc_send = MODULO(para_env_sub%mepos + proc_shift, para_env_sub%num_pe)
+                  proc_recv = MODULO(para_env_sub%mepos - proc_shift, para_env_sub%num_pe)
+
+                  CALL get_group_dist(gd_virtual_sub, proc_send, send_a_start, send_a_end, send_a_size)
+                  CALL get_group_dist(gd_virtual_sub, proc_recv, recv_a_start, recv_a_end, recv_a_size)
 
-               CALL get_group_dist(gd_virtual_sub, proc_send, send_a_start, send_a_end, send_a_size)
-               CALL get_group_dist(gd_virtual_sub, proc_recv, recv_a_start, recv_a_end, recv_a_size)
+                  buffer_send(1:send_a_size, 1:my_B_size) => buffer_send_1D(1:my_B_size*send_a_size)
 
-               buffer_send(:, :send_a_size) = mp2_env%ri_grad%P_ab(ispin)%array(:, send_a_start:send_a_end)
-               CALL mp_sendrecv(buffer_send(:, :send_a_size), proc_send, &
-                                buffer_recv(:, :recv_a_size), proc_recv, para_env_sub%group)
+                  buffer_send(:send_a_size, :) = TRANSPOSE(mp2_env%ri_grad%P_ab(ispin)%array(:, send_a_start:send_a_end))
+                  CALL mp_sendrecv(buffer_send(:send_a_size, :), proc_send, &
+                                   buffer_recv(:, :recv_a_size), proc_recv, para_env_sub%group)
 
-               mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) = &
-                  0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) + buffer_send)
+                  mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) = &
+                     0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array(:, recv_a_start:recv_a_end) + buffer_recv(:, 1:recv_a_size))
 
-            END DO
+               END DO
 
-            DEALLOCATE (buffer_send, buffer_recv)
+               DEALLOCATE (buffer_send_1D, buffer_recv)
+            END BLOCK
          ELSE
             mp2_env%ri_grad%P_ab(ispin)%array(:, :) = 0.5_dp*(mp2_env%ri_grad%P_ab(ispin)%array + &
                                                               TRANSPOSE(mp2_env%ri_grad%P_ab(ispin)%array))