From dc4f4b7f4d5c33b93a5c446c0b9d950476ad4d89 Mon Sep 17 00:00:00 2001
From: Connor Aird <c.aird@ucl.ac.uk>
Date: Thu, 16 May 2024 10:36:37 +0100
Subject: [PATCH] Calculating r higher in the loop nest and dealloc arrays at
 the same level

---
 src/PAO_grid_transform_module.f90 | 11 ++++++-----
 src/exx_evalpao.f90               | 19 ++++++------------
 src/exx_kernel_default.f90        |  6 +++---
 src/ol_ang_coeff_subs.f90         | 33 +++++++++++--------------------
 4 files changed, 27 insertions(+), 42 deletions(-)

diff --git a/src/PAO_grid_transform_module.f90 b/src/PAO_grid_transform_module.f90
index 0d95148cc..e484b6c5c 100644
--- a/src/PAO_grid_transform_module.f90
+++ b/src/PAO_grid_transform_module.f90
@@ -132,7 +132,7 @@ subroutine PAO_or_gradPAO_to_grid(pao_fns, evaluate, direction)
     integer :: npoint ! outputs of check_block
     integer :: count1 ! incremented counter, maps from (l1, acz, m1) to linear index of gridfunctions%griddata
     real(double):: dcellx_block,dcelly_block,dcellz_block ! grid dimensions, should be moved
-    real(double) :: x,y,z ! Temporary variables to reduce indirect accesses
+    real(double) :: x,y,z,r ! Temporary variables to reduce indirect accesses
     real(double) :: rcut ! Input to check_block
     real(double) :: val ! output, written into gridfunctions%griddata
     real(double) :: xblock,yblock,zblock ! inputs to check_block
@@ -145,11 +145,11 @@ subroutine PAO_or_gradPAO_to_grid(pao_fns, evaluate, direction)
        ! Interface to return a value val given arguments
        ! direction,species,l,acz,m,x,y,z. Implemented by
        ! evaluate_pao() and pao_elem_derivative_2().
-       subroutine evaluate(direction,species,l,acz,m,x,y,z,val,sys)
+       subroutine evaluate(direction,species,l,acz,m,x,y,z,r,val,sys)
          use datatypes, only: double
          integer, intent(in) :: species,l,acz,m
          integer, intent(in) :: direction
-         real(kind=double), intent(in) :: x,y,z
+         real(kind=double), intent(in) :: x,y,z,r
          logical, intent(in), optional :: sys 
          real(kind=double), intent(out) :: val
        end subroutine evaluate
@@ -210,7 +210,7 @@ end subroutine evaluate
     !$omp             shared(domain, naba_atoms_of_blocks, offset_position, pao_fns, atomf, &
     !$omp                    dcellx_block, dcelly_block, dcellz_block, dcs_parts, &
     !$omp                    rcut, n_pts_in_block, pao, gridfunctions, direction) &
