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 Shared numerical operations for computing the RI-RS matrix Z_lP.
10 : !> \par History
11 : !> 09.2026 created Jan Wilhelm
12 : ! **************************************************************************************************
13 : MODULE gw_ri_rs_compute_Z_lP_utils
14 : USE cp_blacs_env, ONLY: cp_blacs_env_type
15 : USE cp_dbcsr_api, ONLY: dbcsr_put_block,&
16 : dbcsr_type
17 : USE cp_fm_cholesky, ONLY: cp_fm_cholesky_decompose,&
18 : cp_fm_cholesky_solve
19 : USE cp_fm_struct, ONLY: cp_fm_struct_create,&
20 : cp_fm_struct_release,&
21 : cp_fm_struct_type
22 : USE cp_fm_types, ONLY: cp_fm_create,&
23 : cp_fm_get_info,&
24 : cp_fm_get_submatrix,&
25 : cp_fm_release,&
26 : cp_fm_set_submatrix,&
27 : cp_fm_type
28 : USE kinds, ONLY: dp
29 : USE message_passing, ONLY: mp_para_env_type
30 : #include "./base/base_uses.f90"
31 :
32 : IMPLICIT NONE
33 : PRIVATE
34 :
35 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'gw_ri_rs_compute_Z_lP_utils'
36 : REAL(KIND=dp), PARAMETER, PRIVATE :: jacobi_floor = 1.0E-16_dp
37 :
38 : PUBLIC :: build_gram_jacobi_blas, &
39 : build_jacobi_diag_from_phi, &
40 : scale_rows_by_diag, &
41 : solve_D_lp_distributed, &
42 : store_Z_lP_columns
43 :
44 : CONTAINS
45 :
46 : ! **************************************************************************************************
47 : !> \brief Forms the conditioned dense RI-RS matrix
48 : !>
49 : !> D_ll' = [Σ_μ ϕ_μ(r_l) ϕ_μ(r_l')]²,
50 : !> d_l = 1/sqrt(D_ll),
51 : !> D'_ll' = d_l D_ll' d_l' + λ δ_ll'.
52 : !>
53 : !> Only the dense single-rank solve needs the complete matrix D'.
54 : !> \param phi_local ...
55 : !> \param n_local_grid ...
56 : !> \param n_ao_used ...
57 : !> \param tikhonov ...
58 : !> \param D_local ...
59 : !> \param d_vec_local ...
60 : ! **************************************************************************************************
61 54 : SUBROUTINE build_gram_jacobi_blas(phi_local, n_local_grid, n_ao_used, tikhonov, D_local, &
62 54 : d_vec_local)
63 :
64 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi_local
65 : INTEGER, INTENT(IN) :: n_local_grid, n_ao_used
66 : REAL(KIND=dp), INTENT(IN) :: tikhonov
67 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
68 : INTENT(OUT) :: D_local
69 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: d_vec_local
70 :
71 : CHARACTER(LEN=*), PARAMETER :: routineN = 'build_gram_jacobi_blas'
72 :
73 : INTEGER :: handle, handle_dsyrk, i, j
74 :
75 54 : CALL timeset(routineN, handle)
76 :
77 216 : ALLOCATE (D_local(n_local_grid, n_local_grid))
78 54 : D_local = 0.0_dp
79 :
80 54 : CALL timeset(routineN//'_dsyrk', handle_dsyrk)
81 : CALL dsyrk("L", "N", n_local_grid, n_ao_used, 1.0_dp, phi_local, &
82 54 : n_local_grid, 0.0_dp, D_local, n_local_grid)
83 54 : CALL timestop(handle_dsyrk)
84 :
85 : !$OMP PARALLEL DO DEFAULT(NONE) &
86 : !$OMP SHARED(n_local_grid, D_local, d_vec_local, tikhonov) &
87 : !$OMP PRIVATE(i) &
88 54 : !$OMP SCHEDULE(STATIC)
89 : DO i = 1, n_local_grid
90 : D_local(i, i) = D_local(i, i)**2
91 : d_vec_local(i) = 1.0_dp/SQRT(MAX(D_local(i, i), jacobi_floor))
92 : D_local(i, i) = (D_local(i, i)*d_vec_local(i)**2) + tikhonov
93 : END DO
94 : !$OMP END PARALLEL DO
95 :
96 : !$OMP PARALLEL DO DEFAULT(NONE) &
97 : !$OMP SHARED(n_local_grid, D_local, d_vec_local) &
98 : !$OMP PRIVATE(j, i) &
99 54 : !$OMP SCHEDULE(DYNAMIC)
100 : DO j = 1, n_local_grid
101 : DO i = j + 1, n_local_grid
102 : D_local(i, j) = D_local(i, j)**2
103 : D_local(i, j) = D_local(i, j)*d_vec_local(i)*d_vec_local(j)
104 : D_local(j, i) = D_local(i, j)
105 : END DO
106 : END DO
107 : !$OMP END PARALLEL DO
108 :
109 54 : CALL timestop(handle)
110 :
111 108 : END SUBROUTINE build_gram_jacobi_blas
112 :
113 : ! **************************************************************************************************
114 : !> \brief Computes d_l = 1/sqrt(D_ll) = 1/Σ_μ ϕ_μ(r_l)² without forming D.
115 : !> \param phi_local ...
116 : !> \param n_local_grid ...
117 : !> \param n_ao_used ...
118 : !> \param d_vec_local ...
119 : ! **************************************************************************************************
120 0 : SUBROUTINE build_jacobi_diag_from_phi(phi_local, n_local_grid, n_ao_used, d_vec_local)
121 :
122 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi_local
123 : INTEGER, INTENT(IN) :: n_local_grid, n_ao_used
124 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: d_vec_local
125 :
126 : CHARACTER(LEN=*), PARAMETER :: routineN = 'build_jacobi_diag_from_phi'
127 :
128 : INTEGER :: handle, i, j
129 :
130 0 : CALL timeset(routineN, handle)
131 :
132 : !$OMP PARALLEL DO DEFAULT(NONE) &
133 : !$OMP SHARED(n_local_grid, n_ao_used, phi_local, d_vec_local) &
134 : !$OMP PRIVATE(i, j) &
135 0 : !$OMP SCHEDULE(STATIC)
136 : DO i = 1, n_local_grid
137 : d_vec_local(i) = 0.0_dp
138 : DO j = 1, n_ao_used
139 : d_vec_local(i) = d_vec_local(i) + phi_local(i, j)*phi_local(i, j)
140 : END DO
141 : d_vec_local(i) = 1.0_dp/MAX(d_vec_local(i), jacobi_floor)
142 : END DO
143 : !$OMP END PARALLEL DO
144 :
145 0 : CALL timestop(handle)
146 :
147 0 : END SUBROUTINE build_jacobi_diag_from_phi
148 :
149 : ! **************************************************************************************************
150 : !> \brief Multiplies every matrix row by the corresponding diagonal entry:
151 : !> A(l, :) <- d_l A(l, :).
152 : !> \param matrix ...
153 : !> \param diagonal ...
154 : !> \param nrow ...
155 : !> \param ncol ...
156 : ! **************************************************************************************************
157 108 : SUBROUTINE scale_rows_by_diag(matrix, diagonal, nrow, ncol)
158 :
159 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: matrix
160 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: diagonal
161 : INTEGER, INTENT(IN) :: nrow, ncol
162 :
163 : CHARACTER(LEN=*), PARAMETER :: routineN = 'scale_rows_by_diag'
164 :
165 : INTEGER :: handle, i, j
166 :
167 108 : CALL timeset(routineN, handle)
168 :
169 : !$OMP PARALLEL DO DEFAULT(NONE) &
170 : !$OMP SHARED(ncol, nrow, matrix, diagonal) &
171 : !$OMP PRIVATE(j, i) &
172 108 : !$OMP SCHEDULE(STATIC)
173 : DO j = 1, ncol
174 : DO i = 1, nrow
175 : matrix(i, j) = matrix(i, j)*diagonal(i)
176 : END DO
177 : END DO
178 : !$OMP END PARALLEL DO
179 :
180 108 : CALL timestop(handle)
181 :
182 108 : END SUBROUTINE scale_rows_by_diag
183 :
184 : ! **************************************************************************************************
185 : !> \brief Stores dense Z_lP columns in the distributed block-sparse matrix.
186 : !> \param mat_Z_lP ...
187 : !> \param Z_local ...
188 : !> \param local_grid_idx ...
189 : !> \param n_local_grid ...
190 : !> \param n_loc_ri ...
191 : !> \param atom_P ...
192 : !> \param r_blk_sizes ...
193 : !> \param row_offset ...
194 : !> \param eps_filter ...
195 : ! **************************************************************************************************
196 50 : SUBROUTINE store_Z_lP_columns(mat_Z_lP, Z_local, local_grid_idx, n_local_grid, n_loc_ri, &
197 50 : atom_P, r_blk_sizes, row_offset, eps_filter)
198 :
199 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_Z_lP
200 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: Z_local
201 : INTEGER, DIMENSION(:), INTENT(IN) :: local_grid_idx
202 : INTEGER, INTENT(IN) :: n_local_grid, n_loc_ri, atom_P
203 : INTEGER, DIMENSION(:), INTENT(IN) :: r_blk_sizes, row_offset
204 : REAL(KIND=dp), INTENT(IN) :: eps_filter
205 :
206 : CHARACTER(LEN=*), PARAMETER :: routineN = 'store_Z_lP_columns'
207 :
208 : INTEGER :: current_chunk_size, g_pt, handle, i_blk, &
209 : loc_ptr, r_end, r_start
210 50 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: Z_blk
211 :
212 50 : CALL timeset(routineN, handle)
213 :
214 882 : ALLOCATE (Z_blk(MAXVAL(r_blk_sizes), n_loc_ri))
215 50 : loc_ptr = 1
216 :
217 732 : DO i_blk = 1, SIZE(r_blk_sizes)
218 682 : r_start = row_offset(i_blk) + 1
219 682 : r_end = row_offset(i_blk) + r_blk_sizes(i_blk)
220 682 : current_chunk_size = r_blk_sizes(i_blk)
221 682 : Z_blk = 0.0_dp
222 :
223 18845 : DO WHILE (loc_ptr <= n_local_grid)
224 18791 : g_pt = local_grid_idx(loc_ptr)
225 18791 : IF (g_pt > r_end) EXIT
226 255966 : Z_blk(g_pt - r_start + 1, 1:n_loc_ri) = Z_local(loc_ptr, 1:n_loc_ri)
227 18791 : loc_ptr = loc_ptr + 1
228 : END DO
229 :
230 252928 : IF (MAXVAL(ABS(Z_blk(1:current_chunk_size, 1:n_loc_ri))) > eps_filter) THEN
231 : CALL dbcsr_put_block(mat_Z_lP, row=i_blk, col=atom_P, &
232 670 : block=Z_blk(1:current_chunk_size, 1:n_loc_ri))
233 : END IF
234 : END DO
235 :
236 50 : DEALLOCATE (Z_blk)
237 :
238 50 : CALL timestop(handle)
239 :
240 50 : END SUBROUTINE store_Z_lP_columns
241 :
242 : ! **************************************************************************************************
243 : !> \brief Solves D Z = d with ScaLAPACK for one RI atom and replicated inputs.
244 : !>
245 : !> Each rank builds its block-cyclic part of
246 : !>
247 : !> D'_ll' = d_l [Σ_μ ϕ_μ(r_l) ϕ_μ(r_l')]² d_l' + λ δ_ll'.
248 : !>
249 : !> The Cholesky solution is gathered into the replicated right-hand side.
250 : !> \param phi_local ...
251 : !> \param d_vec ...
252 : !> \param d_lp ...
253 : !> \param n_loc ...
254 : !> \param n_ao ...
255 : !> \param n_rhs ...
256 : !> \param tikhonov ...
257 : !> \param para_env_sub ...
258 : !> \param blacs_env_sub ...
259 : !> \param fm_struct_D ...
260 : !> \param fm_struct_b ...
261 : !> \param fm_D ...
262 : !> \param fm_b ...
263 : !> \param info ...
264 : ! **************************************************************************************************
265 0 : SUBROUTINE solve_D_lp_distributed(phi_local, d_vec, d_lp, n_loc, n_ao, n_rhs, &
266 : tikhonov, para_env_sub, blacs_env_sub, &
267 : fm_struct_D, fm_struct_b, fm_D, fm_b, info)
268 :
269 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi_local
270 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: d_vec
271 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: d_lp
272 : INTEGER, INTENT(IN) :: n_loc, n_ao, n_rhs
273 : REAL(KIND=dp), INTENT(IN) :: tikhonov
274 : TYPE(mp_para_env_type), POINTER :: para_env_sub
275 : TYPE(cp_blacs_env_type), POINTER :: blacs_env_sub
276 : TYPE(cp_fm_struct_type), POINTER :: fm_struct_D, fm_struct_b
277 : TYPE(cp_fm_type), INTENT(INOUT) :: fm_D, fm_b
278 : INTEGER, INTENT(OUT) :: info
279 :
280 : CHARACTER(LEN=*), PARAMETER :: routineN = 'solve_D_lp_distributed'
281 :
282 : INTEGER :: handle, i_loc, ig, j_loc, jg, &
283 : ncol_local, nrow_local
284 0 : INTEGER, DIMENSION(:), POINTER :: col_indices, row_indices
285 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :), &
286 0 : POINTER :: local_data
287 :
288 0 : CALL timeset(routineN, handle)
289 0 : info = 0
290 :
291 0 : NULLIFY (fm_struct_D, fm_struct_b)
292 : CALL cp_fm_struct_create(fm_struct_D, para_env=para_env_sub, &
293 : context=blacs_env_sub, &
294 0 : nrow_global=n_loc, ncol_global=n_loc)
295 : CALL cp_fm_struct_create(fm_struct_b, para_env=para_env_sub, &
296 : context=blacs_env_sub, &
297 0 : nrow_global=n_loc, ncol_global=n_rhs)
298 0 : CALL cp_fm_create(fm_D, fm_struct_D)
299 0 : CALL cp_fm_create(fm_b, fm_struct_b)
300 :
301 : CALL cp_fm_get_info(fm_D, nrow_local=nrow_local, ncol_local=ncol_local, &
302 : row_indices=row_indices, col_indices=col_indices, &
303 0 : local_data=local_data)
304 :
305 0 : IF (nrow_local > 0 .AND. ncol_local > 0) THEN
306 0 : BLOCK
307 : INTEGER, PARAMETER :: ntile = 1024
308 : INTEGER :: ib, ie, jb, je, mb, kb, ti, tj, handle_dgemm
309 0 : REAL(KIND=dp), ALLOCATABLE :: gram_t(:, :), phi_cols_t(:, :), phi_rows_t(:, :)
310 0 : ALLOCATE (phi_rows_t(ntile, n_ao), phi_cols_t(n_ao, ntile), gram_t(ntile, ntile))
311 0 : DO ib = 1, nrow_local, ntile
312 0 : ie = MIN(ib + ntile - 1, nrow_local)
313 0 : mb = ie - ib + 1
314 : !$OMP PARALLEL DO DEFAULT(NONE) &
315 : !$OMP SHARED(mb, n_ao, phi_rows_t, phi_local, row_indices, ib) &
316 0 : !$OMP PRIVATE(ti, j_loc) SCHEDULE(STATIC)
317 : DO j_loc = 1, n_ao
318 : DO ti = 1, mb
319 : phi_rows_t(ti, j_loc) = phi_local(row_indices(ib + ti - 1), j_loc)
320 : END DO
321 : END DO
322 : !$OMP END PARALLEL DO
323 0 : DO jb = 1, ncol_local, ntile
324 0 : je = MIN(jb + ntile - 1, ncol_local)
325 0 : kb = je - jb + 1
326 : !$OMP PARALLEL DO DEFAULT(NONE) &
327 : !$OMP SHARED(kb, n_ao, phi_cols_t, phi_local, col_indices, jb) &
328 0 : !$OMP PRIVATE(tj, i_loc) SCHEDULE(STATIC)
329 : DO tj = 1, kb
330 : DO i_loc = 1, n_ao
331 : phi_cols_t(i_loc, tj) = phi_local(col_indices(jb + tj - 1), i_loc)
332 : END DO
333 : END DO
334 : !$OMP END PARALLEL DO
335 0 : CALL timeset(routineN//'_dgemm', handle_dgemm)
336 : CALL dgemm('N', 'N', mb, kb, n_ao, &
337 : 1.0_dp, phi_rows_t, ntile, phi_cols_t, n_ao, &
338 0 : 0.0_dp, gram_t, ntile)
339 0 : CALL timestop(handle_dgemm)
340 : !$OMP PARALLEL DO DEFAULT(NONE) &
341 : !$OMP SHARED(mb, kb, gram_t, d_vec, row_indices, col_indices, ib, jb) &
342 : !$OMP SHARED(local_data, tikhonov) &
343 0 : !$OMP PRIVATE(ti, tj, ig, jg) SCHEDULE(STATIC)
344 : DO tj = 1, kb
345 : jg = col_indices(jb + tj - 1)
346 : DO ti = 1, mb
347 : ig = row_indices(ib + ti - 1)
348 : local_data(ib + ti - 1, jb + tj - 1) = &
349 : gram_t(ti, tj)*gram_t(ti, tj)*d_vec(ig)*d_vec(jg)
350 : IF (ig == jg) THEN
351 : local_data(ib + ti - 1, jb + tj - 1) = &
352 : local_data(ib + ti - 1, jb + tj - 1) + tikhonov
353 : END IF
354 : END DO
355 : END DO
356 : !$OMP END PARALLEL DO
357 : END DO
358 : END DO
359 0 : DEALLOCATE (phi_rows_t, phi_cols_t, gram_t)
360 : END BLOCK
361 : END IF
362 :
363 0 : CALL cp_fm_set_submatrix(fm_b, d_lp)
364 :
365 0 : CALL cp_fm_cholesky_decompose(fm_D, n=n_loc, info_out=info)
366 0 : IF (info /= 0) CPABORT("pdpotrf failed in solve_D_lp_distributed")
367 :
368 0 : CALL cp_fm_cholesky_solve(fm_D, fm_b, n=n_loc, info_out=info)
369 0 : IF (info /= 0) CPABORT("pdpotrs failed in solve_D_lp_distributed")
370 :
371 0 : CALL cp_fm_get_submatrix(fm_b, d_lp)
372 :
373 0 : CALL cp_fm_release(fm_D)
374 0 : CALL cp_fm_release(fm_b)
375 0 : CALL cp_fm_struct_release(fm_struct_D)
376 0 : CALL cp_fm_struct_release(fm_struct_b)
377 :
378 0 : CALL timestop(handle)
379 :
380 0 : END SUBROUTINE solve_D_lp_distributed
381 :
382 : END MODULE gw_ri_rs_compute_Z_lP_utils
|
