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 Assembly of real and complex k-point operators from real-space DBCSR matrices.
10 : !> Outputs use the distribution of the k-point group; real operators transfer one part only.
11 : !> The module splits this work into a start step and a finish step.
12 : !> One process can start many operators before it finishes any of them.
13 : !> The STANDARD k-point driver uses this order.
14 : !> The routines start, finish and get are collective over the
15 : !> environment-wide communicator. Every rank must call the same
16 : !> sequence of service calls.
17 : !> A context references its host kpoint environment without owning it,
18 : !> so it must not outlive that environment.
19 : !> An explicitly prepared, immutable operator batch can instead be built
20 : !> from bounded context-owned real-space snapshots inside each group.
21 : !> The public round and output-ownership contracts stay the same.
22 : !> Start/finish/get/prepare are not thread safe: serialize them, because their collective
23 : !> communication tolerates no concurrent service call. A separate
24 : !> context per thread only keeps the scratch from aliasing.
25 : ! **************************************************************************************************
26 : MODULE qs_kpoint_operators
27 : USE cp_cfm_basic_linalg, ONLY: cp_cfm_scale_and_add_fm
28 : USE cp_cfm_types, ONLY: cp_cfm_type
29 : USE cp_dbcsr_api, ONLY: &
30 : dbcsr_create, dbcsr_deallocate_matrix, dbcsr_get_info, dbcsr_get_readonly_block_p, &
31 : dbcsr_init_p, dbcsr_p_type, dbcsr_set, dbcsr_type, dbcsr_type_antisymmetric, &
32 : dbcsr_type_symmetric
33 : USE cp_dbcsr_cp2k_link, ONLY: cp_dbcsr_alloc_block_from_nbl
34 : USE cp_dbcsr_operations, ONLY: copy_dbcsr_to_fm,&
35 : copy_fm_to_dbcsr_bc,&
36 : dbcsr_fm_copy_plan_type
37 : USE cp_fm_pool_types, ONLY: cp_fm_pool_type,&
38 : fm_pool_create_fm,&
39 : fm_pool_give_back_fm
40 : USE cp_fm_struct, ONLY: cp_fm_struct_equivalent,&
41 : cp_fm_struct_release,&
42 : cp_fm_struct_retain,&
43 : cp_fm_struct_type
44 : USE cp_fm_types, ONLY: copy_info_type,&
45 : cp_fm_cleanup_copy_general,&
46 : cp_fm_create,&
47 : cp_fm_finish_copy_general,&
48 : cp_fm_get_info,&
49 : cp_fm_release,&
50 : cp_fm_start_copy_general,&
51 : cp_fm_type
52 : USE input_cp2k_kpoints, ONLY: lattice_fft_auto,&
53 : lattice_fft_off,&
54 : lattice_fft_on
55 : USE kinds, ONLY: dp,&
56 : int_8
57 : USE kpoint_lattice_fft, ONLY: cell_to_k_grid_fft,&
58 : lattice_fft_num_threads,&
59 : lattice_fft_shape,&
60 : regular_kpoint_grid
61 : USE kpoint_methods, ONLY: kp_transform_plan_create,&
62 : kp_transform_plan_type,&
63 : rskp_grid_type,&
64 : rskp_transform,&
65 : rskp_transform_grid_extract,&
66 : rskp_transform_grid_prepare,&
67 : rskp_transform_grid_release
68 : USE kpoint_types, ONLY: kpoint_env_type,&
69 : kpoint_type
70 : USE mathconstants, ONLY: gaussi,&
71 : twopi,&
72 : z_one,&
73 : z_zero
74 : USE message_passing, ONLY: mp_para_env_type
75 : USE qs_matrix_pools, ONLY: mpools_get
76 : USE qs_neighbor_list_types, ONLY: neighbor_list_set_p_type
77 :
78 : !$ USE omp_lib, ONLY: omp_get_max_threads, omp_get_num_threads, omp_get_thread_num
79 : #include "./base/base_uses.f90"
80 :
81 : IMPLICIT NONE
82 :
83 : PRIVATE
84 :
85 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'qs_kpoint_operators'
86 :
87 : PUBLIC :: kpoint_operator_context_create, kpoint_operator_context_release
88 : PUBLIC :: kpoint_operator_start, kpoint_operator_finish
89 : PUBLIC :: kpoint_operator_get, kpoint_operator_prepare
90 : PUBLIC :: kpoint_operator_evaluate, kpoint_operator_work_size
91 : PUBLIC :: kpoint_operator_batch_size, kpoint_operator_evaluate_batch
92 : PUBLIC :: kpoint_operator_store, kpoint_operator_get_local, kpoint_operator_get_all
93 :
94 : INTEGER, PARAMETER, PUBLIC :: kpoint_slot_ks = 1, &
95 : kpoint_slot_s = 2, &
96 : kpoint_slot_t = 3
97 :
98 : !> Row of rsmat that carries the spin-free operator image. S and T carry
99 : !> no spin copies: drivers launch one instance per spin of their loop and
100 : !> read this same row in every instance.
101 : INTEGER, PARAMETER, PUBLIC :: kpoint_spin_free = 1
102 :
103 : !> Number of defined slots; grows only inside this module.
104 : INTEGER, PARAMETER, PRIVATE :: kpoint_num_slots = kpoint_slot_t - kpoint_slot_ks + 1
105 :
106 : !> Part indices of one transfer: its real and its imaginary half.
107 : INTEGER, PARAMETER, PRIVATE :: kpoint_part_re = 1, &
108 : kpoint_part_im = 2
109 :
110 : !> Instance states of one (k point, spin, slot) transfer.
111 : !> State transitions happen only inside the service routines.
112 : INTEGER, PARAMETER, PRIVATE :: kpoint_slot_idle = 0, &
113 : kpoint_slot_in_flight = 1, &
114 : kpoint_slot_local = 2
115 :
116 : ! Structural route shared by all captured rows. Whole blocks go to one
117 : ! representative group, then only equivalent tiles are replicated.
118 : TYPE kpoint_operator_rs_route_type
119 : TYPE(kp_transform_plan_type) :: traversal = kp_transform_plan_type()
120 : TYPE(cp_fm_struct_type), POINTER :: matrix_struct => NULL()
121 : INTEGER :: nimage = 0, nvalue = 0, nrow = 0, ncol = 0
122 : INTEGER(KIND=int_8) :: storage = 0
123 : INTEGER, ALLOCATABLE :: row_size(:), col_size(:), ENTRY(:, :), first(:), tile(:, :, :)
124 : INTEGER, ALLOCATABLE :: send_entry(:), send_first(:), scount(:, :), rcount(:, :), &
125 : sdisp(:, :), rdisp(:, :)
126 : REAL(KIND=dp), ALLOCATABLE :: send_values(:)
127 : ! Geometry-only boundaries of contiguous AO-block image groups.
128 : INTEGER, ALLOCATABLE :: block_first(:)
129 : END TYPE kpoint_operator_rs_route_type
130 :
131 : ! Immutable values, not per-k matrices. Metadata belong to the shared route;
132 : ! source identifies starts belonging to the explicitly prepared batch only.
133 : TYPE kpoint_operator_block_type
134 : REAL(KIND=dp), POINTER :: data(:, :) => NULL()
135 : END TYPE kpoint_operator_block_type
136 :
137 : TYPE kpoint_operator_rs_type
138 : INTEGER :: matrix_row = 0
139 : TYPE(dbcsr_p_type), POINTER :: source(:, :) => NULL()
140 : ! Values are owned only for distributed sources. Blocks always borrow:
141 : ! either from values below or, on one world rank, directly from DBCSR.
142 : REAL(KIND=dp), CONTIGUOUS, POINTER :: values(:) => NULL()
143 : TYPE(kpoint_operator_block_type), ALLOCATABLE :: blocks(:)
144 : LOGICAL, ALLOCATABLE :: nonzero(:)
145 : END TYPE kpoint_operator_rs_type
146 :
147 : ! Read-ahead window or complete FFT output per prepared row. Keep output separate
148 : ! from the immutable snapshot; prepare invalidates only the selected slot.
149 : TYPE kpoint_operator_rs_batch_type
150 : INTEGER :: first = 0
151 : COMPLEX(KIND=dp), ALLOCATABLE :: data(:, :, :)
152 : END TYPE kpoint_operator_rs_batch_type
153 :
154 : INTEGER, PARAMETER :: rs_batch_size = 4
155 : INTEGER(KIND=int_8), PARAMETER :: rs_snapshot_bytes = 64_int_8*1024_int_8**2
156 : ! Includes cached tiles and batch phase scratch, divided between all rows/slots.
157 : INTEGER(KIND=int_8), PARAMETER :: rs_batch_bytes = 32_int_8*1024_int_8**2
158 : ! Shared by all prepared spin rows and slots, including transient FFT work.
159 : INTEGER(KIND=int_8), PARAMETER :: rs_fft_bytes = 512_int_8*1024_int_8**2
160 :
161 : ! **************************************************************************************************
162 : !> \brief Caller-owned assembly context for a fixed geometry and parallel layout.
163 : !> OT retains it for one inner SCF loop; other drivers use one call.
164 : !> Release before changing the neighbor lists, matrix pools or k-point groups.
165 : !> Components are private and touched only by this module.
166 : ! **************************************************************************************************
167 : TYPE, PUBLIC :: kpoint_operator_context_type
168 : PRIVATE
169 : ! Instance spin capacity of the owning driver: the bound of the spin
170 : ! loop that start and finish key. Declared once at creation. It never
171 : ! comes from the operators: the first matrix dimension of a spin-free
172 : ! operator such as S counts derivatives, and the images live in the
173 : ! second, so callers address operator rows through matrix_row instead.
174 : INTEGER :: nspin = 0
175 : LOGICAL :: overlap_stored = .FALSE.
176 : ! Borrowed host/distribution references; the context cannot span an energy evaluation.
177 : TYPE(kpoint_type), POINTER :: kpoints => NULL()
178 : TYPE(cp_fm_struct_type), POINTER :: ao_ao_fmstruct => NULL()
179 : TYPE(cp_fm_pool_type), POINTER :: group_fm_pool => NULL()
180 : ! Identity of the neighbor lists used by scratch and prepared snapshots.
181 : TYPE(neighbor_list_set_p_type), DIMENSION(:), &
182 : POINTER :: work_sab => NULL()
183 : ! Owned scratch, created on first use and released with the context.
184 : TYPE(dbcsr_type), POINTER :: rmatrix => NULL()
185 : TYPE(dbcsr_type), POINTER :: cmatrix => NULL()
186 : TYPE(cp_fm_type), ALLOCATABLE, DIMENSION(:) :: fmwork
187 : TYPE(cp_fm_type), ALLOCATABLE :: fmlocal, fmimag
188 : ! Plans share the lifetime and fixed layouts of the two densify workspaces.
189 : TYPE(dbcsr_fm_copy_plan_type), ALLOCATABLE :: copy_plan(:)
190 : ! Immutable snapshots and bounded local output batches. Prepared starts
191 : ! need neither per-k transfers nor outstanding MPI requests.
192 : TYPE(kpoint_operator_rs_type), ALLOCATABLE :: rs(:, :)
193 : TYPE(kpoint_operator_rs_route_type), ALLOCATABLE :: rs_route
194 : TYPE(kpoint_operator_rs_batch_type), ALLOCATABLE :: rs_batch(:, :)
195 : ! Bookkeeping. info is keyed (global k point, spin, slot, part),
196 : ! part kpoint_part_re holding the real half and kpoint_part_im holding
197 : ! the imaginary half. Status is keyed (global round, spin, slot)
198 : ! because start and finish treat all groups of one instance
199 : ! atomically. Allocated on first start.
200 : TYPE(copy_info_type), ALLOCATABLE, DIMENSION(:, :, :, :) :: info
201 : INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: status
202 : END TYPE kpoint_operator_context_type
203 :
204 : CONTAINS
205 :
206 : ! **************************************************************************************************
207 : !> \brief Transform and densify one operator using the context-owned workspaces and routing plans.
208 : !> \param ctx initialized transfer context
209 : !> \param rsmat real-space operator images
210 : !> \param matrix_row operator row, independent of the instance spin
211 : !> \param ik global k-point index
212 : !> \param grid caller-owned reciprocal-grid cache
213 : !> \param use_grid extract from grid instead of a direct phase sum
214 : ! **************************************************************************************************
215 512 : SUBROUTINE kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid, use_grid)
216 :
217 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
218 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
219 : INTEGER, INTENT(IN) :: matrix_row, ik
220 : TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
221 : LOGICAL, INTENT(IN) :: use_grid
222 :
223 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_densify'
224 :
225 : INTEGER :: handle
226 512 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: xkp
227 :
228 512 : CALL timeset(routineN, handle)
229 :
230 : ! Avoid nested component array sections in GCC 16 bounds checks.
231 512 : xkp => ctx%kpoints%xkp
232 512 : IF (ctx%kpoints%use_real_wfn) THEN
233 16 : IF (use_grid) THEN
234 0 : CPASSERT(PRESENT(grid))
235 0 : CALL rskp_transform_grid_extract(grid, ik, ctx%rmatrix)
236 : ELSE
237 16 : CALL dbcsr_set(ctx%rmatrix, 0.0_dp)
238 : CALL rskp_transform(rmatrix=ctx%rmatrix, rsmat=rsmat, ispin=matrix_row, &
239 : xkp=xkp(:, ik), cell_to_index=ctx%kpoints%cell_to_index, &
240 16 : sab_nl=ctx%kpoints%sab_nl)
241 : END IF
242 496 : ELSE IF (use_grid) THEN
243 0 : CPASSERT(PRESENT(grid))
244 0 : CALL rskp_transform_grid_extract(grid, ik, ctx%rmatrix, ctx%cmatrix)
245 : ELSE
246 496 : CALL dbcsr_set(ctx%rmatrix, 0.0_dp)
247 496 : CALL dbcsr_set(ctx%cmatrix, 0.0_dp)
248 : CALL rskp_transform(rmatrix=ctx%rmatrix, cmatrix=ctx%cmatrix, rsmat=rsmat, ispin=matrix_row, &
249 : xkp=xkp(:, ik), cell_to_index=ctx%kpoints%cell_to_index, &
250 496 : sab_nl=ctx%kpoints%sab_nl)
251 : END IF
252 512 : CALL copy_dbcsr_to_fm(ctx%rmatrix, ctx%fmwork(kpoint_part_re), ctx%copy_plan(kpoint_part_re))
253 512 : IF (.NOT. ctx%kpoints%use_real_wfn) THEN
254 496 : CALL copy_dbcsr_to_fm(ctx%cmatrix, ctx%fmwork(kpoint_part_im), ctx%copy_plan(kpoint_part_im))
255 : END IF
256 :
257 512 : CALL timestop(handle)
258 :
259 512 : END SUBROUTINE kpoint_operator_densify
260 :
261 : ! **************************************************************************************************
262 : !> \brief Launch a densified operator to its owning group (only the real part for real wavefunctions).
263 : !> Absent explicit targets, receive both parts through fmlocal. Each copy
264 : !> owns its receive buffer until finish, so the targets can alias.
265 : !> Non-owning groups participate with an empty dummy target.
266 : !> \param kpoints host environment, including group ownership and transfer communicator
267 : !> \param ik global k-point index
268 : !> \param fmwork real/imaginary source matrices on the global distribution
269 : !> \param fmlocal default receiver on the k-point group distribution
270 : !> \param info transfer states, real part first; real wavefunctions allocate only the first part
271 : !> \param target_re optional caller-owned real-part target
272 : !> \param target_im optional imaginary-part target; required with target_re for complex wavefunctions
273 : ! **************************************************************************************************
274 512 : SUBROUTINE kpoint_operator_launch(kpoints, ik, fmwork, fmlocal, info, target_re, target_im)
275 :
276 : TYPE(kpoint_type), POINTER :: kpoints
277 : INTEGER, INTENT(IN) :: ik
278 : TYPE(cp_fm_type), DIMENSION(:), INTENT(IN) :: fmwork
279 : TYPE(cp_fm_type), INTENT(IN) :: fmlocal
280 : TYPE(copy_info_type), DIMENSION(:) :: info
281 : TYPE(cp_fm_type), INTENT(IN), OPTIONAL :: target_re, target_im
282 :
283 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_launch'
284 :
285 : INTEGER :: handle
286 : TYPE(cp_fm_type) :: fmdummy
287 : TYPE(mp_para_env_type), POINTER :: para_env
288 :
289 512 : CALL timeset(routineN, handle)
290 512 : CPASSERT(.NOT. PRESENT(target_im) .OR. PRESENT(target_re))
291 512 : IF (.NOT. kpoints%use_real_wfn) THEN
292 496 : CPASSERT(PRESENT(target_re) .EQV. PRESENT(target_im))
293 : END IF
294 :
295 512 : para_env => kpoints%blacs_env_all%para_env
296 512 : IF (ik >= kpoints%kp_range(1) .AND. ik <= kpoints%kp_range(2)) THEN
297 452 : IF (PRESENT(target_re)) THEN
298 0 : CALL cp_fm_start_copy_general(fmwork(kpoint_part_re), target_re, para_env, info(kpoint_part_re))
299 0 : IF (.NOT. kpoints%use_real_wfn) THEN
300 0 : CALL cp_fm_start_copy_general(fmwork(kpoint_part_im), target_im, para_env, info(kpoint_part_im))
301 : END IF
302 : ELSE
303 452 : CALL cp_fm_start_copy_general(fmwork(kpoint_part_re), fmlocal, para_env, info(kpoint_part_re))
304 452 : IF (.NOT. kpoints%use_real_wfn) THEN
305 436 : CALL cp_fm_start_copy_general(fmwork(kpoint_part_im), fmlocal, para_env, info(kpoint_part_im))
306 : END IF
307 : END IF
308 : ELSE
309 60 : CALL cp_fm_start_copy_general(fmwork(kpoint_part_re), fmdummy, para_env, info(kpoint_part_re))
310 60 : IF (.NOT. kpoints%use_real_wfn) THEN
311 60 : CALL cp_fm_start_copy_general(fmwork(kpoint_part_im), fmdummy, para_env, info(kpoint_part_im))
312 : END IF
313 : END IF
314 :
315 512 : CALL timestop(handle)
316 :
317 512 : END SUBROUTINE kpoint_operator_launch
318 :
319 : ! **************************************************************************************************
320 : !> \brief Receive one group's operator and reclaim every rank's send buffers.
321 : !> CFM/DBCSR outputs consume the context buffer before receiving the next part;
322 : !> FM outputs receive directly. Inactive groups only clean up their sends.
323 : !> \param kpoints parallel environment
324 : !> \param fmlocal shared receive buffer, consumed before receiving the next part
325 : !> \param ik global k-point index
326 : !> \param info real/imaginary transfers started by kpoint_operator_launch
327 : !> \param cmat complex output, real part overwrites and imaginary part adds
328 : !> \param fm_re real FM output
329 : !> \param fm_im optional imaginary FM output
330 : !> \param mat_re real DBCSR output
331 : !> \param mat_im imaginary DBCSR output
332 : ! **************************************************************************************************
333 512 : SUBROUTINE kpoint_operator_receive(kpoints, fmlocal, ik, info, cmat, fm_re, fm_im, mat_re, mat_im)
334 : TYPE(kpoint_type), INTENT(IN), POINTER :: kpoints
335 : TYPE(cp_fm_type) :: fmlocal
336 : INTEGER, INTENT(IN) :: ik
337 : TYPE(copy_info_type), INTENT(INOUT) :: info(:)
338 : TYPE(cp_cfm_type), OPTIONAL :: cmat
339 : TYPE(cp_fm_type), OPTIONAL :: fm_re, fm_im
340 : TYPE(dbcsr_type), INTENT(INOUT), OPTIONAL :: mat_re, mat_im
341 :
342 : INTEGER :: part
343 :
344 512 : IF (ik >= kpoints%kp_range(1) .AND. ik <= kpoints%kp_range(2)) THEN
345 452 : IF (PRESENT(cmat)) THEN
346 436 : CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_re))
347 436 : CALL cp_cfm_scale_and_add_fm(z_zero, cmat, z_one, fmlocal)
348 436 : IF (.NOT. kpoints%use_real_wfn) THEN
349 436 : CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_im))
350 436 : CALL cp_cfm_scale_and_add_fm(z_one, cmat, gaussi, fmlocal)
351 : END IF
352 16 : ELSE IF (PRESENT(fm_re)) THEN
353 16 : CALL cp_fm_finish_copy_general(fm_re, info(kpoint_part_re))
354 16 : IF (kpoints%use_real_wfn) THEN
355 16 : IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
356 : ELSE
357 0 : CALL cp_fm_finish_copy_general(fm_im, info(kpoint_part_im))
358 : END IF
359 : ELSE
360 0 : CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_re))
361 0 : CALL copy_fm_to_dbcsr_bc(fmlocal, mat_re)
362 0 : IF (kpoints%use_real_wfn) THEN
363 0 : fmlocal%local_data(:, :) = 0.0_dp
364 : ELSE
365 0 : CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_im))
366 : END IF
367 0 : CALL copy_fm_to_dbcsr_bc(fmlocal, mat_im)
368 : END IF
369 : END IF
370 2016 : DO part = 1, MERGE(1, 2, kpoints%use_real_wfn)
371 1520 : CALL cp_fm_cleanup_copy_general(info(part))
372 : END DO
373 512 : END SUBROUTINE kpoint_operator_receive
374 :
375 : ! **************************************************************************************************
376 : !> \brief Create an empty operator context.
377 : !> The context owns no buffers after this call; they are created on first
378 : !> use and released with the context.
379 : !> \param ctx context to initialize
380 : !> \param kpoints host k-point environment, referenced, not copied
381 : !> \param ao_ao_fm any square full matrix on the environment-wide communicator;
382 : !> its matrix struct is the template for the context work buffers,
383 : !> which densify on the source side of transfers that cross into the
384 : !> k-point groups
385 : !> \param nspin instance spin space of the owning driver call: the bound of
386 : !> the spin loop its start and finish calls iterate. Fixed for the
387 : !> lifetime of the context
388 : ! **************************************************************************************************
389 39792 : SUBROUTINE kpoint_operator_context_create(ctx, kpoints, ao_ao_fm, nspin)
390 :
391 : TYPE(kpoint_operator_context_type), INTENT(OUT) :: ctx
392 : TYPE(kpoint_type), POINTER :: kpoints
393 : TYPE(cp_fm_type), INTENT(IN) :: ao_ao_fm
394 : INTEGER, INTENT(IN) :: nspin
395 :
396 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_context_create'
397 :
398 : INTEGER :: handle
399 :
400 39792 : CALL timeset(routineN, handle)
401 :
402 39792 : CPASSERT(nspin >= 1)
403 39792 : ctx%nspin = nspin
404 39792 : CPASSERT(ASSOCIATED(kpoints))
405 39792 : ctx%kpoints => kpoints
406 39792 : CALL cp_fm_get_info(ao_ao_fm, matrix_struct=ctx%ao_ao_fmstruct)
407 39792 : CPASSERT(ASSOCIATED(ctx%ao_ao_fmstruct))
408 : ! The template must live on the transfer world. A template from
409 : ! another communicator fails during the transfers, far from this call.
410 39792 : IF (.NOT. ASSOCIATED(ctx%ao_ao_fmstruct%para_env, kpoints%blacs_env_all%para_env)) THEN
411 0 : CALL cp_abort(__LOCATION__, "Operator context requires the global parallel environment")
412 : END IF
413 :
414 39792 : CALL timestop(handle)
415 :
416 39792 : END SUBROUTINE kpoint_operator_context_create
417 :
418 : ! **************************************************************************************************
419 : !> \brief Release the context together with every buffer it created.
420 : !> Call once, at the end of the fixed-layout scope that created the context.
421 : !> Every started instance must be idle at this point. The call aborts
422 : !> when a transfer is still in flight, because nothing would ever wait
423 : !> for its MPI requests.
424 : !> \param ctx context to release
425 : ! **************************************************************************************************
426 39792 : SUBROUTINE kpoint_operator_context_release(ctx)
427 :
428 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
429 :
430 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_context_release'
431 :
432 : INTEGER :: handle, ispin, iwork, slot
433 :
434 39792 : CALL timeset(routineN, handle)
435 :
436 : ! Every started instance must have been finished: an unfinished start
437 : ! holds MPI requests that nothing would ever wait for.
438 39792 : IF (ALLOCATED(ctx%status)) THEN
439 164190 : CPASSERT(ALL(ctx%status == kpoint_slot_idle))
440 : END IF
441 :
442 39792 : IF (ASSOCIATED(ctx%rmatrix)) CALL dbcsr_deallocate_matrix(ctx%rmatrix)
443 39792 : IF (ASSOCIATED(ctx%cmatrix)) CALL dbcsr_deallocate_matrix(ctx%cmatrix)
444 39792 : NULLIFY (ctx%rmatrix, ctx%cmatrix, ctx%work_sab)
445 40376 : IF (ALLOCATED(ctx%copy_plan)) DEALLOCATE (ctx%copy_plan)
446 39792 : IF (ALLOCATED(ctx%rs)) THEN
447 158616 : DO slot = 1, kpoint_num_slots
448 284724 : DO ispin = 1, ctx%nspin
449 245070 : CALL kpoint_operator_rs_clear(ctx%rs(ispin, slot))
450 : END DO
451 : END DO
452 291870 : DEALLOCATE (ctx%rs, ctx%rs_batch)
453 : END IF
454 39792 : IF (ALLOCATED(ctx%rs_route)) THEN
455 39642 : CALL cp_fm_struct_release(ctx%rs_route%matrix_struct)
456 39642 : DEALLOCATE (ctx%rs_route)
457 : END IF
458 39792 : IF (ALLOCATED(ctx%fmwork)) THEN
459 442 : DO iwork = 1, SIZE(ctx%fmwork)
460 442 : CALL cp_fm_release(ctx%fmwork(iwork))
461 : END DO
462 150 : DEALLOCATE (ctx%fmwork)
463 : END IF
464 39792 : IF (ALLOCATED(ctx%fmimag)) THEN
465 56 : CALL cp_fm_release(ctx%fmimag)
466 56 : DEALLOCATE (ctx%fmimag)
467 : END IF
468 39792 : IF (ALLOCATED(ctx%fmlocal)) THEN
469 39792 : CALL fm_pool_give_back_fm(ctx%group_fm_pool, ctx%fmlocal)
470 39792 : DEALLOCATE (ctx%fmlocal)
471 : END IF
472 39912 : IF (ALLOCATED(ctx%info)) DEALLOCATE (ctx%info)
473 39792 : IF (ALLOCATED(ctx%status)) DEALLOCATE (ctx%status)
474 39792 : NULLIFY (ctx%kpoints, ctx%ao_ao_fmstruct, ctx%group_fm_pool)
475 39792 : ctx%overlap_stored = .FALSE.
476 :
477 39792 : CALL timestop(handle)
478 :
479 39792 : END SUBROUTINE kpoint_operator_context_release
480 :
481 : ! **************************************************************************************************
482 : !> \brief Create the context scratch that the first service call needs: the
483 : !> DBCSR work pair with blocks from the current neighbor lists, two
484 : !> densify buffers on the environment-wide struct, and the group
485 : !> transfer target from the group pools.
486 : !> The per-instance slot bookkeeping is not scratch: start creates it
487 : !> on its first call, because the synchronous get call never touches it.
488 : !> The buffers assume a fixed-layout context scope: the
489 : !> host neighbor lists do not change within it, and every buffer use
490 : !> asserts this identity. Buffers from an earlier call stay in place.
491 : !> Called from the service routines only.
492 : !> \param ctx context that receives the buffers
493 : !> \param rsmat real-space image matrices of one operator. The call uses the
494 : !> first matrix as the DBCSR template of the work pair. It reads no
495 : !> matrix data.
496 : ! **************************************************************************************************
497 452 : SUBROUTINE kpoint_operator_ensure(ctx, rsmat)
498 :
499 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
500 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
501 :
502 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_ensure'
503 :
504 : INTEGER :: handle
505 :
506 452 : CALL timeset(routineN, handle)
507 :
508 : ! The pair must still match the host neighbor lists: they are
509 : ! reallocated at every energy evaluation; the context stays within it.
510 452 : IF (ASSOCIATED(ctx%work_sab)) THEN
511 302 : CPASSERT(ASSOCIATED(ctx%work_sab, ctx%kpoints%sab_nl))
512 : END IF
513 452 : IF (.NOT. ASSOCIATED(ctx%rmatrix)) THEN
514 150 : ALLOCATE (ctx%rmatrix)
515 150 : CALL dbcsr_create(ctx%rmatrix, template=rsmat(1, 1)%matrix, matrix_type=dbcsr_type_symmetric)
516 150 : CALL cp_dbcsr_alloc_block_from_nbl(ctx%rmatrix, ctx%kpoints%sab_nl)
517 150 : IF (.NOT. ctx%kpoints%use_real_wfn) THEN
518 142 : ALLOCATE (ctx%cmatrix)
519 142 : CALL dbcsr_create(ctx%cmatrix, template=rsmat(1, 1)%matrix, matrix_type=dbcsr_type_antisymmetric)
520 142 : CALL cp_dbcsr_alloc_block_from_nbl(ctx%cmatrix, ctx%kpoints%sab_nl)
521 : END IF
522 150 : ctx%work_sab => ctx%kpoints%sab_nl
523 : END IF
524 :
525 452 : IF (.NOT. ALLOCATED(ctx%fmwork)) THEN
526 734 : ALLOCATE (ctx%fmwork(MERGE(1, 2, ctx%kpoints%use_real_wfn)))
527 742 : ALLOCATE (ctx%copy_plan(SIZE(ctx%fmwork)))
528 : CALL cp_fm_create(ctx%fmwork(kpoint_part_re), ctx%ao_ao_fmstruct, &
529 150 : name="KPOINT-OPERATOR-WORK-RE")
530 150 : IF (.NOT. ctx%kpoints%use_real_wfn) THEN
531 : CALL cp_fm_create(ctx%fmwork(kpoint_part_im), ctx%ao_ao_fmstruct, &
532 142 : name="KPOINT-OPERATOR-WORK-IM")
533 : END IF
534 : END IF
535 :
536 452 : CALL kpoint_operator_ensure_group(ctx)
537 :
538 452 : CALL timestop(handle)
539 :
540 452 : END SUBROUTINE kpoint_operator_ensure
541 :
542 : ! **************************************************************************************************
543 : !> \brief Acquire only the group-local output workspace; both backends share it.
544 : !> \param ctx operator context
545 : ! **************************************************************************************************
546 48632 : SUBROUTINE kpoint_operator_ensure_group(ctx)
547 :
548 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
549 :
550 48632 : IF (.NOT. ALLOCATED(ctx%fmlocal)) THEN
551 39792 : CALL mpools_get(ctx%kpoints%mpools, ao_ao_fm_pool=ctx%group_fm_pool)
552 39792 : ALLOCATE (ctx%fmlocal)
553 39792 : CALL fm_pool_create_fm(ctx%group_fm_pool, ctx%fmlocal)
554 : END IF
555 :
556 48632 : END SUBROUTINE kpoint_operator_ensure_group
557 :
558 : ! **************************************************************************************************
559 : !> \brief Prepare an immutable operator slot for group-local assembly.
560 : !> Refresh after the source values change, with every instance idle.
561 : !> Other prepared slots and structural routing remain valid. Without
562 : !> matrix_row, capture one row per instance spin; with matrix_row, store
563 : !> one row shared by all instance spins. Failure clears this slot on
564 : !> all ranks and retains the existing transfer fallback.
565 : !> \param ctx context spanning the current initialized k-point environment
566 : !> \param rsmat current real-space operator images
567 : !> \param slot slot to replace (KS, S or T)
568 : !> \param prepared optional collective success flag
569 : !> \param matrix_row optional explicit spin-free operator row
570 : !> \param TARGET optional group-local FM layout for prepared get outputs;
571 : !> otherwise use the host AO/AO pool layout. One layout per context
572 : ! **************************************************************************************************
573 46784 : SUBROUTINE kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row, TARGET)
574 :
575 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
576 : TYPE(dbcsr_p_type), POINTER :: rsmat(:, :)
577 : INTEGER, INTENT(IN) :: slot
578 : LOGICAL, INTENT(OUT), OPTIONAL :: prepared
579 : INTEGER, INTENT(IN), OPTIONAL :: matrix_row
580 : TYPE(cp_fm_type), INTENT(IN), OPTIONAL :: target
581 :
582 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_prepare'
583 :
584 : INTEGER :: compatible, handle, ispin, jslot, nrow, &
585 : row
586 : INTEGER(KIND=int_8) :: retained, storage
587 : LOGICAL :: ready
588 : TYPE(cp_fm_struct_type), POINTER :: fms
589 : TYPE(mp_para_env_type), POINTER :: para_env
590 :
591 46784 : CALL timeset(routineN, handle)
592 46784 : IF (PRESENT(prepared)) prepared = .FALSE.
593 46784 : CPASSERT(ASSOCIATED(ctx%kpoints))
594 46784 : CPASSERT(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
595 46784 : CPASSERT(ASSOCIATED(ctx%kpoints%sab_nl))
596 46784 : IF (ASSOCIATED(ctx%work_sab)) THEN
597 7120 : CPASSERT(ASSOCIATED(ctx%work_sab, ctx%kpoints%sab_nl))
598 : END IF
599 46784 : IF (ALLOCATED(ctx%status)) THEN
600 0 : CPASSERT(ALL(ctx%status == kpoint_slot_idle))
601 : END IF
602 46784 : IF (.NOT. ALLOCATED(ctx%rs)) THEN
603 688410 : ALLOCATE (ctx%rs(ctx%nspin, kpoint_num_slots), ctx%rs_batch(ctx%nspin, kpoint_num_slots))
604 : END IF
605 96858 : DO ispin = 1, ctx%nspin
606 50074 : CALL kpoint_operator_rs_clear(ctx%rs(ispin, slot))
607 50074 : IF (ALLOCATED(ctx%rs_batch(ispin, slot)%data)) DEALLOCATE (ctx%rs_batch(ispin, slot)%data)
608 96858 : ctx%rs_batch(ispin, slot)%first = 0
609 : END DO
610 46784 : nrow = ctx%nspin
611 46784 : IF (PRESENT(matrix_row)) THEN
612 7608 : CPASSERT(matrix_row >= 1 .AND. matrix_row <= SIZE(rsmat, 1))
613 : nrow = 1
614 : ELSE
615 39176 : CPASSERT(SIZE(rsmat, 1) == ctx%nspin)
616 : END IF
617 46784 : para_env => ctx%kpoints%blacs_env_all%para_env
618 46784 : CALL kpoint_operator_ensure_group(ctx)
619 46784 : fms => ctx%fmlocal%matrix_struct
620 46784 : IF (PRESENT(TARGET)) fms => target%matrix_struct
621 46784 : CPASSERT(ASSOCIATED(fms))
622 : ! Routing uses the host group-rank numbering. Different BLACS layouts
623 : ! on that communicator are supported, foreign communicators are not.
624 46784 : compatible = MERGE(1, 0, ASSOCIATED(fms%para_env, ctx%kpoints%para_env_kp))
625 46784 : CALL para_env%min(compatible)
626 46784 : IF (compatible == 0) THEN
627 0 : CALL timestop(handle)
628 22 : RETURN
629 : END IF
630 46784 : IF (.NOT. ALLOCATED(ctx%rs_route)) THEN
631 39664 : ALLOCATE (ctx%rs_route)
632 39664 : CALL kpoint_operator_rs_plan_create(ctx%kpoints, rsmat, fms, rs_snapshot_bytes, ctx%rs_route, ready)
633 39664 : IF (.NOT. ready) THEN
634 22 : DEALLOCATE (ctx%rs_route)
635 22 : CALL timestop(handle)
636 22 : RETURN
637 : END IF
638 39642 : ctx%work_sab => ctx%kpoints%sab_nl
639 : END IF
640 46762 : compatible = MERGE(1, 0, cp_fm_struct_equivalent(fms, ctx%rs_route%matrix_struct))
641 96222 : DO ispin = 1, nrow
642 49460 : row = ispin
643 49460 : IF (PRESENT(matrix_row)) row = matrix_row
644 96222 : IF (.NOT. kpoint_operator_rs_compatible(ctx%rs_route, rsmat, row)) compatible = 0
645 : END DO
646 46762 : CALL para_env%min(compatible)
647 : ! Keep the single-k phase allowance in the snapshot budget for the uncached fallback.
648 46762 : retained = 0_int_8
649 187048 : DO jslot = kpoint_slot_ks, kpoint_slot_t
650 337204 : DO ispin = 1, ctx%nspin
651 150156 : IF (.NOT. ALLOCATED(ctx%rs(ispin, jslot)%blocks)) CYCLE
652 6024 : retained = retained + 128_int_8*SIZE(ctx%rs(ispin, jslot)%blocks, KIND=int_8)
653 146310 : IF (ASSOCIATED(ctx%rs(ispin, jslot)%values)) THEN
654 6024 : retained = retained + 8_int_8*SIZE(ctx%rs(ispin, jslot)%values, KIND=int_8)
655 : END IF
656 : END DO
657 : END DO
658 : storage = ctx%rs_route%storage + retained + INT(nrow, int_8)* &
659 : (MERGE(0_int_8, 8_int_8*ctx%rs_route%nvalue, para_env%num_pe == 1) + &
660 46762 : 128_int_8*SIZE(ctx%rs_route%entry, 2))
661 46762 : CALL para_env%max(storage)
662 46762 : ready = compatible == 1 .AND. storage <= rs_snapshot_bytes
663 46762 : IF (ready) THEN
664 96222 : DO ispin = 1, nrow
665 49460 : row = ispin
666 49460 : IF (PRESENT(matrix_row)) row = matrix_row
667 49460 : CALL kpoint_operator_rs_capture(rsmat, row, ctx%rs_route, ctx%kpoints, ctx%rs(ispin, slot))
668 : CALL kpoint_operator_fft_prepare(ctx%rs(ispin, slot), ctx%rs_route, ctx%kpoints, &
669 96222 : ctx%rs_batch(ispin, slot), rs_fft_bytes/INT(ctx%nspin*kpoint_num_slots, int_8))
670 : END DO
671 : END IF
672 46762 : IF (PRESENT(prepared)) prepared = ready
673 46762 : CALL timestop(handle)
674 :
675 : END SUBROUTINE kpoint_operator_prepare
676 :
677 : ! **************************************************************************************************
678 : !> \brief Release one captured row without invalidating other slots or routing.
679 : !> \param snapshot row to clear
680 : ! **************************************************************************************************
681 178590 : SUBROUTINE kpoint_operator_rs_clear(snapshot)
682 :
683 : TYPE(kpoint_operator_rs_type), INTENT(INOUT) :: snapshot
684 :
685 178590 : IF (ALLOCATED(snapshot%blocks)) DEALLOCATE (snapshot%blocks, snapshot%nonzero)
686 178590 : IF (ASSOCIATED(snapshot%values)) DEALLOCATE (snapshot%values)
687 178590 : NULLIFY (snapshot%source)
688 178590 : snapshot%matrix_row = 0
689 :
690 178590 : END SUBROUTINE kpoint_operator_rs_clear
691 :
692 : ! **************************************************************************************************
693 : !> \brief Transform disjoint AO entries, then exchange only group-owned k-point results.
694 : !> The snapshot stays usable for direct assembly if any rank rejects the FFT.
695 : !> All pool/planner calls and communication precede OpenMP k-point workers.
696 : !> \param snapshot captured real-space blocks
697 : !> \param route prepared group-local layout
698 : !> \param kpoints host sampling and communicators
699 : !> \param work existing matrix-batch cache, filled for all group-owned k points
700 : !> \param max_bytes budget for this row's output, maps and transient FFT/exchange buffers
701 : ! **************************************************************************************************
702 0 : SUBROUTINE kpoint_operator_fft_prepare(snapshot, route, kpoints, work, max_bytes)
703 : TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
704 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
705 : TYPE(kpoint_type), POINTER :: kpoints
706 : TYPE(kpoint_operator_rs_batch_type), INTENT(OUT) :: work
707 : INTEGER(KIND=int_8), INTENT(IN) :: max_bytes
708 :
709 49460 : COMPLEX(KIND=dp), ALLOCATABLE, TARGET :: recvbuf(:), sendbuf(:)
710 49460 : COMPLEX(KIND=dp), CONTIGUOUS, POINTER :: values_k(:, :, :)
711 : INTEGER :: batch, count, first, group, handle, i, &
712 : ik, j, last, local_k, mine, ncell, &
713 : nel, ngroup, nkp, offset, pos, width
714 : INTEGER(KIND=int_8) :: fixed_bytes, per_entry, tile_size
715 49460 : INTEGER, ALLOCATABLE :: cells(:, :), ranges(:, :), rcount(:), &
716 49460 : rdispl(:), scount(:), sdispl(:)
717 : INTEGER, DIMENSION(12) :: layout, layout_max, layout_min
718 : INTEGER, DIMENSION(2) :: choice
719 : INTEGER, DIMENSION(3) :: nfft
720 49460 : REAL(KIND=dp), ALLOCATABLE :: values_rs(:, :, :)
721 : TYPE(cp_fm_struct_type), POINTER :: fms
722 : TYPE(mp_para_env_type), POINTER :: inter, world
723 :
724 49460 : IF (kpoints%lattice_fft == lattice_fft_off) RETURN
725 49460 : nkp = SIZE(kpoints%xkp, 2)
726 49460 : IF (kpoints%lattice_fft == lattice_fft_auto .AND. nkp < 27) RETURN
727 1316 : IF (ANY(kpoints%nkp_grid <= 0)) RETURN
728 1304 : IF (32_int_8*PRODUCT(INT(kpoints%nkp_grid, int_8)) > max_bytes) RETURN
729 326 : IF (.NOT. lattice_fft_shape(kpoints%nkp_grid, nfft, allow_arbitrary=.TRUE.)) RETURN
730 1304 : IF (32_int_8*PRODUCT(INT(nfft, int_8)) > max_bytes) RETURN
731 326 : IF (.NOT. regular_kpoint_grid(kpoints%xkp, kpoints%nkp_grid, allow_incomplete=.TRUE.)) RETURN
732 :
733 326 : world => kpoints%blacs_env_all%para_env
734 326 : inter => kpoints%para_env_inter_kp
735 326 : ngroup = inter%num_pe
736 326 : mine = inter%mepos + 1
737 326 : local_k = kpoints%kp_range(2) - kpoints%kp_range(1) + 1
738 326 : fms => route%matrix_struct
739 : ! Equivalent inter-group ranks must agree on both local indices and global layout.
740 : layout = [fms%nrow_global, fms%ncol_global, fms%nrow_block, fms%ncol_block, &
741 6194 : fms%first_p_pos, fms%context%num_pe, fms%context%mepos, route%nrow, route%ncol]
742 326 : layout_min = layout
743 326 : layout_max = layout
744 326 : CALL inter%min(layout_min)
745 326 : CALL inter%max(layout_max)
746 326 : tile_size = INT(route%nrow, int_8)*route%ncol
747 326 : nel = 0
748 326 : IF (tile_size <= HUGE(nel)) nel = INT(tile_size)
749 : fixed_bytes = 16_int_8*tile_size*local_k + &
750 : 48_int_8*PRODUCT(INT(nfft, int_8))*lattice_fft_num_threads(nfft, nel) + &
751 : 4_int_8*PRODUCT(INT(kpoints%nkp_grid, int_8)) + 12_int_8*nkp + &
752 2282 : 80_int_8*SIZE(route%entry, 2, KIND=int_8) + 24_int_8*ngroup
753 : per_entry = 16_int_8*SIZE(route%entry, 2, KIND=int_8) + &
754 326 : 16_int_8*(INT(nkp, int_8) + INT(ngroup, int_8)*local_k)
755 : width = INT(MAX(0_int_8, MIN(INT((nel - 1)/ngroup + 1, int_8), &
756 326 : (max_bytes - fixed_bytes)/MAX(1_int_8, per_entry))))
757 4890 : choice = [MERGE(1, 0, nel > 0 .AND. ALL(layout_min == layout_max)), width]
758 326 : CALL world%min(choice)
759 326 : IF (choice(1) == 0 .OR. choice(2) == 0) RETURN
760 326 : width = choice(2)
761 326 : IF (kpoints%lattice_fft == lattice_fft_auto .AND. ngroup > 1) THEN
762 468 : IF (16_int_8*width*MINVAL(kpoints%kp_dist(2, :) - kpoints%kp_dist(1, :) + 1) < 32_int_8*1024) RETURN
763 : END IF
764 :
765 190 : CALL timeset("kpoint_operator_fft_prepare", handle)
766 190 : ncell = 2*SIZE(route%entry, 2)
767 950 : ALLOCATE (cells(3, ncell), ranges(2, ngroup))
768 211088 : DO i = 1, SIZE(route%entry, 2)
769 843592 : cells(:, 2*i - 1) = route%entry(5:7, i)
770 843782 : cells(:, 2*i) = -route%entry(5:7, i)
771 : END DO
772 404 : DO group = 1, ngroup
773 214 : ranges(1, group) = INT(INT(nel, int_8)*(group - 1)/ngroup) + 1
774 404 : ranges(2, group) = INT(INT(nel, int_8)*group/ngroup)
775 : END DO
776 950 : ALLOCATE (work%data(local_k, route%nrow, route%ncol))
777 1520 : ALLOCATE (values_rs(width, 1, ncell), sendbuf(width*nkp), recvbuf(width*ngroup*local_k))
778 1140 : ALLOCATE (scount(ngroup), sdispl(ngroup), rcount(ngroup), rdispl(ngroup))
779 671 : DO batch = 1, MAXVAL(ranges(2, :) - ranges(1, :) + 1), width
780 267 : first = ranges(1, mine) + batch - 1
781 267 : last = MIN(first + width - 1, ranges(2, mine))
782 267 : count = MAX(0, last - first + 1)
783 267 : IF (count > 0) THEN
784 267 : CALL kpoint_operator_fft_pack(snapshot, route, first, last, values_rs(1:count, :, :))
785 267 : values_k(1:count, 1:1, 1:nkp) => sendbuf(1:count*nkp)
786 : ! The transform itself supplies a direct result if the FFT backend fails.
787 : CALL cell_to_k_grid_fft(values_rs(1:count, :, :), cells, kpoints%xkp, kpoints%nkp_grid, &
788 267 : values_k)
789 : END IF
790 : offset = 0
791 558 : DO group = 1, ngroup
792 291 : scount(group) = count*(kpoints%kp_dist(2, group) - kpoints%kp_dist(1, group) + 1)
793 291 : sdispl(group) = count*(kpoints%kp_dist(1, group) - 1)
794 291 : first = ranges(1, group) + batch - 1
795 291 : last = MIN(first + width - 1, ranges(2, group))
796 291 : rcount(group) = MAX(0, last - first + 1)*local_k
797 291 : rdispl(group) = offset
798 558 : offset = offset + rcount(group)
799 : END DO
800 267 : CALL inter%alltoall(sendbuf, scount, sdispl, recvbuf, rcount, rdispl)
801 748 : DO group = 1, ngroup
802 291 : first = ranges(1, group) + batch - 1
803 291 : last = MIN(first + width - 1, ranges(2, group))
804 : !$OMP PARALLEL DO DEFAULT(NONE) SCHEDULE(STATIC) &
805 558 : !$OMP SHARED(work,route,recvbuf,rdispl,group,first,last,local_k) PRIVATE(ik,i,j,pos)
806 : DO pos = first, last
807 : i = MODULO(pos - 1, route%nrow) + 1
808 : j = (pos - 1)/route%nrow + 1
809 : DO ik = 1, local_k
810 : work%data(ik, i, j) = recvbuf(rdispl(group) + (ik - 1)*(last - first + 1) + pos - first + 1)
811 : END DO
812 : END DO
813 : !$OMP END PARALLEL DO
814 : END DO
815 : END DO
816 190 : work%first = kpoints%kp_range(1)
817 190 : CALL timestop(handle)
818 99110 : END SUBROUTINE kpoint_operator_fft_prepare
819 :
820 : ! **************************************************************************************************
821 : !> \brief Pack independent AO entries with the same mirror/sign convention as direct assembly.
822 : !> \param snapshot captured real-space blocks
823 : !> \param route prepared group-local layout
824 : !> \param first first flattened local AO entry
825 : !> \param last last flattened local AO entry
826 : !> \param values real-cell batch; odd/even cells hold a block and its Hermitian mirror
827 : ! **************************************************************************************************
828 267 : SUBROUTINE kpoint_operator_fft_pack(snapshot, route, first, last, values)
829 : TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
830 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
831 : INTEGER, INTENT(IN) :: first, last
832 : REAL(KIND=dp), INTENT(OUT) :: values(:, :, :)
833 :
834 : INTEGER :: col, element, iblock, part, row, &
835 : source_col, source_row
836 :
837 861965903 : values(:, :, :) = 0.0_dp
838 : !$OMP PARALLEL DO DEFAULT(NONE) SCHEDULE(STATIC) &
839 267 : !$OMP SHARED(snapshot,route,first,last,values) PRIVATE(element,row,col,iblock,part,source_row,source_col)
840 : DO element = first, last
841 : row = MODULO(element - 1, route%nrow) + 1
842 : col = (element - 1)/route%nrow + 1
843 : DO iblock = 1, SIZE(route%entry, 2)
844 : IF (.NOT. snapshot%nonzero(iblock)) CYCLE
845 : DO part = 1, 2
846 : IF (part == 2 .AND. route%entry(8, iblock) == 0) CYCLE
847 : IF (row < route%tile(1, part, iblock) .OR. row > route%tile(2, part, iblock)) CYCLE
848 : IF (col < route%tile(3, part, iblock) .OR. col > route%tile(4, part, iblock)) CYCLE
849 : source_row = route%matrix_struct%row_indices(row)
850 : source_col = route%matrix_struct%col_indices(col)
851 : IF (part == 2) THEN
852 : source_col = route%matrix_struct%row_indices(row)
853 : source_row = route%matrix_struct%col_indices(col)
854 : END IF
855 : values(element - first + 1, 1, 2*iblock - 2 + part) = &
856 : snapshot%blocks(iblock)%data(source_row - route%entry(1, iblock) + 1, &
857 : source_col - route%entry(2, iblock) + 1)
858 : END DO
859 : END DO
860 : END DO
861 : !$OMP END PARALLEL DO
862 267 : END SUBROUTINE kpoint_operator_fft_pack
863 :
864 : ! **************************************************************************************************
865 : !> \brief Copy a cached matrix into a possibly padded solver matrix, without MPI.
866 : !> \param values cached reciprocal matrix
867 : !> \param TARGET complex output
868 : !> \param fm_re real output
869 : !> \param fm_im optional imaginary output
870 : ! **************************************************************************************************
871 24327 : SUBROUTINE kpoint_operator_rs_copy(values, TARGET, fm_re, fm_im)
872 : COMPLEX(KIND=dp), INTENT(IN) :: values(:, :)
873 : TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
874 : TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
875 :
876 : INTEGER :: ncol, nrow
877 :
878 24327 : nrow = SIZE(values, 1)
879 24327 : ncol = SIZE(values, 2)
880 24327 : CPASSERT(PRESENT(TARGET) .NEQV. PRESENT(fm_re))
881 24327 : IF (PRESENT(TARGET)) THEN
882 29862737 : target%local_data(1:nrow, 1:ncol) = values
883 725250 : target%local_data(nrow + 1:, :) = z_zero
884 23539 : target%local_data(:, ncol + 1:) = z_zero
885 : ELSE
886 89338 : fm_re%local_data(1:nrow, 1:ncol) = REAL(values, dp)
887 8200 : fm_re%local_data(nrow + 1:, :) = 0.0_dp
888 788 : fm_re%local_data(:, ncol + 1:) = 0.0_dp
889 788 : IF (PRESENT(fm_im)) THEN
890 89338 : fm_im%local_data(1:nrow, 1:ncol) = AIMAG(values)
891 8200 : fm_im%local_data(nrow + 1:, :) = 0.0_dp
892 788 : fm_im%local_data(:, ncol + 1:) = 0.0_dp
893 : END IF
894 : END IF
895 24327 : END SUBROUTINE kpoint_operator_rs_copy
896 :
897 : ! **************************************************************************************************
898 : !> \brief Match an explicitly prepared row, sharing spin-free rows across spins.
899 : !> \param ctx operator context
900 : !> \param source live array identity identifying the prepared operation
901 : !> \param ispin instance spin
902 : !> \param slot prepared slot
903 : !> \param matrix_row operator row
904 : !> \return snapshot spin index, or zero for the transfer backend
905 : ! **************************************************************************************************
906 35804 : FUNCTION kpoint_operator_rs_match(ctx, source, ispin, slot, matrix_row) RESULT(irs)
907 :
908 : TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
909 : TYPE(dbcsr_p_type), POINTER :: source(:, :)
910 : INTEGER, INTENT(IN) :: ispin, slot, matrix_row
911 : INTEGER :: irs
912 :
913 : INTEGER :: i
914 :
915 35804 : irs = 0
916 35804 : IF (.NOT. ALLOCATED(ctx%rs)) RETURN
917 35804 : i = ispin
918 35804 : IF (.NOT. ASSOCIATED(ctx%rs(i, slot)%source)) i = 1
919 35804 : IF (ASSOCIATED(ctx%rs(i, slot)%source, source) .AND. ctx%rs(i, slot)%matrix_row == matrix_row) irs = i
920 :
921 : END FUNCTION kpoint_operator_rs_match
922 :
923 : ! **************************************************************************************************
924 : !> \brief Cache whole-block routing to the FM tiles of one representative group.
925 : !> Equivalent ranks in the other groups receive only their tile's blocks.
926 : !> Boundary blocks may overlap several tiles. Missing blocks are packed
927 : !> as zero, permitting different operator sparsities on a fixed AO layout.
928 : !> \param kpoints initialized parallel environment and neighbor lists
929 : !> \param source operator layout template
930 : !> \param matrix_struct target group layout
931 : !> \param max_bytes per-rank budget for routing and capture buffers
932 : !> \param route context-owned structural route
933 : !> \param ready false collectively if unsupported or over budget
934 : ! **************************************************************************************************
935 39664 : SUBROUTINE kpoint_operator_rs_plan_create(kpoints, source, matrix_struct, max_bytes, route, ready)
936 :
937 : TYPE(kpoint_type), POINTER :: kpoints
938 : TYPE(dbcsr_p_type), POINTER :: source(:, :)
939 : TYPE(cp_fm_struct_type), POINTER :: matrix_struct
940 : INTEGER(KIND=int_8), INTENT(IN) :: max_bytes
941 : TYPE(kpoint_operator_rs_route_type), INTENT(OUT) :: route
942 : LOGICAL, INTENT(OUT) :: ready
943 :
944 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_rs_plan_create'
945 :
946 : INTEGER :: compatible, dest, handle, i, iblock, &
947 : igroup, nblock, ncol, ngroup, nrow, p, &
948 : pos, rank
949 : INTEGER(KIND=int_8) :: nr(2), ns(2), storage
950 39664 : INTEGER(KIND=int_8), ALLOCATABLE :: recv_size(:, :), send_size(:, :)
951 39664 : INTEGER, ALLOCATABLE :: identities(:, :), owner(:), &
952 39664 : position(:, :), send_metadata(:, :)
953 39664 : INTEGER, POINTER :: col_size(:), row_size(:)
954 39664 : LOGICAL, ALLOCATABLE :: needed(:)
955 : TYPE(mp_para_env_type), POINTER :: inter, world
956 :
957 39664 : CALL timeset(routineN, handle)
958 39664 : ready = .FALSE.
959 39664 : world => kpoints%blacs_env_all%para_env
960 39664 : inter => kpoints%para_env_inter_kp
961 39664 : CALL dbcsr_get_info(source(1, 1)%matrix, row_blk_size=row_size, col_blk_size=col_size)
962 : CALL kp_transform_plan_create(route%traversal, kpoints%sab_nl, kpoints%cell_to_index, &
963 39664 : SIZE(source, 2), block_template=source(1, 1)%matrix)
964 39664 : compatible = MERGE(1, 0, route%traversal%symmetric)
965 366668 : IF (SUM(row_size) /= matrix_struct%nrow_global .OR. &
966 0 : SUM(col_size) /= matrix_struct%ncol_global) compatible = 0
967 39664 : IF (SIZE(row_size) == SIZE(col_size)) THEN
968 203166 : IF (ANY(row_size /= col_size)) compatible = 0
969 : ELSE
970 0 : compatible = 0
971 : END IF
972 39664 : CALL world%min(compatible)
973 39664 : IF (compatible == 0) THEN
974 0 : CALL timestop(handle)
975 0 : RETURN
976 : END IF
977 39664 : route%nimage = SIZE(source, 2)
978 198320 : ALLOCATE (route%row_size(SIZE(row_size)), route%col_size(SIZE(col_size)))
979 203166 : route%row_size(:) = row_size
980 203166 : route%col_size(:) = col_size
981 : ALLOCATE (identities(2, 0:world%num_pe - 1), owner(0:kpoints%para_env_kp%num_pe - 1), &
982 237984 : needed(0:kpoints%para_env_kp%num_pe - 1))
983 118992 : CALL world%allgather([kpoints%para_env_kp%mepos, inter%mepos], identities)
984 91062 : owner(:) = -1
985 118992 : DO rank = 0, world%num_pe - 1
986 118992 : IF (identities(2, rank) == 0) owner(identities(1, rank)) = rank
987 : END DO
988 91062 : CPASSERT(ALL(owner >= 0))
989 158656 : ALLOCATE (send_size(2, 0:world%num_pe - 1), recv_size(2, 0:world%num_pe - 1))
990 39664 : send_size(:, :) = 0_int_8
991 8505757 : DO i = 1, route%traversal%nentry
992 8466093 : nrow = row_size(route%traversal%row(i))
993 8466093 : ncol = col_size(route%traversal%col(i))
994 8466093 : CALL kpoint_operator_rs_owners(route, matrix_struct, i, needed)
995 20930721 : DO p = 0, SIZE(owner) - 1
996 12424964 : IF (.NOT. needed(p)) CYCLE
997 10316826 : dest = owner(p)
998 10316826 : send_size(1, dest) = send_size(1, dest) + 1_int_8
999 20891057 : send_size(2, dest) = send_size(2, dest) + INT(nrow, int_8)*ncol
1000 : END DO
1001 : END DO
1002 39664 : CALL world%alltoall(send_size, recv_size, 2)
1003 : ! Only inter-rank zero receives from the world. The other groups hold
1004 : ! equivalent FM tiles, as required by kpoint_env_initialize.
1005 277648 : ns = SUM(send_size, DIM=2)
1006 277648 : nr = SUM(recv_size, DIM=2)
1007 39664 : CALL inter%bcast(nr, 0)
1008 : ! Conservative peak: traversal, retained route, temporary metadata and
1009 : ! one snapshot. The full batch budget is checked again by prepare.
1010 : storage = 40_int_8*route%traversal%nentry + &
1011 : 4_int_8*(SIZE(row_size) + SIZE(col_size) + matrix_struct%nrow_global + matrix_struct%ncol_global) + &
1012 : 48_int_8*ns(1) + MERGE(0_int_8, 8_int_8*ns(2), world%num_pe == 1) + 68_int_8*nr(1) + &
1013 39664 : 128_int_8*world%num_pe + 8_int_8*SIZE(owner) + 8_int_8
1014 : ! Singleton groups also retain the AO-block group boundaries.
1015 39664 : IF (SIZE(owner) == 1) storage = storage + 4_int_8*(nr(1) + 1_int_8)
1016 39664 : route%storage = storage
1017 79328 : storage = storage + MERGE(0_int_8, 8_int_8*nr(2), world%num_pe == 1) + 128_int_8*nr(1)
1018 : ! Borrowed values do not consume the snapshot budget; still bound integer route offsets.
1019 198320 : IF (MAXVAL(ns) >= INT(HUGE(0), int_8) .OR. MAXVAL(nr) >= INT(HUGE(0), int_8)) THEN
1020 0 : storage = MAX(storage, max_bytes + 1_int_8)
1021 : END IF
1022 39664 : CALL world%max(storage)
1023 39664 : IF (storage > max_bytes) THEN
1024 22 : CALL timestop(handle)
1025 22 : RETURN
1026 : END IF
1027 : ! The byte bound also keeps all counts/displacements within default INTEGER.
1028 39642 : nblock = INT(nr(1))
1029 39642 : route%nvalue = INT(nr(2))
1030 0 : ALLOCATE (route%scount(2, 0:world%num_pe - 1), route%rcount(2, 0:world%num_pe - 1), &
1031 237852 : route%sdisp(2, 0:world%num_pe - 1), route%rdisp(2, 0:world%num_pe - 1))
1032 277494 : route%scount(:, :) = INT(send_size)
1033 277494 : route%rcount(:, :) = INT(recv_size)
1034 39642 : IF (SIZE(owner) == 1) CALL inter%bcast(route%rcount, 0)
1035 118926 : route%sdisp(:, 0) = 0
1036 118926 : route%rdisp(:, 0) = 0
1037 79284 : DO rank = 1, world%num_pe - 1
1038 118926 : route%sdisp(:, rank) = route%sdisp(:, rank - 1) + route%scount(:, rank - 1)
1039 158568 : route%rdisp(:, rank) = route%rdisp(:, rank - 1) + route%rcount(:, rank - 1)
1040 : END DO
1041 0 : ALLOCATE (route%send_entry(INT(ns(1))), route%send_first(INT(ns(1)) + 1), &
1042 0 : route%send_values(MERGE(0, INT(ns(2)), world%num_pe == 1)), send_metadata(8, INT(ns(1))), &
1043 0 : position(2, 0:world%num_pe - 1), route%entry(8, nblock), &
1044 631200 : route%first(nblock + 1))
1045 277494 : position(:, :) = route%sdisp
1046 8385615 : DO i = 1, route%traversal%nentry
1047 8345973 : nrow = row_size(route%traversal%row(i))
1048 8345973 : ncol = col_size(route%traversal%col(i))
1049 8345973 : CALL kpoint_operator_rs_owners(route, matrix_struct, i, needed)
1050 20570339 : DO p = 0, SIZE(owner) - 1
1051 12184724 : IF (.NOT. needed(p)) CYCLE
1052 10120091 : dest = owner(p)
1053 10120091 : iblock = position(1, dest) + 1
1054 10120091 : pos = position(2, dest) + 1
1055 10120091 : route%send_entry(iblock) = i
1056 10120091 : route%send_first(iblock) = pos
1057 50600455 : send_metadata(1:4, iblock) = [route%traversal%row_offset(i), route%traversal%col_offset(i), nrow, ncol]
1058 40480364 : send_metadata(5:7, iblock) = NINT(route%traversal%symmetry_sign(i))*route%traversal%cell(:, i)
1059 10120091 : send_metadata(8, iblock) = MERGE(1, 0, route%traversal%row(i) /= route%traversal%col(i))
1060 40770879 : position(:, dest) = position(:, dest) + [1, nrow*ncol]
1061 : END DO
1062 : END DO
1063 39642 : route%send_first(INT(ns(1)) + 1) = INT(ns(2)) + 1
1064 39642 : IF (SIZE(owner) == 1) THEN
1065 139650 : CALL world%allgatherv(send_metadata, route%entry, 8*route%rcount(1, :), 8*route%rdisp(1, :))
1066 : ELSE
1067 : CALL world%alltoall(send_metadata, 8*route%scount(1, :), 8*route%sdisp(1, :), &
1068 105408 : route%entry, 8*route%rcount(1, :), 8*route%rdisp(1, :))
1069 11712 : CALL inter%bcast(route%entry, 0)
1070 : END IF
1071 39642 : route%first(1) = 1
1072 14666955 : DO i = 1, nblock
1073 14666955 : route%first(i + 1) = route%first(i) + route%entry(3, i)*route%entry(4, i)
1074 : END DO
1075 39642 : CPASSERT(route%first(nblock + 1) == route%nvalue + 1)
1076 39642 : route%matrix_struct => matrix_struct
1077 39642 : CALL cp_fm_struct_retain(route%matrix_struct)
1078 39642 : CALL kpoint_operator_rs_tile_map(route)
1079 39642 : IF (SIZE(owner) == 1 .AND. nblock > 0) THEN
1080 : ! build_neighbor_lists keeps each symmetric SAB pair's images contiguous
1081 : ! on its unique source rank. Packing preserves this traversal order, so
1082 : ! Allgatherv concatenates whole AO groups, even without global key ordering.
1083 35973986 : ngroup = 1 + COUNT(ANY(route%entry(1:2, 2:) /= route%entry(1:2, :nblock - 1), DIM=1))
1084 83790 : ALLOCATE (route%block_first(ngroup + 1))
1085 27930 : route%block_first(1) = 1
1086 27930 : igroup = 1
1087 9014444 : DO i = 2, nblock
1088 26395976 : IF (ALL(route%entry(1:2, i) == route%entry(1:2, i - 1))) CYCLE
1089 358058 : igroup = igroup + 1
1090 9014444 : route%block_first(igroup) = i
1091 : END DO
1092 27930 : route%block_first(ngroup + 1) = nblock + 1
1093 : END IF
1094 39642 : ready = .TRUE.
1095 39642 : CALL timestop(handle)
1096 :
1097 118992 : END SUBROUTINE kpoint_operator_rs_plan_create
1098 :
1099 : ! **************************************************************************************************
1100 : !> \brief Cache local intervals intersecting each routed block and its Hermitian mirror.
1101 : !> FM indices are ordered globally; dummy indices on empty tiles lie outside the matrix.
1102 : !> \param route fixed routing and target distribution
1103 : ! **************************************************************************************************
1104 39642 : SUBROUTINE kpoint_operator_rs_tile_map(route)
1105 : TYPE(kpoint_operator_rs_route_type), INTENT(INOUT) :: route
1106 :
1107 79284 : INTEGER :: c0, i, nc, nr, r0, rows(0:route%matrix_struct%nrow_global), &
1108 39642 : cols(0:route%matrix_struct%ncol_global)
1109 :
1110 1004388 : rows(:) = 0
1111 1004388 : cols(:) = 0
1112 849280 : DO i = 1, SIZE(route%matrix_struct%row_indices)
1113 809638 : r0 = route%matrix_struct%row_indices(i)
1114 849280 : IF (r0 >= 1 .AND. r0 <= UBOUND(rows, 1)) rows(r0) = 1
1115 : END DO
1116 964746 : DO i = 1, SIZE(route%matrix_struct%col_indices)
1117 925104 : c0 = route%matrix_struct%col_indices(i)
1118 964746 : IF (c0 >= 1 .AND. c0 <= UBOUND(cols, 1)) cols(c0) = 1
1119 : END DO
1120 964746 : DO i = 1, UBOUND(rows, 1)
1121 964746 : rows(i) = rows(i) + rows(i - 1)
1122 : END DO
1123 964746 : DO i = 1, UBOUND(cols, 1)
1124 964746 : cols(i) = cols(i) + cols(i - 1)
1125 : END DO
1126 39642 : route%nrow = rows(UBOUND(rows, 1))
1127 39642 : route%ncol = cols(UBOUND(cols, 1))
1128 118926 : ALLOCATE (route%tile(4, 2, SIZE(route%entry, 2)))
1129 14666955 : DO i = 1, SIZE(route%entry, 2)
1130 14627313 : r0 = route%entry(1, i)
1131 14627313 : c0 = route%entry(2, i)
1132 14627313 : nr = route%entry(3, i)
1133 14627313 : nc = route%entry(4, i)
1134 73136565 : route%tile(:, 1, i) = [rows(r0 - 1) + 1, rows(r0 + nr - 1), cols(c0 - 1) + 1, cols(c0 + nc - 1)]
1135 73176207 : route%tile(:, 2, i) = [rows(c0 - 1) + 1, rows(c0 + nc - 1), cols(r0 - 1) + 1, cols(r0 + nr - 1)]
1136 : END DO
1137 39642 : END SUBROUTINE kpoint_operator_rs_tile_map
1138 :
1139 : ! **************************************************************************************************
1140 : !> \brief Find group-local FM owners intersecting a stored block or its mirror.
1141 : !> \param route source AO block layout and neighbor traversal
1142 : !> \param fms target FM layout within each group
1143 : !> \param i traversal entry
1144 : !> \param needed recipient group-local ranks, without duplicate mirror sends
1145 : ! **************************************************************************************************
1146 16812066 : SUBROUTINE kpoint_operator_rs_owners(route, fms, i, needed)
1147 :
1148 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
1149 : TYPE(cp_fm_struct_type), INTENT(IN) :: fms
1150 : INTEGER, INTENT(IN) :: i
1151 : LOGICAL, INTENT(OUT) :: needed(0:)
1152 :
1153 : INTEGER :: c0, j, nc, nr, pc, pr, r0
1154 33624132 : LOGICAL :: rows(2, 0:fms%context%num_pe(1) - 1), cols(2, 0:fms%context%num_pe(2) - 1)
1155 :
1156 16812066 : r0 = route%traversal%row_offset(i)
1157 16812066 : c0 = route%traversal%col_offset(i)
1158 16812066 : nr = route%row_size(route%traversal%row(i))
1159 16812066 : nc = route%col_size(route%traversal%col(i))
1160 90641130 : rows(:, :) = .FALSE.
1161 67248264 : cols(:, :) = .FALSE.
1162 123750816 : DO j = r0, r0 + nr - 1
1163 106938750 : rows(1, fms%g2p_row(j)) = .TRUE.
1164 123750816 : cols(2, fms%g2p_col(j)) = .TRUE.
1165 : END DO
1166 120217624 : DO j = c0, c0 + nc - 1
1167 103405558 : rows(2, fms%g2p_row(j)) = .TRUE.
1168 120217624 : cols(1, fms%g2p_col(j)) = .TRUE.
1169 : END DO
1170 28342978 : IF (route%traversal%row(i) == route%traversal%col(i)) rows(2, :) = .FALSE.
1171 41421754 : needed(:) = .FALSE.
1172 33624132 : DO pc = 0, fms%context%num_pe(2) - 1
1173 58233820 : DO pr = 0, fms%context%num_pe(1) - 1
1174 : needed(fms%context%blacs2mpi(pr, pc)) = &
1175 45594525 : (rows(1, pr) .AND. cols(1, pc)) .OR. (rows(2, pr) .AND. cols(2, pc))
1176 : END DO
1177 : END DO
1178 :
1179 16812066 : END SUBROUTINE kpoint_operator_rs_owners
1180 :
1181 : ! **************************************************************************************************
1182 : !> \brief Check the source AO layout, not its values or sparse block presence.
1183 : !> \param route fixed context route
1184 : !> \param source current operator images
1185 : !> \param matrix_row row to capture
1186 : !> \return whether the row has the route's image count and AO layout
1187 : ! **************************************************************************************************
1188 49460 : FUNCTION kpoint_operator_rs_compatible(route, source, matrix_row) RESULT(compatible)
1189 :
1190 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
1191 : TYPE(dbcsr_p_type), POINTER :: source(:, :)
1192 : INTEGER, INTENT(IN) :: matrix_row
1193 : LOGICAL :: compatible
1194 :
1195 : INTEGER :: image
1196 49460 : INTEGER, POINTER :: col_size(:), row_size(:)
1197 :
1198 49460 : compatible = .FALSE.
1199 49460 : IF (SIZE(source, 2) /= route%nimage) RETURN
1200 49460 : IF (matrix_row < 1 .OR. matrix_row > SIZE(source, 1)) RETURN
1201 2882174 : DO image = 1, SIZE(source, 2)
1202 2832714 : CALL dbcsr_get_info(source(matrix_row, image)%matrix, row_blk_size=row_size, col_blk_size=col_size)
1203 2832714 : IF (SIZE(row_size) /= SIZE(route%row_size) .OR. SIZE(col_size) /= SIZE(route%col_size)) RETURN
1204 27722122 : IF (ANY(row_size /= route%row_size) .OR. ANY(col_size /= route%col_size)) RETURN
1205 : END DO
1206 49460 : compatible = .TRUE.
1207 :
1208 49460 : END FUNCTION kpoint_operator_rs_compatible
1209 :
1210 : ! **************************************************************************************************
1211 : !> \brief Capture current values using cached routing, with no per-k communication.
1212 : !> \param source current real-space operator images
1213 : !> \param matrix_row row to snapshot
1214 : !> \param route context-owned metadata, counts and reusable packing buffer
1215 : !> \param kpoints initialized host parallel environment
1216 : !> \param snapshot immutable tile-local values
1217 : ! **************************************************************************************************
1218 49460 : SUBROUTINE kpoint_operator_rs_capture(source, matrix_row, route, kpoints, snapshot)
1219 :
1220 : TYPE(dbcsr_p_type), POINTER :: source(:, :)
1221 : INTEGER, INTENT(IN) :: matrix_row
1222 : TYPE(kpoint_operator_rs_route_type), INTENT(INOUT) :: route
1223 : TYPE(kpoint_type), POINTER :: kpoints
1224 : TYPE(kpoint_operator_rs_type), INTENT(OUT) :: snapshot
1225 :
1226 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_rs_capture'
1227 :
1228 : INTEGER :: col, handle, i, iblock, ncol, nrow, pos, &
1229 : row
1230 : LOGICAL :: found
1231 49460 : REAL(KIND=dp), POINTER :: block(:, :)
1232 : TYPE(mp_para_env_type), POINTER :: inter, world
1233 :
1234 49460 : CALL timeset(routineN, handle)
1235 49460 : world => kpoints%blacs_env_all%para_env
1236 49460 : inter => kpoints%para_env_inter_kp
1237 19403947 : ALLOCATE (snapshot%blocks(SIZE(route%entry, 2)), snapshot%nonzero(SIZE(route%entry, 2)))
1238 49460 : snapshot%source => source
1239 49460 : snapshot%matrix_row = matrix_row
1240 49460 : IF (world%num_pe == 1) THEN
1241 : ! No value copy on a single world rank. Lookups happen before workers start.
1242 0 : DO iblock = 1, SIZE(route%send_entry)
1243 0 : i = route%send_entry(iblock)
1244 : CALL dbcsr_get_readonly_block_p(source(matrix_row, route%traversal%image(i))%matrix, &
1245 0 : route%traversal%row(i), route%traversal%col(i), block, found)
1246 0 : snapshot%nonzero(iblock) = found
1247 0 : IF (found) THEN
1248 0 : CPASSERT(SIZE(block, 1) == route%entry(3, iblock))
1249 0 : CPASSERT(SIZE(block, 2) == route%entry(4, iblock))
1250 0 : snapshot%blocks(iblock)%data => block
1251 : END IF
1252 : END DO
1253 0 : CALL timestop(handle)
1254 0 : RETURN
1255 : END IF
1256 148380 : ALLOCATE (snapshot%values(route%nvalue))
1257 742191540 : route%send_values(:) = 0.0_dp
1258 : !$OMP PARALLEL DO DEFAULT(NONE) SHARED(source, matrix_row, route) &
1259 49460 : !$OMP PRIVATE(iblock, i, block, found, nrow, ncol, pos, row, col) SCHEDULE(STATIC)
1260 : DO iblock = 1, SIZE(route%send_entry)
1261 : i = route%send_entry(iblock)
1262 : CALL dbcsr_get_readonly_block_p(source(matrix_row, route%traversal%image(i))%matrix, &
1263 : route%traversal%row(i), route%traversal%col(i), block, found)
1264 : IF (.NOT. found) CYCLE
1265 : nrow = route%row_size(route%traversal%row(i))
1266 : ncol = route%col_size(route%traversal%col(i))
1267 : CPASSERT(SIZE(block, 1) == nrow .AND. SIZE(block, 2) == ncol)
1268 : pos = route%send_first(iblock)
1269 : DO col = 1, ncol
1270 : DO row = 1, nrow
1271 : route%send_values(pos) = block(row, col)
1272 : pos = pos + 1
1273 : END DO
1274 : END DO
1275 : END DO
1276 : !$OMP END PARALLEL DO
1277 49460 : IF (kpoints%para_env_kp%num_pe == 1) THEN
1278 170830 : CALL world%allgatherv(route%send_values, snapshot%values, route%rcount(2, :), route%rdisp(2, :))
1279 : ELSE
1280 : CALL world%alltoall(route%send_values, route%scount(2, :), route%sdisp(2, :), &
1281 137646 : snapshot%values, route%rcount(2, :), route%rdisp(2, :))
1282 15294 : CALL inter%bcast(snapshot%values, 0)
1283 : END IF
1284 19206107 : DO iblock = 1, SIZE(route%entry, 2)
1285 84110427 : snapshot%nonzero(iblock) = ANY(snapshot%values(route%first(iblock):route%first(iblock + 1) - 1) /= 0.0_dp)
1286 : snapshot%blocks(iblock)%data(1:route%entry(3, iblock), 1:route%entry(4, iblock)) => &
1287 19206107 : snapshot%values(route%first(iblock):route%first(iblock + 1) - 1)
1288 : END DO
1289 49460 : snapshot%source => source
1290 49460 : snapshot%matrix_row = matrix_row
1291 49460 : CALL timestop(handle)
1292 :
1293 98920 : END SUBROUTINE kpoint_operator_rs_capture
1294 :
1295 : ! **************************************************************************************************
1296 : !> \brief Deliver one prepared operator, building a bounded window of local k points on a miss.
1297 : !> The callers retain their round/spin/slot ordering. Each row has its own
1298 : !> window so interleaved spins and H/S/T do not evict one another.
1299 : !> Only prepared values are read ahead; transfer and grid paths are unchanged.
1300 : !> \param ctx serialized operator context
1301 : !> \param irs captured row, possibly shared by spin-free instances
1302 : !> \param slot prepared operator slot
1303 : !> \param ik global k-point index owned by this group
1304 : !> \param TARGET optional complex output on the prepared distribution
1305 : !> \param fm_re optional real output instead of target; fm_im may be omitted for a real operator
1306 : !> \param fm_im optional imaginary output on the same distribution
1307 : ! **************************************************************************************************
1308 35758 : SUBROUTINE kpoint_operator_rs_build(ctx, irs, slot, ik, TARGET, fm_re, fm_im)
1309 :
1310 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1311 : INTEGER, INTENT(IN) :: irs, slot, ik
1312 : TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
1313 : TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1314 :
1315 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_rs_build'
1316 :
1317 : INTEGER :: first, handle, k, last, ncol, nrow, width
1318 : INTEGER(KIND=int_8) :: bytes_per_k
1319 : LOGICAL :: layout_ok
1320 35758 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: xkp
1321 :
1322 35758 : CALL timeset(routineN, handle)
1323 35758 : IF (PRESENT(TARGET)) THEN
1324 34380 : layout_ok = cp_fm_struct_equivalent(target%matrix_struct, ctx%rs_route%matrix_struct)
1325 : ELSE
1326 1378 : CPASSERT(PRESENT(fm_re))
1327 1378 : layout_ok = cp_fm_struct_equivalent(fm_re%matrix_struct, ctx%rs_route%matrix_struct)
1328 1378 : IF (PRESENT(fm_im)) THEN
1329 1080 : layout_ok = layout_ok .AND. cp_fm_struct_equivalent(fm_im%matrix_struct, ctx%rs_route%matrix_struct)
1330 : END IF
1331 : END IF
1332 34678 : CPASSERT(layout_ok)
1333 : ! Avoid nested component array sections in GCC 16 bounds checks.
1334 35758 : xkp => ctx%kpoints%xkp
1335 35758 : first = ctx%kpoints%kp_range(1)
1336 35758 : last = ctx%kpoints%kp_range(2)
1337 35758 : CPASSERT(ik >= first .AND. ik <= last)
1338 35758 : nrow = ctx%rs_route%nrow
1339 35758 : ncol = ctx%rs_route%ncol
1340 35758 : bytes_per_k = 16_int_8*(INT(nrow, int_8)*ncol + SIZE(ctx%rs_route%entry, 2, KIND=int_8))
1341 : width = MIN(rs_batch_size, last - first + 1, &
1342 35758 : INT(rs_batch_bytes/(INT(ctx%nspin*kpoint_num_slots, int_8)*MAX(1_int_8, bytes_per_k))))
1343 35758 : IF (ALLOCATED(ctx%rs_batch(irs, slot)%data)) width = SIZE(ctx%rs_batch(irs, slot)%data, 1)
1344 35758 : IF ((width < 2 .AND. .NOT. ALLOCATED(ctx%rs_batch(irs, slot)%data)) .OR. nrow == 0 .OR. ncol == 0) THEN
1345 : ! No extra dense tile for a large/empty layout or a single local k point.
1346 : CALL kpoint_operator_rs_build_batch(ctx%rs(irs, slot), ctx%rs_route, xkp(:, ik:ik), &
1347 11431 : TARGET=TARGET, fm_re=fm_re, fm_im=fm_im)
1348 : ELSE
1349 : ASSOCIATE (work => ctx%rs_batch(irs, slot))
1350 26715 : IF (.NOT. ALLOCATED(work%data)) ALLOCATE (work%data(width, nrow, ncol))
1351 24327 : IF (work%first == 0 .OR. ik < work%first .OR. ik >= work%first + width) THEN
1352 5687 : work%first = first + ((ik - first)/width)*width
1353 : CALL kpoint_operator_rs_build_batch(ctx%rs(irs, slot), ctx%rs_route, &
1354 5687 : xkp(:, work%first:MIN(work%first + width - 1, last)), batch=work%data)
1355 : END IF
1356 24327 : k = ik - work%first + 1
1357 : ! Deliver a work copy: eigensolvers may overwrite H/S without corrupting the window.
1358 48654 : CALL kpoint_operator_rs_copy(work%data(k, :, :), TARGET, fm_re, fm_im)
1359 : END ASSOCIATE
1360 : END IF
1361 35758 : IF (ctx%kpoints%use_real_wfn) THEN
1362 298 : IF (PRESENT(TARGET)) target%local_data(:, :) = CMPLX(REAL(target%local_data, dp), 0.0_dp, dp)
1363 298 : IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
1364 : END IF
1365 35758 : CALL timestop(handle)
1366 :
1367 35758 : END SUBROUTINE kpoint_operator_rs_build
1368 :
1369 : ! **************************************************************************************************
1370 : !> \brief Accumulate several k points in one route/value traversal, without MPI or BLACS.
1371 : !> The leading batch dimension makes updates to different k points contiguous.
1372 : !> Columns remain independent OpenMP work items; each target retains the
1373 : !> original normal/mirrored contribution order. Phase scratch is call-local.
1374 : !> The singleton fallback writes directly to the supplied CFM or FM pair.
1375 : !> \param snapshot immutable real-space contributions
1376 : !> \param route structural block routing and target tile map
1377 : !> \param xkp actual coordinates, one column per k point
1378 : !> \param batch optional packed complex output (k point, local row, local column)
1379 : !> \param TARGET optional single complex output instead of batch
1380 : !> \param fm_re optional single real output instead of batch/target
1381 : !> \param fm_im optional single imaginary output
1382 : !> \param phase_work optional private phase storage supplied by a worker
1383 : !> \param parallel_columns split columns across threads, disabled inside k-point workers
1384 : ! **************************************************************************************************
1385 20660 : SUBROUTINE kpoint_operator_rs_build_batch(snapshot, route, xkp, batch, TARGET, fm_re, fm_im, phase_work, parallel_columns)
1386 :
1387 : TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
1388 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
1389 : REAL(KIND=dp), INTENT(IN) :: xkp(:, :)
1390 : COMPLEX(KIND=dp), CONTIGUOUS, INTENT(INOUT), &
1391 : OPTIONAL, TARGET :: batch(:, :, :)
1392 : TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
1393 : TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1394 : COMPLEX(KIND=dp), INTENT(INOUT), OPTIONAL, TARGET :: phase_work(:, :)
1395 : LOGICAL, INTENT(IN), OPTIONAL :: parallel_columns
1396 :
1397 20660 : COMPLEX(KIND=dp), ALLOCATABLE, TARGET :: owned_phase(:, :)
1398 20660 : COMPLEX(KIND=dp), CONTIGUOUS, POINTER :: output(:, :, :)
1399 20660 : COMPLEX(KIND=dp), POINTER :: phase(:, :)
1400 : INTEGER :: first, iblock, k, last, nbatch, ncol, &
1401 : nthreads, previous, thread
1402 : LOGICAL :: parallel
1403 : REAL(KIND=dp) :: arg
1404 :
1405 20660 : nbatch = SIZE(xkp, 2)
1406 20660 : CPASSERT(nbatch >= 1 .AND. nbatch <= rs_batch_size .AND. SIZE(xkp, 1) == 3)
1407 82640 : CPASSERT(COUNT([PRESENT(batch), PRESENT(TARGET), PRESENT(fm_re)]) == 1)
1408 20660 : NULLIFY (output)
1409 20660 : IF (PRESENT(batch)) THEN
1410 5687 : CPASSERT(SIZE(batch, 1) >= nbatch)
1411 5687 : output => batch
1412 : ELSE
1413 14973 : CPASSERT(nbatch == 1)
1414 14973 : IF (PRESENT(TARGET)) THEN
1415 14383 : output(1:1, 1:SIZE(target%local_data, 1), 1:SIZE(target%local_data, 2)) => target%local_data
1416 : ELSE
1417 590 : CPASSERT(PRESENT(fm_re))
1418 1213655 : fm_re%local_data(:, :) = 0.0_dp
1419 106129 : IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
1420 : END IF
1421 : END IF
1422 20660 : IF (ASSOCIATED(output)) output(:, :, :) = z_zero
1423 20660 : IF (route%nrow == 0 .OR. route%ncol == 0) RETURN
1424 20660 : IF (PRESENT(phase_work)) THEN
1425 3542 : CPASSERT(SIZE(phase_work, 1) >= nbatch)
1426 3542 : CPASSERT(SIZE(phase_work, 2) == SIZE(route%entry, 2))
1427 3542 : phase => phase_work(1:nbatch, :)
1428 : ELSE
1429 68472 : ALLOCATE (owned_phase(nbatch, SIZE(route%entry, 2)))
1430 17118 : phase => owned_phase
1431 : END IF
1432 20660 : previous = 0
1433 9754413 : DO iblock = 1, SIZE(route%entry, 2)
1434 9733753 : IF (.NOT. snapshot%nonzero(iblock)) CYCLE
1435 : ! Neighboring blocks often share a lattice translation, including its symmetry sign.
1436 9714507 : IF (previous > 0) THEN
1437 10095442 : IF (ALL(route%entry(5:7, iblock) == route%entry(5:7, previous))) THEN
1438 81947 : phase(:, iblock) = phase(:, previous)
1439 : CYCLE
1440 : END IF
1441 : END IF
1442 29087036 : DO k = 1, nbatch
1443 77637172 : arg = twopi*DOT_PRODUCT(REAL(route%entry(5:7, iblock), dp), xkp(:, k))
1444 29087036 : phase(k, iblock) = CMPLX(COS(arg), SIN(arg), KIND=dp)
1445 : END DO
1446 39906 : previous = iblock
1447 : END DO
1448 20660 : ncol = route%ncol
1449 20660 : parallel = .TRUE.
1450 20660 : IF (PRESENT(parallel_columns)) parallel = parallel_columns
1451 20660 : nthreads = 1
1452 20660 : !$ IF (parallel) nthreads = MIN(omp_get_max_threads(), ncol)
1453 17118 : IF (nthreads > 1) THEN
1454 : !$OMP PARALLEL DEFAULT(NONE) NUM_THREADS(nthreads) SHARED(snapshot,route,phase,output,fm_re,fm_im,ncol) &
1455 : !$OMP PRIVATE(first,last,thread,nthreads)
1456 : thread = 0
1457 : nthreads = 1
1458 : !$ thread = omp_get_thread_num()
1459 : !$ nthreads = omp_get_num_threads()
1460 : first = ncol*thread/nthreads + 1
1461 : last = ncol*(thread + 1)/nthreads
1462 0 : CALL kpoint_operator_rs_accumulate(snapshot, route, phase, first, last, output, fm_re, fm_im)
1463 : !$OMP END PARALLEL
1464 : ELSE
1465 : ! Called inside a k-point worker: no nested parallel region, timer or allocation.
1466 20660 : CALL kpoint_operator_rs_accumulate(snapshot, route, phase, 1, ncol, output, fm_re, fm_im)
1467 : END IF
1468 :
1469 41320 : END SUBROUTINE kpoint_operator_rs_build_batch
1470 :
1471 : ! **************************************************************************************************
1472 : !> \brief Shared borrowed/packed phase-sum kernel; each caller owns a disjoint column interval.
1473 : !> \param snapshot read-only block views
1474 : !> \param route structural layout
1475 : !> \param phase precomputed phases (k, block)
1476 : !> \param first first local column owned by this call
1477 : !> \param last last local column owned by this call
1478 : !> \param output packed complex output, or disassociated for the FM pair
1479 : !> \param fm_re real output
1480 : !> \param fm_im imaginary output
1481 : ! **************************************************************************************************
1482 20660 : SUBROUTINE kpoint_operator_rs_accumulate(snapshot, route, phase, first, last, output, fm_re, fm_im)
1483 : TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
1484 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
1485 : COMPLEX(KIND=dp), INTENT(IN) :: phase(:, :)
1486 : INTEGER, INTENT(IN) :: first, last
1487 : COMPLEX(KIND=dp), CONTIGUOUS, INTENT(IN), POINTER :: output(:, :, :)
1488 : TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1489 :
1490 : COMPLEX(KIND=dp) :: factors(rs_batch_size), value
1491 20660 : COMPLEX(KIND=dp), CONTIGUOUS, POINTER :: full(:, :)
1492 : INTEGER :: c0, col, hi, iblock, j, k, lo, nc, nr, &
1493 : part, r0, row, source_col, source_row
1494 : LOGICAL :: full_fm
1495 : REAL(KIND=dp) :: element, phase_im, phase_re
1496 :
1497 0 : IF (first > last) RETURN
1498 20660 : NULLIFY (full)
1499 : full_fm = .NOT. ASSOCIATED(output) .AND. &
1500 : route%nrow == route%matrix_struct%nrow_global .AND. &
1501 20660 : route%ncol == route%matrix_struct%ncol_global
1502 20660 : IF (ASSOCIATED(output)) THEN
1503 : IF (SIZE(output, 1) == 1 .AND. &
1504 20070 : route%nrow == route%matrix_struct%nrow_global .AND. &
1505 : route%ncol == route%matrix_struct%ncol_global) THEN
1506 3719 : full(1:SIZE(output, 2), 1:SIZE(output, 3)) => output
1507 : END IF
1508 : END IF
1509 9754413 : DO iblock = 1, SIZE(route%entry, 2)
1510 9733753 : IF (.NOT. snapshot%nonzero(iblock)) CYCLE
1511 9714507 : r0 = route%entry(1, iblock)
1512 9714507 : c0 = route%entry(2, iblock)
1513 9714507 : nr = route%entry(3, iblock)
1514 9714507 : nc = route%entry(4, iblock)
1515 9714507 : IF (full_fm) THEN
1516 : ! Single-rank split-FM caches have contiguous global rows and columns.
1517 : ! Keep each block's direct/mirrored order, without per-element index lookup.
1518 7402 : phase_re = REAL(phase(1, iblock), dp)
1519 7402 : phase_im = AIMAG(phase(1, iblock))
1520 7402 : lo = MAX(first, c0)
1521 7402 : hi = MIN(last, c0 + nc - 1)
1522 103628 : DO j = lo, hi
1523 103628 : !$OMP SIMD PRIVATE(element)
1524 : DO row = 1, nr
1525 1250938 : element = snapshot%blocks(iblock)%data(row, j - c0 + 1)
1526 1250938 : fm_re%local_data(r0 + row - 1, j) = fm_re%local_data(r0 + row - 1, j) + phase_re*element
1527 1250938 : IF (PRESENT(fm_im)) THEN
1528 1250938 : fm_im%local_data(r0 + row - 1, j) = fm_im%local_data(r0 + row - 1, j) + phase_im*element
1529 : END IF
1530 : END DO
1531 : END DO
1532 7402 : IF (route%entry(8, iblock) /= 0) THEN
1533 3190 : lo = MAX(first, r0)
1534 3190 : hi = MIN(last, r0 + nr - 1)
1535 44660 : DO j = lo, hi
1536 44660 : !$OMP SIMD PRIVATE(element)
1537 : DO row = 1, nc
1538 539110 : element = snapshot%blocks(iblock)%data(j - r0 + 1, row)
1539 539110 : fm_re%local_data(c0 + row - 1, j) = fm_re%local_data(c0 + row - 1, j) + phase_re*element
1540 539110 : IF (PRESENT(fm_im)) THEN
1541 539110 : fm_im%local_data(c0 + row - 1, j) = fm_im%local_data(c0 + row - 1, j) - phase_im*element
1542 : END IF
1543 : END DO
1544 : END DO
1545 : END IF
1546 : CYCLE
1547 : END IF
1548 9707105 : IF (ASSOCIATED(full)) THEN
1549 : ! Update whole contiguous blocks, clipped to this worker's columns.
1550 1763319 : lo = MAX(first, c0)
1551 1763319 : hi = MIN(last, c0 + nc - 1)
1552 1763319 : IF (lo <= hi) THEN
1553 : full(r0:r0 + nr - 1, lo:hi) = full(r0:r0 + nr - 1, lo:hi) + &
1554 147298673 : phase(1, iblock)*snapshot%blocks(iblock)%data(:, lo - c0 + 1:hi - c0 + 1)
1555 : END IF
1556 1763319 : IF (route%entry(8, iblock) /= 0) THEN
1557 1067738 : lo = MAX(first, r0)
1558 1067738 : hi = MIN(last, r0 + nr - 1)
1559 1067738 : IF (lo <= hi) THEN
1560 : full(c0:c0 + nc - 1, lo:hi) = full(c0:c0 + nc - 1, lo:hi) + &
1561 : CONJG(phase(1, iblock))* &
1562 91307294 : TRANSPOSE(snapshot%blocks(iblock)%data(lo - r0 + 1:hi - r0 + 1, :))
1563 : END IF
1564 : END IF
1565 : CYCLE
1566 : END IF
1567 23852018 : DO part = 1, 2
1568 15887572 : IF (part == 2 .AND. route%entry(8, iblock) == 0) CYCLE
1569 14130865 : lo = MAX(first, route%tile(3, part, iblock))
1570 14130865 : hi = MIN(last, route%tile(4, part, iblock))
1571 14130865 : IF (lo > hi) CYCLE
1572 45641866 : factors(1:SIZE(phase, 1)) = phase(:, iblock)
1573 27958111 : IF (part == 2) factors(1:SIZE(phase, 1)) = CONJG(factors(1:SIZE(phase, 1)))
1574 109776484 : DO j = lo, hi
1575 87682587 : col = route%matrix_struct%col_indices(j)
1576 615599898 : DO row = route%tile(1, part, iblock), route%tile(2, part, iblock)
1577 512029739 : source_row = route%matrix_struct%row_indices(row)
1578 512029739 : source_col = col
1579 512029739 : IF (part == 2) THEN
1580 213867430 : source_col = source_row
1581 213867430 : source_row = col
1582 : END IF
1583 512029739 : element = snapshot%blocks(iblock)%data(source_row - r0 + 1, source_col - c0 + 1)
1584 599712326 : IF (ASSOCIATED(output)) THEN
1585 : !$OMP SIMD
1586 : DO k = 1, SIZE(phase, 1)
1587 1038477198 : output(k, row, j) = output(k, row, j) + factors(k)*element
1588 : END DO
1589 : ELSE
1590 135186823 : value = factors(1)*element
1591 135186823 : fm_re%local_data(row, j) = fm_re%local_data(row, j) + REAL(value, dp)
1592 135186823 : IF (PRESENT(fm_im)) fm_im%local_data(row, j) = fm_im%local_data(row, j) + AIMAG(value)
1593 : END IF
1594 : END DO
1595 : END DO
1596 : END DO
1597 : END DO
1598 20660 : END SUBROUTINE kpoint_operator_rs_accumulate
1599 :
1600 : ! **************************************************************************************************
1601 : !> \brief Size per-worker phase scratch for direct assembly into a solver matrix.
1602 : !> \param ctx prepared context
1603 : !> \param nentry number of route entries
1604 : ! **************************************************************************************************
1605 1492 : SUBROUTINE kpoint_operator_work_size(ctx, nentry)
1606 : TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
1607 : INTEGER, INTENT(OUT) :: nentry
1608 :
1609 1492 : CPASSERT(ALLOCATED(ctx%rs_route))
1610 1492 : nentry = SIZE(ctx%rs_route%entry, 2)
1611 1492 : END SUBROUTINE kpoint_operator_work_size
1612 :
1613 : ! **************************************************************************************************
1614 : !> \brief Evaluate an immutable prepared operator with exclusively caller-owned output and scratch.
1615 : !> Prepare/refresh/release remain collective and serialized. Evaluate contains no MPI,
1616 : !> DBCSR lookup, window update or nested OpenMP region. Source blocks must stay unchanged.
1617 : !> \param ctx prepared read-only context
1618 : !> \param ispin instance spin
1619 : !> \param slot operator slot
1620 : !> \param ik global k index owned by this group
1621 : !> \param phase private phase scratch, sized with work_size
1622 : !> \param TARGET optional private complex singleton output
1623 : !> \param fm_re optional private real output instead of TARGET
1624 : !> \param fm_im optional private imaginary output, accompanying fm_re for a complex operator
1625 : ! **************************************************************************************************
1626 3542 : SUBROUTINE kpoint_operator_evaluate(ctx, ispin, slot, ik, phase, TARGET, fm_re, fm_im)
1627 : TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
1628 : INTEGER, INTENT(IN) :: ispin, slot, ik
1629 : TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
1630 : COMPLEX(KIND=dp), INTENT(INOUT), TARGET :: phase(:, :)
1631 : TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1632 :
1633 3542 : REAL(KIND=dp), POINTER :: xkp(:, :)
1634 : LOGICAL :: cached
1635 : INTEGER :: k, row
1636 :
1637 0 : CPASSERT(ik >= ctx%kpoints%kp_range(1))
1638 3542 : CPASSERT(ik <= ctx%kpoints%kp_range(2))
1639 3542 : row = ispin
1640 3542 : IF (.NOT. ASSOCIATED(ctx%rs(row, slot)%source)) row = 1
1641 3542 : CPASSERT(ASSOCIATED(ctx%rs(row, slot)%source))
1642 3542 : xkp => ctx%kpoints%xkp
1643 3542 : k = ik - ctx%rs_batch(row, slot)%first + 1
1644 3542 : cached = .FALSE.
1645 3542 : IF (ALLOCATED(ctx%rs_batch(row, slot)%data)) cached = k >= 1 .AND. k <= SIZE(ctx%rs_batch(row, slot)%data, 1)
1646 : IF (cached) THEN
1647 0 : CALL kpoint_operator_rs_copy(ctx%rs_batch(row, slot)%data(k, :, :), TARGET, fm_re, fm_im)
1648 : ELSE
1649 : CALL kpoint_operator_rs_build_batch(ctx%rs(row, slot), ctx%rs_route, xkp(:, ik:ik), &
1650 : TARGET=TARGET, fm_re=fm_re, fm_im=fm_im, &
1651 3542 : phase_work=phase, parallel_columns=.FALSE.)
1652 : END IF
1653 3542 : IF (ctx%kpoints%use_real_wfn) THEN
1654 0 : IF (PRESENT(TARGET)) target%local_data(:, :) = CMPLX(REAL(target%local_data, dp), 0.0_dp, dp)
1655 0 : IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
1656 : END IF
1657 3542 : END SUBROUTINE kpoint_operator_evaluate
1658 :
1659 : ! **************************************************************************************************
1660 : !> \brief Start one operator for one global round and spin. The service walks
1661 : !> all groups with a k point in this round; groups without one still
1662 : !> participate in the transfers. Every rank must iterate through the
1663 : !> maximum local k-point count. Finish every started instance with
1664 : !> kpoint_operator_finish. The first start also creates the slot
1665 : !> bookkeeping of the context.
1666 : !> Collective over the environment-wide communicator: every rank must
1667 : !> call the same sequence of service calls.
1668 : !> \param ctx context created by kpoint_operator_context_create
1669 : !> \param ikp global round, up to the maximum k-point count across groups
1670 : !> \param ispin instance spin: keys the bookkeeping
1671 : !> \param rsmat real-space image matrices of the operator. A prepared operator
1672 : !> batch uses the values captured by prepare; otherwise start reads
1673 : !> the current values. The matching finish never reads rsmat
1674 : !> \param slot instance slot, optionally backed by a prepared snapshot
1675 : !> \param grid prepared reciprocal-grid cache of the caller
1676 : !> \param use_grid extract from grid instead of a direct phase sum
1677 : !> \param matrix_row row of rsmat to transform. Spin-free operators such
1678 : !> as S and T must pass kpoint_spin_free: their first matrix
1679 : !> dimension counts derivatives, not spins. A spin index is out of
1680 : !> bounds in plain SCF, but with derivatives stored it reads another
1681 : !> valid row silently
1682 : ! **************************************************************************************************
1683 30564 : SUBROUTINE kpoint_operator_start(ctx, ikp, ispin, rsmat, slot, grid, use_grid, matrix_row)
1684 :
1685 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1686 : INTEGER, INTENT(IN) :: ikp, ispin
1687 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1688 : INTEGER, INTENT(IN) :: slot
1689 : TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
1690 : LOGICAL, INTENT(IN) :: use_grid
1691 : INTEGER, INTENT(IN) :: matrix_row
1692 :
1693 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_start'
1694 :
1695 : INTEGER :: handle, igroup, ik, irs, kpmax
1696 30564 : INTEGER, DIMENSION(:, :), POINTER :: kp_dist
1697 :
1698 30564 : CALL timeset(routineN, handle)
1699 :
1700 30564 : CPASSERT(ASSOCIATED(ctx%kpoints))
1701 30564 : CPASSERT(ASSOCIATED(ctx%kpoints%sab_nl))
1702 : ! The instance spin is bounded by the spin space declared at creation;
1703 : ! the matrix row by this operator's own rows. The two bounds differ for
1704 : ! spin-free operators such as S, whose first dimension counts
1705 : ! real-space blocks.
1706 30564 : CPASSERT(ispin >= 1 .AND. ispin <= ctx%nspin)
1707 30564 : CPASSERT(matrix_row >= 1 .AND. matrix_row <= SIZE(rsmat, 1))
1708 62066 : kpmax = MAXVAL(ctx%kpoints%kp_dist(2, :) - ctx%kpoints%kp_dist(1, :) + 1)
1709 30564 : CPASSERT(ikp >= 1 .AND. ikp <= kpmax)
1710 30564 : CPASSERT(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
1711 : ! State is per round; allocate expensive copy descriptors only when
1712 : ! an actual distributed transfer is launched.
1713 30564 : IF (.NOT. ALLOCATED(ctx%status)) THEN
1714 56850 : ALLOCATE (ctx%status(kpmax, ctx%nspin, kpoint_num_slots))
1715 164190 : ctx%status(:, :, :) = kpoint_slot_idle
1716 : END IF
1717 30564 : CPASSERT(ctx%status(ikp, ispin, slot) == kpoint_slot_idle)
1718 30564 : IF (.NOT. use_grid) THEN
1719 30564 : irs = kpoint_operator_rs_match(ctx, rsmat, ispin, slot, matrix_row)
1720 30564 : IF (irs > 0) THEN
1721 30524 : CPASSERT(ASSOCIATED(ctx%work_sab, ctx%kpoints%sab_nl))
1722 30524 : ctx%status(ikp, ispin, slot) = kpoint_slot_local
1723 30524 : CALL timestop(handle)
1724 30524 : RETURN
1725 : END IF
1726 : END IF
1727 40 : CALL kpoint_operator_ensure(ctx, rsmat)
1728 :
1729 : ! Only the transfer backend needs copy descriptors. A prepared-only
1730 : ! or get-only context never carries them. Size by the declared spin
1731 : ! space, not by the possibly spin-free operator's row count.
1732 40 : IF (.NOT. ALLOCATED(ctx%info)) THEN
1733 : ! One copy_info per in-flight part: start takes it INTENT(OUT),
1734 : ! so two concurrent starts cannot share one instance.
1735 0 : ALLOCATE (ctx%info(SIZE(ctx%kpoints%xkp, 2), ctx%nspin, &
1736 380 : kpoint_num_slots, MERGE(1, 2, ctx%kpoints%use_real_wfn)))
1737 : END IF
1738 :
1739 40 : kp_dist => ctx%kpoints%kp_dist
1740 80 : DO igroup = 1, ctx%kpoints%nkp_groups
1741 40 : ik = kp_dist(1, igroup) + ikp - 1
1742 40 : IF (ik > kp_dist(2, igroup)) CYCLE
1743 40 : CALL kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid, use_grid)
1744 : CALL kpoint_operator_launch(ctx%kpoints, ik, ctx%fmwork, ctx%fmlocal, &
1745 80 : ctx%info(ik, ispin, slot, :))
1746 : END DO
1747 40 : ctx%status(ikp, ispin, slot) = kpoint_slot_in_flight
1748 :
1749 40 : CALL timestop(handle)
1750 :
1751 30564 : END SUBROUTINE kpoint_operator_start
1752 :
1753 : ! **************************************************************************************************
1754 : !> \brief Finish one started instance: build its prepared group-local operator,
1755 : !> or merge received parts into cmat (real part first, overwrite then add)
1756 : !> and reclaim the send side on every group.
1757 : !> The slot returns to idle.
1758 : !> Collective over the environment-wide communicator: every rank must
1759 : !> call the same sequence of service calls.
1760 : !> \param ctx context that started the instance
1761 : !> \param ikp global round of the start call
1762 : !> \param ispin spin component of the start call
1763 : !> \param slot slot of the start call; the slot must be in flight
1764 : !> \param cmat caller-allocated complex matrix that receives the operator.
1765 : !> It must share the distribution of the k-point group. Unchanged on
1766 : !> groups without a k point in this round
1767 : !> \param fm_re optional real FM output instead of cmat for real wavefunctions, on the group layout
1768 : ! **************************************************************************************************
1769 30564 : SUBROUTINE kpoint_operator_finish(ctx, ikp, ispin, slot, cmat, fm_re)
1770 :
1771 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1772 : INTEGER, INTENT(IN) :: ikp, ispin, slot
1773 : TYPE(cp_cfm_type), OPTIONAL :: cmat
1774 : TYPE(cp_fm_type), OPTIONAL :: fm_re
1775 :
1776 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_finish'
1777 :
1778 : INTEGER :: handle, igroup, ik, irs, kpmax
1779 : INTEGER, DIMENSION(2) :: kp_range
1780 30564 : INTEGER, DIMENSION(:, :), POINTER :: kp_dist
1781 : LOGICAL :: layout_ok
1782 :
1783 30564 : CALL timeset(routineN, handle)
1784 :
1785 30564 : CPASSERT(ASSOCIATED(ctx%kpoints))
1786 30564 : CPASSERT(ALLOCATED(ctx%status))
1787 30564 : kpmax = SIZE(ctx%status, 1)
1788 30564 : CPASSERT(ikp >= 1 .AND. ikp <= kpmax)
1789 30564 : CPASSERT(ispin >= 1 .AND. ispin <= ctx%nspin)
1790 30564 : CPASSERT(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
1791 30564 : CPASSERT(ctx%status(ikp, ispin, slot) /= kpoint_slot_idle)
1792 : ! The merge writes cmat in place. A foreign distribution would fail
1793 : ! inside the BLACS calls, far from this argument.
1794 :
1795 30564 : CPASSERT(PRESENT(cmat) .NEQV. PRESENT(fm_re))
1796 30564 : CPASSERT(.NOT. PRESENT(fm_re) .OR. ctx%kpoints%use_real_wfn)
1797 91692 : kp_range = ctx%kpoints%kp_range
1798 30564 : IF (ctx%status(ikp, ispin, slot) == kpoint_slot_local) THEN
1799 30524 : irs = ispin
1800 30524 : IF (.NOT. ASSOCIATED(ctx%rs(irs, slot)%source)) irs = 1
1801 30524 : ik = kp_range(1) + ikp - 1
1802 30524 : IF (ik <= kp_range(2)) THEN
1803 30520 : CALL kpoint_operator_rs_build(ctx, irs, slot, ik, cmat, fm_re=fm_re)
1804 : END IF
1805 30524 : ctx%status(ikp, ispin, slot) = kpoint_slot_idle
1806 30524 : CALL timestop(handle)
1807 30524 : RETURN
1808 : END IF
1809 40 : CPASSERT(ctx%status(ikp, ispin, slot) == kpoint_slot_in_flight)
1810 40 : IF (PRESENT(cmat)) THEN
1811 24 : layout_ok = cp_fm_struct_equivalent(cmat%matrix_struct, ctx%fmlocal%matrix_struct)
1812 : ELSE
1813 16 : layout_ok = cp_fm_struct_equivalent(fm_re%matrix_struct, ctx%fmlocal%matrix_struct)
1814 : END IF
1815 40 : CPASSERT(layout_ok)
1816 40 : kp_dist => ctx%kpoints%kp_dist
1817 80 : DO igroup = 1, ctx%kpoints%nkp_groups
1818 40 : ik = kp_dist(1, igroup) + ikp - 1
1819 40 : IF (ik > kp_dist(2, igroup)) CYCLE
1820 : CALL kpoint_operator_receive(ctx%kpoints, ctx%fmlocal, ik, ctx%info(ik, ispin, slot, :), &
1821 80 : cmat=cmat, fm_re=fm_re)
1822 : END DO
1823 40 : ctx%status(ikp, ispin, slot) = kpoint_slot_idle
1824 :
1825 40 : CALL timestop(handle)
1826 :
1827 30564 : END SUBROUTINE kpoint_operator_finish
1828 :
1829 : ! **************************************************************************************************
1830 : !> \brief Assemble one operator for one global round and spin in one call:
1831 : !> run the single-group transfer of the service for every group and
1832 : !> finish it immediately, through transfer states local to this call.
1833 : !> The persistent slot machinery stays with start and finish.
1834 : !> Exactly one output flavor must be present: cmat receives the complex
1835 : !> operator (real part overwrites, imaginary part adds) through the
1836 : !> context buffer; the fm pair receives the two parts separately,
1837 : !> transferred directly into fm_re and fm_im on the owning group (any
1838 : !> distribution: the general copy redistributes); or the dbcsr pair
1839 : !> receives the two parts as block-cyclic matrices built by the service,
1840 : !> replacing the caller-side copy_fm_to_dbcsr_bc detour. The dbcsr
1841 : !> matrices must be freshly initialized (dbcsr_init_p) and the cmat
1842 : !> flavor must share the distribution of the k-point group.
1843 : !> Collective over the environment-wide communicator: every rank must
1844 : !> call the same sequence of service calls.
1845 : !> \param ctx context created by kpoint_operator_context_create
1846 : !> \param ikp global round, up to the maximum k-point count across groups
1847 : !> \param ispin instance spin: bounds-checked against the context spin space
1848 : !> \param rsmat real-space image matrices of the operator
1849 : !> \param cmat complex matrix output
1850 : !> \param fm_re real-part matrix output
1851 : !> \param fm_im imaginary-part output; may be omitted for real wavefunctions
1852 : !> \param mat_re real-part block-cyclic dbcsr output
1853 : !> \param mat_im imaginary-part block-cyclic dbcsr output
1854 : !> \param matrix_row row of rsmat to transform. Spin-free operators such
1855 : !> as S and T must pass kpoint_spin_free: their first matrix
1856 : !> dimension counts derivatives, not spins. A spin index is out of
1857 : !> bounds in plain SCF, but with derivatives stored it reads another
1858 : !> valid row silently
1859 : !> \param slot optional prepared slot: opt in to its immutable snapshot.
1860 : !> Without slot, get always reads current source values. A missing or
1861 : !> mismatched snapshot retains the original transfer path. Prepared FM
1862 : !> targets must match the layout selected at prepare; inactive groups
1863 : !> do not inspect or modify their output arguments.
1864 : !> \param grid optional lattice-FFT grid, overriding a prepared snapshot for this request
1865 : ! **************************************************************************************************
1866 5652 : SUBROUTINE kpoint_operator_get(ctx, ikp, ispin, rsmat, cmat, fm_re, fm_im, &
1867 : mat_re, mat_im, matrix_row, slot, grid)
1868 :
1869 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1870 : INTEGER, INTENT(IN) :: ikp, ispin
1871 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1872 : TYPE(cp_cfm_type), OPTIONAL :: cmat
1873 : TYPE(cp_fm_type), OPTIONAL :: fm_re, fm_im
1874 : TYPE(dbcsr_type), INTENT(INOUT), OPTIONAL :: mat_re, mat_im
1875 : INTEGER, INTENT(IN) :: matrix_row
1876 : INTEGER, INTENT(IN), OPTIONAL :: slot
1877 : TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
1878 :
1879 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_get'
1880 :
1881 : INTEGER :: handle, igroup, ik, irs, kpmax
1882 5652 : INTEGER, DIMENSION(:, :), POINTER :: kp_dist
1883 : LOGICAL :: layout_ok
1884 67824 : TYPE(copy_info_type) :: info(2)
1885 :
1886 5652 : CALL timeset(routineN, handle)
1887 :
1888 22608 : CPASSERT(COUNT([PRESENT(cmat), PRESENT(fm_re), PRESENT(mat_re)]) == 1)
1889 5652 : IF (PRESENT(fm_re) .OR. PRESENT(fm_im)) THEN
1890 852 : CPASSERT(PRESENT(fm_re))
1891 852 : CPASSERT(PRESENT(fm_im) .OR. ctx%kpoints%use_real_wfn)
1892 : END IF
1893 5652 : IF (PRESENT(mat_re) .OR. PRESENT(mat_im)) THEN
1894 228 : CPASSERT(PRESENT(mat_re) .AND. PRESENT(mat_im))
1895 : END IF
1896 :
1897 5652 : CPASSERT(ASSOCIATED(ctx%kpoints))
1898 5652 : CPASSERT(ASSOCIATED(ctx%kpoints%sab_nl))
1899 5652 : CPASSERT(ispin >= 1 .AND. ispin <= ctx%nspin)
1900 5652 : CPASSERT(matrix_row >= 1 .AND. matrix_row <= SIZE(rsmat, 1))
1901 13780 : kpmax = MAXVAL(ctx%kpoints%kp_dist(2, :) - ctx%kpoints%kp_dist(1, :) + 1)
1902 5652 : CPASSERT(ikp >= 1 .AND. ikp <= kpmax)
1903 5652 : IF (PRESENT(slot) .AND. .NOT. PRESENT(grid)) THEN
1904 5240 : CPASSERT(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
1905 5240 : irs = kpoint_operator_rs_match(ctx, rsmat, ispin, slot, matrix_row)
1906 5240 : IF (irs > 0) THEN
1907 5240 : ik = ctx%kpoints%kp_range(1) + ikp - 1
1908 5240 : IF (ik <= ctx%kpoints%kp_range(2)) THEN
1909 5238 : IF (PRESENT(mat_re)) THEN
1910 228 : IF (.NOT. ALLOCATED(ctx%fmimag)) THEN
1911 56 : ALLOCATE (ctx%fmimag)
1912 56 : CALL cp_fm_create(ctx%fmimag, ctx%fmlocal%matrix_struct)
1913 : END IF
1914 228 : CALL kpoint_operator_rs_build(ctx, irs, slot, ik, fm_re=ctx%fmlocal, fm_im=ctx%fmimag)
1915 228 : CALL copy_fm_to_dbcsr_bc(ctx%fmlocal, mat_re)
1916 228 : CALL copy_fm_to_dbcsr_bc(ctx%fmimag, mat_im)
1917 : ELSE
1918 5010 : CALL kpoint_operator_rs_build(ctx, irs, slot, ik, TARGET=cmat, fm_re=fm_re, fm_im=fm_im)
1919 : END IF
1920 : END IF
1921 5240 : CALL timestop(handle)
1922 5240 : RETURN
1923 : END IF
1924 : END IF
1925 412 : CALL kpoint_operator_ensure(ctx, rsmat)
1926 412 : IF (PRESENT(cmat)) THEN
1927 : ! The merge writes cmat in place. A foreign distribution would fail
1928 : ! inside the BLACS calls, far from this argument.
1929 412 : layout_ok = cp_fm_struct_equivalent(cmat%matrix_struct, ctx%fmlocal%matrix_struct)
1930 412 : CPASSERT(layout_ok)
1931 : END IF
1932 :
1933 412 : kp_dist => ctx%kpoints%kp_dist
1934 884 : DO igroup = 1, ctx%kpoints%nkp_groups
1935 472 : ik = kp_dist(1, igroup) + ikp - 1
1936 472 : IF (ik > kp_dist(2, igroup)) CYCLE
1937 : ! Absent fm targets are forwarded as absent optional arguments. Both
1938 : ! CFM and DBCSR outputs use the same context-owned receive buffer.
1939 472 : CALL kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid=grid, use_grid=PRESENT(grid))
1940 : CALL kpoint_operator_launch(ctx%kpoints, ik, ctx%fmwork, ctx%fmlocal, info, &
1941 472 : target_re=fm_re, target_im=fm_im)
1942 : CALL kpoint_operator_receive(ctx%kpoints, ctx%fmlocal, ik, info, cmat=cmat, fm_re=fm_re, fm_im=fm_im, &
1943 884 : mat_re=mat_re, mat_im=mat_im)
1944 : END DO
1945 :
1946 412 : CALL timestop(handle)
1947 :
1948 5652 : END SUBROUTINE kpoint_operator_get
1949 :
1950 : ! **************************************************************************************************
1951 : !> \brief Read an operator inside a group-local OT loop, after collective cache preparation.
1952 : !> A single distributed group retains the sparse transform and distribution.
1953 : !> One-rank or multiple groups read their owned cache, without world collectives.
1954 : !> \param matrix_rs current real-space operator
1955 : !> \param kpoints host environment
1956 : !> \param kp local k-point state
1957 : !> \param ispin source row (one for spin-free S/T)
1958 : !> \param cache_re prepared real part on the k-point group
1959 : !> \param cache_im prepared imaginary part on the k-point group
1960 : !> \param matrix_re newly initialized sparse real output
1961 : !> \param matrix_im newly initialized sparse imaginary output
1962 : ! **************************************************************************************************
1963 6423 : SUBROUTINE kpoint_operator_get_local(matrix_rs, kpoints, kp, ispin, cache_re, cache_im, &
1964 : matrix_re, matrix_im)
1965 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: matrix_rs
1966 : TYPE(kpoint_type), POINTER :: kpoints
1967 : TYPE(kpoint_env_type), POINTER :: kp
1968 : INTEGER, INTENT(IN) :: ispin
1969 : TYPE(cp_fm_type), INTENT(IN) :: cache_re, cache_im
1970 : TYPE(dbcsr_type), POINTER :: matrix_re, matrix_im
1971 :
1972 : INTEGER :: nkp_groups
1973 6423 : INTEGER, DIMENSION(:, :, :), POINTER :: cell_to_index
1974 6423 : REAL(dp), DIMENSION(:, :), POINTER :: xkp
1975 : TYPE(neighbor_list_set_p_type), DIMENSION(:), &
1976 6423 : POINTER :: sab_nl
1977 :
1978 6423 : NULLIFY (cell_to_index, sab_nl, xkp)
1979 0 : CPASSERT(ASSOCIATED(matrix_rs))
1980 6423 : CPASSERT(ispin >= 1 .AND. ispin <= SIZE(matrix_rs, 1))
1981 6423 : CPASSERT(ASSOCIATED(kpoints))
1982 6423 : CPASSERT(ASSOCIATED(kp))
1983 :
1984 6423 : CALL dbcsr_init_p(matrix_re)
1985 6423 : CALL dbcsr_init_p(matrix_im)
1986 6423 : nkp_groups = kpoints%nkp_groups
1987 6423 : xkp => kpoints%xkp
1988 6423 : cell_to_index => kpoints%cell_to_index
1989 6423 : sab_nl => kpoints%sab_nl
1990 6423 : IF (nkp_groups == 1 .AND. kpoints%para_env_kp%num_pe > 1) THEN
1991 5360 : CPASSERT(ASSOCIATED(xkp))
1992 5360 : CPASSERT(ASSOCIATED(cell_to_index))
1993 5360 : CPASSERT(ASSOCIATED(sab_nl))
1994 : CALL dbcsr_create(matrix_re, template=matrix_rs(1, 1)%matrix, &
1995 5360 : matrix_type=dbcsr_type_symmetric)
1996 : CALL dbcsr_create(matrix_im, template=matrix_rs(1, 1)%matrix, &
1997 5360 : matrix_type=dbcsr_type_antisymmetric)
1998 5360 : CALL cp_dbcsr_alloc_block_from_nbl(matrix_re, sab_nl)
1999 5360 : CALL cp_dbcsr_alloc_block_from_nbl(matrix_im, sab_nl)
2000 5360 : CALL dbcsr_set(matrix_re, 0.0_dp)
2001 5360 : CALL dbcsr_set(matrix_im, 0.0_dp)
2002 : CALL rskp_transform(rmatrix=matrix_re, cmatrix=matrix_im, rsmat=matrix_rs, ispin=ispin, &
2003 5360 : xkp=xkp(1:3, kp%nkpoint), cell_to_index=cell_to_index, sab_nl=sab_nl)
2004 : ELSE
2005 1063 : CALL copy_fm_to_dbcsr_bc(cache_re, matrix_re)
2006 1063 : CALL copy_fm_to_dbcsr_bc(cache_im, matrix_im)
2007 : END IF
2008 6423 : END SUBROUTINE kpoint_operator_get_local
2009 :
2010 : ! **************************************************************************************************
2011 : !> \brief Fill caller-owned matrices for one operator row, including independent history snapshots.
2012 : !> Prepare and fallback transfers are collective. Prepared one-rank groups evaluate k points
2013 : !> concurrently, with no MPI or matrix allocation/refcounts inside the worker loop.
2014 : !> \param kpoints host environment
2015 : !> \param ao_ao_fm global-layout template
2016 : !> \param rsmat current real-space images
2017 : !> \param slot operator slot
2018 : !> \param cmat allocated group-local outputs, one per locally owned k point
2019 : !> \param matrix_row source row; S/T use kpoint_spin_free
2020 : ! **************************************************************************************************
2021 1190 : SUBROUTINE kpoint_operator_get_all(kpoints, ao_ao_fm, rsmat, slot, cmat, matrix_row)
2022 : TYPE(kpoint_type), POINTER :: kpoints
2023 : TYPE(cp_fm_type), INTENT(IN) :: ao_ao_fm
2024 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
2025 : INTEGER, INTENT(IN) :: slot
2026 : TYPE(cp_cfm_type), DIMENSION(:), INTENT(INOUT) :: cmat
2027 : INTEGER, INTENT(IN) :: matrix_row
2028 :
2029 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_get_all'
2030 :
2031 1190 : COMPLEX(KIND=dp), ALLOCATABLE :: phase(:, :)
2032 : INTEGER :: handle, ik, ikp, kplocal, kpmax, nentry, &
2033 : nworkers
2034 : LOGICAL :: layout_ok, prepared
2035 1190 : TYPE(kpoint_operator_context_type) :: ctx
2036 :
2037 1190 : CALL timeset(routineN, handle)
2038 1190 : kplocal = SIZE(kpoints%kp_env)
2039 1190 : CPASSERT(SIZE(cmat) == kplocal .AND. kplocal > 0)
2040 1190 : CALL kpoint_operator_context_create(ctx, kpoints, ao_ao_fm, 1)
2041 1190 : CALL kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row=matrix_row)
2042 1190 : IF (prepared .AND. kpoints%para_env_kp%num_pe == 1) THEN
2043 3372 : DO ikp = 1, kplocal
2044 2336 : layout_ok = cp_fm_struct_equivalent(cmat(ikp)%matrix_struct, ctx%rs_route%matrix_struct)
2045 3372 : CPASSERT(layout_ok)
2046 : END DO
2047 1036 : CALL kpoint_operator_work_size(ctx, nentry)
2048 1036 : nworkers = 1
2049 1036 : !$ nworkers = MIN(omp_get_max_threads(), kplocal)
2050 : !$OMP PARALLEL DEFAULT(NONE) NUM_THREADS(nworkers) SHARED(ctx,kpoints,cmat,slot,kplocal,nentry) &
2051 1036 : !$OMP PRIVATE(phase,ikp,ik)
2052 : ALLOCATE (phase(1, nentry))
2053 : !$OMP DO SCHEDULE(DYNAMIC, 1)
2054 : DO ikp = 1, kplocal
2055 : ik = kpoints%kp_range(1) + ikp - 1
2056 : CALL kpoint_operator_evaluate(ctx, 1, slot, ik, phase, cmat(ikp))
2057 : END DO
2058 : !$OMP END DO
2059 1036 : DEALLOCATE (phase)
2060 : !$OMP END PARALLEL
2061 : ELSE
2062 308 : kpmax = MAXVAL(kpoints%kp_dist(2, :) - kpoints%kp_dist(1, :) + 1)
2063 398 : DO ikp = 1, kpmax
2064 : CALL kpoint_operator_get(ctx, ikp, 1, rsmat, cmat=cmat(MIN(ikp, kplocal)), &
2065 398 : matrix_row=matrix_row, slot=slot)
2066 : END DO
2067 : END IF
2068 1190 : CALL kpoint_operator_context_release(ctx)
2069 1190 : CALL timestop(handle)
2070 1190 : END SUBROUTINE kpoint_operator_get_all
2071 :
2072 : ! **************************************************************************************************
2073 : !> \brief Borrow the group-local OT cache for one operator slot.
2074 : !> \param kp owning k-point environment
2075 : !> \param slot KS, overlap or kinetic operator slot
2076 : !> \param ispin spin index, used only for the KS slot
2077 : !> \return non-owning real/imaginary FM array
2078 : ! **************************************************************************************************
2079 37021 : FUNCTION kpoint_operator_cache(kp, slot, ispin) RESULT(cache)
2080 : TYPE(kpoint_env_type), INTENT(IN) :: kp
2081 : INTEGER, INTENT(IN) :: slot, ispin
2082 : TYPE(cp_fm_type), POINTER :: cache(:)
2083 :
2084 68371 : SELECT CASE (slot)
2085 : CASE (kpoint_slot_ks)
2086 31350 : cache => kp%ot_hmat(:, ispin)
2087 : CASE (kpoint_slot_s)
2088 3581 : cache => kp%ot_smat
2089 : CASE (kpoint_slot_t)
2090 37021 : cache => kp%ot_tmat
2091 : END SELECT
2092 37021 : END FUNCTION kpoint_operator_cache
2093 :
2094 : ! **************************************************************************************************
2095 : !> \brief Sum images with block GEMMs over local k points, then fill each Hermitian mirror once.
2096 : !> Each worker owns one AO block and spin. Scratch uses the same packing
2097 : !> and k-tile upper bounds as the worker-count and tile-size choices.
2098 : !> \param ctx prepared single-rank-group context and owned output caches
2099 : !> \param slot operator slot
2100 : !> \param nrows number of captured spin rows
2101 : ! **************************************************************************************************
2102 1548 : SUBROUTINE kpoint_operator_store_blocks(ctx, slot, nrows)
2103 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
2104 : INTEGER, INTENT(IN) :: slot, nrows
2105 :
2106 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_store_blocks'
2107 :
2108 : INTEGER :: c0, col, first, handle, ic, igroup, ikp, &
2109 : ispin, j, k, last, nbatch, nc, nelem, &
2110 : ngroup, nimages, nk, nr, nworkers, r0, &
2111 : row, width
2112 : INTEGER(KIND=int_8) :: nfactor, npack, work_size
2113 1548 : REAL(KIND=dp), ALLOCATABLE, TARGET :: work(:)
2114 1548 : REAL(KIND=dp), CONTIGUOUS, POINTER :: factors(:, :, :), RESULT(:, :, :), &
2115 1548 : values(:, :)
2116 1548 : REAL(KIND=dp), POINTER :: xkp(:, :)
2117 1548 : TYPE(cp_fm_type), POINTER :: cache(:)
2118 : TYPE(kpoint_env_type), POINTER :: kp
2119 :
2120 1548 : CALL timeset(routineN, handle)
2121 1548 : nk = SIZE(ctx%kpoints%kp_env)
2122 1548 : nc = MERGE(1, 2, ctx%kpoints%use_real_wfn)
2123 1548 : xkp => ctx%kpoints%xkp
2124 : ASSOCIATE (kpoints => ctx%kpoints)
2125 : ! Clear padding as well as physical rows/columns, as in the per-k-point path.
2126 4355 : DO ikp = 1, nk
2127 2807 : kp => kpoints%kp_env(ikp)%kpoint_env
2128 7734 : DO ispin = 1, nrows
2129 3379 : cache => kpoint_operator_cache(kp, slot, ispin)
2130 12944 : DO ic = 1, nc
2131 8142281 : cache(ic)%local_data(:, :) = 0.0_dp
2132 : END DO
2133 : END DO
2134 : END DO
2135 : END ASSOCIATE
2136 1548 : IF (nk == 0 .OR. SIZE(ctx%rs_route%entry, 2) == 0) THEN
2137 0 : CALL timestop(handle)
2138 : RETURN
2139 : END IF
2140 1548 : ngroup = SIZE(ctx%rs_route%block_first) - 1
2141 : nworkers = 1
2142 1548 : !$ nworkers = omp_get_max_threads()
2143 1548 : nworkers = MIN(nworkers, ngroup*nrows)
2144 1548 : CALL kpoint_operator_batch_size(ctx, nworkers, 0, width, work_size)
2145 : !$OMP PARALLEL DEFAULT(NONE) NUM_THREADS(nworkers) &
2146 : !$OMP SHARED(ctx,slot,nrows,ngroup,nk,nc,width,work_size,xkp) &
2147 : !$OMP PRIVATE(work,values,factors,result,igroup,ispin,first,last,nr,nelem,nimages,r0,c0,j,col,row, &
2148 1548 : !$OMP k,ikp,nbatch,ic,kp,cache,npack,nfactor)
2149 : ALLOCATE (work(work_size))
2150 : !$OMP DO COLLAPSE(2) SCHEDULE(DYNAMIC, 1)
2151 : DO igroup = 1, ngroup
2152 : DO ispin = 1, nrows
2153 : first = ctx%rs_route%block_first(igroup)
2154 : last = ctx%rs_route%block_first(igroup + 1) - 1
2155 : r0 = ctx%rs_route%entry(1, first)
2156 : c0 = ctx%rs_route%entry(2, first)
2157 : nr = ctx%rs_route%entry(3, first)
2158 : nelem = nr*ctx%rs_route%entry(4, first)
2159 : IF (nelem == 0) CYCLE
2160 : IF (.NOT. ANY(ctx%rs(ispin, slot)%nonzero(first:last))) CYCLE
2161 : nimages = last - first + 1
2162 : npack = INT(nelem, int_8)*nimages
2163 : nfactor = INT(nimages, int_8)*width*nc
2164 : values(1:nelem, 1:nimages) => work(1:npack)
2165 : factors(1:nimages, 1:width, 1:nc) => work(npack + 1:npack + nfactor)
2166 : RESULT(1:nelem, 1:width, 1:nc) => work(npack + nfactor + 1:npack + nfactor + INT(nelem, int_8)*width*nc)
2167 : DO k = 1, nk, width
2168 : nbatch = MIN(width, nk - k + 1)
2169 : CALL kpoint_operator_contract_block(ctx%rs(ispin, slot), ctx%rs_route, first, last, &
2170 : xkp(:, ctx%kpoints%kp_range(1) + k - 1: &
2171 : ctx%kpoints%kp_range(1) + k + nbatch - 2), &
2172 : values, factors, RESULT, pack=k == 1)
2173 : DO j = 1, nbatch
2174 : ikp = k + j - 1
2175 : kp => ctx%kpoints%kp_env(ikp)%kpoint_env
2176 : cache => kpoint_operator_cache(kp, slot, ispin)
2177 : DO ic = 1, nc
2178 : DO col = 1, ctx%rs_route%entry(4, first)
2179 : !$OMP SIMD
2180 : DO row = 1, nr
2181 : cache(ic)%local_data(r0 + row - 1, c0 + col - 1) = RESULT(row + (col - 1)*nr, j, ic)
2182 : END DO
2183 : END DO
2184 : IF (ctx%rs_route%entry(8, first) == 0) CYCLE
2185 : DO row = 1, nr
2186 : !$OMP SIMD
2187 : DO col = 1, ctx%rs_route%entry(4, first)
2188 : cache(ic)%local_data(c0 + col - 1, r0 + row - 1) = &
2189 : MERGE(1.0_dp, -1.0_dp, ic == kpoint_part_re)*RESULT(row + (col - 1)*nr, j, ic)
2190 : END DO
2191 : END DO
2192 : END DO
2193 : END DO
2194 : END DO
2195 : END DO
2196 : END DO
2197 : !$OMP END DO
2198 1548 : DEALLOCATE (work)
2199 : !$OMP END PARALLEL
2200 1548 : CALL timestop(handle)
2201 3096 : END SUBROUTINE kpoint_operator_store_blocks
2202 :
2203 : ! **************************************************************************************************
2204 : !> \brief Size local k tiles, budgeting operator outputs and block scratch over all workers.
2205 : !> A block exceeding the budget is processed by one worker with one k point.
2206 : !> OT writes existing caches and distributes block/spin jobs; dense callers distribute k tiles.
2207 : !> \param ctx prepared single-rank-group context
2208 : !> \param nworkers requested workers, reduced to fit the minimum working set
2209 : !> \param nmatrices simultaneously retained complex operator tiles; zero writes existing split-FM caches
2210 : !> \param width number of k points per worker
2211 : !> \param work_size real block scratch elements per worker
2212 : ! **************************************************************************************************
2213 26858 : SUBROUTINE kpoint_operator_batch_size(ctx, nworkers, nmatrices, width, work_size)
2214 : TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
2215 : INTEGER, INTENT(INOUT) :: nworkers
2216 : INTEGER, INTENT(IN) :: nmatrices
2217 : INTEGER, INTENT(OUT) :: width
2218 : INTEGER(KIND=int_8), INTENT(OUT) :: work_size
2219 :
2220 : INTEGER :: first, igroup, nc, nk
2221 : INTEGER(KIND=int_8) :: max_pack, max_per_k, nelem, nimages, &
2222 : output_bytes, per_k_bytes
2223 :
2224 26858 : CPASSERT(ctx%kpoints%para_env_kp%num_pe == 1)
2225 26858 : CPASSERT(nmatrices >= 0 .AND. nworkers > 0)
2226 26858 : max_pack = 0_int_8
2227 26858 : max_per_k = 0_int_8
2228 26858 : IF (SIZE(ctx%rs_route%entry, 2) > 0) THEN
2229 410476 : DO igroup = 1, SIZE(ctx%rs_route%block_first) - 1
2230 383618 : first = ctx%rs_route%block_first(igroup)
2231 383618 : nelem = INT(ctx%rs_route%entry(3, first), int_8)*ctx%rs_route%entry(4, first)
2232 383618 : nimages = ctx%rs_route%block_first(igroup + 1) - first
2233 383618 : max_pack = MAX(max_pack, nelem*nimages)
2234 410476 : max_per_k = MAX(max_per_k, nelem + nimages)
2235 : END DO
2236 : END IF
2237 : output_bytes = 16_int_8*nmatrices*ctx%rs_route%matrix_struct%local_leading_dimension* &
2238 26858 : MAX(1, ctx%rs_route%ncol)
2239 26858 : nc = 2
2240 26858 : IF (nmatrices == 0) nc = MERGE(1, 2, ctx%kpoints%use_real_wfn)
2241 26858 : per_k_bytes = 8_int_8*nc*max_per_k + output_bytes
2242 26858 : nk = SIZE(ctx%kpoints%kp_env)
2243 26858 : IF (nmatrices > 0) nworkers = MIN(nworkers, nk)
2244 26858 : nworkers = MIN(nworkers, INT(MAX(1_int_8, rs_batch_bytes/(8_int_8*max_pack + per_k_bytes))))
2245 26858 : width = INT(MIN(INT(nk, int_8), MAX(1_int_8, (rs_batch_bytes/nworkers - 8_int_8*max_pack)/per_k_bytes)))
2246 26858 : IF (nmatrices > 0) width = MIN(width, (nk + nworkers - 1)/nworkers)
2247 26858 : work_size = max_pack + INT(nc, int_8)*width*max_per_k
2248 26858 : END SUBROUTINE kpoint_operator_batch_size
2249 :
2250 : ! **************************************************************************************************
2251 : !> \brief Assemble a caller-owned complex k tile using real block/image contractions.
2252 : !> The prepared route and snapshots are immutable; workers share no output or scratch.
2253 : !> \param ctx prepared single-rank-group context
2254 : !> \param ispin operator spin
2255 : !> \param slot prepared operator slot
2256 : !> \param first_k first global k index
2257 : !> \param batch complex output (local row, local column, k point), including padding
2258 : !> \param work private real scratch sized by kpoint_operator_batch_size
2259 : ! **************************************************************************************************
2260 28762 : SUBROUTINE kpoint_operator_evaluate_batch(ctx, ispin, slot, first_k, batch, work)
2261 : TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
2262 : INTEGER, INTENT(IN) :: ispin, slot, first_k
2263 : COMPLEX(KIND=dp), CONTIGUOUS, INTENT(OUT) :: batch(:, :, :)
2264 : REAL(KIND=dp), CONTIGUOUS, INTENT(INOUT), TARGET :: work(:)
2265 :
2266 : CHARACTER(len=*), PARAMETER :: routineN = 'kpoint_operator_evaluate_batch'
2267 :
2268 : INTEGER :: c0, col, first, handle, igroup, j, last, &
2269 : nc, nelem, nimages, nr, r0, row, spin, &
2270 : width
2271 : INTEGER(KIND=int_8) :: nfactor, npack
2272 28762 : REAL(KIND=dp), CONTIGUOUS, POINTER :: factors(:, :, :), RESULT(:, :, :), &
2273 28762 : values(:, :)
2274 28762 : REAL(KIND=dp), POINTER :: xkp(:, :)
2275 :
2276 28762 : CALL timeset(routineN, handle)
2277 28762 : CPASSERT(ctx%kpoints%para_env_kp%num_pe == 1)
2278 28762 : CPASSERT(SIZE(batch, 1) >= ctx%rs_route%nrow)
2279 28762 : CPASSERT(SIZE(batch, 2) >= ctx%rs_route%ncol)
2280 28762 : width = SIZE(batch, 3)
2281 28762 : CPASSERT(first_k >= ctx%kpoints%kp_range(1))
2282 28762 : CPASSERT(first_k + width - 1 <= ctx%kpoints%kp_range(2))
2283 28762 : spin = ispin
2284 28762 : IF (.NOT. ASSOCIATED(ctx%rs(spin, slot)%source)) spin = 1
2285 28762 : CPASSERT(ASSOCIATED(ctx%rs(spin, slot)%source))
2286 28762 : xkp => ctx%kpoints%xkp
2287 61252339 : batch(:, :, :) = z_zero
2288 28762 : IF (SIZE(ctx%rs_route%entry, 2) == 0) THEN
2289 0 : CALL timestop(handle)
2290 0 : RETURN
2291 : END IF
2292 441074 : DO igroup = 1, SIZE(ctx%rs_route%block_first) - 1
2293 412312 : first = ctx%rs_route%block_first(igroup)
2294 412312 : last = ctx%rs_route%block_first(igroup + 1) - 1
2295 412920 : IF (.NOT. ANY(ctx%rs(spin, slot)%nonzero(first:last))) CYCLE
2296 412312 : r0 = ctx%rs_route%entry(1, first)
2297 412312 : c0 = ctx%rs_route%entry(2, first)
2298 412312 : nr = ctx%rs_route%entry(3, first)
2299 412312 : nc = ctx%rs_route%entry(4, first)
2300 412312 : nelem = nr*nc
2301 412312 : IF (nelem == 0) CYCLE
2302 412312 : nimages = last - first + 1
2303 412312 : npack = INT(nelem, int_8)*nimages
2304 412312 : nfactor = 2_int_8*nimages*width
2305 412312 : CPASSERT(SIZE(work, KIND=int_8) >= npack + nfactor + 2_int_8*nelem*width)
2306 412312 : values(1:nelem, 1:nimages) => work(1:npack)
2307 412312 : factors(1:nimages, 1:width, 1:2) => work(npack + 1:npack + nfactor)
2308 412312 : RESULT(1:nelem, 1:width, 1:2) => work(npack + nfactor + 1:npack + nfactor + 2_int_8*nelem*width)
2309 : CALL kpoint_operator_contract_block(ctx%rs(spin, slot), ctx%rs_route, first, last, &
2310 : xkp(:, first_k:first_k + width - 1), &
2311 412312 : values, factors, RESULT, pack=.TRUE.)
2312 1355586 : DO j = 1, width
2313 5574466 : DO col = 1, nc
2314 5574466 : !$OMP SIMD
2315 : DO row = 1, nr
2316 : batch(r0 + row - 1, c0 + col - 1, j) = &
2317 35375450 : CMPLX(RESULT(row + (col - 1)*nr, j, 1), RESULT(row + (col - 1)*nr, j, 2), KIND=dp)
2318 : END DO
2319 : END DO
2320 914512 : IF (ctx%rs_route%entry(8, first) == 0) CYCLE
2321 4508432 : DO row = 1, nr
2322 4324581 : !$OMP SIMD
2323 : DO col = 1, nc
2324 24445017 : batch(c0 + col - 1, r0 + row - 1, j) = CONJG(batch(r0 + row - 1, c0 + col - 1, j))
2325 : END DO
2326 : END DO
2327 : END DO
2328 : END DO
2329 28762 : CALL timestop(handle)
2330 28762 : END SUBROUTINE kpoint_operator_evaluate_batch
2331 :
2332 : ! **************************************************************************************************
2333 : !> \brief Contract one AO block's images over a k tile for real or complex destinations.
2334 : !> \param snapshot immutable real-space values
2335 : !> \param route shared AO-block geometry
2336 : !> \param first first image entry
2337 : !> \param last last image entry
2338 : !> \param xkp k coordinates for this tile
2339 : !> \param values packed real source block columns
2340 : !> \param factors real cosine and optional sine factors
2341 : !> \param RESULT real and optional imaginary output blocks
2342 : !> \param pack pack source values on the first tile; later tiles reuse them
2343 : ! **************************************************************************************************
2344 429738 : SUBROUTINE kpoint_operator_contract_block(snapshot, route, first, last, xkp, values, factors, RESULT, pack)
2345 : TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
2346 : TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
2347 : INTEGER, INTENT(IN) :: first, last
2348 : REAL(KIND=dp), INTENT(IN) :: xkp(:, :)
2349 : REAL(KIND=dp), CONTIGUOUS, INTENT(INOUT) :: values(:, :), factors(:, :, :)
2350 : REAL(KIND=dp), CONTIGUOUS, INTENT(OUT) :: RESULT(:, :, :)
2351 : LOGICAL, INTENT(IN) :: pack
2352 :
2353 : INTEGER :: col, iblock, ic, ik, nactive, nelem, nr, &
2354 : row
2355 : REAL(KIND=dp) :: arg
2356 :
2357 429738 : nelem = SIZE(values, 1)
2358 429738 : nr = route%entry(3, first)
2359 429738 : nactive = 0
2360 11015152 : DO iblock = first, last
2361 10585414 : IF (.NOT. snapshot%nonzero(iblock)) CYCLE
2362 10560302 : nactive = nactive + 1
2363 10560302 : IF (pack) THEN
2364 79397728 : DO col = 1, route%entry(4, first)
2365 79397728 : !$OMP SIMD
2366 : DO row = 1, nr
2367 583746658 : values(row + (col - 1)*nr, nactive) = snapshot%blocks(iblock)%data(row, col)
2368 : END DO
2369 : END DO
2370 : END IF
2371 36083169 : DO ik = 1, SIZE(xkp, 2)
2372 100272068 : arg = twopi*DOT_PRODUCT(REAL(route%entry(5:7, iblock), dp), xkp(:, ik))
2373 25068017 : factors(nactive, ik, 1) = COS(arg)
2374 35628319 : IF (SIZE(factors, 3) == 2) factors(nactive, ik, 2) = SIN(arg)
2375 : END DO
2376 : END DO
2377 1289214 : DO ic = 1, SIZE(factors, 3)
2378 : CALL dgemm('N', 'N', nelem, SIZE(xkp, 2), nactive, 1.0_dp, values, nelem, &
2379 1289214 : factors(:, :, ic), SIZE(factors, 1), 0.0_dp, RESULT(:, :, ic), nelem)
2380 : END DO
2381 429738 : END SUBROUTINE kpoint_operator_contract_block
2382 :
2383 : ! **************************************************************************************************
2384 : !> \brief Prepare owned split-FM operators for OT without entering an SCF eigensolver.
2385 : !> H (or the initialization density) retains all source spins and is refreshed on each call.
2386 : !> S is stored once per context; the caller controls when to prepare T.
2387 : !> Outputs survive the context. Only routing and scratch persist between calls.
2388 : !> \param ctx context with fixed geometry, AO layout and spin capacity
2389 : !> \param matrix_ks Hamiltonian or initialization density
2390 : !> \param matrix_s spin-free overlap
2391 : !> \param matrix_t optional spin-free kinetic operator; absent leaves the existing kinetic cache alone
2392 : ! **************************************************************************************************
2393 1396 : SUBROUTINE kpoint_operator_store(ctx, matrix_ks, matrix_s, matrix_t)
2394 : TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
2395 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: matrix_ks, matrix_s
2396 : TYPE(dbcsr_p_type), DIMENSION(:, :), OPTIONAL, &
2397 : POINTER :: matrix_t
2398 :
2399 : CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_store'
2400 : INTEGER :: handle, ic, ikp, ispin, ncomp, nspin
2401 : LOGICAL :: kinetic, reuse
2402 : TYPE(cp_fm_struct_type), POINTER :: matrix_struct
2403 : TYPE(kpoint_env_type), POINTER :: kp
2404 : TYPE(kpoint_type), POINTER :: kpoints
2405 :
2406 1396 : CALL timeset(routineN, handle)
2407 1396 : kpoints => ctx%kpoints
2408 1396 : nspin = SIZE(matrix_ks, 1)
2409 1396 : CPASSERT(nspin == 1 .OR. nspin == ctx%nspin)
2410 1396 : ncomp = MERGE(1, 2, kpoints%use_real_wfn)
2411 1396 : kinetic = PRESENT(matrix_t)
2412 1396 : IF (kinetic) kinetic = ASSOCIATED(matrix_t)
2413 1396 : CALL kpoint_operator_ensure_group(ctx)
2414 1396 : matrix_struct => ctx%fmlocal%matrix_struct
2415 3924 : DO ikp = 1, SIZE(kpoints%kp_env)
2416 2528 : kp => kpoints%kp_env(ikp)%kpoint_env
2417 2528 : reuse = .FALSE.
2418 2528 : IF (ASSOCIATED(kp%ot_hmat)) THEN
2419 6915 : IF (ALL(SHAPE(kp%ot_hmat) == [ncomp, nspin])) THEN
2420 2305 : reuse = cp_fm_struct_equivalent(kp%ot_hmat(1, 1)%matrix_struct, matrix_struct)
2421 : END IF
2422 : END IF
2423 2305 : IF (.NOT. reuse) THEN
2424 223 : CALL cp_fm_release(kp%ot_hmat)
2425 1696 : ALLOCATE (kp%ot_hmat(ncomp, nspin))
2426 491 : DO ispin = 1, nspin
2427 1027 : DO ic = 1, ncomp
2428 804 : CALL cp_fm_create(kp%ot_hmat(ic, ispin), matrix_struct)
2429 : END DO
2430 : END DO
2431 : END IF
2432 2528 : CALL allocate_spin_free(kp%ot_smat)
2433 3924 : IF (kinetic) CALL allocate_spin_free(kp%ot_tmat)
2434 : END DO
2435 1396 : CALL store_slot(matrix_ks, kpoint_slot_ks, nspin)
2436 1396 : IF (.NOT. ctx%overlap_stored) THEN
2437 368 : CALL store_slot(matrix_s, kpoint_slot_s, 1)
2438 368 : ctx%overlap_stored = .TRUE.
2439 : END IF
2440 1396 : IF (kinetic) CALL store_slot(matrix_t, kpoint_slot_t, 1)
2441 1396 : CALL timestop(handle)
2442 :
2443 : CONTAINS
2444 :
2445 : ! **************************************************************************************************
2446 : !> \brief ...
2447 : !> \param cache ...
2448 : ! **************************************************************************************************
2449 2868 : SUBROUTINE allocate_spin_free(cache)
2450 : TYPE(cp_fm_type), DIMENSION(:), POINTER :: cache
2451 :
2452 : INTEGER :: ic
2453 : LOGICAL :: reuse
2454 :
2455 2868 : reuse = .FALSE.
2456 2868 : IF (ASSOCIATED(cache)) THEN
2457 2571 : IF (SIZE(cache) == ncomp) reuse = cp_fm_struct_equivalent(cache(1)%matrix_struct, matrix_struct)
2458 : END IF
2459 2571 : IF (reuse) RETURN
2460 297 : CALL cp_fm_release(cache)
2461 1485 : ALLOCATE (cache(ncomp))
2462 891 : DO ic = 1, ncomp
2463 891 : CALL cp_fm_create(cache(ic), matrix_struct)
2464 : END DO
2465 : END SUBROUTINE allocate_spin_free
2466 :
2467 : ! **************************************************************************************************
2468 : !> \brief ...
2469 : !> \param rsmat ...
2470 : !> \param slot ...
2471 : !> \param nrows ...
2472 : ! **************************************************************************************************
2473 1946 : SUBROUTINE store_slot(rsmat, slot, nrows)
2474 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
2475 : INTEGER, INTENT(IN) :: slot, nrows
2476 :
2477 : INTEGER :: fft_ready, ikp, ispin, kpmax
2478 : LOGICAL :: prepared, ready, use_blocks, use_grid
2479 1946 : TYPE(rskp_grid_type), ALLOCATABLE :: grids(:)
2480 :
2481 : use_grid = kpoints%lattice_fft == lattice_fft_on .OR. &
2482 1946 : (kpoints%lattice_fft == lattice_fft_auto .AND. SIZE(kpoints%xkp, 2) >= 27)
2483 1946 : IF (slot == kpoint_slot_ks .AND. nrows == ctx%nspin) THEN
2484 1396 : CALL kpoint_operator_prepare(ctx, rsmat, slot, prepared)
2485 : ELSE
2486 550 : CALL kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row=kpoint_spin_free)
2487 : END IF
2488 : ! Keep block GEMMs for direct assembly; get copies any prepared FFT output.
2489 1946 : use_blocks = prepared .AND. kpoints%para_env_kp%num_pe == 1
2490 1946 : IF (prepared) THEN
2491 4260 : DO ispin = 1, nrows
2492 4704 : use_blocks = use_blocks .AND. .NOT. ALLOCATED(ctx%rs_batch(ispin, slot)%data)
2493 : END DO
2494 : END IF
2495 1946 : IF (use_blocks) THEN
2496 1548 : CALL kpoint_operator_store_blocks(ctx, slot, nrows)
2497 : ELSE
2498 796 : kpmax = MAXVAL(kpoints%kp_dist(2, :) - kpoints%kp_dist(1, :) + 1)
2499 : ! Keep the existing forward lattice-FFT choice when distributed snapshots cannot be prepared.
2500 398 : use_grid = use_grid .AND. .NOT. prepared
2501 : IF (use_grid) THEN
2502 0 : CALL kpoint_operator_ensure(ctx, rsmat)
2503 0 : ALLOCATE (grids(nrows))
2504 0 : DO ispin = 1, nrows
2505 : CALL rskp_transform_grid_prepare(grids(ispin), ctx%rmatrix, rsmat, ispin, &
2506 : kpoints%xkp, kpoints%nkp_grid, kpoints%cell_to_index, &
2507 : kpoints%sab_nl, ready, &
2508 0 : max_storage_bytes=512_int_8*1024_int_8**2/INT(nspin + 1, int_8))
2509 0 : use_grid = use_grid .AND. ready
2510 : END DO
2511 0 : fft_ready = MERGE(1, 0, use_grid)
2512 0 : CALL kpoints%blacs_env_all%para_env%min(fft_ready)
2513 0 : use_grid = fft_ready == 1
2514 : END IF
2515 1142 : DO ikp = 1, kpmax
2516 1932 : DO ispin = 1, nrows
2517 1534 : IF (use_grid) THEN
2518 0 : CALL store_round(rsmat, slot, ikp, ispin, grid=grids(ispin))
2519 : ELSE
2520 790 : CALL store_round(rsmat, slot, ikp, ispin)
2521 : END IF
2522 : END DO
2523 : END DO
2524 398 : IF (ALLOCATED(grids)) THEN
2525 0 : DO ispin = 1, nrows
2526 0 : CALL rskp_transform_grid_release(grids(ispin))
2527 : END DO
2528 0 : DEALLOCATE (grids)
2529 : END IF
2530 : END IF
2531 : ! The split-FM caches now own the results. Do not retain source block
2532 : ! pointers or duplicate snapshots across Hamiltonian rebuilds.
2533 1946 : IF (ALLOCATED(ctx%rs)) THEN
2534 4354 : DO ispin = 1, ctx%nspin
2535 2408 : CALL kpoint_operator_rs_clear(ctx%rs(ispin, slot))
2536 2408 : IF (ALLOCATED(ctx%rs_batch(ispin, slot)%data)) DEALLOCATE (ctx%rs_batch(ispin, slot)%data)
2537 4354 : ctx%rs_batch(ispin, slot)%first = 0
2538 : END DO
2539 : END IF
2540 1946 : END SUBROUTINE store_slot
2541 :
2542 : ! **************************************************************************************************
2543 : !> \brief ...
2544 : !> \param rsmat ...
2545 : !> \param slot ...
2546 : !> \param ikp ...
2547 : !> \param ispin ...
2548 : !> \param grid ...
2549 : ! **************************************************************************************************
2550 790 : SUBROUTINE store_round(rsmat, slot, ikp, ispin, grid)
2551 : TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
2552 : INTEGER, INTENT(IN) :: slot, ikp, ispin
2553 : TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
2554 :
2555 790 : TYPE(cp_fm_type), DIMENSION(:), POINTER :: cache
2556 : TYPE(cp_fm_type), POINTER :: imag
2557 : TYPE(kpoint_env_type), POINTER :: kp
2558 :
2559 : ! A valid dummy cache is passed on inactive tail groups; get leaves it untouched.
2560 790 : kp => kpoints%kp_env(MIN(ikp, SIZE(kpoints%kp_env)))%kpoint_env
2561 1580 : cache => kpoint_operator_cache(kp, slot, ispin)
2562 790 : NULLIFY (imag)
2563 790 : IF (SIZE(cache) == 2) imag => cache(2)
2564 : CALL kpoint_operator_get(ctx, ikp, ispin, rsmat, fm_re=cache(1), fm_im=imag, &
2565 790 : matrix_row=ispin, slot=slot, grid=grid)
2566 790 : END SUBROUTINE store_round
2567 : END SUBROUTINE kpoint_operator_store
2568 :
2569 0 : END MODULE qs_kpoint_operators
|