-    !$omp             private(ia, ipart, iblock, l1, acz, m1, count1, x, y, z, val, position, &
+    !$omp             private(ia, ipart, iblock, l1, acz, m1, count1, x, y, z, r, val, position, &
     !$omp                     npoint, r_store, ip_store, x_store, y_store, z_store, my_species, &
     !$omp                     xblock, yblock, zblock, iatom, xatom, yatom, zatom, naba_part_label, ind_part, icover)
     blocks_loop_omp: do iblock = 1, domain%groups_on_node ! primary set of blocks
@@ -240,11 +240,12 @@ end subroutine evaluate
                 x = x_store(ip)
                 y = y_store(ip)
                 z = z_store(ip)
+                r = sqrt(x*x+y*y+z*z)
                 ! For this point-atom offset, we accumulate the PAO on the grid
                 l_loop: do l1 = 0,pao(my_species)%greatest_angmom
                    z_loop: do acz = 1,pao(my_species)%angmom(l1)%n_zeta_in_angmom
                       m_loop: do m1=-l1,l1
-                         call evaluate(direction,my_species,l1,acz,m1,x,y,z,val)
+                         call evaluate(direction,my_species,l1,acz,m1,x,y,z,r,val)
                          gridfunctions(pao_fns)%griddata(position) = val
                          position = position + n_pts_in_block
                       end do m_loop
diff --git a/src/exx_evalpao.f90 b/src/exx_evalpao.f90
index a54627ea2..4e4248dc0 100644
--- a/src/exx_evalpao.f90
+++ b/src/exx_evalpao.f90
@@ -64,18 +64,13 @@ subroutine exx_phi_on_grid(inode,atom,spec,extent,xyz,nsuppfuncs,phi_on_grid,r_i
     real(double)        :: int, n, rest
     real(double)        :: xyz_delta(3), xyz_offset(3)
 
-    integer             :: count1, nsf1
-    integer             :: ierr, stat
-    integer             :: i, j, ngrid   
+    integer             :: count1
+    integer             :: i, ngrid   
     integer             :: ijk_delta(3), ijk(3), kji(3), njk(3), nji(3)
-    integer             :: n1, n2, npts
-
-    real(double)        :: xj1, xj2, a, b, c, d, del_r, alpha
-    real(double)        :: a_table, b_table, c_table, d_table
     integer             :: i_dummy = 0
     
     ! << Called subroutine variables >>
-    real(double)        :: pao_val, y_val    
+    real(double)        :: pao_val
 
     real(double),  parameter :: a1g_norm  = sqrt(one/(four*pi))
     real(double),  parameter :: t1u_norm  = sqrt(three/(four*pi))
@@ -144,9 +139,6 @@ subroutine exx_phi_on_grid(inode,atom,spec,extent,xyz,nsuppfuncs,phi_on_grid,r_i
                 nji(i) = ngrid
 
              end if
-             !write(*,'(5F12.6,4I6)') xyz(i), n, int, rest, &
-             !     xyz_delta(i), ijk(i), njk(i), kji(i), nji(i)
-             !write(*,'(4I6)') ijk(i), njk(i), kji(i), nji(i)
 
           end do
           mx = mx +kji(1)-1
@@ -161,13 +153,14 @@ subroutine exx_phi_on_grid(inode,atom,spec,extent,xyz,nsuppfuncs,phi_on_grid,r_i
        !$omp parallel do collapse(3) schedule(runtime) default(none)   & 
        !$omp     shared(mx,my,mz,px,py,pz,grid_spacing,xyz_offset,pao, &
        !$omp            spec,phi_on_grid,i_dummy,exx_cartesian,extent) &
-       !$omp    private(nx,ny,nz,x,y,z,count1,l1,acz,m1,pao_val)
+       !$omp    private(nx,ny,nz,x,y,z,r,count1,l1,acz,m1,pao_val)
        grid_x_loop: do nx = mx, px
           grid_y_loop: do ny = my, py
              grid_z_loop: do nz = mz, pz
                 x = nx*grid_spacing + xyz_offset(1)
                 y = ny*grid_spacing + xyz_offset(2)
                 z = nz*grid_spacing + xyz_offset(3)
+                r = sqrt(x*x+y*y+z*z)
 
                 count1  = 1
                 angu_loop: do l1 = 0, pao(spec)%greatest_angmom
@@ -176,7 +169,7 @@ subroutine exx_phi_on_grid(inode,atom,spec,extent,xyz,nsuppfuncs,phi_on_grid,r_i
 
                       magn_loop: do m1 = -l1, l1                      
                       
-                         call evaluate_pao(i_dummy,spec,l1,acz,m1,x,y,z,pao_val,exx_cartesian)
+                         call evaluate_pao(i_dummy,spec,l1,acz,m1,x,y,z,r,pao_val,exx_cartesian)
 
                          ! Put pao_val directly into phi_on_grid
                          ! (only for primitive PAOs and not for blips)
diff --git a/src/exx_kernel_default.f90 b/src/exx_kernel_default.f90
index 517dddc7f..4eeff51d9 100644
--- a/src/exx_kernel_default.f90
+++ b/src/exx_kernel_default.f90
@@ -1531,13 +1531,13 @@ subroutine m_kern_exx_eri(k_off, kpart, ib_nd_acc, ibaddr, nbnab, &
                                !
                             end do jb_loop
                             !
+                            if ( exx_alloc ) call exx_mem_alloc(extent,0,0,'Ome_kj_1d_buffer','dealloc')
+                            if ( exx_alloc ) call exx_mem_alloc(extent,jb%nsup,0,'phi_j','dealloc')
+                            !
                          end if
                          !
                       end if
                       !
-                      if ( exx_alloc ) call exx_mem_alloc(extent,0,0,'Ome_kj_1d_buffer','dealloc')
-                      if ( exx_alloc ) call exx_mem_alloc(extent,jb%nsup,0,'phi_j','dealloc')                                
-                      !
 !!$
 !!$ ****[ j end loop ]****
 !!$                      
diff --git a/src/ol_ang_coeff_subs.f90 b/src/ol_ang_coeff_subs.f90
index 04dc44318..88e73348f 100644
--- a/src/ol_ang_coeff_subs.f90
+++ b/src/ol_ang_coeff_subs.f90
@@ -769,17 +769,14 @@ subroutine convert_basis(x,y,z,r,thet,phi)
     implicit none
     !routine to convert Cartesian displacements into displacements
     !in spherical polar coordinates
-    real(double), intent(in) :: x,y,z
-    real(double), intent(inout) :: r,thet,phi
-    real(double) :: x2,y2,z2,mod_r,mod_r_plane,num
+    real(double), intent(in) :: x,y,z,r
+    real(double), intent(inout) :: thet,phi
+    real(double) :: x2,y2,mod_r_plane
 
     x2 = x*x
     y2 = y*y
-    z2 = z*z
-    mod_r = sqrt(x2+y2+z2)
     mod_r_plane = sqrt(x2+y2)
-    if(mod_r<very_small) then
-       r = zero
+    if(r<very_small) then
        thet = zero
        phi = zero
        return
@@ -788,7 +785,7 @@ subroutine convert_basis(x,y,z,r,thet,phi)
     if(abs(z)<very_small) then
        thet = half*pi
     else
-       thet = acos(z/mod_r) !need to make sure this ratio doesn't disappear as well..
+       thet = acos(z/r) !need to make sure this ratio doesn't disappear as well..
     endif
 
     if(abs(x2)<very_small.and.abs(y2)<very_small) then
@@ -812,9 +809,6 @@ subroutine convert_basis(x,y,z,r,thet,phi)
           endif
        endif
     endif
-   
-    r = mod_r
-
   end subroutine convert_basis
 !!***
 
@@ -845,7 +839,7 @@ end subroutine convert_basis
 !!    ie. stay in Cartesian ones (should be more efficient)   
 !!  SOURCE
 !!
-  subroutine evaluate_pao(i_vector,sp,l,nz,m,x,y,z,pao_val,system)
+  subroutine evaluate_pao(i_vector,sp,l,nz,m,x,y,z,r,pao_val,system)
 
     use datatypes
     use numbers
@@ -855,11 +849,11 @@ subroutine evaluate_pao(i_vector,sp,l,nz,m,x,y,z,pao_val,system)
 
     integer, intent(in) :: i_vector ! dummy argument, included to satisfy interface in PAO_grid_transform_module
     integer, intent(in) :: sp,l,nz,m
-    real(double), intent(in) :: x,y,z
+    real(double), intent(in) :: x,y,z,r
     logical,      intent(in), optional :: system 
     real(double), intent(out) :: pao_val
     !
-    real(double) :: r,theta,phi,del_r,y_val,val
+    real(double) :: theta,phi,del_r,y_val
     real(double) :: a, b, c, d, r1, r2, r3, r4, rr
     integer :: npts, j
     logical :: cartesian
@@ -871,9 +865,6 @@ subroutine evaluate_pao(i_vector,sp,l,nz,m,x,y,z,pao_val,system)
        cartesian = .false.
     end if
     !
-    !compute radius
-    r = sqrt(x*x+y*y+z*z) 
-    !
     !interpolate for required value of radial function
     npts  = pao(sp)%angmom(l)%zeta(nz)%length
     del_r = (pao(sp)%angmom(l)%zeta(nz)%cutoff/&
@@ -1003,7 +994,7 @@ end subroutine pp_elem
 !!    Added dummy argument to satisfy interface in PAO_grid_transform_module
 !!  SOURCE
 !!
-  subroutine pao_elem_derivative_2(i_vector,spec,l,nzeta,m,x_i,y_i,z_i,drvtv_val,sys_dummy)
+  subroutine pao_elem_derivative_2(i_vector,spec,l,nzeta,m,x_i,y_i,z_i,r_my,drvtv_val,sys_dummy)
 
     use datatypes
     use numbers
@@ -1012,18 +1003,18 @@ subroutine pao_elem_derivative_2(i_vector,spec,l,nzeta,m,x_i,y_i,z_i,drvtv_val,s
     
     !RC 09/11/03 using (now debugged) routine pp_elem_derivative (see 
     ! above) as template for this sbrt pao_elem_derivative
-    real(double), intent(in) :: x_i,y_i,z_i
+    real(double), intent(in) :: x_i,y_i,z_i,r_my
     logical,      intent(in), optional :: sys_dummy
     real(double), intent(out) :: drvtv_val
     integer, intent(in) :: i_vector, l, m, spec, nzeta
     integer :: n1,n2
     real(double) :: del_x,xj1,xj2,f_r,df_r,a,b,c,d,alpha,beta,gamma,delta
-    real(double) :: x,y,z,r_my,theta,phi,c_r,c_theta,c_phi,x_n,y_n,z_n
+    real(double) :: x,y,z,theta,phi,c_r,c_theta,c_phi,x_n,y_n,z_n
     real(double) :: fm_phi, dfm_phi, val1, val2, rl, rl1, tmpr, tmpdr
 
     !routine to calculate value of gradient of f(r)*Re(Y_lm(x,y,z))
     !along specified Cartesian unit vector
-    r_my = sqrt((x_i*x_i)+(y_i*y_i)+(z_i*z_i))
+    !r_my = sqrt((x_i*x_i)+(y_i*y_i)+(z_i*z_i))
     call pao_rad_vals_fetch(spec,l,nzeta,m,r_my,f_r,df_r)
     if(r_my>very_small) then 
        x_n = x_i/r_my; y_n = y_i/r_my; z_n = z_i/r_my