Skip to content

Commit

Permalink
RI HFX: optimized process grid for (MO RI AO) tensor
Browse files Browse the repository at this point in the history 
Distribute over RI index only since MO index is much smaller than the
other indices
  • Loading branch information
@pseewald
pseewald committed Mar 7, 2020
1 parent 3c30079 commit f7de199
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 26 deletions.
24 changes: 2 additions & 22 deletions src/hfx_ri.F
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,8 +165,6 @@ SUBROUTINE hfx_ri_pre_scf_mo(qs_env, ri_data)
CALL dbcsr_t_clear(ri_data%t_3c_int_ctr_3(1, 1))
CALL dbcsr_t_destroy(t_2c_int(1))

IF (ASSOCIATED(ri_data%pgrid_2)) CALL tensor_change_pgrid(ri_data%t_3c_int_ctr_2(1, 1), ri_data%pgrid_2, &
nodata=.TRUE., unit_nr=unit_nr)
CALL dbcsr_t_copy(t_3c_int_2_ctr(1, 1), ri_data%t_3c_int_ctr_2(1, 1), order=[2, 1, 3], move_data=.TRUE.)
CALL dbcsr_t_destroy(t_3c_int_2_ctr(1, 1))
CALL dbcsr_t_filter(ri_data%t_3c_int_ctr_2(1, 1), ri_data%filter_eps)
Expand Down Expand Up @@ -592,25 +590,20 @@ SUBROUTINE hfx_ri_update_ks_mo(qs_env, ri_data, ks_matrix, mo_coeff, &
INTEGER, DIMENSION(2) :: pdims_2d
TYPE(dbcsr_distribution_type) :: ks_dist
TYPE(dbcsr_t_pgrid_type) :: pgrid_2d
TYPE(dbcsr_t_pgrid_type), POINTER :: pgrid, pgrid_opt
TYPE(dbcsr_t_pgrid_type), POINTER :: pgrid
TYPE(dbcsr_t_type) :: ks_t, ks_t_split, mo_coeff_t, &
mo_coeff_t_split
TYPE(dbcsr_t_type), DIMENSION(1, 1) :: t_3c_int_ctr, t_3c_int_ctr_opt, &
t_3c_int_ctr_opt_copy

CALL timeset(routineN, handle)

NULLIFY (pgrid_opt)

CPASSERT(SIZE(ks_matrix, 2) == 1)

unit_nr = ri_data%unit_nr

IF (geometry_did_change) THEN
CALL hfx_ri_pre_scf_mo(qs_env, ri_data)
ELSE
IF (ASSOCIATED(ri_data%pgrid_2)) CALL tensor_change_pgrid(ri_data%t_3c_int_ctr_2(1, 1), &
ri_data%pgrid_2, unit_nr=unit_nr)
ENDIF

ALLOCATE (bounds(2, 1))
Expand Down Expand Up @@ -694,8 +687,7 @@ SUBROUTINE hfx_ri_update_ks_mo(qs_env, ri_data, ks_matrix, mo_coeff, &
contract_1=[1, 2], notcontract_1=[3], &
contract_2=[1, 2], notcontract_2=[3], &
map_1=[1], map_2=[2], filter_eps=ri_data%filter_eps, &
unit_nr=unit_nr, move_data=.TRUE., &
pgrid_opt_1=pgrid_opt)
unit_nr=unit_nr, move_data=.TRUE.)
CALL timestop(handle2)

