Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Bethe-Salpeter equation on the RI-AO GW, the GW of PROPERTIES%BANDSTRUCTURE%GW with explicit
10 : !> three-center RI integrals over atomic orbitals (no RI_RS): reads PROPERTIES%BSE, builds the
11 : !> window-sized RI slabs and the effective inverse dielectric matrix from bs_env, and hands
12 : !> them to the solver core through prepare_bse_env_from_bandstructure_gw.
13 : !> Non-periodic systems only. Reads no real-time propagation input.
14 : !> \par History
15 : !> 09.2026 created [Maximilian Graml]
16 : !> \author Maximilian Graml
17 : ! **************************************************************************************************
18 : MODULE bse_bs_env_link
19 : USE bse_input, ONLY: read_bse_section
20 : USE bse_main, ONLY: bse_solve,&
21 : prepare_bse_env_from_bandstructure_gw
22 : USE bse_matvec, ONLY: bse_memory_left_GB,&
23 : bse_resolve_memory_per_proc
24 : USE bse_types, ONLY: bse_env_release,&
25 : bse_env_type
26 : USE bse_util, ONLY: bse_window_from_dft
27 : USE cp_blacs_env, ONLY: cp_blacs_env_type
28 : USE cp_dbcsr_api, ONLY: &
29 : dbcsr_create, dbcsr_distribution_get, dbcsr_distribution_new, dbcsr_distribution_release, &
30 : dbcsr_distribution_type, dbcsr_finalize, dbcsr_get_info, dbcsr_get_stored_coordinates, &
31 : dbcsr_put_block, dbcsr_release, dbcsr_type, dbcsr_type_no_symmetry, dbcsr_work_create
32 : USE cp_dbcsr_operations, ONLY: copy_dbcsr_to_fm
33 : USE cp_dbcsr_redistribution, ONLY: cp_dbcsr_from_subgroups
34 : USE cp_fm_basic_linalg, ONLY: cp_fm_scale_and_add
35 : USE cp_fm_diag, ONLY: cp_fm_power
36 : USE cp_fm_struct, ONLY: cp_fm_struct_create,&
37 : cp_fm_struct_release,&
38 : cp_fm_struct_type
39 : USE cp_fm_types, ONLY: cp_fm_create,&
40 : cp_fm_release,&
41 : cp_fm_set_all,&
42 : cp_fm_to_fm,&
43 : cp_fm_to_fm_submat,&
44 : cp_fm_type
45 : USE cp_log_handling, ONLY: cp_get_default_logger,&
46 : cp_logger_get_default_io_unit,&
47 : cp_logger_type
48 : USE dbt_api, ONLY: &
49 : dbt_batched_contract_finalize, dbt_batched_contract_init, dbt_clear, dbt_contract, &
50 : dbt_create, dbt_destroy, dbt_get_block, dbt_get_stored_coordinates, &
51 : dbt_iterator_blocks_left, dbt_iterator_next_block, dbt_iterator_start, dbt_iterator_stop, &
52 : dbt_iterator_type, dbt_pgrid_create, dbt_pgrid_destroy, dbt_pgrid_type, dbt_type
53 : USE dbt_matrix_redistribution, ONLY: cp_fm_to_dbt
54 : USE gw_tensor_large_cell_Gamma, ONLY: compute_3c_integrals
55 : USE gw_utils_fm, ONLY: fm_contract_ABA
56 : USE input_section_types, ONLY: section_vals_type
57 : USE kinds, ONLY: dp,&
58 : int_8
59 : USE libint_2c_3c, ONLY: libint_potential_type
60 : USE mo_window, ONLY: mo_window_type
61 : USE mp2_ri_2c, ONLY: RI_2c_integral_mat
62 : USE parallel_gemm_api, ONLY: parallel_gemm
63 : USE post_scf_bandstructure_types, ONLY: post_scf_bandstructure_type
64 : USE qs_environment_types, ONLY: get_qs_env,&
65 : qs_environment_type
66 : USE qs_tensors, ONLY: neighbor_list_3c_destroy
67 : USE qs_tensors_types, ONLY: create_2c_tensor,&
68 : create_3c_tensor
69 : #include "./base/base_uses.f90"
70 :
71 : IMPLICIT NONE
72 :
73 : PRIVATE
74 :
75 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'bse_bs_env_link'
76 :
77 : ! block size of the MO index p in the tensors: the AO and RI indices come in atomic blocks, the MO
78 : ! index has none. A block is the unit the tensor grid distributes and filters, so the size bounds
79 : ! the largest block of the MO product, (RI atom) x (q sector) x mo_block_size
80 : INTEGER, PARAMETER, PRIVATE :: mo_block_size = 64
81 :
82 : PUBLIC :: bse_from_bs_env
83 :
84 : CONTAINS
85 :
86 : ! **************************************************************************************************
87 : !> \brief Entry from the properties driver after the bandstructure block. Hands the solver core what
88 : !> its kernels are built from, for p, q in the active window:
89 : !> B^P_pq = sum_QR (pq|Q) [M^-1]_QR [V^1/2]_RP,
90 : !> ε^-1_eff,PQ = δ_PQ + sum_RS [V^-1/2]_PR W^c_RS(iω=0) [V^-1/2]_SQ,
91 : !> and the quasiparticle energies ε_p of the GW, so that, with the core's
92 : !> \bar{B}^P_ij = sum_Q ε^-1_eff,PQ B^Q_ij,
93 : !> A_ia,jb = δ_ij δ_ab (ε_a-ε_i) + α sum_P B^P_ia B^P_jb - w sum_P \bar{B}^P_ij B^P_ab
94 : !> B_ia,jb = α sum_P B^P_ia B^P_jb - w sum_P \bar{B}^P_ib B^P_ja
95 : !> V_PQ: bare Coulomb and M_PQ: RI metric in the RI basis, W^c = W - V: the correlation part
96 : !> of the GW's static screened interaction. Order: checks, section, window, screening,
97 : !> per-spin slabs, seam, solve, release
98 : !> \param qs_env source of bs_env; the slab build reads its blacs context and matrix distribution
99 : !> \param bse_section PROPERTIES%BSE
100 : ! **************************************************************************************************
101 6 : SUBROUTINE bse_from_bs_env(qs_env, bse_section)
102 : TYPE(qs_environment_type), POINTER :: qs_env
103 : TYPE(section_vals_type), POINTER :: bse_section
104 :
105 : CHARACTER(LEN=*), PARAMETER :: routineN = 'bse_from_bs_env'
106 :
107 : INTEGER :: handle, ispin, n_mo_real, n_RI, n_spin, &
108 : n_window, unit_nr
109 6 : INTEGER, ALLOCATABLE, DIMENSION(:) :: homo_full, homo_red, virt_red
110 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: eigenval_dft, eigenval_qp
111 24 : TYPE(bse_env_type) :: bse_env
112 : TYPE(cp_fm_type) :: fm_eps_inv, fm_M_inv_V_sqrt
113 6 : TYPE(cp_fm_type), ALLOCATABLE, DIMENSION(:) :: fm_S_ab, fm_S_ia, fm_S_ij
114 : TYPE(cp_logger_type), POINTER :: logger
115 : TYPE(mo_window_type) :: window
116 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
117 :
118 6 : CALL timeset(routineN, handle)
119 :
120 6 : logger => cp_get_default_logger()
121 6 : unit_nr = cp_logger_get_default_io_unit(logger)
122 6 : NULLIFY (bs_env)
123 6 : CALL get_qs_env(qs_env, bs_env=bs_env)
124 : ! the driver calls this after the bandstructure block, which creates bs_env
125 6 : IF (.NOT. ASSOCIATED(bs_env)) CPABORT("PROPERTIES%BSE found no GW environment (bs_env)")
126 : ! check_cp2k_input refused each of these at input parsing with a message that says what to change;
127 : ! Gamma is the one k-point that carries the MOs and the energies read below
128 6 : IF (.NOT. bs_env%do_gw .OR. bs_env%do_periodic .OR. bs_env%do_gw_ri_rs .OR. bs_env%nkp_bs_and_DOS /= 1) THEN
129 0 : CPABORT("PROPERTIES%BSE needs a non-periodic, Gamma-only GW without RI_RS")
130 : END IF
131 : ! read_bse_input told the GW at its init to keep its static W and the 3c neighbour list
132 6 : IF (.NOT. (bs_env%keep_W_freq_zero .AND. bs_env%keep_3c_nl) .OR. &
133 : .NOT. ASSOCIATED(bs_env%fm_W_MIC_freq_zero%matrix_struct)) THEN
134 0 : CPABORT("PROPERTIES%BSE: the GW kept no static W(0) or no 3c neighbour list")
135 : END IF
136 : ! every RI batch of the slabs goes to one tensor group; unset groups would skip them all silently
137 6 : IF (bs_env%num_tensor_groups < 1 .OR. bs_env%tensor_group_color < 0) THEN
138 0 : CPABORT("PROPERTIES%BSE: the tensor groups of the GW are not set")
139 : END IF
140 :
141 6 : n_spin = bs_env%n_spin
142 6 : n_RI = bs_env%n_RI
143 : ! the canonical orthogonalisation parks the linearly dependent directions above n_mo_retained
144 6 : n_mo_real = MIN(bs_env%n_mo_retained, bs_env%n_ao)
145 :
146 6 : CALL read_bse_section(bse_section, bse_env)
147 6 : CALL bse_resolve_memory_per_proc(bse_env%memory_per_proc_gb, bs_env%para_env, unit_nr)
148 :
149 6 : CALL link_window(bs_env, bse_env, n_mo_real, window, homo_red, virt_red, unit_nr)
150 6 : n_window = window%last_mo - window%first_mo + 1
151 :
152 6 : CALL link_screening(qs_env, bs_env, unit_nr, fm_M_inv_V_sqrt, fm_eps_inv)
153 : ! the GW kept W(0) for the screening and nl_3c for the slabs; each goes once read
154 6 : CALL cp_fm_release(bs_env%fm_W_MIC_freq_zero)
155 :
156 54 : ALLOCATE (fm_S_ia(n_spin), fm_S_ij(n_spin), fm_S_ab(n_spin))
157 48 : ALLOCATE (eigenval_qp(n_window, n_spin), eigenval_dft(n_window, n_spin), homo_full(n_spin))
158 14 : DO ispin = 1, n_spin
159 : CALL link_slabs(qs_env, bs_env, bse_env, ispin, window, homo_red(ispin), virt_red(ispin), fm_M_inv_V_sqrt, &
160 8 : fm_S_ij(ispin), fm_S_ia(ispin), fm_S_ab(ispin), unit_nr)
161 136 : eigenval_qp(:, ispin) = bs_env%eigenval_GW(window%first_mo:window%last_mo, 1, ispin)
162 136 : eigenval_dft(:, ispin) = bs_env%eigenval_scf_Gamma(window%first_mo:window%last_mo, ispin)
163 14 : homo_full(ispin) = bs_env%n_occ(ispin)
164 : END DO
165 6 : CALL cp_fm_release(fm_M_inv_V_sqrt)
166 6 : CALL neighbor_list_3c_destroy(bs_env%nl_3c)
167 :
168 : CALL prepare_bse_env_from_bandstructure_gw(bse_env, fm_S_ia, fm_S_ij, fm_S_ab, fm_eps_inv, &
169 : eigenval_qp, eigenval_dft, homo_full, homo_red, virt_red, &
170 6 : n_RI, qs_env, unit_nr)
171 6 : CALL cp_fm_release(fm_eps_inv)
172 6 : CALL bse_solve(bse_env, bs_env%fm_mo_coeff_Gamma(1:n_spin), unit_nr)
173 6 : CALL bse_env_release(bse_env)
174 :
175 6 : CALL timestop(handle)
176 :
177 18 : END SUBROUTINE bse_from_bs_env
178 :
179 : ! **************************************************************************************************
180 : !> \brief The active window on the DFT axis over the real (retained) MOs, one absolute range for
181 : !> every spin, and the levels inside it per spin: homo_red = n_occ - first_mo + 1 occupied,
182 : !> virt_red = last_mo - n_occ virtual, with n_occ the occupied MOs of the spin
183 : !> \param bs_env DFT eigenvalues and occupations at Gamma
184 : !> \param bse_env its cutoffs, read from the section already
185 : !> \param n_mo_real MOs that are states, not linear-dependency placeholders
186 : !> \param window first and last MO of the window
187 : !> \param homo_red occupied levels in the window per spin, allocated here
188 : !> \param virt_red virtual levels in the window per spin, allocated here
189 : !> \param unit_nr output unit, positive on the writing rank only
190 : ! **************************************************************************************************
191 6 : SUBROUTINE link_window(bs_env, bse_env, n_mo_real, window, homo_red, virt_red, unit_nr)
192 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
193 : TYPE(bse_env_type), INTENT(IN) :: bse_env
194 : INTEGER, INTENT(IN) :: n_mo_real
195 : TYPE(mo_window_type), INTENT(OUT) :: window
196 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: homo_red, virt_red
197 : INTEGER, INTENT(IN) :: unit_nr
198 :
199 : CHARACTER(LEN=*), PARAMETER :: routineN = 'link_window'
200 :
201 : INTEGER :: handle, ispin, n_spin
202 : INTEGER, ALLOCATABLE, DIMENSION(:) :: homo, n_mo
203 :
204 6 : CALL timeset(routineN, handle)
205 :
206 6 : n_spin = bs_env%n_spin
207 36 : ALLOCATE (n_mo(n_spin), homo(n_spin), homo_red(n_spin), virt_red(n_spin))
208 14 : n_mo(:) = n_mo_real
209 14 : homo(:) = bs_env%n_occ(1:n_spin)
210 : CALL bse_window_from_dft(bs_env%eigenval_scf_Gamma(1:n_mo_real, 1:n_spin), n_mo, homo, &
211 6 : bse_env%bse_cutoff_occ, bse_env%bse_cutoff_empty, window, warn=.TRUE.)
212 14 : DO ispin = 1, n_spin
213 8 : homo_red(ispin) = homo(ispin) - window%first_mo + 1
214 8 : virt_red(ispin) = window%last_mo - homo(ispin)
215 8 : IF (homo_red(ispin) < 1) CPABORT("PROPERTIES%BSE: ENERGY_CUTOFF_OCC leaves no occupied level")
216 14 : IF (virt_red(ispin) < 1) CPABORT("PROPERTIES%BSE: ENERGY_CUTOFF_EMPTY leaves no empty level")
217 : END DO
218 9 : IF (unit_nr > 0) WRITE (unit_nr, '(T2,A4,T7,A,T61,2I10)') 'BSE|', &
219 6 : 'Active window on the DFT axis (first, last MO)', window%first_mo, window%last_mo
220 6 : DEALLOCATE (n_mo, homo)
221 :
222 6 : CALL timestop(handle)
223 :
224 6 : END SUBROUTINE link_window
225 :
226 : ! **************************************************************************************************
227 : !> \brief The RI transform and the effective inverse dielectric matrix from the GW's static W,
228 : !> with V_PQ the bare Coulomb and M_PQ the RI metric in the RI basis:
229 : !> [M^-1 V^1/2]_QP = sum_R [M^-1]_QR [V^1/2]_RP, so that the slabs
230 : !> B^P_pq = sum_Q (pq|Q) [M^-1 V^1/2]_QP give
231 : !> sum_P B^P_pq B^P_rs = sum_QR (pq|Q) [M^-1 V M^-1]_QR (R|rs);
232 : !> with the stored W~_PQ = sum_RS [M^-1]_PR W^c_RS(iω=0) [M^-1]_SQ,
233 : !> ε^-1_eff,PQ = δ_PQ + sum_RS [V^-1/2]_PR W^c_RS(iω=0) [V^-1/2]_SQ, so that
234 : !> sum_PQ B^P_pq ε^-1_eff,PQ B^Q_rs = sum_RS (pq|R) [M^-1 (V + W^c(iω=0)) M^-1]_RS (S|rs),
235 : !> the full static W of the GW. The GW builds the same factor M^-1 V^1/2 per frequency from a
236 : !> Cholesky factor of V (compute_MinvVsqrt_Vsqrt); the symmetric roots here give V^-1/2 as
237 : !> well, which the stored W~ needs, with the eigenvalue threshold of the GW (eps_eigval_mat_RI).
238 : !> Two diagonalisations of an RI x RI matrix, once per run
239 : !> \param qs_env for the 2c RI integrals
240 : !> \param bs_env RI metric and its inverse, the stored static W~, the RI x RI matrix structure
241 : !> \param unit_nr output unit, positive on the writing rank only
242 : !> \param fm_M_inv_V_sqrt [M^-1 V^1/2]_QP, RI x RI, created here
243 : !> \param fm_eps_inv ε^-1_eff,PQ, RI x RI, created here
244 : ! **************************************************************************************************
245 6 : SUBROUTINE link_screening(qs_env, bs_env, unit_nr, fm_M_inv_V_sqrt, fm_eps_inv)
246 : TYPE(qs_environment_type), POINTER :: qs_env
247 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
248 : INTEGER, INTENT(IN) :: unit_nr
249 : TYPE(cp_fm_type), INTENT(OUT) :: fm_M_inv_V_sqrt, fm_eps_inv
250 :
251 : CHARACTER(LEN=*), PARAMETER :: routineN = 'link_screening'
252 :
253 : INTEGER :: handle, n_dependent, n_RI
254 : TYPE(cp_fm_type) :: fm_V_inv_sqrt, fm_V_sqrt, fm_Wc, fm_work
255 6 : TYPE(cp_fm_type), ALLOCATABLE, DIMENSION(:, :) :: fm_M, fm_V
256 : TYPE(libint_potential_type) :: coulomb_op
257 :
258 6 : CALL timeset(routineN, handle)
259 :
260 6 : n_RI = bs_env%n_RI
261 : ! bare Coulomb V_PQ (coulomb_op keeps its default, the full 1/r) and the RI metric M_PQ
262 6 : CALL RI_2c_integral_mat(qs_env, fm_V, bs_env%fm_RI_RI%matrix_struct, n_RI, coulomb_op)
263 6 : CALL RI_2c_integral_mat(qs_env, fm_M, bs_env%fm_RI_RI%matrix_struct, n_RI, bs_env%ri_metric)
264 :
265 : ! [V^1/2]_PQ and [V^-1/2]_PQ by symmetric roots; eigenvalues of V under the GW's threshold are dropped
266 6 : CALL cp_fm_create(fm_V_sqrt, bs_env%fm_RI_RI%matrix_struct)
267 6 : CALL cp_fm_create(fm_V_inv_sqrt, bs_env%fm_RI_RI%matrix_struct)
268 6 : CALL cp_fm_create(fm_work, bs_env%fm_RI_RI%matrix_struct)
269 6 : CALL cp_fm_to_fm(fm_V(1, 1), fm_V_sqrt)
270 6 : CALL cp_fm_power(fm_V_sqrt, fm_work, 0.5_dp, bs_env%eps_eigval_mat_RI, n_dependent)
271 6 : CALL cp_fm_to_fm(fm_V(1, 1), fm_V_inv_sqrt)
272 6 : CALL cp_fm_power(fm_V_inv_sqrt, fm_work, -0.5_dp, bs_env%eps_eigval_mat_RI, n_dependent)
273 6 : IF (unit_nr > 0 .AND. n_dependent > 0) WRITE (unit_nr, '(T2,A4,T7,A,T71,I10)') 'BSE|', &
274 0 : 'RI Coulomb directions dropped by the eigenvalue threshold', n_dependent
275 :
276 : ! [M^-1 V^1/2]_QP = sum_R [M^-1]_QR [V^1/2]_RP
277 6 : CALL cp_fm_create(fm_M_inv_V_sqrt, bs_env%fm_RI_RI%matrix_struct)
278 : CALL parallel_gemm("N", "N", n_RI, n_RI, n_RI, 1.0_dp, bs_env%fm_Minv_Gamma, fm_V_sqrt, &
279 6 : 0.0_dp, fm_M_inv_V_sqrt)
280 :
281 : ! W^c_PQ(iω=0) = sum_RS M_PR W~_RS M_SQ, in place on a copy of the stored W~
282 6 : CALL cp_fm_create(fm_Wc, bs_env%fm_RI_RI%matrix_struct)
283 6 : CALL cp_fm_to_fm(bs_env%fm_W_MIC_freq_zero, fm_Wc)
284 6 : CALL fm_contract_ABA(fm_M(1, 1), fm_Wc)
285 : ! ε^-1_eff,PQ = δ_PQ + sum_RS [V^-1/2]_PR W^c_RS(iω=0) [V^-1/2]_SQ, the sum in the work matrix
286 6 : CALL fm_contract_ABA(fm_V_inv_sqrt, fm_Wc, fm_work)
287 6 : CALL cp_fm_create(fm_eps_inv, bs_env%fm_RI_RI%matrix_struct)
288 6 : CALL cp_fm_set_all(fm_eps_inv, 0.0_dp, 1.0_dp)
289 6 : CALL cp_fm_scale_and_add(1.0_dp, fm_eps_inv, 1.0_dp, fm_work)
290 :
291 6 : CALL cp_fm_release(fm_V)
292 6 : CALL cp_fm_release(fm_M)
293 6 : CALL cp_fm_release(fm_V_sqrt)
294 6 : CALL cp_fm_release(fm_V_inv_sqrt)
295 6 : CALL cp_fm_release(fm_Wc)
296 6 : CALL cp_fm_release(fm_work)
297 :
298 6 : CALL timestop(handle)
299 :
300 12 : END SUBROUTINE link_screening
301 :
302 : ! **************************************************************************************************
303 : !> \brief The window columns of the MO coefficients, C_μp for p in the window, as a 2-index
304 : !> tensor (AO atom | MO block) on the tensor group's grid, the MO blocking the caller's.
305 : !> The coefficients are a full matrix on all ranks, the tensor lives on the ranks of one
306 : !> tensor group, and no copy between the two exists. The routine takes the route of the GW's
307 : !> own matrices (cp_fm_to_dbt: full matrix, DBCSR matrix on all ranks, DBCSR matrix on
308 : !> the group, tensor) and creates the two DBCSR matrices of that route in the
309 : !> (AO atom | MO block) blocking, for which the GW holds no template. A small object,
310 : !> n_ao x n_window: the steps are there for the change of communicator, not for memory
311 : !> \param bs_env MO coefficients at Gamma, AO block sizes, the tensor group
312 : !> \param dbcsr_dist distribution of the SCF matrices, the template of the global matrix
313 : !> \param ispin spin channel
314 : !> \param window first and last MO of the active window
315 : !> \param sizes_mo MO block sizes of the tensor, summing to the window
316 : !> \param t_C the tensor, created here
317 : ! **************************************************************************************************
318 176 : SUBROUTINE link_mo_tensor(bs_env, dbcsr_dist, ispin, window, sizes_mo, t_C)
319 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
320 : TYPE(dbcsr_distribution_type), POINTER :: dbcsr_dist
321 : INTEGER, INTENT(IN) :: ispin
322 : TYPE(mo_window_type), INTENT(IN) :: window
323 : INTEGER, DIMENSION(:), INTENT(IN) :: sizes_mo
324 : TYPE(dbt_type), INTENT(OUT) :: t_C
325 :
326 : CHARACTER(LEN=*), PARAMETER :: routineN = 'link_mo_tensor'
327 :
328 : INTEGER :: handle, n_window
329 16 : INTEGER, ALLOCATABLE, DIMENSION(:) :: dist_1, dist_2
330 : INTEGER, DIMENSION(2) :: pdims_2d
331 : TYPE(cp_fm_struct_type), POINTER :: struct_C_window
332 : TYPE(cp_fm_type) :: fm_C_window
333 : TYPE(dbcsr_distribution_type) :: dist_global, dist_group, dist_tmpl
334 : TYPE(dbcsr_type) :: mat_global, mat_group
335 48 : TYPE(dbt_pgrid_type) :: pgrid_2d
336 :
337 16 : CALL timeset(routineN, handle)
338 :
339 16 : n_window = window%last_mo - window%first_mo + 1
340 : ! the MO blocks cover the window exactly
341 48 : IF (SUM(sizes_mo) /= n_window) CPABORT("MO block sizes do not sum to the active window")
342 :
343 16 : pdims_2d = 0
344 16 : CALL dbt_pgrid_create(bs_env%para_env_tensor, pdims_2d, pgrid_2d)
345 16 : CALL create_2c_tensor(t_C, dist_1, dist_2, pgrid_2d, bs_env%sizes_AO, sizes_mo, name="(AO | MO)")
346 16 : DEALLOCATE (dist_1, dist_2)
347 16 : CALL dbt_pgrid_destroy(pgrid_2d)
348 :
349 : ! the window columns as a full matrix, then the route of the GW's own matrices (cp_fm_to_dbt):
350 : ! global DBCSR, group DBCSR, tensor; the coefficients are never replicated
351 16 : NULLIFY (struct_C_window)
352 : CALL cp_fm_struct_create(struct_C_window, template_fmstruct=bs_env%fm_mo_coeff_Gamma(ispin)%matrix_struct, &
353 16 : ncol_global=n_window)
354 16 : CALL cp_fm_create(fm_C_window, struct_C_window)
355 16 : CALL cp_fm_struct_release(struct_C_window)
356 16 : CALL cp_fm_to_fm_submat(bs_env%fm_mo_coeff_Gamma(ispin), fm_C_window, bs_env%n_ao, n_window, 1, window%first_mo, 1, 1)
357 16 : CALL link_matrix_round_robin(dbcsr_dist, bs_env%sizes_AO, sizes_mo, "C window", dist_global, mat_global)
358 : ! the group matrix on the grid of the GW's AO matrices of the tensor group
359 16 : CALL dbcsr_get_info(bs_env%mat_ao_ao_tensor%matrix, distribution=dist_tmpl)
360 16 : CALL link_matrix_round_robin(dist_tmpl, bs_env%sizes_AO, sizes_mo, "C window, group", dist_group, mat_group)
361 : CALL cp_fm_to_dbt(fm_C_window, mat_global, mat_group, t_C, bs_env%para_env, &
362 16 : bs_env%para_env_tensor%num_pe, bs_env%eps_filter)
363 :
364 16 : CALL cp_fm_release(fm_C_window)
365 16 : CALL dbcsr_release(mat_global)
366 16 : CALL dbcsr_release(mat_group)
367 16 : CALL dbcsr_distribution_release(dist_global)
368 16 : CALL dbcsr_distribution_release(dist_group)
369 :
370 16 : CALL timestop(handle)
371 :
372 32 : END SUBROUTINE link_mo_tensor
373 :
374 : ! **************************************************************************************************
375 : !> \brief A DBCSR matrix with the given blocking on the communicator and process grid of a template
376 : !> distribution, the blocks round robin over the grid rows and columns
377 : !> \param dist_tmpl the template
378 : !> \param row_sizes row block sizes
379 : !> \param col_sizes column block sizes
380 : !> \param matrix_name name of the matrix
381 : !> \param matrix_dist the distribution of the matrix, created here, released after the matrix
382 : !> \param new_matrix the matrix, created here
383 : ! **************************************************************************************************
384 56 : SUBROUTINE link_matrix_round_robin(dist_tmpl, row_sizes, col_sizes, matrix_name, matrix_dist, new_matrix)
385 : TYPE(dbcsr_distribution_type), INTENT(IN) :: dist_tmpl
386 : INTEGER, DIMENSION(:), INTENT(IN) :: row_sizes, col_sizes
387 : CHARACTER(LEN=*), INTENT(IN) :: matrix_name
388 : TYPE(dbcsr_distribution_type), INTENT(OUT) :: matrix_dist
389 : TYPE(dbcsr_type), INTENT(OUT) :: new_matrix
390 :
391 : INTEGER :: i_block, npcols, nprows
392 56 : INTEGER, DIMENSION(:), POINTER :: col_blk, col_dist, row_blk, row_dist
393 :
394 56 : NULLIFY (col_blk, col_dist, row_blk, row_dist)
395 56 : CALL dbcsr_distribution_get(dist_tmpl, nprows=nprows, npcols=npcols)
396 392 : ALLOCATE (row_dist(SIZE(row_sizes)), col_dist(SIZE(col_sizes)), row_blk(SIZE(row_sizes)), col_blk(SIZE(col_sizes)))
397 448 : row_dist(:) = [(MOD(i_block - 1, nprows), i_block=1, SIZE(row_sizes))]
398 232 : col_dist(:) = [(MOD(i_block - 1, npcols), i_block=1, SIZE(col_sizes))]
399 224 : row_blk(:) = row_sizes(:)
400 144 : col_blk(:) = col_sizes(:)
401 56 : CALL dbcsr_distribution_new(matrix_dist, template=dist_tmpl, row_dist=row_dist, col_dist=col_dist)
402 : CALL dbcsr_create(new_matrix, name=matrix_name, dist=matrix_dist, matrix_type=dbcsr_type_no_symmetry, &
403 56 : row_blk_size=row_blk, col_blk_size=col_blk)
404 56 : DEALLOCATE (row_dist, col_dist, row_blk, col_blk)
405 :
406 112 : END SUBROUTINE link_matrix_round_robin
407 :
408 : ! **************************************************************************************************
409 : !> \brief Block sizes of the window index p: blocks of at most mo_block_size levels inside the
410 : !> occupied and inside the virtual sector, so that no block straddles the two
411 : !> \param n_occ occupied levels in the window
412 : !> \param n_virt virtual levels in the window
413 : !> \param block_sizes levels per block, occupied blocks first, allocated here
414 : !> \param n_block_occ occupied blocks
415 : ! **************************************************************************************************
416 8 : SUBROUTINE link_sector_blocks(n_occ, n_virt, block_sizes, n_block_occ)
417 : INTEGER, INTENT(IN) :: n_occ, n_virt
418 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: block_sizes
419 : INTEGER, INTENT(OUT) :: n_block_occ
420 :
421 : INTEGER :: n_block_virt
422 :
423 8 : n_block_occ = (n_occ - 1)/mo_block_size + 1
424 8 : n_block_virt = (n_virt - 1)/mo_block_size + 1
425 24 : ALLOCATE (block_sizes(n_block_occ + n_block_virt))
426 24 : block_sizes(:) = mo_block_size
427 8 : block_sizes(n_block_occ) = n_occ - (n_block_occ - 1)*mo_block_size
428 8 : block_sizes(n_block_occ + n_block_virt) = n_virt - (n_block_virt - 1)*mo_block_size
429 :
430 8 : END SUBROUTINE link_sector_blocks
431 :
432 : ! **************************************************************************************************
433 : !> \brief Contiguous RI atom ranges of at most max_RI_per_batch RI functions; an atom larger than
434 : !> that forms its own batch
435 : !> \param sizes_RI RI functions per atom
436 : !> \param max_RI_per_batch RI functions per batch the memory left allows
437 : !> \param batch_first first atom per batch, allocated here
438 : !> \param batch_last last atom per batch, allocated here
439 : ! **************************************************************************************************
440 8 : SUBROUTINE link_ri_batches(sizes_RI, max_RI_per_batch, batch_first, batch_last)
441 : INTEGER, DIMENSION(:), INTENT(IN) :: sizes_RI
442 : INTEGER, INTENT(IN) :: max_RI_per_batch
443 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: batch_first, batch_last
444 :
445 : INTEGER :: i_atom, n_atom, n_batch, n_RI_in_batch
446 8 : INTEGER, ALLOCATABLE, DIMENSION(:) :: first_atom, last_atom
447 :
448 8 : n_atom = SIZE(sizes_RI)
449 32 : ALLOCATE (first_atom(n_atom), last_atom(n_atom))
450 32 : n_batch = 0
451 32 : n_RI_in_batch = 0
452 32 : DO i_atom = 1, n_atom
453 24 : IF (n_RI_in_batch > 0 .AND. n_RI_in_batch + sizes_RI(i_atom) > max_RI_per_batch) THEN
454 0 : last_atom(n_batch) = i_atom - 1
455 : n_RI_in_batch = 0
456 : END IF
457 24 : IF (n_RI_in_batch == 0) THEN
458 8 : n_batch = n_batch + 1
459 8 : first_atom(n_batch) = i_atom
460 : END IF
461 32 : n_RI_in_batch = n_RI_in_batch + sizes_RI(i_atom)
462 : END DO
463 8 : last_atom(n_batch) = n_atom
464 32 : ALLOCATE (batch_first(n_batch), batch_last(n_batch))
465 16 : batch_first(:) = first_atom(1:n_batch)
466 16 : batch_last(:) = last_atom(1:n_batch)
467 8 : DEALLOCATE (first_atom, last_atom)
468 :
469 8 : END SUBROUTINE link_ri_batches
470 :
471 : ! **************************************************************************************************
472 : !> \brief What MEMORY_PER_PROC leaves to the slab build per rank and the largest RI batch it allows:
473 : !> per RI function of a batch, on the ranks of its tensor group the 3c integrals (P|μν), the
474 : !> half-transform (P μ|q) and the product (P|qp), and on all ranks the block matrices of the
475 : !> batch, their global copies and the batch's full matrix; one atom per batch with a warning
476 : !> when one RI atom alone exceeds what is left
477 : !> \param bs_env sizes and communicators
478 : !> \param bse_env the resolved MEMORY_PER_PROC
479 : !> \param n_window levels in the active window
480 : !> \param unit_nr output unit, positive on the writing rank only
481 : !> \param budget_bytes what is left, in bytes; the RI transform chunks its work matrix by it as well
482 : !> \param max_RI_per_batch RI functions per batch within what is left, at least one
483 : ! **************************************************************************************************
484 8 : SUBROUTINE link_batch_budget(bs_env, bse_env, n_window, unit_nr, budget_bytes, max_RI_per_batch)
485 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
486 : TYPE(bse_env_type), INTENT(IN) :: bse_env
487 : INTEGER, INTENT(IN) :: n_window, unit_nr
488 : INTEGER(KIND=int_8), INTENT(OUT) :: budget_bytes
489 : INTEGER, INTENT(OUT) :: max_RI_per_batch
490 :
491 : CHARACTER(LEN=16) :: left_str, need_str
492 : INTEGER(KIND=int_8) :: bytes_per_RI, largest_batch_bytes, &
493 : one_atom_bytes
494 : REAL(KIND=dp) :: mem_left_GB
495 :
496 8 : mem_left_GB = bse_memory_left_GB(bse_env%memory_per_proc_gb, bs_env%para_env)
497 8 : budget_bytes = INT(mem_left_GB*1.0E9_dp, KIND=int_8)
498 : ! per RI function of a batch, all distributed: on the ranks of its tensor group the 3c
499 : ! integrals (P|μν), the half-transform (P μ|q) and the product (P|qp); on all ranks the block
500 : ! matrices of the batch, their global copies (two inside the group-to-global hop) and the
501 : ! batch's full matrix, four (P|qp); the contractions' own scratch is not counted
502 : bytes_per_RI = 8_int_8*((INT(bs_env%n_ao, int_8)**2 + INT(bs_env%n_ao, int_8)*INT(n_window, int_8) + &
503 : INT(n_window, int_8)**2)/INT(bs_env%para_env_tensor%num_pe, int_8) + &
504 8 : 4_int_8*INT(n_window, int_8)**2/INT(bs_env%para_env%num_pe, int_8))
505 32 : max_RI_per_batch = INT(MAX(1_int_8, MIN(INT(SUM(bs_env%sizes_RI), int_8), budget_bytes/bytes_per_RI)))
506 32 : one_atom_bytes = INT(MAXVAL(bs_env%sizes_RI), int_8)*bytes_per_RI
507 8 : largest_batch_bytes = MAX(INT(max_RI_per_batch, int_8)*bytes_per_RI, one_atom_bytes)
508 8 : IF (unit_nr > 0) THEN
509 4 : WRITE (unit_nr, '(T2,A4,T7,A,T67,F14.3)') 'BSE|', 'Memory per MPI rank for the BSE, MEMORY_PER_PROC (GB)', &
510 8 : bse_env%memory_per_proc_gb
511 4 : WRITE (unit_nr, '(T2,A4,T7,A,T67,F14.3)') 'BSE|', 'Left for the 3c integrals after the memory held (GB)', &
512 8 : mem_left_GB
513 4 : WRITE (unit_nr, '(T2,A4,T7,A,T67,F14.3)') 'BSE|', 'Largest RI batch per MPI rank (GB)', &
514 8 : REAL(largest_batch_bytes, dp)*1.0E-9_dp
515 : END IF
516 : ! nothing smaller than one atom per batch exists: run with it and say so
517 8 : IF (one_atom_bytes > budget_bytes) THEN
518 0 : WRITE (need_str, '(F16.3)') REAL(one_atom_bytes, dp)*1.0E-9_dp
519 0 : WRITE (left_str, '(F16.3)') mem_left_GB
520 : CALL cp_warn(__LOCATION__, &
521 : "PROPERTIES%BSE: the 3c integral batch of one RI atom needs "//TRIM(ADJUSTL(need_str))// &
522 : " GB per MPI rank, "//TRIM(ADJUSTL(left_str))//" GB are left of MEMORY_PER_PROC; one atom "// &
523 0 : "per batch. Raise MEMORY_PER_PROC or use more MPI ranks.")
524 : END IF
525 :
526 8 : END SUBROUTINE link_batch_budget
527 :
528 : ! **************************************************************************************************
529 : !> \brief The three window-sized slabs of one spin, B^P_pq = sum_Q (pq|Q) [M^-1 V^1/2]_QP with
530 : !> (Q|pq) = sum_μν C_μp (Q|μν) C_νq, batched over RI atoms: the 3c integrals of the batch
531 : !> are rebuilt and transformed over ν, then over μ, by two tensor contractions on the batch's
532 : !> tensor group; the product goes to all ranks and its rows are copied into the slabs; the
533 : !> RI transform by M^-1 V^1/2 follows in column chunks. Column conventions of the solver core:
534 : !> ia = (i-1) n_virt + a, ij = (i-1) n_occ + j, ab = (a-1) n_virt + b
535 : !> \param qs_env blacs context, distribution of the SCF matrices, environment of the 3c integrals
536 : !> \param bs_env RI and AO sizes, the 3c neighbour list and tensor template, the tensor groups
537 : !> \param bse_env its memory budget and memory check
538 : !> \param ispin spin channel
539 : !> \param window first and last MO of the active window
540 : !> \param n_occ occupied levels in the window
541 : !> \param n_virt virtual levels in the window
542 : !> \param fm_M_inv_V_sqrt the RI transform [M^-1 V^1/2]_QP, RI x RI
543 : !> \param fm_S_ij B^P_ij, RI x n_occ*n_occ, created here
544 : !> \param fm_S_ia B^P_ia, RI x n_occ*n_virt, created here
545 : !> \param fm_S_ab B^P_ab, RI x n_virt*n_virt, created here
546 : !> \param unit_nr output unit, positive on the writing rank only
547 : ! **************************************************************************************************
548 16 : SUBROUTINE link_slabs(qs_env, bs_env, bse_env, ispin, window, n_occ, n_virt, fm_M_inv_V_sqrt, fm_S_ij, fm_S_ia, &
549 : fm_S_ab, unit_nr)
550 : TYPE(qs_environment_type), POINTER :: qs_env
551 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
552 : TYPE(bse_env_type), INTENT(IN) :: bse_env
553 : INTEGER, INTENT(IN) :: ispin
554 : TYPE(mo_window_type), INTENT(IN) :: window
555 : INTEGER, INTENT(IN) :: n_occ, n_virt
556 : TYPE(cp_fm_type), INTENT(IN) :: fm_M_inv_V_sqrt
557 : TYPE(cp_fm_type), INTENT(OUT) :: fm_S_ij, fm_S_ia, fm_S_ab
558 : INTEGER, INTENT(IN) :: unit_nr
559 :
560 : CHARACTER(LEN=*), PARAMETER :: routineN = 'link_slabs'
561 :
562 : INTEGER :: atom_first, atom_last, handle, i_batch, &
563 : max_RI_per_batch, n_batch, &
564 : n_block_occ, n_RI, n_window
565 : INTEGER(KIND=int_8) :: budget_bytes
566 8 : INTEGER, ALLOCATABLE, DIMENSION(:) :: batch_first, batch_last, dist_1, dist_2, &
567 8 : dist_3, sizes_p, sizes_q
568 : INTEGER, DIMENSION(3) :: pdims_3d
569 : TYPE(cp_blacs_env_type), POINTER :: blacs_env
570 : TYPE(cp_fm_struct_type), POINTER :: struct_ab, struct_ia, struct_ij
571 : TYPE(dbcsr_distribution_type), POINTER :: dbcsr_dist
572 24 : TYPE(dbt_pgrid_type) :: pgrid_3d
573 328 : TYPE(dbt_type) :: t_3c, t_3c_half, t_3c_mo, t_C_p, t_C_q
574 :
575 8 : CALL timeset(routineN, handle)
576 :
577 8 : NULLIFY (blacs_env, dbcsr_dist, struct_ab, struct_ia, struct_ij)
578 8 : CALL get_qs_env(qs_env, blacs_env=blacs_env, dbcsr_dist=dbcsr_dist)
579 8 : n_RI = bs_env%n_RI
580 8 : n_window = n_occ + n_virt
581 : CALL cp_fm_struct_create(struct_ij, para_env=bs_env%para_env, context=blacs_env, &
582 8 : nrow_global=n_RI, ncol_global=n_occ*n_occ)
583 : CALL cp_fm_struct_create(struct_ia, para_env=bs_env%para_env, context=blacs_env, &
584 8 : nrow_global=n_RI, ncol_global=n_occ*n_virt)
585 : CALL cp_fm_struct_create(struct_ab, para_env=bs_env%para_env, context=blacs_env, &
586 8 : nrow_global=n_RI, ncol_global=n_virt*n_virt)
587 : ! the slabs collect (Q|pq) batch by batch; the RI transform at the end makes them B^P_pq in place
588 8 : CALL cp_fm_create(fm_S_ij, struct_ij, name="fm_S_ij", set_zero=.TRUE.)
589 8 : CALL cp_fm_create(fm_S_ia, struct_ia, name="fm_S_ia", set_zero=.TRUE.)
590 8 : CALL cp_fm_create(fm_S_ab, struct_ab, name="fm_S_ab", set_zero=.TRUE.)
591 8 : CALL cp_fm_struct_release(struct_ij)
592 8 : CALL cp_fm_struct_release(struct_ia)
593 8 : CALL cp_fm_struct_release(struct_ab)
594 :
595 : ! the window index q in two blocks, occupied and virtual, so that every (P, q, p) block of the
596 : ! product covers the whole q range of one pair type; p in blocks inside each sector
597 8 : ALLOCATE (sizes_q(2))
598 24 : sizes_q(:) = [n_occ, n_virt]
599 8 : CALL link_sector_blocks(n_occ, n_virt, sizes_p, n_block_occ)
600 8 : CALL link_mo_tensor(bs_env, dbcsr_dist, ispin, window, sizes_q, t_C_q)
601 8 : CALL link_mo_tensor(bs_env, dbcsr_dist, ispin, window, sizes_p, t_C_p)
602 :
603 8 : CALL link_batch_budget(bs_env, bse_env, n_window, unit_nr, budget_bytes, max_RI_per_batch)
604 8 : CALL link_ri_batches(bs_env%sizes_RI, max_RI_per_batch, batch_first, batch_last)
605 8 : n_batch = SIZE(batch_first)
606 8 : IF (unit_nr > 0) WRITE (unit_nr, '(T2,A4,T7,A,T71,I10)') 'BSE|', 'RI atom batches of the 3c integral transformation', n_batch
607 :
608 : ! (RI AO | MO) of the half-transform and (RI | MO MO) of the product on the tensor group's grid
609 8 : pdims_3d = 0
610 8 : CALL dbt_pgrid_create(bs_env%para_env_tensor, pdims_3d, pgrid_3d)
611 : CALL create_3c_tensor(t_3c_half, dist_1, dist_2, dist_3, pgrid_3d, bs_env%sizes_RI, bs_env%sizes_AO, &
612 8 : sizes_q, map1=[1, 2], map2=[3], name="(RI AO | MO)")
613 8 : DEALLOCATE (dist_1, dist_2, dist_3)
614 8 : CALL dbt_pgrid_destroy(pgrid_3d)
615 8 : pdims_3d = 0
616 : CALL dbt_pgrid_create(bs_env%para_env_tensor, pdims_3d, pgrid_3d, &
617 32 : tensor_dims=[bs_env%n_atom, SIZE(sizes_q), SIZE(sizes_p)])
618 : CALL create_3c_tensor(t_3c_mo, dist_1, dist_2, dist_3, pgrid_3d, bs_env%sizes_RI, sizes_q, sizes_p, &
619 8 : map1=[1], map2=[2, 3], name="(RI | MO MO)")
620 8 : DEALLOCATE (dist_1, dist_2, dist_3)
621 8 : CALL dbt_pgrid_destroy(pgrid_3d)
622 8 : CALL dbt_create(bs_env%t_RI_AO__AO, t_3c)
623 :
624 : ! C_q contracts once per batch in a layout-compatible contraction: the batched context keeps its replicated
625 : ! copy across the batches; on the small operand only, so that no grid changes under the neighbour list
626 : ! of the 3c integrals. No context on C_p: no layout of (Pμ|q) fits both contractions, so the second remaps it
627 8 : CALL dbt_batched_contract_init(t_C_q)
628 :
629 16 : DO i_batch = 1, n_batch
630 8 : atom_first = batch_first(i_batch)
631 8 : atom_last = batch_last(i_batch)
632 : ! one tensor group per batch, round robin; every rank joins the hop into the slabs
633 8 : IF (MODULO(i_batch - 1, bs_env%num_tensor_groups) == bs_env%tensor_group_color) THEN
634 24 : CALL compute_3c_integrals(qs_env, bs_env, t_3c, atoms_RI=[atom_first, atom_last])
635 : ! (Pμ|q) = sum_ν (P|μν) C_νq
636 : CALL dbt_contract(alpha=1.0_dp, tensor_1=t_3c, tensor_2=t_C_q, beta=0.0_dp, tensor_3=t_3c_half, &
637 : contract_1=[3], notcontract_1=[1, 2], map_1=[1, 2], &
638 : contract_2=[1], notcontract_2=[2], map_2=[3], &
639 8 : filter_eps=bs_env%eps_filter)
640 8 : CALL dbt_clear(t_3c)
641 : ! (P|qp) = sum_μ (Pμ|q) C_μp, distributed over the group
642 : CALL dbt_contract(alpha=1.0_dp, tensor_1=t_3c_half, tensor_2=t_C_p, beta=0.0_dp, tensor_3=t_3c_mo, &
643 : contract_1=[2], notcontract_1=[1, 3], map_1=[1, 2], &
644 : contract_2=[1], notcontract_2=[2], map_2=[3], &
645 8 : filter_eps=bs_env%eps_filter)
646 8 : CALL dbt_clear(t_3c_half)
647 : END IF
648 : CALL link_slab_rows(bs_env, blacs_env, dbcsr_dist, t_3c_mo, n_occ, n_virt, sizes_p, n_block_occ, atom_first, atom_last, &
649 8 : fm_S_ij, fm_S_ia, fm_S_ab)
650 16 : CALL dbt_clear(t_3c_mo)
651 : END DO
652 :
653 8 : CALL dbt_batched_contract_finalize(t_C_q)
654 8 : CALL dbt_destroy(t_3c)
655 8 : CALL dbt_destroy(t_3c_half)
656 8 : CALL dbt_destroy(t_3c_mo)
657 8 : CALL dbt_destroy(t_C_q)
658 8 : CALL dbt_destroy(t_C_p)
659 8 : DEALLOCATE (sizes_q, sizes_p, batch_first, batch_last)
660 :
661 : ! B^P_pq = sum_Q [M^-1 V^1/2]_QP (Q|pq), in column chunks, in place
662 8 : CALL apply_ri_transform(fm_M_inv_V_sqrt, fm_S_ij, n_RI, n_occ*n_occ, budget_bytes)
663 8 : CALL apply_ri_transform(fm_M_inv_V_sqrt, fm_S_ia, n_RI, n_occ*n_virt, budget_bytes)
664 8 : CALL apply_ri_transform(fm_M_inv_V_sqrt, fm_S_ab, n_RI, n_virt*n_virt, budget_bytes)
665 :
666 8 : CALL timestop(handle)
667 :
668 24 : END SUBROUTINE link_slabs
669 :
670 : ! **************************************************************************************************
671 : !> \brief The batch product (P|qp) into the rows of its RI functions in the three slabs: the local
672 : !> (P, q, p) blocks into one block matrix per pair type on the tensor group's grid, the
673 : !> group-to-global hop of the GW, the block-cyclic copy, the row copy into the slab; every
674 : !> rank takes part, a rank outside the batch's group with no blocks. The pair (q, p) lands
675 : !> in column (p-1) n_q + q of its slab, n_q = n_occ for ij and n_virt for ia and ab:
676 : !> ij = (i-1) n_occ + j at (q, p) = (j, i), ia = (i-1) n_virt + a at (a, i),
677 : !> ab = (a-1) n_virt + b at (b, a)
678 : !> \param bs_env RI offsets and sizes per atom, the communicator of all ranks
679 : !> \param blacs_env context of the slabs
680 : !> \param dbcsr_dist a distribution over all ranks, the template of the global block matrices
681 : !> \param t_3c_mo (P|qp) of the RI atoms atom_first..atom_last, q in two blocks (occupied, virtual),
682 : !> p in sector blocks
683 : !> \param n_occ occupied levels in the window
684 : !> \param n_virt virtual levels in the window
685 : !> \param sizes_p the p blocks, occupied blocks first
686 : !> \param n_block_occ occupied blocks of p
687 : !> \param atom_first first RI atom of the batch
688 : !> \param atom_last last RI atom of the batch
689 : !> \param fm_S_ij slab of the ij pairs; the rows of the batch are written
690 : !> \param fm_S_ia slab of the ia pairs; the rows of the batch are written
691 : !> \param fm_S_ab slab of the ab pairs; the rows of the batch are written
692 : ! **************************************************************************************************
693 8 : SUBROUTINE link_slab_rows(bs_env, blacs_env, dbcsr_dist, t_3c_mo, n_occ, n_virt, sizes_p, n_block_occ, atom_first, atom_last, &
694 : fm_S_ij, fm_S_ia, fm_S_ab)
695 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
696 : TYPE(cp_blacs_env_type), POINTER :: blacs_env
697 : TYPE(dbcsr_distribution_type), POINTER :: dbcsr_dist
698 : TYPE(dbt_type), INTENT(INOUT) :: t_3c_mo
699 : INTEGER, INTENT(IN) :: n_occ, n_virt
700 : INTEGER, DIMENSION(:), INTENT(IN) :: sizes_p
701 : INTEGER, INTENT(IN) :: n_block_occ, atom_first, atom_last
702 : TYPE(cp_fm_type), INTENT(IN) :: fm_S_ij, fm_S_ia, fm_S_ab
703 :
704 : CHARACTER(LEN=*), PARAMETER :: routineN = 'link_slab_rows'
705 : INTEGER, PARAMETER :: pair_ab = 3, pair_ia = 2, pair_ij = 1
706 :
707 : INTEGER :: col_block, group_handle, handle, i_pair, my_rank, n_atom_batch, n_block_p, &
708 : n_RI_batch, npcol, nprow, owner_rank, p_block, q_block, rank, RI_first, row_block
709 8 : INTEGER, ALLOCATABLE, DIMENSION(:, :) :: cart_coords
710 : INTEGER, DIMENSION(3) :: blk_index, blk_size, n_col_blocks, &
711 : n_cols_pair, p_block_offset, &
712 : q_block_of_pair
713 16 : INTEGER, DIMENSION(:), POINTER :: col_dist, col_sizes, row_dist, row_sizes
714 8 : INTEGER, DIMENSION(:, :), POINTER :: pgrid
715 : LOGICAL :: found
716 8 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: block_Pqp
717 : TYPE(cp_fm_struct_type), POINTER :: struct_batch
718 : TYPE(cp_fm_type) :: fm_batch
719 56 : TYPE(dbcsr_distribution_type), DIMENSION(3) :: dist_global, dist_group
720 56 : TYPE(dbcsr_type), DIMENSION(3) :: mat_global, mat_group
721 : TYPE(dbt_iterator_type) :: iter
722 :
723 8 : CALL timeset(routineN, handle)
724 :
725 8 : NULLIFY (col_dist, col_sizes, pgrid, row_dist, row_sizes, struct_batch)
726 8 : RI_first = bs_env%i_RI_start_from_atom(atom_first)
727 8 : n_RI_batch = bs_env%i_RI_end_from_atom(atom_last) - RI_first + 1
728 8 : n_atom_batch = atom_last - atom_first + 1
729 8 : n_block_p = SIZE(sizes_p)
730 32 : n_cols_pair = [n_occ*n_occ, n_occ*n_virt, n_virt*n_virt]
731 32 : n_col_blocks = [n_block_occ, n_block_occ, n_block_p - n_block_occ]
732 8 : q_block_of_pair = [1, 2, 2]
733 : ! column block col_block of a pair type is the p block p_block = p_block_offset + col_block
734 32 : p_block_offset = [0, 0, n_block_occ]
735 :
736 : ! the group grid of the tensor as the process grid of the local matrices, its ranks placed by
737 : ! their coordinates, so that every block the tensor holds is local to its matrix
738 : ASSOCIATE (cart => t_3c_mo%pgrid%mp_comm_2d)
739 8 : my_rank = cart%mepos
740 8 : group_handle = cart%get_handle()
741 8 : nprow = cart%num_pe_cart(1)
742 8 : npcol = cart%num_pe_cart(2)
743 48 : ALLOCATE (cart_coords(2, 0:cart%num_pe - 1), pgrid(0:nprow - 1, 0:npcol - 1))
744 32 : DO rank = 0, cart%num_pe - 1
745 16 : CALL cart%coords(rank, cart_coords(:, rank))
746 24 : pgrid(cart_coords(1, rank), cart_coords(2, rank)) = rank
747 : END DO
748 : END ASSOCIATE
749 32 : ALLOCATE (row_dist(n_atom_batch), row_sizes(n_atom_batch))
750 32 : DO row_block = 1, n_atom_batch
751 96 : CALL dbt_get_stored_coordinates(t_3c_mo, [atom_first + row_block - 1, 1, 1], owner_rank)
752 24 : row_dist(row_block) = cart_coords(1, owner_rank)
753 32 : row_sizes(row_block) = bs_env%sizes_RI(atom_first + row_block - 1)
754 : END DO
755 32 : DO i_pair = 1, 3
756 96 : ALLOCATE (col_dist(n_col_blocks(i_pair)), col_sizes(n_col_blocks(i_pair)))
757 48 : DO col_block = 1, n_col_blocks(i_pair)
758 24 : p_block = p_block_offset(i_pair) + col_block
759 96 : CALL dbt_get_stored_coordinates(t_3c_mo, [atom_first, q_block_of_pair(i_pair), p_block], owner_rank)
760 24 : col_dist(col_block) = cart_coords(2, owner_rank)
761 : ! the q range of the pair type times the p block
762 64 : col_sizes(col_block) = MERGE(n_occ, n_virt, i_pair == pair_ij)*sizes_p(p_block)
763 : END DO
764 : CALL dbcsr_distribution_new(dist_group(i_pair), group=group_handle, pgrid=pgrid, &
765 24 : row_dist=row_dist, col_dist=col_dist)
766 : CALL dbcsr_create(mat_group(i_pair), name="(P|qp) batch, group", dist=dist_group(i_pair), &
767 24 : matrix_type=dbcsr_type_no_symmetry, row_blk_size=row_sizes, col_blk_size=col_sizes)
768 : ! the threaded put_block below cannot create the work matrix itself
769 24 : CALL dbcsr_work_create(mat_group(i_pair), work_mutable=.TRUE.)
770 : ! the global matrix of the batch, round robin over the process grid of the slabs
771 24 : CALL link_matrix_round_robin(dbcsr_dist, row_sizes, col_sizes, "(P|qp) batch", dist_global(i_pair), mat_global(i_pair))
772 32 : DEALLOCATE (col_dist, col_sizes)
773 : END DO
774 8 : DEALLOCATE (pgrid, cart_coords, row_dist, row_sizes)
775 :
776 : ! the local (P, q, p) blocks as (P, qp) blocks, q fastest, the column order of the slabs; the
777 : ! iterator needs the full thread team
778 : !$OMP PARALLEL DEFAULT(NONE) SHARED(t_3c_mo, mat_group, n_block_occ, atom_first, p_block_offset, my_rank) &
779 8 : !$OMP PRIVATE(iter, blk_index, blk_size, block_Pqp, found, q_block, p_block, i_pair, row_block, col_block, owner_rank)
780 : CALL dbt_iterator_start(iter, t_3c_mo)
781 : DO WHILE (dbt_iterator_blocks_left(iter))
782 : CALL dbt_iterator_next_block(iter, blk_index, blk_size=blk_size)
783 : q_block = blk_index(2)
784 : p_block = blk_index(3)
785 : ! pair types of the (q block, p block) blocks: (occ, occ) ij, (virt, occ) ia, (virt, virt) ab;
786 : ! (occ, virt) is the ai block, which no slab holds
787 : IF (q_block == 1 .AND. p_block > n_block_occ) CYCLE
788 : IF (q_block == 1) THEN
789 : i_pair = pair_ij
790 : ELSE IF (p_block <= n_block_occ) THEN
791 : i_pair = pair_ia
792 : ELSE
793 : i_pair = pair_ab
794 : END IF
795 : ALLOCATE (block_Pqp(blk_size(1), blk_size(2), blk_size(3)))
796 : CALL dbt_get_block(t_3c_mo, blk_index, blk_size, block_Pqp, found)
797 : IF (found) THEN
798 : row_block = blk_index(1) - atom_first + 1
799 : col_block = p_block - p_block_offset(i_pair)
800 : CALL dbcsr_get_stored_coordinates(mat_group(i_pair), row_block, col_block, owner_rank)
801 : ! the group matrix was laid out by the tensor's coordinates: the rank holding the block owns it
802 : IF (owner_rank /= my_rank) CPABORT("(P|qp) block is not local to its group matrix")
803 : CALL dbcsr_put_block(mat_group(i_pair), row_block, col_block, &
804 : RESHAPE(block_Pqp, [blk_size(1), blk_size(2)*blk_size(3)]))
805 : END IF
806 : DEALLOCATE (block_Pqp)
807 : END DO
808 : CALL dbt_iterator_stop(iter)
809 : !$OMP END PARALLEL
810 :
811 32 : DO i_pair = 1, 3
812 24 : CALL dbcsr_finalize(mat_group(i_pair))
813 24 : CALL cp_dbcsr_from_subgroups(mat_group(i_pair), mat_global(i_pair), bs_env%para_env)
814 : CALL cp_fm_struct_create(struct_batch, para_env=bs_env%para_env, context=blacs_env, &
815 24 : nrow_global=n_RI_batch, ncol_global=n_cols_pair(i_pair))
816 24 : CALL cp_fm_create(fm_batch, struct_batch)
817 24 : CALL cp_fm_struct_release(struct_batch)
818 24 : CALL copy_dbcsr_to_fm(mat_global(i_pair), fm_batch)
819 : ! known bottleneck: cp_fm_to_fm_submat copies column by column, one pdcopy per pair column; the
820 : ! memory it takes on single ranks can stay until the end of the run and set their peak. A block
821 : ! redistribution of fm_batch into the rows of the slab would replace it
822 8 : SELECT CASE (i_pair)
823 : CASE (pair_ij)
824 8 : CALL cp_fm_to_fm_submat(fm_batch, fm_S_ij, n_RI_batch, n_cols_pair(i_pair), 1, 1, RI_first, 1)
825 : CASE (pair_ia)
826 8 : CALL cp_fm_to_fm_submat(fm_batch, fm_S_ia, n_RI_batch, n_cols_pair(i_pair), 1, 1, RI_first, 1)
827 : CASE (pair_ab)
828 24 : CALL cp_fm_to_fm_submat(fm_batch, fm_S_ab, n_RI_batch, n_cols_pair(i_pair), 1, 1, RI_first, 1)
829 : END SELECT
830 24 : CALL cp_fm_release(fm_batch)
831 24 : CALL dbcsr_release(mat_group(i_pair))
832 24 : CALL dbcsr_release(mat_global(i_pair))
833 24 : CALL dbcsr_distribution_release(dist_group(i_pair))
834 56 : CALL dbcsr_distribution_release(dist_global(i_pair))
835 : END DO
836 :
837 8 : CALL timestop(handle)
838 :
839 32 : END SUBROUTINE link_slab_rows
840 :
841 : ! **************************************************************************************************
842 : !> \brief B^P_c = sum_Q [M^-1 V^1/2]_QP (Q|c) for every column c of a slab, in place: the columns pass in
843 : !> chunks through one RI x chunk work matrix of at most budget_bytes per rank,
844 : !> max_cols_chunk = budget_bytes N_all/(8 n_RI) with N_all the ranks of the slab
845 : !> \param fm_M_inv_V_sqrt the RI transform [M^-1 V^1/2]_QP, RI x RI
846 : !> \param fm_slab (Q|c) on entry, B^P_c on exit
847 : !> \param n_RI RI functions, the rows of the slab
848 : !> \param n_cols columns of the slab
849 : !> \param budget_bytes per-rank budget of the work matrix, which is distributed over the ranks
850 : ! **************************************************************************************************
851 48 : SUBROUTINE apply_ri_transform(fm_M_inv_V_sqrt, fm_slab, n_RI, n_cols, budget_bytes)
852 : TYPE(cp_fm_type), INTENT(IN) :: fm_M_inv_V_sqrt, fm_slab
853 : INTEGER, INTENT(IN) :: n_RI, n_cols
854 : INTEGER(KIND=int_8), INTENT(IN) :: budget_bytes
855 :
856 : CHARACTER(LEN=*), PARAMETER :: routineN = 'apply_ri_transform'
857 :
858 : INTEGER :: col_first, handle, max_cols_chunk, &
859 : n_cols_chunk
860 : TYPE(cp_fm_struct_type), POINTER :: struct_work
861 : TYPE(cp_fm_type) :: fm_work
862 :
863 24 : CALL timeset(routineN, handle)
864 :
865 : max_cols_chunk = INT(MAX(1_int_8, MIN(INT(n_cols, int_8), &
866 : budget_bytes*INT(fm_slab%matrix_struct%para_env%num_pe, int_8) &
867 24 : /(8_int_8*INT(n_RI, int_8)))))
868 24 : NULLIFY (struct_work)
869 : CALL cp_fm_struct_create(struct_work, para_env=fm_slab%matrix_struct%para_env, &
870 24 : context=fm_slab%matrix_struct%context, nrow_global=n_RI, ncol_global=max_cols_chunk)
871 24 : CALL cp_fm_create(fm_work, struct_work)
872 24 : CALL cp_fm_struct_release(struct_work)
873 48 : DO col_first = 1, n_cols, max_cols_chunk
874 24 : n_cols_chunk = MIN(max_cols_chunk, n_cols - col_first + 1)
875 : ! work_P,c = sum_Q [M^-1 V^1/2]_QP slab_Q,c for the columns c of this chunk, then back into the
876 : ! slab; the copy back is column by column (one pdcopy each), see link_slab_rows
877 : CALL parallel_gemm("T", "N", n_RI, n_cols_chunk, n_RI, 1.0_dp, fm_M_inv_V_sqrt, fm_slab, 0.0_dp, fm_work, &
878 24 : b_first_col=col_first, c_first_col=1)
879 : CALL cp_fm_to_fm_submat(fm_work, fm_slab, nrow=n_RI, ncol=n_cols_chunk, s_firstrow=1, s_firstcol=1, &
880 48 : t_firstrow=1, t_firstcol=col_first)
881 : END DO
882 24 : CALL cp_fm_release(fm_work)
883 :
884 24 : CALL timestop(handle)
885 :
886 24 : END SUBROUTINE apply_ri_transform
887 :
888 : END MODULE bse_bs_env_link
|
