Skip to content

Commit

Permalink
task_list: Compute matrix buffer offsets during generate_qs_task_list
Browse files Browse the repository at this point in the history
  • Loading branch information
@oschuett
oschuett committed Jul 18, 2020
1 parent 460a82c commit d15f8a6
Show file tree
Hide file tree
Showing 3 changed files with 183 additions and 208 deletions.
63 changes: 34 additions & 29 deletions src/qs_collocate_density.F
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,10 +106,8 @@ MODULE qs_collocate_density
rs_pw_transfer
USE rs_pw_interface, ONLY: density_rs2pw,&
density_rs2pw_basic
USE task_list_methods, ONLY: rs_blocks_type,&
rs_copy_matrix,&
USE task_list_methods, ONLY: rs_copy_matrix,&
rs_distribute_matrix,&
rs_get_block_p,&
rs_scatter_matrix
USE task_list_types, ONLY: atom_pair_type,&
task_list_type,&
Expand Down Expand Up @@ -1426,13 +1424,13 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
routineP = moduleN//':'//routineN

CHARACTER(LEN=default_string_length) :: my_basis_type
INTEGER :: ga_gb_function, handle, iatom, iatom_old, igrid_level, ikind, ikind_old, img, &
img_old, ipair, ipgf, iset, iset_old, itask, ithread, jatom, jatom_old, jkind, jkind_old, &
jpgf, jset, jset_old, lb, lbr, lbw, maxco, maxsgf_set, n, na1, na2, natoms, nb1, nb2, &
nblock, ncoa, ncob, nimages, nr, nrlevel, nseta, nsetb, ntasks, nthread, nw, nxy, nz, &
nzsize, sgfa, sgfb, ub
INTEGER :: block_size, ga_gb_function, handle, iatom, iatom_old, igrid_level, ikind, &
ikind_old, img, img_old, ipair, ipgf, iset, iset_old, itask, ithread, jatom, jatom_old, &
jkind, jkind_old, jpgf, jset, jset_old, lb, lbr, lbw, maxco, maxsgf_set, n, na1, na2, &
natoms, nb1, nb2, nblock, ncoa, ncob, nimages, nr, nrlevel, nseta, nsetb, nsgfa, nsgfb, &
ntasks, nthread, nw, nxy, nz, nzsize, offset, pair_index, sgfa, sgfb, ub
INTEGER, DIMENSION(:), POINTER :: la_max, la_min, lb_max, lb_min, npgfa, &
npgfb, nsgfa, nsgfb
npgfb, nsgf_seta, nsgf_setb
INTEGER, DIMENSION(:, :), POINTER :: first_sgfa, first_sgfb
LOGICAL :: atom_pair_changed, distributed_rs_grids, &
do_kp, my_compute_grad, &
Expand All @@ -1443,7 +1441,6 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
REAL(KIND=dp), DIMENSION(:, :), POINTER :: p_block, pab, sphi_a, sphi_b, work, &
zeta, zetb
REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: pabt, workt
TYPE(atom_pair_type), DIMENSION(:), POINTER :: atom_pair_recv, atom_pair_send
TYPE(cell_type), POINTER :: cell
TYPE(cube_info_type), DIMENSION(:), POINTER :: cube_info
TYPE(dbcsr_p_type), DIMENSION(:), POINTER :: deltap
Expand All @@ -1459,7 +1456,6 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
POINTER :: rs_descs
TYPE(realspace_grid_p_type), DIMENSION(:), POINTER :: rs_rho
TYPE(realspace_grid_type) :: threadlocal_rsgrid
TYPE(rs_blocks_type) :: blocks_recv
TYPE(task_list_type), POINTER :: task_list, task_list_soft
TYPE(task_type), DIMENSION(:), POINTER :: tasks

Expand All @@ -1470,7 +1466,7 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,

NULLIFY (qs_kind, cell, dft_control, orb_basis_set, deltap, &
qs_kind_set, particle_set, rs_rho, pw_env, rs_descs, la_max, &
la_min, lb_max, lb_min, npgfa, npgfb, nsgfa, nsgfb, p_block, &
la_min, lb_max, lb_min, npgfa, npgfb, nsgf_seta, nsgf_setb, p_block, &
sphi_a, sphi_b, zeta, zetb, first_sgfa, first_sgfb, tasks, pabt, &
workt, lgrid)

Expand Down Expand Up @@ -1585,8 +1581,6 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
IF (my_soft) task_list => task_list_soft
CPASSERT(ASSOCIATED(task_list))
tasks => task_list%tasks
atom_pair_send => task_list%atom_pair_send
atom_pair_recv => task_list%atom_pair_recv
ntasks = task_list%ntasks

! *** set up the pw multi-grids
Expand Down Expand Up @@ -1641,10 +1635,9 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,

! distribute the matrix
IF (distributed_rs_grids) THEN
CALL rs_scatter_matrix(rs_descs=rs_descs, pmats=deltap, blocks_recv=blocks_recv, &
atom_pair_send=atom_pair_send, atom_pair_recv=atom_pair_recv)
CALL rs_scatter_matrix(task_list, rs_descs, deltap)
ELSE
CALL rs_copy_matrix(pmats=deltap, blocks_recv=blocks_recv, atom_pair_recv=atom_pair_recv)
CALL rs_copy_matrix(task_list, deltap)
ENDIF

