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