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 Matrix-free application of the BSE matrices A and B to trial vectors from RI slabs that
10 : !> are sliced along the RI index over all MPI ranks
11 : !> \par History
12 : !> 09.2026 created [Maximilian Graml]
13 : ! **************************************************************************************************
14 : MODULE bse_matvec
15 :
16 : USE bse_util, ONLY: ia_of_occ_virt,&
17 : occ_of_ia,&
18 : virt_of_ia
19 : USE cp_blacs_env, ONLY: cp_blacs_env_create,&
20 : cp_blacs_env_release,&
21 : cp_blacs_env_type
22 : USE cp_fm_struct, ONLY: cp_fm_struct_create,&
23 : cp_fm_struct_release,&
24 : cp_fm_struct_type
25 : USE cp_fm_types, ONLY: cp_fm_create,&
26 : cp_fm_get_info,&
27 : cp_fm_get_submatrix,&
28 : cp_fm_release,&
29 : cp_fm_set_all,&
30 : cp_fm_to_fm_submat_general,&
31 : cp_fm_type
32 : USE input_constants, ONLY: bse_precond_full_diag
33 : USE kinds, ONLY: dp
34 : USE message_passing, ONLY: mp_mem_avail_per_rank_GB,&
35 : mp_para_env_type
36 : #include "./base/base_uses.f90"
37 :
38 : IMPLICIT NONE
39 :
40 : PRIVATE
41 :
42 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'bse_matvec'
43 :
44 : ! share of the free memory per rank the iterative solver may use, for the pass width of
45 : ! bse_matvec_apply and the default budget of the subspace ceiling: half, since the probe counts
46 : ! reclaimable cache and is optimistic
47 : REAL(KIND=dp), PARAMETER, PUBLIC :: mem_fraction = 0.5_dp
48 :
49 : PUBLIC :: bse_matvec_env_type, bse_matvec_create, bse_matvec_release, bse_matvec_apply, &
50 : bse_matvec_diagonal, bse_matvec_subblock, bse_matvec_vector_struct, bse_matvec_selfcheck
51 :
52 : ! **************************************************************************************************
53 : !> \brief RI slabs of one spin channel, every rank holding n_ri_loc whole RI slices
54 : !> \param alpha prefactor of the exchange term (2 singlet, 0 triplet)
55 : !> \param w_fac prefactor of the screened term (0 switches it off)
56 : !> \param eps_diff ε_a - ε_i at the compound index ia = (i-1)*virt + a of bse_util's ia_of_occ_virt
57 : !> \param B_ia B^P_ia at (a, i, P)
58 : !> \param B_bar_ia \bar{B}^P_ia at (a, i, P), only for do_abba
59 : !> \param B_bar_ij \bar{B}^P_ij at (j, i, P)
60 : !> \param B_ab B^P_ab at (b, a, P)
61 : !> \param row_count rows of a trial vector owned by each rank, in rank order
62 : !> \param row_displ first row of each rank minus one, so that row_count and row_displ describe
63 : !> the contiguous ia range per rank that allgatherv and sum_scatter need
64 : !> \param blacs_env npe x 1 process grid of the slabs and of all trial vectors
65 : !> \param block_cols trial vectors per pass in bse_matvec_apply, -1 to size it from free memory
66 : ! **************************************************************************************************
67 : TYPE bse_matvec_env_type
68 : INTEGER :: homo = 0, virt = 0, n_ov = 0, &
69 : n_ri_loc = 0, block_cols = -1
70 : REAL(KIND=dp) :: alpha = 0.0_dp, w_fac = 0.0_dp
71 : LOGICAL :: do_abba = .FALSE.
72 : INTEGER, ALLOCATABLE, DIMENSION(:) :: row_count, row_displ
73 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: eps_diff
74 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: B_ia, B_bar_ia, B_bar_ij, B_ab
75 : TYPE(cp_blacs_env_type), POINTER :: blacs_env => NULL()
76 : TYPE(mp_para_env_type), POINTER :: para_env => NULL()
77 : END TYPE bse_matvec_env_type
78 :
79 : CONTAINS
80 :
81 : ! **************************************************************************************************
82 : !> \brief Moves the RI slabs onto an npe x 1 process grid, such that every rank owns whole RI
83 : !> slices, and stores them as local 3-index arrays: the contraction over each P in
84 : !> bse_matvec_apply is then a local DGEMM and only the final sum over P crosses the ranks
85 : !> \param mv_env the environment created here
86 : !> \param fm_S_ia B^P_ia, N_RI x homo*virt
87 : !> \param fm_S_bar_ij \bar{B}^P_ij, the slab that enters W_ij,ab, N_RI x homo*homo
88 : !> \param fm_S_ab B^P_ab, N_RI x virt*virt
89 : !> \param fm_S_bar_ia \bar{B}^P_ia, the slab that enters W_ib,aj, N_RI x homo*virt
90 : !> \param eps_reduced quasiparticle energies of the homo+virt active levels
91 : !> \param homo occupied levels of the active window
92 : !> \param virt virtual levels of the active window
93 : !> \param alpha prefactor of the exchange term, 2 singlet, 0 triplet
94 : !> \param w_fac prefactor of the screened term, 0 switches it off
95 : !> \param do_abba slice \bar{B}^P_ia as well, for the application of B
96 : !> \param unit_nr output unit, positive on the writing rank only
97 : !> \param block_cols trial vectors per pass of bse_matvec_apply; absent or -1 sizes it from the free memory
98 : ! **************************************************************************************************
99 6 : SUBROUTINE bse_matvec_create(mv_env, fm_S_ia, fm_S_bar_ij, fm_S_ab, fm_S_bar_ia, eps_reduced, &
100 : homo, virt, alpha, w_fac, do_abba, unit_nr, block_cols)
101 :
102 : TYPE(bse_matvec_env_type), INTENT(OUT) :: mv_env
103 : TYPE(cp_fm_type), INTENT(IN) :: fm_S_ia, fm_S_bar_ij, fm_S_ab, &
104 : fm_S_bar_ia
105 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: eps_reduced
106 : INTEGER, INTENT(IN) :: homo, virt
107 : REAL(KIND=dp), INTENT(IN) :: alpha, w_fac
108 : LOGICAL, INTENT(IN) :: do_abba
109 : INTEGER, INTENT(IN) :: unit_nr
110 : INTEGER, INTENT(IN), OPTIONAL :: block_cols
111 :
112 : CHARACTER(LEN=*), PARAMETER :: routineN = 'bse_matvec_create'
113 :
114 : CHARACTER(LEN=12) :: n_idle_str, n_ri_str, npe_str
115 : INTEGER :: a, handle, i, n_ri, n_ri_slab, &
116 : n_row_block
117 : REAL(KIND=dp) :: mem_slabs_GB, n_slab_entries
118 :
119 6 : CALL timeset(routineN, handle)
120 :
121 6 : mv_env%homo = homo
122 6 : mv_env%virt = virt
123 6 : IF (PRESENT(block_cols)) mv_env%block_cols = block_cols
124 6 : mv_env%n_ov = homo*virt
125 6 : mv_env%alpha = alpha
126 6 : mv_env%w_fac = w_fac
127 6 : mv_env%do_abba = do_abba
128 6 : mv_env%para_env => fm_S_ia%matrix_struct%para_env
129 :
130 : CALL cp_blacs_env_create(mv_env%blacs_env, mv_env%para_env, &
131 18 : grid_2d=[mv_env%para_env%num_pe, 1])
132 :
133 : ! one contiguous ia range per rank, so that a trial vector is gathered with allgatherv and
134 : ! its image scattered back with sum_scatter, rather than replicated by two allreduces
135 6 : n_row_block = contiguous_row_block(mv_env%n_ov, mv_env%para_env%num_pe)
136 30 : ALLOCATE (mv_env%row_count(mv_env%para_env%num_pe), mv_env%row_displ(mv_env%para_env%num_pe))
137 18 : DO i = 1, mv_env%para_env%num_pe
138 12 : mv_env%row_displ(i) = MIN(mv_env%n_ov, (i - 1)*n_row_block)
139 18 : mv_env%row_count(i) = MIN(mv_env%n_ov, i*n_row_block) - mv_env%row_displ(i)
140 : END DO
141 :
142 : ! eps_diff_ia = ε_a - ε_i at the compound index ia = (i-1)*virt + a, the row index of every trial vector
143 18 : ALLOCATE (mv_env%eps_diff(mv_env%n_ov))
144 30 : DO i = 1, homo
145 318 : DO a = 1, virt
146 312 : mv_env%eps_diff(ia_of_occ_virt(i, a, virt)) = eps_reduced(homo + a) - eps_reduced(i)
147 : END DO
148 : END DO
149 :
150 : ! every slab is sliced on the same grid, so every rank owns the same RI indices of each slab
151 6 : CALL slice_slab(mv_env, fm_S_ia, virt, homo, mv_env%B_ia, mv_env%n_ri_loc)
152 6 : IF (w_fac /= 0.0_dp) THEN
153 6 : CALL slice_slab(mv_env, fm_S_bar_ij, homo, homo, mv_env%B_bar_ij, n_ri_slab)
154 6 : CPASSERT(n_ri_slab == mv_env%n_ri_loc)
155 6 : CALL slice_slab(mv_env, fm_S_ab, virt, virt, mv_env%B_ab, n_ri_slab)
156 6 : CPASSERT(n_ri_slab == mv_env%n_ri_loc)
157 6 : IF (do_abba) THEN
158 4 : CALL slice_slab(mv_env, fm_S_bar_ia, virt, homo, mv_env%B_bar_ia, n_ri_slab)
159 4 : CPASSERT(n_ri_slab == mv_env%n_ri_loc)
160 : END IF
161 : END IF
162 :
163 6 : CALL cp_fm_get_info(fm_S_ia, nrow_global=n_ri)
164 6 : IF (mv_env%para_env%num_pe > n_ri .AND. unit_nr > 0) THEN
165 0 : WRITE (npe_str, '(I0)') mv_env%para_env%num_pe
166 0 : WRITE (n_ri_str, '(I0)') n_ri
167 0 : WRITE (n_idle_str, '(I0)') mv_env%para_env%num_pe - n_ri
168 : CALL cp_warn(__LOCATION__, &
169 : "BSE iterative solver: more MPI ranks ("//TRIM(npe_str)// &
170 : ") than RI basis functions ("//TRIM(n_ri_str)//"); "// &
171 0 : TRIM(n_idle_str)//" ranks hold no RI slice and idle in the kernel application.")
172 : END IF
173 :
174 6 : IF (unit_nr > 0) THEN
175 : ! pair entries of the sliced slabs per RI index
176 3 : n_slab_entries = REAL(homo*virt, dp)
177 3 : IF (w_fac /= 0.0_dp) THEN
178 3 : n_slab_entries = n_slab_entries + REAL(homo, dp)**2 + REAL(virt, dp)**2
179 3 : IF (do_abba) n_slab_entries = n_slab_entries + REAL(homo*virt, dp)
180 : END IF
181 3 : mem_slabs_GB = n_slab_entries*REAL(mv_env%n_ri_loc, dp)*8.0E-9_dp
182 3 : WRITE (unit_nr, '(T2,A4,T7,A,T71,I10)') 'BSE|', &
183 6 : 'Max. number of RI functions per MPI rank', mv_env%n_ri_loc
184 3 : WRITE (unit_nr, '(T2,A4,T7,A,T67,F14.3)') 'BSE|', &
185 6 : 'Memory of the RI slabs per MPI rank (GB)', mem_slabs_GB
186 : END IF
187 :
188 6 : CALL timestop(handle)
189 :
190 12 : END SUBROUTINE bse_matvec_create
191 :
192 : ! **************************************************************************************************
193 : !> \brief Rows of a trial vector per rank, the smallest block that leaves no rank beyond the last
194 : !> \param n_ov number of transitions, the global row count
195 : !> \param npe ranks of the npe x 1 grid
196 : !> \return rows of every full block; the last rank takes what remains, possibly none
197 : ! **************************************************************************************************
198 28 : PURE FUNCTION contiguous_row_block(n_ov, npe) RESULT(n_row_block)
199 :
200 : INTEGER, INTENT(IN) :: n_ov, npe
201 : INTEGER :: n_row_block
202 :
203 28 : n_row_block = (n_ov + npe - 1)/npe
204 :
205 28 : END FUNCTION contiguous_row_block
206 :
207 : ! **************************************************************************************************
208 : !> \brief Redistributes one N_RI x n_fast*n_slow slab to whole RI rows per rank and copies the owned
209 : !> rows into a contiguous (n_fast, n_slow, n_ri_loc) array, slab_loc(x, y, p) = slab(P_p, (y-1)*n_fast + x)
210 : !> for the p-th owned RI index P_p, so that every slice slab_loc(:, :, p) is a n_fast x n_slow DGEMM operand
211 : !> \param mv_env process grid and communicator of the slices
212 : !> \param fm_slab the slab on the grid of the caller, N_RI x n_fast*n_slow
213 : !> \param n_fast fast index of the pair index of the slab
214 : !> \param n_slow slow index of the pair index of the slab
215 : !> \param slab_loc the owned slices, allocated here
216 : !> \param n_ri_loc RI indices owned by this rank, the third extent of slab_loc
217 : ! **************************************************************************************************
218 22 : SUBROUTINE slice_slab(mv_env, fm_slab, n_fast, n_slow, slab_loc, n_ri_loc)
219 :
220 : TYPE(bse_matvec_env_type), INTENT(IN) :: mv_env
221 : TYPE(cp_fm_type), INTENT(IN) :: fm_slab
222 : INTEGER, INTENT(IN) :: n_fast, n_slow
223 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), &
224 : INTENT(OUT) :: slab_loc
225 : INTEGER, INTENT(OUT) :: n_ri_loc
226 :
227 : CHARACTER(LEN=*), PARAMETER :: routineN = 'slice_slab'
228 :
229 : INTEGER :: handle, n_ri, npair, p
230 : TYPE(cp_fm_struct_type), POINTER :: fm_struct
231 : TYPE(cp_fm_type) :: fm_sliced
232 :
233 22 : CALL timeset(routineN, handle)
234 :
235 : ! row count from the struct: the slabs may carry fewer rows than the full RI basis
236 22 : CALL cp_fm_get_info(fm_slab, nrow_global=n_ri, ncol_global=npair)
237 22 : CPASSERT(npair == n_fast*n_slow)
238 :
239 : ! one-row blocks on the npe x 1 grid spread the RI rows cyclically over the ranks; the copy
240 : ! below turns each strided local_data row into a contiguous (n_fast, n_slow) slice
241 22 : NULLIFY (fm_struct)
242 : CALL cp_fm_struct_create(fm_struct, para_env=mv_env%para_env, context=mv_env%blacs_env, &
243 : nrow_global=n_ri, ncol_global=npair, nrow_block=1, &
244 22 : force_block=.TRUE.)
245 22 : CALL cp_fm_create(fm_sliced, fm_struct, name="fm_slab_ri_sliced")
246 22 : CALL cp_fm_struct_release(fm_struct)
247 : CALL cp_fm_to_fm_submat_general(fm_slab, fm_sliced, n_ri, npair, 1, 1, 1, 1, &
248 22 : fm_slab%matrix_struct%context)
249 :
250 22 : CALL cp_fm_get_info(fm_sliced, nrow_local=n_ri_loc)
251 110 : ALLOCATE (slab_loc(n_fast, n_slow, n_ri_loc))
252 935 : DO p = 1, n_ri_loc
253 2761 : slab_loc(:, :, p) = RESHAPE(fm_sliced%local_data(p, 1:npair), [n_fast, n_slow])
254 : END DO
255 22 : CALL cp_fm_release(fm_sliced)
256 :
257 22 : CALL timestop(handle)
258 :
259 66 : END SUBROUTINE slice_slab
260 :
261 : ! **************************************************************************************************
262 : !> \brief Frees the sliced slabs, the transition energies and the process grid of mv_env
263 : !> \param mv_env the environment released
264 : ! **************************************************************************************************
265 6 : SUBROUTINE bse_matvec_release(mv_env)
266 :
267 : TYPE(bse_matvec_env_type), INTENT(INOUT) :: mv_env
268 :
269 6 : IF (ALLOCATED(mv_env%row_count)) DEALLOCATE (mv_env%row_count)
270 6 : IF (ALLOCATED(mv_env%row_displ)) DEALLOCATE (mv_env%row_displ)
271 6 : IF (ALLOCATED(mv_env%eps_diff)) DEALLOCATE (mv_env%eps_diff)
272 6 : IF (ALLOCATED(mv_env%B_ia)) DEALLOCATE (mv_env%B_ia)
273 6 : IF (ALLOCATED(mv_env%B_bar_ia)) DEALLOCATE (mv_env%B_bar_ia)
274 6 : IF (ALLOCATED(mv_env%B_bar_ij)) DEALLOCATE (mv_env%B_bar_ij)
275 6 : IF (ALLOCATED(mv_env%B_ab)) DEALLOCATE (mv_env%B_ab)
276 6 : IF (ASSOCIATED(mv_env%blacs_env)) CALL cp_blacs_env_release(mv_env%blacs_env)
277 6 : NULLIFY (mv_env%para_env)
278 :
279 6 : END SUBROUTINE bse_matvec_release
280 :
281 : ! **************************************************************************************************
282 : !> \brief Matrix structure of a block of trial vectors: rows ia distributed, all columns local
283 : !> \param mv_env process grid and communicator of the trial vectors
284 : !> \param ncol_global number of trial vectors
285 : !> \param fm_struct created here, released by the caller
286 : ! **************************************************************************************************
287 22 : SUBROUTINE bse_matvec_vector_struct(mv_env, ncol_global, fm_struct)
288 :
289 : TYPE(bse_matvec_env_type), INTENT(IN) :: mv_env
290 : INTEGER, INTENT(IN) :: ncol_global
291 : TYPE(cp_fm_struct_type), POINTER :: fm_struct
292 :
293 22 : NULLIFY (fm_struct)
294 : ! the row block matches mv_env%row_count, which bse_matvec_apply's collectives rely on
295 : CALL cp_fm_struct_create(fm_struct, para_env=mv_env%para_env, context=mv_env%blacs_env, &
296 : nrow_global=mv_env%n_ov, ncol_global=ncol_global, &
297 : nrow_block=contiguous_row_block(mv_env%n_ov, mv_env%para_env%num_pe), &
298 22 : force_block=.TRUE.)
299 :
300 22 : END SUBROUTINE bse_matvec_vector_struct
301 :
302 : ! **************************************************************************************************
303 : !> \brief Applies A (and B) to ncol trial vectors without forming an N_ov x N_ov object,
304 : !> (A Z)_ia = (ε_a-ε_i) Z_ia + α sum_P B^P_ia sum_jb B^P_jb Z_jb - w sum_Pjb \bar{B}^P_ij B^P_ab Z_jb
305 : !> (B Z)_ia = α sum_P B^P_ia sum_jb B^P_jb Z_jb - w sum_Pjb \bar{B}^P_ib B^P_ja Z_jb
306 : !> α: exchange prefactor (2 singlet, 0 triplet), w: prefactor of the screened term,
307 : !> \bar{B}: slab bound by the caller of bse_matvec_create (B itself for TDHF and ALPHA
308 : !> screening, sum_Q [1+Q(0)]^-1_PQ B^Q otherwise); the sums over P run over the local RI
309 : !> slices and are completed by a sum over the ranks
310 : !> \param mv_env slabs, prefactors and transition energies
311 : !> \param fm_Z trial vectors, columns first_col .. first_col+ncol-1 are read
312 : !> \param first_col first trial vector read, and first column written
313 : !> \param ncol number of trial vectors
314 : !> \param fm_AZ receives A Z in the same columns
315 : !> \param fm_BZ receives B Z in the same columns, or from first_col_BZ on
316 : !> \param first_col_BZ first column of fm_BZ written, first_col by default
317 : ! **************************************************************************************************
318 82 : SUBROUTINE bse_matvec_apply(mv_env, fm_Z, first_col, ncol, fm_AZ, fm_BZ, first_col_BZ)
319 :
320 : TYPE(bse_matvec_env_type), INTENT(IN), TARGET :: mv_env
321 : TYPE(cp_fm_type), INTENT(IN) :: fm_Z
322 : INTEGER, INTENT(IN) :: first_col, ncol
323 : TYPE(cp_fm_type), INTENT(IN) :: fm_AZ
324 : TYPE(cp_fm_type), INTENT(IN), OPTIONAL :: fm_BZ
325 : INTEGER, INTENT(IN), OPTIONAL :: first_col_BZ
326 :
327 : CHARACTER(LEN=*), PARAMETER :: routineN = 'bse_matvec_apply'
328 :
329 : INTEGER :: col, col_shift_B, handle, handle2, ia, &
330 : iloc, k, me, n_work_bufs, nb, nb_max, &
331 : nrow_local, p
332 82 : INTEGER, DIMENSION(:), POINTER :: row_indices
333 : LOGICAL :: do_B
334 : REAL(KIND=dp) :: mem_avail_GB
335 82 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: R_loc
336 82 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:), TARGET :: BabZ_buf, RA_buf, RB_buf, Z_buf
337 82 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: BjaZ_ab, T_Pk
338 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :), &
339 82 : POINTER :: B_ia_2d, BabZ_wide, RA, RB, Z, Z_wide
340 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :, :), &
341 82 : POINTER :: BabZ_3d, RA_3d, RB_3d, Z_3d
342 :
343 82 : CALL timeset(routineN, handle)
344 :
345 82 : do_B = PRESENT(fm_BZ)
346 82 : IF (do_B) THEN
347 : ! B Z needs the \bar{B}^P_ia slab, which is sliced for do_abba only
348 66 : CPASSERT(mv_env%do_abba)
349 : END IF
350 82 : col_shift_B = 0
351 82 : IF (PRESENT(first_col_BZ)) col_shift_B = first_col_BZ - first_col
352 :
353 82 : CALL cp_fm_get_info(fm_Z, nrow_local=nrow_local, row_indices=row_indices)
354 : ! the collectives below read the rows of one rank as one range, as bse_matvec_vector_struct
355 : ! lays them out; a vector built elsewhere would scatter into the wrong rows
356 82 : me = mv_env%para_env%mepos + 1
357 82 : CPASSERT(nrow_local == mv_env%row_count(me))
358 82 : CPASSERT(nrow_local == 0 .OR. row_indices(1) == mv_env%row_displ(me) + 1)
359 246 : ALLOCATE (R_loc(nrow_local))
360 :
361 : ! columns per pass from the memory of the replicated work arrays Z, B_ab Z, R^A (and R^B)
362 82 : IF (do_B) THEN
363 : n_work_bufs = 4
364 : ELSE
365 16 : n_work_bufs = 3
366 : END IF
367 82 : IF (mv_env%block_cols > 0) THEN
368 : ! pinned by input: the replicated work arrays are then a known n_work_bufs*n_ov*nb doubles
369 82 : nb_max = MIN(ncol, mv_env%block_cols)
370 : ELSE
371 0 : CALL mp_mem_avail_per_rank_GB(mv_env%para_env, mem_avail_GB)
372 0 : nb_max = ncol
373 0 : IF (mem_avail_GB > 0.0_dp) THEN
374 : nb_max = INT(MIN(REAL(ncol, dp), &
375 0 : mem_fraction*mem_avail_GB*1.0E9_dp/(8.0_dp*REAL(n_work_bufs, dp)*REAL(mv_env%n_ov, dp))))
376 : END IF
377 : END IF
378 82 : nb_max = MAX(nb_max, 1)
379 :
380 410 : ALLOCATE (Z_buf(mv_env%n_ov*nb_max), RA_buf(mv_env%n_ov*nb_max), BabZ_buf(mv_env%n_ov*nb_max))
381 82 : IF (do_B) THEN
382 330 : ALLOCATE (RB_buf(mv_env%n_ov*nb_max), BjaZ_ab(mv_env%virt, mv_env%virt))
383 : END IF
384 328 : ALLOCATE (T_Pk(mv_env%n_ri_loc, nb_max))
385 82 : B_ia_2d(1:mv_env%n_ov, 1:mv_env%n_ri_loc) => mv_env%B_ia
386 :
387 164 : DO col = first_col, first_col + ncol - 1, nb_max
388 82 : nb = MIN(nb_max, first_col + ncol - col)
389 : ! three views of one contiguous buffer, bounds-remapping pointer assignment (Fortran 2003):
390 : ! Z_ia,k for the exchange DGEMM, Z_3d(a, i, k) per column, Z_wide(a, (i k)) for the wide product
391 82 : Z(1:mv_env%n_ov, 1:nb) => Z_buf(1:mv_env%n_ov*nb)
392 82 : Z_3d(1:mv_env%virt, 1:mv_env%homo, 1:nb) => Z_buf(1:mv_env%n_ov*nb)
393 82 : Z_wide(1:mv_env%virt, 1:mv_env%homo*nb) => Z_buf(1:mv_env%n_ov*nb)
394 82 : RA(1:mv_env%n_ov, 1:nb) => RA_buf(1:mv_env%n_ov*nb)
395 82 : RA_3d(1:mv_env%virt, 1:mv_env%homo, 1:nb) => RA_buf(1:mv_env%n_ov*nb)
396 82 : BabZ_3d(1:mv_env%virt, 1:mv_env%homo, 1:nb) => BabZ_buf(1:mv_env%n_ov*nb)
397 82 : BabZ_wide(1:mv_env%virt, 1:mv_env%homo*nb) => BabZ_buf(1:mv_env%n_ov*nb)
398 82 : IF (do_B) THEN
399 66 : RB(1:mv_env%n_ov, 1:nb) => RB_buf(1:mv_env%n_ov*nb)
400 66 : RB_3d(1:mv_env%virt, 1:mv_env%homo, 1:nb) => RB_buf(1:mv_env%n_ov*nb)
401 : END IF
402 :
403 : ! Z_ia,k on every rank: every rank owns a contiguous ia range, so the column arrives sorted
404 82 : CALL timeset(routineN//"_gather", handle2)
405 330 : DO k = 1, nb
406 : CALL mv_env%para_env%allgatherv(fm_Z%local_data(1:nrow_local, col + k - 1), Z(:, k), &
407 330 : mv_env%row_count, mv_env%row_displ)
408 : END DO
409 82 : CALL timestop(handle2)
410 :
411 : ! exchange, identical in A Z and B Z: R_ia,k = α sum_P B^P_ia t_Pk, t_Pk = sum_jb B^P_jb Z_jb,k
412 82 : CALL timeset(routineN//"_exchange", handle2)
413 82 : RA(:, :) = 0.0_dp
414 82 : IF (mv_env%alpha /= 0.0_dp .AND. mv_env%n_ri_loc > 0) THEN
415 : CALL DGEMM('T', 'N', mv_env%n_ri_loc, nb, mv_env%n_ov, 1.0_dp, B_ia_2d, mv_env%n_ov, Z, mv_env%n_ov, &
416 82 : 0.0_dp, T_Pk, mv_env%n_ri_loc)
417 : CALL DGEMM('N', 'N', mv_env%n_ov, nb, mv_env%n_ri_loc, mv_env%alpha, B_ia_2d, mv_env%n_ov, T_Pk, mv_env%n_ri_loc, &
418 82 : 0.0_dp, RA, mv_env%n_ov)
419 : END IF
420 20336 : IF (do_B) RB(:, :) = RA(:, :)
421 82 : CALL timestop(handle2)
422 :
423 82 : CALL timeset(routineN//"_W", handle2)
424 82 : IF (mv_env%w_fac /= 0.0_dp) THEN
425 3485 : DO p = 1, mv_env%n_ri_loc
426 : ! A: R_ia,k -= w sum_jb \bar{B}^P_ij B^P_ab Z_jb,k, BabZ_ib,k = sum_b B^P_ab Z_jb,k first
427 : CALL DGEMM('T', 'N', mv_env%virt, mv_env%homo*nb, mv_env%virt, 1.0_dp, mv_env%B_ab(:, :, p), mv_env%virt, &
428 3403 : Z_wide, mv_env%virt, 0.0_dp, BabZ_wide, mv_env%virt)
429 13695 : DO k = 1, nb
430 : CALL DGEMM('N', 'N', mv_env%virt, mv_env%homo, mv_env%homo, -mv_env%w_fac, BabZ_3d(:, :, k), mv_env%virt, &
431 13695 : mv_env%B_bar_ij(:, :, p), mv_env%homo, 1.0_dp, RA_3d(:, :, k), mv_env%virt)
432 : END DO
433 : ! B: R_ia,k -= w sum_jb \bar{B}^P_ib B^P_ja Z_jb,k, BjaZ_ab = sum_j B^P_ja Z_jb,k first
434 3485 : IF (do_B) THEN
435 11288 : DO k = 1, nb
436 : CALL DGEMM('N', 'T', mv_env%virt, mv_env%virt, mv_env%homo, 1.0_dp, mv_env%B_ia(:, :, p), mv_env%virt, &
437 8549 : Z_3d(:, :, k), mv_env%virt, 0.0_dp, BjaZ_ab, mv_env%virt)
438 : CALL DGEMM('N', 'N', mv_env%virt, mv_env%homo, mv_env%virt, -mv_env%w_fac, BjaZ_ab, mv_env%virt, &
439 11288 : mv_env%B_bar_ia(:, :, p), mv_env%virt, 1.0_dp, RB_3d(:, :, k), mv_env%virt)
440 : END DO
441 : END IF
442 : END DO
443 : END IF
444 82 : CALL timestop(handle2)
445 :
446 : ! the sum over the ranks lands on the owner of each row, never replicated
447 82 : CALL timeset(routineN//"_reduce", handle2)
448 : ! (A Z)_ia,k = (ε_a-ε_i) Z_ia,k + R^A_ia,k and (B Z)_ia,k = R^B_ia,k on the local rows
449 330 : DO k = 1, nb
450 248 : CALL mv_env%para_env%sum_scatter(RA(:, k:k), R_loc, mv_env%row_count)
451 6200 : DO iloc = 1, nrow_local
452 5952 : ia = row_indices(iloc)
453 6200 : fm_AZ%local_data(iloc, col + k - 1) = R_loc(iloc) + mv_env%eps_diff(ia)*Z(ia, k)
454 : END DO
455 330 : IF (do_B) THEN
456 206 : CALL mv_env%para_env%sum_scatter(RB(:, k:k), R_loc, mv_env%row_count)
457 5150 : fm_BZ%local_data(1:nrow_local, col + col_shift_B + k - 1) = R_loc(1:nrow_local)
458 : END IF
459 : END DO
460 492 : CALL timestop(handle2)
461 : END DO
462 :
463 82 : DEALLOCATE (Z_buf, RA_buf, BabZ_buf, T_Pk)
464 82 : IF (do_B) THEN
465 66 : DEALLOCATE (RB_buf, BjaZ_ab)
466 : END IF
467 :
468 82 : CALL timestop(handle)
469 :
470 246 : END SUBROUTINE bse_matvec_apply
471 :
472 : ! **************************************************************************************************
473 : !> \brief Diagonal used by the Davidson correction, either ε_a-ε_i or the full diagonal
474 : !> A_ia,ia = ε_a-ε_i + α sum_P (B^P_ia)^2 - w sum_P \bar{B}^P_ii B^P_aa
475 : !> with α, w and \bar{B} as in bse_matvec_apply
476 : !> \param mv_env slabs, prefactors and transition energies
477 : !> \param precond_kind bse_precond_full_diag or bse_precond_qp_diff
478 : !> \param diag replicated, size N_ov
479 : ! **************************************************************************************************
480 6 : SUBROUTINE bse_matvec_diagonal(mv_env, precond_kind, diag)
481 :
482 : TYPE(bse_matvec_env_type), INTENT(IN) :: mv_env
483 : INTEGER, INTENT(IN) :: precond_kind
484 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: diag
485 :
486 : CHARACTER(LEN=*), PARAMETER :: routineN = 'bse_matvec_diagonal'
487 :
488 : INTEGER :: a, handle, i, ia, p
489 :
490 6 : CALL timeset(routineN, handle)
491 :
492 294 : diag(:) = 0.0_dp
493 6 : IF (precond_kind == bse_precond_full_diag) THEN
494 255 : DO p = 1, mv_env%n_ri_loc
495 1251 : DO i = 1, mv_env%homo
496 13197 : DO a = 1, mv_env%virt
497 11952 : ia = ia_of_occ_virt(i, a, mv_env%virt)
498 11952 : diag(ia) = diag(ia) + mv_env%alpha*mv_env%B_ia(a, i, p)**2
499 12948 : IF (mv_env%w_fac /= 0.0_dp) THEN
500 11952 : diag(ia) = diag(ia) - mv_env%w_fac*mv_env%B_bar_ij(i, i, p)*mv_env%B_ab(a, a, p)
501 : END IF
502 : END DO
503 : END DO
504 : END DO
505 582 : CALL mv_env%para_env%sum(diag)
506 : END IF
507 294 : diag(:) = diag(:) + mv_env%eps_diff(:)
508 :
509 6 : CALL timestop(handle)
510 :
511 6 : END SUBROUTINE bse_matvec_diagonal
512 :
513 : ! **************************************************************************************************
514 : !> \brief Exact A (and B) on a list of transitions, replicated on every rank,
515 : !> A_kl = δ_kl (ε_a-ε_i) + α sum_P B^P_ia B^P_jb - w sum_P \bar{B}^P_ij B^P_ab
516 : !> B_kl = α sum_P B^P_ia B^P_jb - w sum_P \bar{B}^P_ib B^P_ja
517 : !> with k = (i,a), l = (j,b) and α, w, \bar{B} as in bse_matvec_apply
518 : !> \param mv_env slabs, prefactors and transition energies
519 : !> \param ia_list global transition indices ia = (i-1)*virt + a of the block
520 : !> \param A_sub SIZE(ia_list) x SIZE(ia_list)
521 : !> \param B_sub same, only formed when present
522 : ! **************************************************************************************************
523 0 : SUBROUTINE bse_matvec_subblock(mv_env, ia_list, A_sub, B_sub)
524 :
525 : TYPE(bse_matvec_env_type), INTENT(IN) :: mv_env
526 : INTEGER, DIMENSION(:), INTENT(IN) :: ia_list
527 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT) :: A_sub
528 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT), &
529 : OPTIONAL :: B_sub
530 :
531 : CHARACTER(LEN=*), PARAMETER :: routineN = 'bse_matvec_subblock'
532 :
533 : INTEGER :: handle, k, l, n_sub, p
534 0 : INTEGER, ALLOCATABLE, DIMENSION(:) :: a_virt, i_occ
535 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: B_Pk
536 :
537 0 : CALL timeset(routineN, handle)
538 :
539 0 : n_sub = SIZE(ia_list)
540 0 : ALLOCATE (i_occ(n_sub), a_virt(n_sub), B_Pk(MAX(mv_env%n_ri_loc, 1), n_sub))
541 0 : DO k = 1, n_sub
542 0 : i_occ(k) = occ_of_ia(ia_list(k), mv_env%virt)
543 0 : a_virt(k) = virt_of_ia(ia_list(k), mv_env%virt)
544 : END DO
545 :
546 : ! exchange, identical in A and B: α sum_P B_Pk B_Pl with B_Pk = B^P_ia at ia = ia_list(k)
547 0 : A_sub(:, :) = 0.0_dp
548 0 : IF (mv_env%alpha /= 0.0_dp .AND. mv_env%n_ri_loc > 0) THEN
549 : ! P outermost: the gathered entries of one slice lie within one contiguous (a, i) block
550 0 : DO p = 1, mv_env%n_ri_loc
551 0 : DO k = 1, n_sub
552 0 : B_Pk(p, k) = mv_env%B_ia(a_virt(k), i_occ(k), p)
553 : END DO
554 : END DO
555 : CALL DGEMM('T', 'N', n_sub, n_sub, mv_env%n_ri_loc, mv_env%alpha, B_Pk, SIZE(B_Pk, 1), B_Pk, SIZE(B_Pk, 1), &
556 0 : 0.0_dp, A_sub, n_sub)
557 : END IF
558 0 : IF (PRESENT(B_sub)) B_sub(:, :) = A_sub(:, :)
559 :
560 0 : IF (mv_env%w_fac /= 0.0_dp) THEN
561 0 : DO p = 1, mv_env%n_ri_loc
562 0 : DO l = 1, n_sub
563 0 : DO k = 1, n_sub
564 0 : A_sub(k, l) = A_sub(k, l) - mv_env%w_fac*mv_env%B_bar_ij(i_occ(l), i_occ(k), p)*mv_env%B_ab(a_virt(l), a_virt(k), p)
565 : END DO
566 : END DO
567 0 : IF (PRESENT(B_sub)) THEN
568 0 : DO l = 1, n_sub
569 0 : DO k = 1, n_sub
570 0 : B_sub(k, l) = B_sub(k, l) - mv_env%w_fac*mv_env%B_ia(a_virt(k), i_occ(l), p)*mv_env%B_bar_ia(a_virt(l), i_occ(k), p)
571 : END DO
572 : END DO
573 : END IF
574 : END DO
575 : END IF
576 :
577 0 : CALL mv_env%para_env%sum(A_sub)
578 0 : IF (PRESENT(B_sub)) CALL mv_env%para_env%sum(B_sub)
579 0 : DO k = 1, n_sub
580 0 : A_sub(k, k) = A_sub(k, k) + mv_env%eps_diff(ia_list(k))
581 : END DO
582 :
583 0 : DEALLOCATE (i_occ, a_virt, B_Pk)
584 :
585 0 : CALL timestop(handle)
586 :
587 0 : END SUBROUTINE bse_matvec_subblock
588 :
589 : ! **************************************************************************************************
590 : !> \brief Debug check of the matrix-free application against the explicit matrices A (and B),
591 : !> dev = max_ia,k |(A Z)_ia,k - sum_jb A_ia,jb Z_jb,k| over up to eight unit vectors Z_ia,k = δ_ia,k
592 : !> and one dense vector Z_ia = sin(ia), the same for B, and max_ia |d_ia - A_ia,ia| for the diagonal
593 : !> \param mv_env the environment under test
594 : !> \param fm_A_explicit A from create_A_and_B, N_ov x N_ov on the grid of the caller
595 : !> \param unit_nr output unit, positive on the writing rank only
596 : !> \param fm_B_explicit B, present for an ABBA run
597 : ! **************************************************************************************************
598 0 : SUBROUTINE bse_matvec_selfcheck(mv_env, fm_A_explicit, unit_nr, fm_B_explicit)
599 :
600 : TYPE(bse_matvec_env_type), INTENT(IN) :: mv_env
601 : TYPE(cp_fm_type), INTENT(IN) :: fm_A_explicit
602 : INTEGER, INTENT(IN) :: unit_nr
603 : TYPE(cp_fm_type), INTENT(IN), OPTIONAL :: fm_B_explicit
604 :
605 : CHARACTER(LEN=*), PARAMETER :: routineN = 'bse_matvec_selfcheck'
606 :
607 : INTEGER :: handle, ia, iloc, k, n_ov, n_unit, ncol, &
608 : nrow_local
609 0 : INTEGER, DIMENSION(:), POINTER :: row_indices
610 : REAL(KIND=dp) :: dev_A, dev_B, dev_diag
611 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: diag
612 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: mv_result, ref_matrix, ref_result, Z
613 : TYPE(cp_fm_struct_type), POINTER :: fm_struct
614 : TYPE(cp_fm_type) :: fm_AZ, fm_BZ, fm_Z
615 :
616 0 : CALL timeset(routineN, handle)
617 :
618 0 : n_ov = mv_env%n_ov
619 : ! eight unit vectors read eight columns of the explicit matrix exactly; the dense column covers the whole contraction
620 0 : n_unit = MIN(n_ov, 8)
621 0 : dev_B = 0.0_dp
622 0 : ncol = n_unit + 1
623 :
624 : ! trial vectors: e_k for k <= n_unit, then one dense column Z_ia = sin(ia)
625 0 : ALLOCATE (Z(n_ov, ncol))
626 0 : Z(:, :) = 0.0_dp
627 0 : DO k = 1, n_unit
628 0 : Z(k, k) = 1.0_dp
629 : END DO
630 0 : DO ia = 1, n_ov
631 0 : Z(ia, ncol) = SIN(REAL(ia, dp))
632 : END DO
633 :
634 0 : CALL bse_matvec_vector_struct(mv_env, ncol, fm_struct)
635 0 : CALL cp_fm_create(fm_Z, fm_struct, name="fm_Z_selfcheck")
636 0 : CALL cp_fm_create(fm_AZ, fm_struct, name="fm_AZ_selfcheck")
637 0 : CALL cp_fm_set_all(fm_AZ, 0.0_dp)
638 0 : IF (PRESENT(fm_B_explicit)) THEN
639 0 : CALL cp_fm_create(fm_BZ, fm_struct, name="fm_BZ_selfcheck")
640 0 : CALL cp_fm_set_all(fm_BZ, 0.0_dp)
641 : END IF
642 0 : CALL cp_fm_struct_release(fm_struct)
643 :
644 0 : CALL cp_fm_get_info(fm_Z, nrow_local=nrow_local, row_indices=row_indices)
645 0 : DO k = 1, ncol
646 0 : DO iloc = 1, nrow_local
647 0 : fm_Z%local_data(iloc, k) = Z(row_indices(iloc), k)
648 : END DO
649 : END DO
650 :
651 0 : IF (PRESENT(fm_B_explicit)) THEN
652 0 : CALL bse_matvec_apply(mv_env, fm_Z, 1, ncol, fm_AZ, fm_BZ)
653 : ELSE
654 0 : CALL bse_matvec_apply(mv_env, fm_Z, 1, ncol, fm_AZ)
655 : END IF
656 :
657 0 : ALLOCATE (ref_matrix(n_ov, n_ov), mv_result(n_ov, ncol), ref_result(n_ov, ncol), diag(n_ov))
658 :
659 0 : CALL cp_fm_get_submatrix(fm_A_explicit, ref_matrix)
660 0 : CALL cp_fm_get_submatrix(fm_AZ, mv_result)
661 0 : ref_result(:, :) = MATMUL(ref_matrix, Z)
662 0 : dev_A = MAXVAL(ABS(mv_result - ref_result))
663 0 : CALL bse_matvec_diagonal(mv_env, bse_precond_full_diag, diag)
664 0 : dev_diag = 0.0_dp
665 0 : DO ia = 1, n_ov
666 0 : dev_diag = MAX(dev_diag, ABS(diag(ia) - ref_matrix(ia, ia)))
667 : END DO
668 :
669 0 : IF (PRESENT(fm_B_explicit)) THEN
670 0 : CALL cp_fm_get_submatrix(fm_B_explicit, ref_matrix)
671 0 : CALL cp_fm_get_submatrix(fm_BZ, mv_result)
672 0 : ref_result(:, :) = MATMUL(ref_matrix, Z)
673 0 : dev_B = MAXVAL(ABS(mv_result - ref_result))
674 : END IF
675 :
676 0 : IF (unit_nr > 0) THEN
677 0 : WRITE (unit_nr, '(T2,A10,T13,A,T59,ES22.6)') 'BSE|DEBUG|', &
678 0 : 'Max deviation matvec vs explicit A', dev_A
679 0 : IF (PRESENT(fm_B_explicit)) THEN
680 0 : WRITE (unit_nr, '(T2,A10,T13,A,T59,ES22.6)') 'BSE|DEBUG|', &
681 0 : 'Max deviation matvec vs explicit B', dev_B
682 : END IF
683 0 : WRITE (unit_nr, '(T2,A10,T13,A,T59,ES22.6)') 'BSE|DEBUG|', &
684 0 : 'Max deviation diagonal vs explicit A', dev_diag
685 : END IF
686 :
687 0 : DEALLOCATE (Z, ref_matrix, mv_result, ref_result, diag)
688 0 : CALL cp_fm_release(fm_Z)
689 0 : CALL cp_fm_release(fm_AZ)
690 0 : IF (PRESENT(fm_B_explicit)) CALL cp_fm_release(fm_BZ)
691 :
692 0 : CALL timestop(handle)
693 :
694 0 : END SUBROUTINE bse_matvec_selfcheck
695 :
696 0 : END MODULE bse_matvec
|