! map all tasks on the grids
Expand All @@ -1653,14 +1646,15 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
!$OMP SHARED(ntasks,tasks,nimages,natoms,particle_set,pabt,workt), &
!$OMP SHARED(my_basis_type,my_soft,maxco,ncoset,nthread), &
!$OMP SHARED(cell,cube_info,ga_gb_function), &
!$OMP SHARED(rs_rho,lgrid,gridlevel_info,task_list,qs_kind_set, blocks_recv), &
!$OMP SHARED(rs_rho,lgrid,gridlevel_info,task_list,qs_kind_set), &
!$OMP PRIVATE(igrid_level,iatom,jatom,iset,jset,ipgf,jpgf,ikind,jkind,pab,work), &
!$OMP PRIVATE(img,img_old,iatom_old,jatom_old,iset_old,jset_old,ikind_old,jkind_old), &
!$OMP PRIVATE(qs_kind,orb_basis_set,first_sgfa,la_max,la_min), &
!$OMP PRIVATE(npgfa,nseta,nsgfa,sphi_a,zeta,first_sgfb,lb_max,lb_min,npgfb), &
!$OMP PRIVATE(nsetb,nsgfb,sphi_b,zetb,p_block, p_block_set), &
!$OMP PRIVATE(npgfa,nseta,nsgf_seta,sphi_a,zeta,first_sgfb,lb_max,lb_min,npgfb), &
!$OMP PRIVATE(nsetb,nsgf_setb,sphi_b,zetb,p_block, p_block_set), &
!$OMP PRIVATE(atom_pair_changed,ncoa,sgfa,ncob,sgfb,rab,ra,rb), &
!$OMP PRIVATE(na1,na2,nb1,nb2,scale,use_subpatch,ithread,lb,ub,n,nw), &
!$OMP PRIVATE(nsgfa,nsgfb,pair_index,block_size,offset), &
!$OMP PRIVATE(itask,nz,nxy,nzsize,nrlevel,nblock,lbw,lbr,nr,threadlocal_rsgrid)

ithread = 0
Expand Down Expand Up @@ -1697,6 +1691,7 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
jset = tasks(itask)%jset
ipgf = tasks(itask)%ipgf
jpgf = tasks(itask)%jpgf
pair_index = tasks(itask)%pair_index
ikind = particle_set(iatom)%atomic_kind%kind_number
jkind = particle_set(jatom)%atomic_kind%kind_number
Expand All @@ -1713,7 +1708,8 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
lmin=la_min, &
npgf=npgfa, &
nset=nseta, &
nsgf_set=nsgfa, &
nsgf_set=nsgf_seta, &
nsgf=nsgfa, &
sphi=sphi_a, &
zet=zeta)
ENDIF
Expand All @@ -1727,12 +1723,21 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
lmin=lb_min, &
npgf=npgfb, &
nset=nsetb, &
nsgf_set=nsgfb, &
nsgf_set=nsgf_setb, &
nsgf=nsgfb, &
sphi=sphi_b, &
zet=zetb)
ENDIF
CALL rs_get_block_p(blocks_recv, tasks(itask), symmetric=.TRUE., block=p_block)
! fetch block from buffer
block_size = nsgfa*nsgfb
offset = task_list%pair_offsets_recv(pair_index)
IF (iatom <= jatom) THEN
p_block(1:nsgfa, 1:nsgfb) => task_list%buffer_recv(offset + 1:offset + block_size)
ELSE
!iatom / jatom swapped
p_block(1:nsgfb, 1:nsgfa) => task_list%buffer_recv(offset + 1:offset + block_size)
ENDIF
iatom_old = iatom
jatom_old = jatom
Expand All @@ -1754,16 +1759,16 @@ SUBROUTINE calculate_rho_elec(matrix_p, matrix_p_kp, rho, rho_gspace, total_rho,
sgfb = first_sgfb(1, jset)
IF (ALLOCATED(p_block_set)) DEALLOCATE (p_block_set)
ALLOCATE (p_block_set(nsgfa(iset), nsgfb(jset)))
ALLOCATE (p_block_set(nsgf_seta(iset), nsgf_setb(jset)))
IF (iatom <= jatom) THEN
p_block_set(:, :) = p_block(sgfa:sgfa + nsgfa(iset) - 1, sgfb:sgfb + nsgfb(jset) - 1)
p_block_set(:, :) = p_block(sgfa:sgfa + nsgf_seta(iset) - 1, sgfb:sgfb + nsgf_setb(jset) - 1)
ELSE
p_block_set(:, :) = TRANSPOSE(p_block(sgfb:sgfb + nsgfb(jset) - 1, sgfa:sgfa + nsgfa(iset) - 1))
p_block_set(:, :) = TRANSPOSE(p_block(sgfb:sgfb + nsgf_setb(jset) - 1, sgfa:sgfa + nsgf_seta(iset) - 1))
ENDIF
work(1:ncoa, 1:nsgfb(jset)) = MATMUL(sphi_a(1:ncoa, sgfa:sgfa + nsgfa(iset) - 1), p_block_set)
pab(1:ncoa, 1:ncob) = MATMUL(work(1:ncoa, 1:nsgfb(jset)), &
TRANSPOSE(sphi_b(1:ncob, sgfb:sgfb + nsgfb(jset) - 1)))
work(1:ncoa, 1:nsgf_setb(jset)) = MATMUL(sphi_a(1:ncoa, sgfa:sgfa + nsgf_seta(iset) - 1), p_block_set)
pab(1:ncoa, 1:ncob) = MATMUL(work(1:ncoa, 1:nsgf_setb(jset)), &
TRANSPOSE(sphi_b(1:ncob, sgfb:sgfb + nsgf_setb(jset) - 1)))
iset_old = iset
jset_old = jset
Expand Down

0 comments on commit d15f8a6

Please sign in to comment.