CALL dbcsr_t_destroy(t_3c_int_ctr_opt(1, 1))
Expand All @@ -707,18 +699,6 @@ SUBROUTINE hfx_ri_update_ks_mo(qs_env, ri_data, ks_matrix, mo_coeff, &
CALL dbcsr_t_copy_tensor_to_matrix(ks_t, ks_matrix(ispin, 1)%matrix, summation=.TRUE.)
CALL dbcsr_t_destroy(ks_t)

IF (ispin == nspins) THEN
CPASSERT(ASSOCIATED(pgrid_opt))
IF (ASSOCIATED(ri_data%pgrid_2)) THEN
CALL dbcsr_t_pgrid_destroy(ri_data%pgrid_2)
DEALLOCATE (ri_data%pgrid_2)
ENDIF
ri_data%pgrid_2 => pgrid_opt
ELSE
CALL dbcsr_t_pgrid_destroy(pgrid_opt)
DEALLOCATE (pgrid_opt)
ENDIF

ENDDO

CALL timestop(handle)
Expand Down
26 changes: 22 additions & 4 deletions src/hfx_types.F
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,8 @@ MODULE hfx_types
USE machine, ONLY: m_chdir,&
m_getcwd
USE mathlib, ONLY: erfc_cutoff
USE message_passing, ONLY: mp_cart_create
USE message_passing, ONLY: mp_cart_create,&
mp_environ
USE orbital_pointers, ONLY: nco,&
ncoset,&
nso
Expand Down Expand Up @@ -392,7 +393,7 @@ MODULE hfx_types
TYPE(dbcsr_t_type), DIMENSION(:, :), ALLOCATABLE :: t_3c_int_ctr_1
TYPE(dbcsr_t_pgrid_type), POINTER :: pgrid_1 => NULL()

! 3c integral tensor in ( AO | RI AO) format for contraction
! 3c integral tensor in ( AO | RI AO) format for contraction (also used for (MO RI AO) tensor)
TYPE(dbcsr_t_type), DIMENSION(:, :), ALLOCATABLE :: t_3c_int_ctr_2
TYPE(dbcsr_t_pgrid_type), POINTER :: pgrid_2 => NULL()

Expand Down Expand Up @@ -1063,10 +1064,12 @@ SUBROUTINE hfx_ri_init(ri_data, qs_kind_set, particle_set, para_env, basis_cntrl
TYPE(cp_para_env_type) :: para_env
INTEGER, INTENT(IN) :: basis_cntrl

INTEGER :: ikind, mp_comm_3d, natom, nkind
INTEGER :: ikind, iproc, mp_comm_3d, natom, nkind, &
nproc
INTEGER, ALLOCATABLE, DIMENSION(:) :: dist_AO_1, dist_AO_2, dist_RI
INTEGER, DIMENSION(3) :: pcoord, pdims
LOGICAL :: same_op
TYPE(dbcsr_t_pgrid_type), POINTER :: pgrid
TYPE(distribution_3d_type) :: dist_3d
TYPE(gto_basis_set_p_type), ALLOCATABLE, &
DIMENSION(:) :: basis_set_AO, basis_set_RI
Expand Down Expand Up @@ -1154,8 +1157,23 @@ SUBROUTINE hfx_ri_init(ri_data, qs_kind_set, particle_set, para_env, basis_cntrl
CALL split_block_sizes(ri_data%bsizes_AO, ri_data%bsizes_AO_split, ri_data%max_bsize)
CALL split_block_sizes(ri_data%bsizes_RI, ri_data%bsizes_RI_split, ri_data%max_bsize)

IF (ri_data%flavor == ri_mo) THEN
CALL mp_environ(nproc, iproc, para_env%group)
ALLOCATE (ri_data%pgrid_2)
! we distribute (MO RI AO) tensor over RI index only due to 2 reasons:
! - MO index is small and should not be distributed, RI index is the largest
! - we change representation from (MO | RI AO) -> (MO RI | AO) and both tensors should have
! same process grid (communication-less copy). Also tensor contraction algorithm requires that
! tensors are distributed mainly in the combined index.
pdims = [1, nproc, 1]
CALL dbcsr_t_pgrid_create(para_env%group, pdims, ri_data%pgrid_2)
pgrid => ri_data%pgrid_2
ELSE
pgrid => ri_data%pgrid
ENDIF

CALL create_3c_tensor(ri_data%t_3c_int_ctr_2(1, 1), ri_data%dist2_ao_1, ri_data%dist2_ri, ri_data%dist2_ao_2, &
ri_data%pgrid, ri_data%bsizes_AO, ri_data%bsizes_RI_split, ri_data%bsizes_AO_split, &
pgrid, ri_data%bsizes_AO, ri_data%bsizes_RI_split, ri_data%bsizes_AO_split, &
[1], [2, 3], name="(AO | RI AO)")

ALLOCATE (ri_data%t_3c_int_ctr_3(1, 1))
Expand Down

0 comments on commit f7de199

Please sign in to comment.