(git:6344619)
Loading...
Searching...
No Matches
qs_kpoint_operators.F
Go to the documentation of this file.
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! **************************************************************************************************
28 USE cp_cfm_types, ONLY: cp_cfm_type
29 USE cp_dbcsr_api, ONLY: &
31 dbcsr_init_p, dbcsr_p_type, dbcsr_set, dbcsr_type, dbcsr_type_antisymmetric, &
32 dbcsr_type_symmetric
44 USE cp_fm_types, ONLY: copy_info_type,&
54 USE kinds, ONLY: dp,&
55 int_8
63 USE kpoint_types, ONLY: kpoint_env_type,&
65 USE mathconstants, ONLY: gaussi,&
66 twopi,&
67 z_one,&
68 z_zero
72
73!$ USE omp_lib, ONLY: omp_get_max_threads, omp_get_num_threads, omp_get_thread_num
74#include "./base/base_uses.f90"
75
76 IMPLICIT NONE
77
78 PRIVATE
79
80 CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'qs_kpoint_operators'
81
87
88 INTEGER, PARAMETER, PUBLIC :: kpoint_slot_ks = 1, &
89 kpoint_slot_s = 2, &
91
92 !> Row of rsmat that carries the spin-free operator image. S and T carry
93 !> no spin copies: drivers launch one instance per spin of their loop and
94 !> read this same row in every instance.
95 INTEGER, PARAMETER, PUBLIC :: kpoint_spin_free = 1
96
97 !> Number of defined slots; grows only inside this module.
98 INTEGER, PARAMETER, PRIVATE :: kpoint_num_slots = kpoint_slot_t - kpoint_slot_ks + 1
99
100 !> Part indices of one transfer: its real and its imaginary half.
101 INTEGER, PARAMETER, PRIVATE :: kpoint_part_re = 1, &
102 kpoint_part_im = 2
103
104 !> Instance states of one (k point, spin, slot) transfer.
105 !> State transitions happen only inside the service routines.
106 INTEGER, PARAMETER, PRIVATE :: kpoint_slot_idle = 0, &
107 kpoint_slot_in_flight = 1, &
108 kpoint_slot_local = 2
109
110 ! Structural route shared by all captured rows. Whole blocks go to one
111 ! representative group, then only equivalent tiles are replicated.
112 TYPE kpoint_operator_rs_route_type
114 TYPE(cp_fm_struct_type), POINTER :: matrix_struct => null()
115 INTEGER :: nimage = 0, nvalue = 0, nrow = 0, ncol = 0
116 INTEGER(KIND=int_8) :: storage = 0
117 INTEGER, ALLOCATABLE :: row_size(:), col_size(:), entry(:, :), first(:), tile(:, :, :)
118 INTEGER, ALLOCATABLE :: send_entry(:), send_first(:), scount(:, :), rcount(:, :), &
119 sdisp(:, :), rdisp(:, :)
120 REAL(kind=dp), ALLOCATABLE :: send_values(:)
121 END TYPE kpoint_operator_rs_route_type
122
123 ! Immutable values, not per-k matrices. Metadata belong to the shared route;
124 ! source identifies starts belonging to the explicitly prepared batch only.
125 TYPE kpoint_operator_block_type
126 REAL(kind=dp), POINTER :: data(:, :) => null()
127 END TYPE kpoint_operator_block_type
128
129 TYPE kpoint_operator_rs_type
130 INTEGER :: matrix_row = 0
131 TYPE(dbcsr_p_type), POINTER :: source(:, :) => null()
132 ! Values are owned only for distributed sources. Blocks always borrow:
133 ! either from values below or, on one world rank, directly from DBCSR.
134 REAL(kind=dp), CONTIGUOUS, POINTER :: values(:) => null()
135 TYPE(kpoint_operator_block_type), ALLOCATABLE :: blocks(:)
136 LOGICAL, ALLOCATABLE :: nonzero(:)
137 END TYPE kpoint_operator_rs_type
138
139 ! Small read-ahead window per prepared row. Keep mutable output separate
140 ! from the immutable snapshot; prepare invalidates only the selected slot.
141 TYPE kpoint_operator_rs_batch_type
142 INTEGER :: first = 0
143 COMPLEX(KIND=dp), ALLOCATABLE :: data(:, :, :)
144 END TYPE kpoint_operator_rs_batch_type
145
146 INTEGER, PARAMETER :: rs_batch_size = 4
147 ! Includes cached tiles and batch phase scratch, divided between all rows/slots.
148 INTEGER(KIND=int_8), PARAMETER :: rs_batch_bytes = 32_int_8*1024_int_8**2
149
150! **************************************************************************************************
151!> \brief Caller-owned assembly context for one driver call.
152!> Created and released in the same driver routine and passed to helpers
153!> as an argument; never stored on the host kpoint_type.
154!> Components are private and touched only by this module.
155! **************************************************************************************************
157 PRIVATE
158 ! Instance spin space of the owning driver call: the bound of the spin
159 ! loop that start and finish key. Declared once at creation. It never
160 ! comes from the operators: the first matrix dimension of a spin-free
161 ! operator such as S counts derivatives, and the images live in the
162 ! second, so callers address operator rows through matrix_row instead.
163 INTEGER :: nspin = 0
164 ! Borrowed host/distribution references; the context cannot span an energy evaluation.
165 TYPE(kpoint_type), POINTER :: kpoints => null()
166 TYPE(cp_fm_struct_type), POINTER :: ao_ao_fmstruct => null()
167 TYPE(cp_fm_pool_type), POINTER :: group_fm_pools => null()
168 ! Identity of the neighbor lists used by scratch and prepared snapshots.
169 TYPE(neighbor_list_set_p_type), DIMENSION(:), &
170 POINTER :: work_sab => null()
171 ! Owned scratch, created on first use and released with the context.
172 TYPE(dbcsr_type), POINTER :: rmatrix => null()
173 TYPE(dbcsr_type), POINTER :: cmatrix => null()
174 TYPE(cp_fm_type), ALLOCATABLE, DIMENSION(:) :: fmwork
175 TYPE(cp_fm_type), ALLOCATABLE :: fmlocal, fmimag
176 ! Plans share the lifetime and fixed layouts of the two densify workspaces.
177 TYPE(dbcsr_fm_copy_plan_type), ALLOCATABLE :: copy_plan(:)
178 ! Immutable snapshots and bounded local output batches. Prepared starts
179 ! need neither per-k transfers nor outstanding MPI requests.
180 TYPE(kpoint_operator_rs_type), ALLOCATABLE :: rs(:, :)
181 TYPE(kpoint_operator_rs_route_type), ALLOCATABLE :: rs_route
182 TYPE(kpoint_operator_rs_batch_type), ALLOCATABLE :: rs_batch(:, :)
183 ! Bookkeeping. info is keyed (global k point, spin, slot, part),
184 ! part kpoint_part_re holding the real half and kpoint_part_im holding
185 ! the imaginary half. Status is keyed (global round, spin, slot)
186 ! because start and finish treat all groups of one instance
187 ! atomically. Allocated on first start.
188 TYPE(copy_info_type), ALLOCATABLE, DIMENSION(:, :, :, :) :: info
189 INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: status
191
192CONTAINS
193
194! **************************************************************************************************
195!> \brief Transform and densify one operator using the context-owned workspaces and routing plans.
196!> \param ctx initialized transfer context
197!> \param rsmat real-space operator images
198!> \param matrix_row operator row, independent of the instance spin
199!> \param ik global k-point index
200!> \param grid caller-owned reciprocal-grid cache
201!> \param use_grid extract from grid instead of a direct phase sum
202! **************************************************************************************************
203 SUBROUTINE kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid, use_grid)
204
205 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
206 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
207 INTEGER, INTENT(IN) :: matrix_row, ik
208 TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
209 LOGICAL, INTENT(IN) :: use_grid
210
211 CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_densify'
212
213 INTEGER :: handle
214 REAL(KIND=dp), DIMENSION(:, :), POINTER :: xkp
215
216 CALL timeset(routinen, handle)
217
218 ! Avoid nested component array sections in GCC 16 bounds checks.
219 xkp => ctx%kpoints%xkp
220 IF (ctx%kpoints%use_real_wfn) THEN
221 IF (use_grid) THEN
222 cpassert(PRESENT(grid))
223 CALL rskp_transform_grid_extract(grid, ik, ctx%rmatrix)
224 ELSE
225 CALL dbcsr_set(ctx%rmatrix, 0.0_dp)
226 CALL rskp_transform(rmatrix=ctx%rmatrix, rsmat=rsmat, ispin=matrix_row, &
227 xkp=xkp(:, ik), cell_to_index=ctx%kpoints%cell_to_index, &
228 sab_nl=ctx%kpoints%sab_nl)
229 END IF
230 ELSE IF (use_grid) THEN
231 cpassert(PRESENT(grid))
232 CALL rskp_transform_grid_extract(grid, ik, ctx%rmatrix, ctx%cmatrix)
233 ELSE
234 CALL dbcsr_set(ctx%rmatrix, 0.0_dp)
235 CALL dbcsr_set(ctx%cmatrix, 0.0_dp)
236 CALL rskp_transform(rmatrix=ctx%rmatrix, cmatrix=ctx%cmatrix, rsmat=rsmat, ispin=matrix_row, &
237 xkp=xkp(:, ik), cell_to_index=ctx%kpoints%cell_to_index, &
238 sab_nl=ctx%kpoints%sab_nl)
239 END IF
240 CALL copy_dbcsr_to_fm(ctx%rmatrix, ctx%fmwork(kpoint_part_re), ctx%copy_plan(kpoint_part_re))
241 IF (.NOT. ctx%kpoints%use_real_wfn) THEN
242 CALL copy_dbcsr_to_fm(ctx%cmatrix, ctx%fmwork(kpoint_part_im), ctx%copy_plan(kpoint_part_im))
243 END IF
244
245 CALL timestop(handle)
246
247 END SUBROUTINE kpoint_operator_densify
248
249! **************************************************************************************************
250!> \brief Launch a densified operator to its owning group (only the real part for real wavefunctions).
251!> Absent explicit targets, receive both parts through fmlocal. Each copy
252!> owns its receive buffer until finish, so the targets can alias.
253!> Non-owning groups participate with an empty dummy target.
254!> \param kpoints host environment, including group ownership and transfer communicator
255!> \param ik global k-point index
256!> \param fmwork real/imaginary source matrices on the global distribution
257!> \param fmlocal default receiver on the k-point group distribution
258!> \param info transfer states, real part first; real wavefunctions allocate only the first part
259!> \param target_re optional caller-owned real-part target
260!> \param target_im optional imaginary-part target; required with target_re for complex wavefunctions
261! **************************************************************************************************
262 SUBROUTINE kpoint_operator_launch(kpoints, ik, fmwork, fmlocal, info, target_re, target_im)
263
264 TYPE(kpoint_type), POINTER :: kpoints
265 INTEGER, INTENT(IN) :: ik
266 TYPE(cp_fm_type), DIMENSION(:), INTENT(IN) :: fmwork
267 TYPE(cp_fm_type), INTENT(IN) :: fmlocal
268 TYPE(copy_info_type), DIMENSION(:) :: info
269 TYPE(cp_fm_type), INTENT(IN), OPTIONAL :: target_re, target_im
270
271 CHARACTER(LEN=*), PARAMETER :: routineN = 'kpoint_operator_launch'
272
273 INTEGER :: handle
274 TYPE(cp_fm_type) :: fmdummy
275 TYPE(mp_para_env_type), POINTER :: para_env
276
277 CALL timeset(routinen, handle)
278 cpassert(.NOT. PRESENT(target_im) .OR. PRESENT(target_re))
279 IF (.NOT. kpoints%use_real_wfn) THEN
280 cpassert(PRESENT(target_re) .EQV. PRESENT(target_im))
281 END IF
282
283 para_env => kpoints%blacs_env_all%para_env
284 IF (ik >= kpoints%kp_range(1) .AND. ik <= kpoints%kp_range(2)) THEN
285 IF (PRESENT(target_re)) THEN
286 CALL cp_fm_start_copy_general(fmwork(kpoint_part_re), target_re, para_env, info(kpoint_part_re))
287 IF (.NOT. kpoints%use_real_wfn) THEN
288 CALL cp_fm_start_copy_general(fmwork(kpoint_part_im), target_im, para_env, info(kpoint_part_im))
289 END IF
290 ELSE
291 CALL cp_fm_start_copy_general(fmwork(kpoint_part_re), fmlocal, para_env, info(kpoint_part_re))
292 IF (.NOT. kpoints%use_real_wfn) THEN
293 CALL cp_fm_start_copy_general(fmwork(kpoint_part_im), fmlocal, para_env, info(kpoint_part_im))
294 END IF
295 END IF
296 ELSE
297 CALL cp_fm_start_copy_general(fmwork(kpoint_part_re), fmdummy, para_env, info(kpoint_part_re))
298 IF (.NOT. kpoints%use_real_wfn) THEN
299 CALL cp_fm_start_copy_general(fmwork(kpoint_part_im), fmdummy, para_env, info(kpoint_part_im))
300 END IF
301 END IF
302
303 CALL timestop(handle)
304
305 END SUBROUTINE kpoint_operator_launch
306
307! **************************************************************************************************
308!> \brief Receive one group's operator and reclaim every rank's send buffers.
309!> CFM/DBCSR outputs consume the context buffer before receiving the next part;
310!> FM outputs receive directly. Inactive groups only clean up their sends.
311!> \param kpoints parallel environment
312!> \param fmlocal shared receive buffer, consumed before receiving the next part
313!> \param ik global k-point index
314!> \param info real/imaginary transfers started by kpoint_operator_launch
315!> \param cmat complex output, real part overwrites and imaginary part adds
316!> \param fm_re real FM output
317!> \param fm_im optional imaginary FM output
318!> \param mat_re real DBCSR output
319!> \param mat_im imaginary DBCSR output
320! **************************************************************************************************
321 SUBROUTINE kpoint_operator_receive(kpoints, fmlocal, ik, info, cmat, fm_re, fm_im, mat_re, mat_im)
322 TYPE(kpoint_type), INTENT(IN), POINTER :: kpoints
323 TYPE(cp_fm_type) :: fmlocal
324 INTEGER, INTENT(IN) :: ik
325 TYPE(copy_info_type), INTENT(INOUT) :: info(:)
326 TYPE(cp_cfm_type), OPTIONAL :: cmat
327 TYPE(cp_fm_type), OPTIONAL :: fm_re, fm_im
328 TYPE(dbcsr_type), INTENT(INOUT), OPTIONAL :: mat_re, mat_im
329
330 INTEGER :: part
331
332 IF (ik >= kpoints%kp_range(1) .AND. ik <= kpoints%kp_range(2)) THEN
333 IF (PRESENT(cmat)) THEN
334 CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_re))
335 CALL cp_cfm_scale_and_add_fm(z_zero, cmat, z_one, fmlocal)
336 IF (.NOT. kpoints%use_real_wfn) THEN
337 CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_im))
338 CALL cp_cfm_scale_and_add_fm(z_one, cmat, gaussi, fmlocal)
339 END IF
340 ELSE IF (PRESENT(fm_re)) THEN
341 CALL cp_fm_finish_copy_general(fm_re, info(kpoint_part_re))
342 IF (kpoints%use_real_wfn) THEN
343 IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
344 ELSE
345 CALL cp_fm_finish_copy_general(fm_im, info(kpoint_part_im))
346 END IF
347 ELSE
348 CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_re))
349 CALL copy_fm_to_dbcsr_bc(fmlocal, mat_re)
350 IF (kpoints%use_real_wfn) THEN
351 fmlocal%local_data(:, :) = 0.0_dp
352 ELSE
353 CALL cp_fm_finish_copy_general(fmlocal, info(kpoint_part_im))
354 END IF
355 CALL copy_fm_to_dbcsr_bc(fmlocal, mat_im)
356 END IF
357 END IF
358 DO part = 1, merge(1, 2, kpoints%use_real_wfn)
359 CALL cp_fm_cleanup_copy_general(info(part))
360 END DO
361 END SUBROUTINE kpoint_operator_receive
362
363! **************************************************************************************************
364!> \brief Create an empty operator context.
365!> The context owns no buffers after this call; they are created on first
366!> use and released with the context.
367!> \param ctx context to initialize
368!> \param kpoints host k-point environment, referenced, not copied
369!> \param ao_ao_fm any square full matrix on the environment-wide communicator;
370!> its matrix struct is the template for the context work buffers,
371!> which densify on the source side of transfers that cross into the
372!> k-point groups
373!> \param nspin instance spin space of the owning driver call: the bound of
374!> the spin loop its start and finish calls iterate. Fixed for the
375!> lifetime of the context
376! **************************************************************************************************
377 SUBROUTINE kpoint_operator_context_create(ctx, kpoints, ao_ao_fm, nspin)
378
379 TYPE(kpoint_operator_context_type), INTENT(OUT) :: ctx
380 TYPE(kpoint_type), POINTER :: kpoints
381 TYPE(cp_fm_type), INTENT(IN) :: ao_ao_fm
382 INTEGER, INTENT(IN) :: nspin
383
384 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_context_create'
385
386 INTEGER :: handle
387
388 CALL timeset(routinen, handle)
389
390 cpassert(nspin >= 1)
391 ctx%nspin = nspin
392 cpassert(ASSOCIATED(kpoints))
393 ctx%kpoints => kpoints
394 CALL cp_fm_get_info(ao_ao_fm, matrix_struct=ctx%ao_ao_fmstruct)
395 cpassert(ASSOCIATED(ctx%ao_ao_fmstruct))
396 ! The template must live on the transfer world. A template from
397 ! another communicator fails during the transfers, far from this call.
398 IF (.NOT. ASSOCIATED(ctx%ao_ao_fmstruct%para_env, kpoints%blacs_env_all%para_env)) THEN
399 CALL cp_abort(__location__, "Operator context requires the global parallel environment")
400 END IF
401
402 CALL timestop(handle)
403
404 END SUBROUTINE kpoint_operator_context_create
405
406! **************************************************************************************************
407!> \brief Release the context together with every buffer it created.
408!> Call once, at the exit of the routine that created the context.
409!> Every started instance must be idle at this point. The call aborts
410!> when a transfer is still in flight, because nothing would ever wait
411!> for its MPI requests.
412!> \param ctx context to release
413! **************************************************************************************************
415
416 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
417
418 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_context_release'
419
420 INTEGER :: handle, ispin, iwork, slot
421
422 CALL timeset(routinen, handle)
423
424 ! Every started instance must have been finished: an unfinished start
425 ! holds MPI requests that nothing would ever wait for.
426 IF (ALLOCATED(ctx%status)) THEN
427 cpassert(all(ctx%status == kpoint_slot_idle))
428 END IF
429
430 IF (ASSOCIATED(ctx%rmatrix)) CALL dbcsr_deallocate_matrix(ctx%rmatrix)
431 IF (ASSOCIATED(ctx%cmatrix)) CALL dbcsr_deallocate_matrix(ctx%cmatrix)
432 NULLIFY (ctx%rmatrix, ctx%cmatrix, ctx%work_sab)
433 IF (ALLOCATED(ctx%copy_plan)) DEALLOCATE (ctx%copy_plan)
434 IF (ALLOCATED(ctx%rs)) THEN
435 DO slot = 1, kpoint_num_slots
436 DO ispin = 1, ctx%nspin
437 CALL kpoint_operator_rs_clear(ctx%rs(ispin, slot))
438 END DO
439 END DO
440 DEALLOCATE (ctx%rs, ctx%rs_batch)
441 END IF
442 IF (ALLOCATED(ctx%rs_route)) THEN
443 CALL cp_fm_struct_release(ctx%rs_route%matrix_struct)
444 DEALLOCATE (ctx%rs_route)
445 END IF
446 IF (ALLOCATED(ctx%fmwork)) THEN
447 DO iwork = 1, SIZE(ctx%fmwork)
448 CALL cp_fm_release(ctx%fmwork(iwork))
449 END DO
450 DEALLOCATE (ctx%fmwork)
451 END IF
452 IF (ALLOCATED(ctx%fmimag)) THEN
453 CALL cp_fm_release(ctx%fmimag)
454 DEALLOCATE (ctx%fmimag)
455 END IF
456 IF (ALLOCATED(ctx%fmlocal)) THEN
457 CALL fm_pool_give_back_fm(ctx%group_fm_pools, ctx%fmlocal)
458 DEALLOCATE (ctx%fmlocal)
459 END IF
460 IF (ALLOCATED(ctx%info)) DEALLOCATE (ctx%info)
461 IF (ALLOCATED(ctx%status)) DEALLOCATE (ctx%status)
462 NULLIFY (ctx%kpoints, ctx%ao_ao_fmstruct, ctx%group_fm_pools)
463
464 CALL timestop(handle)
465
467
468! **************************************************************************************************
469!> \brief Create the context scratch that the first service call needs: the
470!> DBCSR work pair with blocks from the current neighbor lists, two
471!> densify buffers on the environment-wide struct, and the group
472!> transfer target from the group pools.
473!> The per-instance slot bookkeeping is not scratch: start creates it
474!> on its first call, because the synchronous get call never touches it.
475!> The buffers assume the context lifetime of one driver call: the
476!> host neighbor lists do not change within it, and every buffer use
477!> asserts this identity. Buffers from an earlier call stay in place.
478!> Called from the service routines only.
479!> \param ctx context that receives the buffers
480!> \param rsmat real-space image matrices of one operator. The call uses the
481!> first matrix as the DBCSR template of the work pair. It reads no
482!> matrix data.
483! **************************************************************************************************
484 SUBROUTINE kpoint_operator_ensure(ctx, rsmat)
485
486 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
487 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
488
489 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_ensure'
490
491 INTEGER :: handle
492
493 CALL timeset(routinen, handle)
494
495 ! The pair must still match the host neighbor lists: they are
496 ! reallocated at every energy evaluation, while the context spans one
497 ! driver call within it.
498 IF (ASSOCIATED(ctx%work_sab)) THEN
499 cpassert(ASSOCIATED(ctx%work_sab, ctx%kpoints%sab_nl))
500 END IF
501 IF (.NOT. ASSOCIATED(ctx%rmatrix)) THEN
502 ALLOCATE (ctx%rmatrix)
503 CALL dbcsr_create(ctx%rmatrix, template=rsmat(1, 1)%matrix, matrix_type=dbcsr_type_symmetric)
504 CALL cp_dbcsr_alloc_block_from_nbl(ctx%rmatrix, ctx%kpoints%sab_nl)
505 IF (.NOT. ctx%kpoints%use_real_wfn) THEN
506 ALLOCATE (ctx%cmatrix)
507 CALL dbcsr_create(ctx%cmatrix, template=rsmat(1, 1)%matrix, matrix_type=dbcsr_type_antisymmetric)
508 CALL cp_dbcsr_alloc_block_from_nbl(ctx%cmatrix, ctx%kpoints%sab_nl)
509 END IF
510 ctx%work_sab => ctx%kpoints%sab_nl
511 END IF
512
513 IF (.NOT. ALLOCATED(ctx%fmwork)) THEN
514 ALLOCATE (ctx%fmwork(merge(1, 2, ctx%kpoints%use_real_wfn)))
515 ALLOCATE (ctx%copy_plan(SIZE(ctx%fmwork)))
516 CALL cp_fm_create(ctx%fmwork(kpoint_part_re), ctx%ao_ao_fmstruct, &
517 name="KPOINT-OPERATOR-WORK-RE")
518 IF (.NOT. ctx%kpoints%use_real_wfn) THEN
519 CALL cp_fm_create(ctx%fmwork(kpoint_part_im), ctx%ao_ao_fmstruct, &
520 name="KPOINT-OPERATOR-WORK-IM")
521 END IF
522 END IF
523
524 CALL kpoint_operator_ensure_group(ctx)
525
526 CALL timestop(handle)
527
528 END SUBROUTINE kpoint_operator_ensure
529
530! **************************************************************************************************
531!> \brief Acquire only the group-local output workspace; both backends share it.
532!> \param ctx operator context
533! **************************************************************************************************
534 SUBROUTINE kpoint_operator_ensure_group(ctx)
535
536 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
537
538 IF (.NOT. ALLOCATED(ctx%fmlocal)) THEN
539 CALL mpools_get(ctx%kpoints%mpools, ao_ao_fm_pools=ctx%group_fm_pools)
540 ALLOCATE (ctx%fmlocal)
541 CALL fm_pool_create_fm(ctx%group_fm_pools, ctx%fmlocal)
542 END IF
543
544 END SUBROUTINE kpoint_operator_ensure_group
545
546! **************************************************************************************************
547!> \brief Prepare an immutable operator slot for group-local assembly.
548!> Refresh after the source values change, with every instance idle.
549!> Other prepared slots and structural routing remain valid. Without
550!> matrix_row, capture one row per instance spin; with matrix_row, store
551!> one row shared by all instance spins. Failure clears this slot on
552!> all ranks and retains the existing transfer fallback.
553!> \param ctx context spanning the current initialized k-point environment
554!> \param rsmat current real-space operator images
555!> \param slot slot to replace (KS, S or T)
556!> \param prepared optional collective success flag
557!> \param matrix_row optional explicit spin-free operator row
558!> \param TARGET optional group-local FM layout for prepared get outputs;
559!> otherwise use the host AO/AO pool layout. One layout per context
560! **************************************************************************************************
561 SUBROUTINE kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row, TARGET)
562
563 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
564 TYPE(dbcsr_p_type), POINTER :: rsmat(:, :)
565 INTEGER, INTENT(IN) :: slot
566 LOGICAL, INTENT(OUT), OPTIONAL :: prepared
567 INTEGER, INTENT(IN), OPTIONAL :: matrix_row
568 TYPE(cp_fm_type), INTENT(IN), OPTIONAL :: target
569
570 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_prepare'
571 INTEGER(KIND=int_8), PARAMETER :: max_bytes = 64_int_8*1024_int_8**2
572
573 INTEGER :: compatible, handle, ispin, jslot, nrow, &
574 row
575 INTEGER(KIND=int_8) :: retained, storage
576 LOGICAL :: ready
577 TYPE(cp_fm_struct_type), POINTER :: fms
578 TYPE(mp_para_env_type), POINTER :: para_env
579
580 CALL timeset(routinen, handle)
581 IF (PRESENT(prepared)) prepared = .false.
582 cpassert(ASSOCIATED(ctx%kpoints))
583 cpassert(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
584 cpassert(ASSOCIATED(ctx%kpoints%sab_nl))
585 IF (ASSOCIATED(ctx%work_sab)) THEN
586 cpassert(ASSOCIATED(ctx%work_sab, ctx%kpoints%sab_nl))
587 END IF
588 IF (ALLOCATED(ctx%status)) THEN
589 cpassert(all(ctx%status == kpoint_slot_idle))
590 END IF
591 IF (.NOT. ALLOCATED(ctx%rs)) THEN
592 ALLOCATE (ctx%rs(ctx%nspin, kpoint_num_slots), ctx%rs_batch(ctx%nspin, kpoint_num_slots))
593 END IF
594 DO ispin = 1, ctx%nspin
595 CALL kpoint_operator_rs_clear(ctx%rs(ispin, slot))
596 IF (ALLOCATED(ctx%rs_batch(ispin, slot)%data)) DEALLOCATE (ctx%rs_batch(ispin, slot)%data)
597 ctx%rs_batch(ispin, slot)%first = 0
598 END DO
599 nrow = ctx%nspin
600 IF (PRESENT(matrix_row)) THEN
601 cpassert(matrix_row >= 1 .AND. matrix_row <= SIZE(rsmat, 1))
602 nrow = 1
603 ELSE
604 cpassert(SIZE(rsmat, 1) == ctx%nspin)
605 END IF
606 para_env => ctx%kpoints%blacs_env_all%para_env
607 CALL kpoint_operator_ensure_group(ctx)
608 fms => ctx%fmlocal%matrix_struct
609 IF (PRESENT(TARGET)) fms => target%matrix_struct
610 cpassert(ASSOCIATED(fms))
611 ! Routing uses the host group-rank numbering. Different BLACS layouts
612 ! on that communicator are supported, foreign communicators are not.
613 compatible = merge(1, 0, ASSOCIATED(fms%para_env, ctx%kpoints%para_env_kp))
614 CALL para_env%min(compatible)
615 IF (compatible == 0) THEN
616 CALL timestop(handle)
617 RETURN
618 END IF
619 IF (.NOT. ALLOCATED(ctx%rs_route)) THEN
620 ALLOCATE (ctx%rs_route)
621 CALL kpoint_operator_rs_plan_create(ctx%kpoints, rsmat, fms, max_bytes, ctx%rs_route, ready)
622 IF (.NOT. ready) THEN
623 DEALLOCATE (ctx%rs_route)
624 CALL timestop(handle)
625 RETURN
626 END IF
627 ctx%work_sab => ctx%kpoints%sab_nl
628 END IF
629 compatible = merge(1, 0, cp_fm_struct_equivalent(fms, ctx%rs_route%matrix_struct))
630 DO ispin = 1, nrow
631 row = ispin
632 IF (PRESENT(matrix_row)) row = matrix_row
633 IF (.NOT. kpoint_operator_rs_compatible(ctx%rs_route, rsmat, row)) compatible = 0
634 END DO
635 CALL para_env%min(compatible)
636 ! Keep the single-k phase allowance in the snapshot budget for the uncached fallback.
637 retained = 0_int_8
638 DO jslot = kpoint_slot_ks, kpoint_slot_t
639 DO ispin = 1, ctx%nspin
640 IF (.NOT. ALLOCATED(ctx%rs(ispin, jslot)%blocks)) cycle
641 retained = retained + 128_int_8*SIZE(ctx%rs(ispin, jslot)%blocks, kind=int_8)
642 IF (ASSOCIATED(ctx%rs(ispin, jslot)%values)) THEN
643 retained = retained + 8_int_8*SIZE(ctx%rs(ispin, jslot)%values, kind=int_8)
644 END IF
645 END DO
646 END DO
647 storage = ctx%rs_route%storage + retained + int(nrow, int_8)* &
648 (merge(0_int_8, 8_int_8*ctx%rs_route%nvalue, para_env%num_pe == 1) + &
649 128_int_8*SIZE(ctx%rs_route%entry, 2))
650 CALL para_env%max(storage)
651 ready = compatible == 1 .AND. storage <= max_bytes
652 IF (ready) THEN
653 DO ispin = 1, nrow
654 row = ispin
655 IF (PRESENT(matrix_row)) row = matrix_row
656 CALL kpoint_operator_rs_capture(rsmat, row, ctx%rs_route, ctx%kpoints, ctx%rs(ispin, slot))
657 END DO
658 END IF
659 IF (PRESENT(prepared)) prepared = ready
660 CALL timestop(handle)
661
662 END SUBROUTINE kpoint_operator_prepare
663
664! **************************************************************************************************
665!> \brief Release one captured row without invalidating other slots or routing.
666!> \param snapshot row to clear
667! **************************************************************************************************
668 SUBROUTINE kpoint_operator_rs_clear(snapshot)
669
670 TYPE(kpoint_operator_rs_type), INTENT(INOUT) :: snapshot
671
672 IF (ALLOCATED(snapshot%blocks)) DEALLOCATE (snapshot%blocks, snapshot%nonzero)
673 IF (ASSOCIATED(snapshot%values)) DEALLOCATE (snapshot%values)
674 NULLIFY (snapshot%source)
675 snapshot%matrix_row = 0
676
677 END SUBROUTINE kpoint_operator_rs_clear
678
679! **************************************************************************************************
680!> \brief Match an explicitly prepared row, sharing spin-free rows across spins.
681!> \param ctx operator context
682!> \param source live array identity identifying the prepared operation
683!> \param ispin instance spin
684!> \param slot prepared slot
685!> \param matrix_row operator row
686!> \return snapshot spin index, or zero for the transfer backend
687! **************************************************************************************************
688 FUNCTION kpoint_operator_rs_match(ctx, source, ispin, slot, matrix_row) RESULT(irs)
689
690 TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
691 TYPE(dbcsr_p_type), POINTER :: source(:, :)
692 INTEGER, INTENT(IN) :: ispin, slot, matrix_row
693 INTEGER :: irs
694
695 INTEGER :: i
696
697 irs = 0
698 IF (.NOT. ALLOCATED(ctx%rs)) RETURN
699 i = ispin
700 IF (.NOT. ASSOCIATED(ctx%rs(i, slot)%source)) i = 1
701 IF (ASSOCIATED(ctx%rs(i, slot)%source, source) .AND. ctx%rs(i, slot)%matrix_row == matrix_row) irs = i
702
703 END FUNCTION kpoint_operator_rs_match
704
705! **************************************************************************************************
706!> \brief Cache whole-block routing to the FM tiles of one representative group.
707!> Equivalent ranks in the other groups receive only their tile's blocks.
708!> Boundary blocks may overlap several tiles. Missing blocks are packed
709!> as zero, permitting different operator sparsities on a fixed AO layout.
710!> \param kpoints initialized parallel environment and neighbor lists
711!> \param source operator layout template
712!> \param matrix_struct target group layout
713!> \param max_bytes per-rank budget for routing and capture buffers
714!> \param route context-owned structural route
715!> \param ready false collectively if unsupported or over budget
716! **************************************************************************************************
717 SUBROUTINE kpoint_operator_rs_plan_create(kpoints, source, matrix_struct, max_bytes, route, ready)
718
719 TYPE(kpoint_type), POINTER :: kpoints
720 TYPE(dbcsr_p_type), POINTER :: source(:, :)
721 TYPE(cp_fm_struct_type), POINTER :: matrix_struct
722 INTEGER(KIND=int_8), INTENT(IN) :: max_bytes
723 TYPE(kpoint_operator_rs_route_type), INTENT(OUT) :: route
724 LOGICAL, INTENT(OUT) :: ready
725
726 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_rs_plan_create'
727
728 INTEGER :: compatible, dest, handle, i, iblock, &
729 nblock, ncol, nrow, p, pos, rank
730 INTEGER(KIND=int_8) :: nr(2), ns(2), storage
731 INTEGER(KIND=int_8), ALLOCATABLE :: recv_size(:, :), send_size(:, :)
732 INTEGER, ALLOCATABLE :: identities(:, :), owner(:), &
733 position(:, :), send_metadata(:, :)
734 INTEGER, POINTER :: col_size(:), row_size(:)
735 LOGICAL, ALLOCATABLE :: needed(:)
736 TYPE(mp_para_env_type), POINTER :: inter, world
737
738 CALL timeset(routinen, handle)
739 ready = .false.
740 world => kpoints%blacs_env_all%para_env
741 inter => kpoints%para_env_inter_kp
742 CALL dbcsr_get_info(source(1, 1)%matrix, row_blk_size=row_size, col_blk_size=col_size)
743 CALL kp_transform_plan_create(route%traversal, kpoints%sab_nl, kpoints%cell_to_index, &
744 SIZE(source, 2), block_template=source(1, 1)%matrix)
745 compatible = merge(1, 0, route%traversal%symmetric)
746 IF (sum(row_size) /= matrix_struct%nrow_global .OR. &
747 sum(col_size) /= matrix_struct%ncol_global) compatible = 0
748 IF (SIZE(row_size) == SIZE(col_size)) THEN
749 IF (any(row_size /= col_size)) compatible = 0
750 ELSE
751 compatible = 0
752 END IF
753 CALL world%min(compatible)
754 IF (compatible == 0) THEN
755 CALL timestop(handle)
756 RETURN
757 END IF
758 route%nimage = SIZE(source, 2)
759 ALLOCATE (route%row_size(SIZE(row_size)), route%col_size(SIZE(col_size)))
760 route%row_size(:) = row_size
761 route%col_size(:) = col_size
762 ALLOCATE (identities(2, 0:world%num_pe - 1), owner(0:kpoints%para_env_kp%num_pe - 1), &
763 needed(0:kpoints%para_env_kp%num_pe - 1))
764 CALL world%allgather([kpoints%para_env_kp%mepos, inter%mepos], identities)
765 owner(:) = -1
766 DO rank = 0, world%num_pe - 1
767 IF (identities(2, rank) == 0) owner(identities(1, rank)) = rank
768 END DO
769 cpassert(all(owner >= 0))
770 ALLOCATE (send_size(2, 0:world%num_pe - 1), recv_size(2, 0:world%num_pe - 1))
771 send_size(:, :) = 0_int_8
772 DO i = 1, route%traversal%nentry
773 nrow = row_size(route%traversal%row(i))
774 ncol = col_size(route%traversal%col(i))
775 CALL kpoint_operator_rs_owners(route, matrix_struct, i, needed)
776 DO p = 0, SIZE(owner) - 1
777 IF (.NOT. needed(p)) cycle
778 dest = owner(p)
779 send_size(1, dest) = send_size(1, dest) + 1_int_8
780 send_size(2, dest) = send_size(2, dest) + int(nrow, int_8)*ncol
781 END DO
782 END DO
783 CALL world%alltoall(send_size, recv_size, 2)
784 ! Only inter-rank zero receives from the world. The other groups hold
785 ! equivalent FM tiles, as required by kpoint_env_initialize.
786 ns = sum(send_size, dim=2)
787 nr = sum(recv_size, dim=2)
788 CALL inter%bcast(nr, 0)
789 ! Conservative peak: traversal, retained route, temporary metadata and
790 ! one snapshot. The full batch budget is checked again by prepare.
791 storage = 40_int_8*route%traversal%nentry + &
792 4_int_8*(SIZE(row_size) + SIZE(col_size) + matrix_struct%nrow_global + matrix_struct%ncol_global) + &
793 48_int_8*ns(1) + merge(0_int_8, 8_int_8*ns(2), world%num_pe == 1) + 68_int_8*nr(1) + &
794 128_int_8*world%num_pe + 8_int_8*SIZE(owner) + 8_int_8
795 route%storage = storage
796 storage = storage + merge(0_int_8, 8_int_8*nr(2), world%num_pe == 1) + 128_int_8*nr(1)
797 ! Borrowed values do not consume the snapshot budget; still bound integer route offsets.
798 IF (maxval(ns) >= int(huge(0), int_8) .OR. maxval(nr) >= int(huge(0), int_8)) THEN
799 storage = max(storage, max_bytes + 1_int_8)
800 END IF
801 CALL world%max(storage)
802 IF (storage > max_bytes) THEN
803 CALL timestop(handle)
804 RETURN
805 END IF
806 ! The byte bound also keeps all counts/displacements within default INTEGER.
807 nblock = int(nr(1))
808 route%nvalue = int(nr(2))
809 ALLOCATE (route%scount(2, 0:world%num_pe - 1), route%rcount(2, 0:world%num_pe - 1), &
810 route%sdisp(2, 0:world%num_pe - 1), route%rdisp(2, 0:world%num_pe - 1))
811 route%scount(:, :) = int(send_size)
812 route%rcount(:, :) = int(recv_size)
813 IF (SIZE(owner) == 1) CALL inter%bcast(route%rcount, 0)
814 route%sdisp(:, 0) = 0
815 route%rdisp(:, 0) = 0
816 DO rank = 1, world%num_pe - 1
817 route%sdisp(:, rank) = route%sdisp(:, rank - 1) + route%scount(:, rank - 1)
818 route%rdisp(:, rank) = route%rdisp(:, rank - 1) + route%rcount(:, rank - 1)
819 END DO
820 ALLOCATE (route%send_entry(int(ns(1))), route%send_first(int(ns(1)) + 1), &
821 route%send_values(merge(0, int(ns(2)), world%num_pe == 1)), send_metadata(8, int(ns(1))), &
822 position(2, 0:world%num_pe - 1), route%entry(8, nblock), &
823 route%first(nblock + 1))
824 position(:, :) = route%sdisp
825 DO i = 1, route%traversal%nentry
826 nrow = row_size(route%traversal%row(i))
827 ncol = col_size(route%traversal%col(i))
828 CALL kpoint_operator_rs_owners(route, matrix_struct, i, needed)
829 DO p = 0, SIZE(owner) - 1
830 IF (.NOT. needed(p)) cycle
831 dest = owner(p)
832 iblock = position(1, dest) + 1
833 pos = position(2, dest) + 1
834 route%send_entry(iblock) = i
835 route%send_first(iblock) = pos
836 send_metadata(1:4, iblock) = [route%traversal%row_offset(i), route%traversal%col_offset(i), nrow, ncol]
837 send_metadata(5:7, iblock) = nint(route%traversal%symmetry_sign(i))*route%traversal%cell(:, i)
838 send_metadata(8, iblock) = merge(1, 0, route%traversal%row(i) /= route%traversal%col(i))
839 position(:, dest) = position(:, dest) + [1, nrow*ncol]
840 END DO
841 END DO
842 route%send_first(int(ns(1)) + 1) = int(ns(2)) + 1
843 IF (SIZE(owner) == 1) THEN
844 CALL world%allgatherv(send_metadata, route%entry, 8*route%rcount(1, :), 8*route%rdisp(1, :))
845 ELSE
846 CALL world%alltoall(send_metadata, 8*route%scount(1, :), 8*route%sdisp(1, :), &
847 route%entry, 8*route%rcount(1, :), 8*route%rdisp(1, :))
848 CALL inter%bcast(route%entry, 0)
849 END IF
850 route%first(1) = 1
851 DO i = 1, nblock
852 route%first(i + 1) = route%first(i) + route%entry(3, i)*route%entry(4, i)
853 END DO
854 cpassert(route%first(nblock + 1) == route%nvalue + 1)
855 route%matrix_struct => matrix_struct
856 CALL cp_fm_struct_retain(route%matrix_struct)
857 CALL kpoint_operator_rs_tile_map(route)
858 ready = .true.
859 CALL timestop(handle)
860
861 END SUBROUTINE kpoint_operator_rs_plan_create
862
863! **************************************************************************************************
864!> \brief Cache local intervals intersecting each routed block and its Hermitian mirror.
865!> FM indices are ordered globally; dummy indices on empty tiles lie outside the matrix.
866!> \param route fixed routing and target distribution
867! **************************************************************************************************
868 SUBROUTINE kpoint_operator_rs_tile_map(route)
869 TYPE(kpoint_operator_rs_route_type), INTENT(INOUT) :: route
870
871 INTEGER :: c0, i, nc, nr, r0, rows(0:route%matrix_struct%nrow_global), &
872 cols(0:route%matrix_struct%ncol_global)
873
874 rows(:) = 0
875 cols(:) = 0
876 DO i = 1, SIZE(route%matrix_struct%row_indices)
877 r0 = route%matrix_struct%row_indices(i)
878 IF (r0 >= 1 .AND. r0 <= ubound(rows, 1)) rows(r0) = 1
879 END DO
880 DO i = 1, SIZE(route%matrix_struct%col_indices)
881 c0 = route%matrix_struct%col_indices(i)
882 IF (c0 >= 1 .AND. c0 <= ubound(cols, 1)) cols(c0) = 1
883 END DO
884 DO i = 1, ubound(rows, 1)
885 rows(i) = rows(i) + rows(i - 1)
886 END DO
887 DO i = 1, ubound(cols, 1)
888 cols(i) = cols(i) + cols(i - 1)
889 END DO
890 route%nrow = rows(ubound(rows, 1))
891 route%ncol = cols(ubound(cols, 1))
892 ALLOCATE (route%tile(4, 2, SIZE(route%entry, 2)))
893 DO i = 1, SIZE(route%entry, 2)
894 r0 = route%entry(1, i)
895 c0 = route%entry(2, i)
896 nr = route%entry(3, i)
897 nc = route%entry(4, i)
898 route%tile(:, 1, i) = [rows(r0 - 1) + 1, rows(r0 + nr - 1), cols(c0 - 1) + 1, cols(c0 + nc - 1)]
899 route%tile(:, 2, i) = [rows(c0 - 1) + 1, rows(c0 + nc - 1), cols(r0 - 1) + 1, cols(r0 + nr - 1)]
900 END DO
901 END SUBROUTINE kpoint_operator_rs_tile_map
902
903! **************************************************************************************************
904!> \brief Find group-local FM owners intersecting a stored block or its mirror.
905!> \param route source AO block layout and neighbor traversal
906!> \param fms target FM layout within each group
907!> \param i traversal entry
908!> \param needed recipient group-local ranks, without duplicate mirror sends
909! **************************************************************************************************
910 SUBROUTINE kpoint_operator_rs_owners(route, fms, i, needed)
911
912 TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
913 TYPE(cp_fm_struct_type), INTENT(IN) :: fms
914 INTEGER, INTENT(IN) :: i
915 LOGICAL, INTENT(OUT) :: needed(0:)
916
917 INTEGER :: c0, j, nc, nr, pc, pr, r0
918 LOGICAL :: rows(2, 0:fms%context%num_pe(1) - 1), cols(2, 0:fms%context%num_pe(2) - 1)
919
920 r0 = route%traversal%row_offset(i)
921 c0 = route%traversal%col_offset(i)
922 nr = route%row_size(route%traversal%row(i))
923 nc = route%col_size(route%traversal%col(i))
924 rows(:, :) = .false.
925 cols(:, :) = .false.
926 DO j = r0, r0 + nr - 1
927 rows(1, fms%g2p_row(j)) = .true.
928 cols(2, fms%g2p_col(j)) = .true.
929 END DO
930 DO j = c0, c0 + nc - 1
931 rows(2, fms%g2p_row(j)) = .true.
932 cols(1, fms%g2p_col(j)) = .true.
933 END DO
934 IF (route%traversal%row(i) == route%traversal%col(i)) rows(2, :) = .false.
935 needed(:) = .false.
936 DO pc = 0, fms%context%num_pe(2) - 1
937 DO pr = 0, fms%context%num_pe(1) - 1
938 needed(fms%context%blacs2mpi(pr, pc)) = &
939 (rows(1, pr) .AND. cols(1, pc)) .OR. (rows(2, pr) .AND. cols(2, pc))
940 END DO
941 END DO
942
943 END SUBROUTINE kpoint_operator_rs_owners
944
945! **************************************************************************************************
946!> \brief Check the source AO layout, not its values or sparse block presence.
947!> \param route fixed context route
948!> \param source current operator images
949!> \param matrix_row row to capture
950!> \return whether the row has the route's image count and AO layout
951! **************************************************************************************************
952 FUNCTION kpoint_operator_rs_compatible(route, source, matrix_row) RESULT(compatible)
953
954 TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
955 TYPE(dbcsr_p_type), POINTER :: source(:, :)
956 INTEGER, INTENT(IN) :: matrix_row
957 LOGICAL :: compatible
958
959 INTEGER :: image
960 INTEGER, POINTER :: col_size(:), row_size(:)
961
962 compatible = .false.
963 IF (SIZE(source, 2) /= route%nimage) RETURN
964 IF (matrix_row < 1 .OR. matrix_row > SIZE(source, 1)) RETURN
965 DO image = 1, SIZE(source, 2)
966 CALL dbcsr_get_info(source(matrix_row, image)%matrix, row_blk_size=row_size, col_blk_size=col_size)
967 IF (SIZE(row_size) /= SIZE(route%row_size) .OR. SIZE(col_size) /= SIZE(route%col_size)) RETURN
968 IF (any(row_size /= route%row_size) .OR. any(col_size /= route%col_size)) RETURN
969 END DO
970 compatible = .true.
971
972 END FUNCTION kpoint_operator_rs_compatible
973
974! **************************************************************************************************
975!> \brief Capture current values using cached routing, with no per-k communication.
976!> \param source current real-space operator images
977!> \param matrix_row row to snapshot
978!> \param route context-owned metadata, counts and reusable packing buffer
979!> \param kpoints initialized host parallel environment
980!> \param snapshot immutable tile-local values
981! **************************************************************************************************
982 SUBROUTINE kpoint_operator_rs_capture(source, matrix_row, route, kpoints, snapshot)
983
984 TYPE(dbcsr_p_type), POINTER :: source(:, :)
985 INTEGER, INTENT(IN) :: matrix_row
986 TYPE(kpoint_operator_rs_route_type), INTENT(INOUT) :: route
987 TYPE(kpoint_type), POINTER :: kpoints
988 TYPE(kpoint_operator_rs_type), INTENT(OUT) :: snapshot
989
990 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_rs_capture'
991
992 INTEGER :: col, handle, i, iblock, ncol, nrow, pos, &
993 row
994 LOGICAL :: found
995 REAL(kind=dp), POINTER :: block(:, :)
996 TYPE(mp_para_env_type), POINTER :: inter, world
997
998 CALL timeset(routinen, handle)
999 world => kpoints%blacs_env_all%para_env
1000 inter => kpoints%para_env_inter_kp
1001 ALLOCATE (snapshot%blocks(SIZE(route%entry, 2)), snapshot%nonzero(SIZE(route%entry, 2)))
1002 snapshot%source => source
1003 snapshot%matrix_row = matrix_row
1004 IF (world%num_pe == 1) THEN
1005 ! No value copy on a single world rank. Lookups happen before workers start.
1006 DO iblock = 1, SIZE(route%send_entry)
1007 i = route%send_entry(iblock)
1008 CALL dbcsr_get_readonly_block_p(source(matrix_row, route%traversal%image(i))%matrix, &
1009 route%traversal%row(i), route%traversal%col(i), block, found)
1010 snapshot%nonzero(iblock) = found
1011 IF (found) THEN
1012 cpassert(SIZE(block, 1) == route%entry(3, iblock))
1013 cpassert(SIZE(block, 2) == route%entry(4, iblock))
1014 snapshot%blocks(iblock)%data => block
1015 END IF
1016 END DO
1017 CALL timestop(handle)
1018 RETURN
1019 END IF
1020 ALLOCATE (snapshot%values(route%nvalue))
1021 route%send_values(:) = 0.0_dp
1022!$OMP PARALLEL DO DEFAULT(NONE) SHARED(source, matrix_row, route) &
1023!$OMP PRIVATE(iblock, i, block, found, nrow, ncol, pos, row, col) SCHEDULE(STATIC)
1024 DO iblock = 1, SIZE(route%send_entry)
1025 i = route%send_entry(iblock)
1026 CALL dbcsr_get_readonly_block_p(source(matrix_row, route%traversal%image(i))%matrix, &
1027 route%traversal%row(i), route%traversal%col(i), block, found)
1028 IF (.NOT. found) cycle
1029 nrow = route%row_size(route%traversal%row(i))
1030 ncol = route%col_size(route%traversal%col(i))
1031 cpassert(SIZE(block, 1) == nrow .AND. SIZE(block, 2) == ncol)
1032 pos = route%send_first(iblock)
1033 DO col = 1, ncol
1034 DO row = 1, nrow
1035 route%send_values(pos) = block(row, col)
1036 pos = pos + 1
1037 END DO
1038 END DO
1039 END DO
1040!$OMP END PARALLEL DO
1041 IF (kpoints%para_env_kp%num_pe == 1) THEN
1042 CALL world%allgatherv(route%send_values, snapshot%values, route%rcount(2, :), route%rdisp(2, :))
1043 ELSE
1044 CALL world%alltoall(route%send_values, route%scount(2, :), route%sdisp(2, :), &
1045 snapshot%values, route%rcount(2, :), route%rdisp(2, :))
1046 CALL inter%bcast(snapshot%values, 0)
1047 END IF
1048 DO iblock = 1, SIZE(route%entry, 2)
1049 snapshot%nonzero(iblock) = any(snapshot%values(route%first(iblock):route%first(iblock + 1) - 1) /= 0.0_dp)
1050 snapshot%blocks(iblock)%data(1:route%entry(3, iblock), 1:route%entry(4, iblock)) => &
1051 snapshot%values(route%first(iblock):route%first(iblock + 1) - 1)
1052 END DO
1053 snapshot%source => source
1054 snapshot%matrix_row = matrix_row
1055 CALL timestop(handle)
1056
1057 END SUBROUTINE kpoint_operator_rs_capture
1058
1059! **************************************************************************************************
1060!> \brief Deliver one prepared operator, building a bounded window of local k points on a miss.
1061!> The callers retain their round/spin/slot ordering. Each row has its own
1062!> window so interleaved spins and H/S/T do not evict one another.
1063!> Only prepared values are read ahead; transfer and grid paths are unchanged.
1064!> \param ctx serialized operator context
1065!> \param irs captured row, possibly shared by spin-free instances
1066!> \param slot prepared operator slot
1067!> \param ik global k-point index owned by this group
1068!> \param TARGET optional complex output on the prepared distribution
1069!> \param fm_re optional real output instead of target; fm_im may be omitted for a real operator
1070!> \param fm_im optional imaginary output on the same distribution
1071! **************************************************************************************************
1072 SUBROUTINE kpoint_operator_rs_build(ctx, irs, slot, ik, TARGET, fm_re, fm_im)
1073
1074 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1075 INTEGER, INTENT(IN) :: irs, slot, ik
1076 TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
1077 TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1078
1079 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_rs_build'
1080
1081 INTEGER :: first, handle, k, last, ncol, nrow, width
1082 INTEGER(KIND=int_8) :: bytes_per_k
1083 LOGICAL :: layout_ok
1084 REAL(kind=dp), DIMENSION(:, :), POINTER :: xkp
1085
1086 CALL timeset(routinen, handle)
1087 IF (PRESENT(TARGET)) THEN
1088 layout_ok = cp_fm_struct_equivalent(target%matrix_struct, ctx%rs_route%matrix_struct)
1089 ELSE
1090 cpassert(PRESENT(fm_re))
1091 layout_ok = cp_fm_struct_equivalent(fm_re%matrix_struct, ctx%rs_route%matrix_struct)
1092 IF (PRESENT(fm_im)) THEN
1093 layout_ok = layout_ok .AND. cp_fm_struct_equivalent(fm_im%matrix_struct, ctx%rs_route%matrix_struct)
1094 END IF
1095 END IF
1096 cpassert(layout_ok)
1097 ! Avoid nested component array sections in GCC 16 bounds checks.
1098 xkp => ctx%kpoints%xkp
1099 first = ctx%kpoints%kp_range(1)
1100 last = ctx%kpoints%kp_range(2)
1101 cpassert(ik >= first .AND. ik <= last)
1102 nrow = ctx%rs_route%nrow
1103 ncol = ctx%rs_route%ncol
1104 bytes_per_k = 16_int_8*(int(nrow, int_8)*ncol + SIZE(ctx%rs_route%entry, 2, kind=int_8))
1105 width = min(rs_batch_size, last - first + 1, &
1106 int(rs_batch_bytes/(int(ctx%nspin*kpoint_num_slots, int_8)*max(1_int_8, bytes_per_k))))
1107 IF (width < 2 .OR. nrow == 0 .OR. ncol == 0) THEN
1108 ! No extra dense tile for a large/empty layout or a single local k point.
1109 CALL kpoint_operator_rs_build_batch(ctx%rs(irs, slot), ctx%rs_route, xkp(:, ik:ik), &
1110 TARGET=TARGET, fm_re=fm_re, fm_im=fm_im)
1111 ELSE
1112 associate(work => ctx%rs_batch(irs, slot))
1113 IF (.NOT. ALLOCATED(work%data)) ALLOCATE (work%data(width, nrow, ncol))
1114 IF (work%first == 0 .OR. ik < work%first .OR. ik >= work%first + width) THEN
1115 work%first = first + ((ik - first)/width)*width
1116 CALL kpoint_operator_rs_build_batch(ctx%rs(irs, slot), ctx%rs_route, &
1117 xkp(:, work%first:min(work%first + width - 1, last)), batch=work%data)
1118 END IF
1119 k = ik - work%first + 1
1120 ! Deliver a work copy: eigensolvers may overwrite H/S without corrupting the window.
1121 IF (PRESENT(TARGET)) THEN
1122 target%local_data(1:nrow, 1:ncol) = work%data(k, :, :)
1123 target%local_data(nrow + 1:, :) = z_zero
1124 target%local_data(:, ncol + 1:) = z_zero
1125 ELSE
1126 cpassert(PRESENT(fm_re))
1127 fm_re%local_data(1:nrow, 1:ncol) = real(work%data(k, :, :), dp)
1128 IF (PRESENT(fm_im)) fm_im%local_data(1:nrow, 1:ncol) = aimag(work%data(k, :, :))
1129 fm_re%local_data(nrow + 1:, :) = 0.0_dp
1130 fm_re%local_data(:, ncol + 1:) = 0.0_dp
1131 IF (PRESENT(fm_im)) fm_im%local_data(nrow + 1:, :) = 0.0_dp
1132 IF (PRESENT(fm_im)) fm_im%local_data(:, ncol + 1:) = 0.0_dp
1133 END IF
1134 END associate
1135 END IF
1136 IF (ctx%kpoints%use_real_wfn) THEN
1137 IF (PRESENT(TARGET)) target%local_data(:, :) = cmplx(real(target%local_data, dp), 0.0_dp, dp)
1138 IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
1139 END IF
1140 CALL timestop(handle)
1141
1142 END SUBROUTINE kpoint_operator_rs_build
1143
1144! **************************************************************************************************
1145!> \brief Accumulate several k points in one route/value traversal, without MPI or BLACS.
1146!> The leading batch dimension makes updates to different k points contiguous.
1147!> Columns remain independent OpenMP work items; each target retains the
1148!> original normal/mirrored contribution order. Phase scratch is call-local.
1149!> The singleton fallback writes directly to the supplied CFM or FM pair.
1150!> \param snapshot immutable real-space contributions
1151!> \param route structural block routing and target tile map
1152!> \param xkp actual coordinates, one column per k point
1153!> \param batch optional packed complex output (k point, local row, local column)
1154!> \param TARGET optional single complex output instead of batch
1155!> \param fm_re optional single real output instead of batch/target
1156!> \param fm_im optional single imaginary output
1157!> \param phase_work optional private phase storage supplied by a worker
1158!> \param parallel_columns split columns across threads, disabled inside k-point workers
1159! **************************************************************************************************
1160 SUBROUTINE kpoint_operator_rs_build_batch(snapshot, route, xkp, batch, TARGET, fm_re, fm_im, phase_work, parallel_columns)
1161
1162 TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
1163 TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
1164 REAL(kind=dp), INTENT(IN) :: xkp(:, :)
1165 COMPLEX(KIND=dp), CONTIGUOUS, INTENT(INOUT), &
1166 OPTIONAL, TARGET :: batch(:, :, :)
1167 TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
1168 TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1169 COMPLEX(KIND=dp), INTENT(INOUT), OPTIONAL, TARGET :: phase_work(:, :)
1170 LOGICAL, INTENT(IN), OPTIONAL :: parallel_columns
1171
1172 COMPLEX(KIND=dp), ALLOCATABLE, TARGET :: owned_phase(:, :)
1173 COMPLEX(KIND=dp), CONTIGUOUS, POINTER :: output(:, :, :)
1174 COMPLEX(KIND=dp), POINTER :: phase(:, :)
1175 INTEGER :: first, iblock, k, last, nbatch, ncol, &
1176 nthreads, thread
1177 LOGICAL :: parallel
1178 REAL(kind=dp) :: arg
1179
1180 nbatch = SIZE(xkp, 2)
1181 cpassert(nbatch >= 1 .AND. nbatch <= rs_batch_size .AND. SIZE(xkp, 1) == 3)
1182 cpassert(count([PRESENT(batch), PRESENT(TARGET), PRESENT(fm_re)]) == 1)
1183 NULLIFY (output)
1184 IF (PRESENT(batch)) THEN
1185 cpassert(SIZE(batch, 1) >= nbatch)
1186 output => batch
1187 ELSE
1188 cpassert(nbatch == 1)
1189 IF (PRESENT(TARGET)) THEN
1190 output(1:1, 1:SIZE(target%local_data, 1), 1:SIZE(target%local_data, 2)) => target%local_data
1191 ELSE
1192 cpassert(PRESENT(fm_re))
1193 fm_re%local_data(:, :) = 0.0_dp
1194 IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
1195 END IF
1196 END IF
1197 IF (ASSOCIATED(output)) output(:, :, :) = z_zero
1198 IF (route%nrow == 0 .OR. route%ncol == 0) RETURN
1199 IF (PRESENT(phase_work)) THEN
1200 cpassert(SIZE(phase_work, 1) >= nbatch)
1201 cpassert(SIZE(phase_work, 2) == SIZE(route%entry, 2))
1202 phase => phase_work(1:nbatch, :)
1203 ELSE
1204 ALLOCATE (owned_phase(nbatch, SIZE(route%entry, 2)))
1205 phase => owned_phase
1206 END IF
1207 DO iblock = 1, SIZE(route%entry, 2)
1208 IF (.NOT. snapshot%nonzero(iblock)) cycle
1209 DO k = 1, nbatch
1210 arg = twopi*dot_product(real(route%entry(5:7, iblock), dp), xkp(:, k))
1211 phase(k, iblock) = cmplx(cos(arg), sin(arg), kind=dp)
1212 END DO
1213 END DO
1214 ncol = route%ncol
1215 parallel = .true.
1216 IF (PRESENT(parallel_columns)) parallel = parallel_columns
1217 nthreads = 1
1218!$ IF (parallel) nthreads = MIN(omp_get_max_threads(), ncol)
1219 IF (nthreads > 1) THEN
1220!$OMP PARALLEL DEFAULT(NONE) NUM_THREADS(nthreads) SHARED(snapshot,route,phase,output,fm_re,fm_im,ncol) &
1221!$OMP PRIVATE(first,last,thread,nthreads)
1222 thread = 0
1223 nthreads = 1
1224!$ thread = omp_get_thread_num()
1225!$ nthreads = omp_get_num_threads()
1226 first = ncol*thread/nthreads + 1
1227 last = ncol*(thread + 1)/nthreads
1228 CALL kpoint_operator_rs_accumulate(snapshot, route, phase, first, last, output, fm_re, fm_im)
1229!$OMP END PARALLEL
1230 ELSE
1231 ! Called inside a k-point worker: no nested parallel region, timer or allocation.
1232 CALL kpoint_operator_rs_accumulate(snapshot, route, phase, 1, ncol, output, fm_re, fm_im)
1233 END IF
1234
1235 END SUBROUTINE kpoint_operator_rs_build_batch
1236
1237! **************************************************************************************************
1238!> \brief Shared borrowed/packed phase-sum kernel; each caller owns a disjoint column interval.
1239!> \param snapshot read-only block views
1240!> \param route structural layout
1241!> \param phase precomputed phases (k, block)
1242!> \param first first local column owned by this call
1243!> \param last last local column owned by this call
1244!> \param output packed complex output, or disassociated for the FM pair
1245!> \param fm_re real output
1246!> \param fm_im imaginary output
1247! **************************************************************************************************
1248 SUBROUTINE kpoint_operator_rs_accumulate(snapshot, route, phase, first, last, output, fm_re, fm_im)
1249 TYPE(kpoint_operator_rs_type), INTENT(IN) :: snapshot
1250 TYPE(kpoint_operator_rs_route_type), INTENT(IN) :: route
1251 COMPLEX(KIND=dp), INTENT(IN) :: phase(:, :)
1252 INTEGER, INTENT(IN) :: first, last
1253 COMPLEX(KIND=dp), CONTIGUOUS, INTENT(IN), POINTER :: output(:, :, :)
1254 TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1255
1256 COMPLEX(KIND=dp) :: factors(rs_batch_size), value
1257 COMPLEX(KIND=dp), CONTIGUOUS, POINTER :: full(:, :)
1258 INTEGER :: c0, col, hi, iblock, j, k, lo, nc, nr, &
1259 part, r0, row, source_col, source_row
1260 REAL(kind=dp) :: element
1261
1262 IF (first > last) RETURN
1263 NULLIFY (full)
1264 IF (ASSOCIATED(output)) THEN
1265 IF (SIZE(output, 1) == 1 .AND. &
1266 route%nrow == route%matrix_struct%nrow_global .AND. &
1267 route%ncol == route%matrix_struct%ncol_global) THEN
1268 full(1:SIZE(output, 2), 1:SIZE(output, 3)) => output
1269 END IF
1270 END IF
1271 DO iblock = 1, SIZE(route%entry, 2)
1272 IF (.NOT. snapshot%nonzero(iblock)) cycle
1273 r0 = route%entry(1, iblock)
1274 c0 = route%entry(2, iblock)
1275 nr = route%entry(3, iblock)
1276 nc = route%entry(4, iblock)
1277 IF (ASSOCIATED(full)) THEN
1278 ! Update whole contiguous blocks, clipped to this worker's columns.
1279 lo = max(first, c0)
1280 hi = min(last, c0 + nc - 1)
1281 IF (lo <= hi) THEN
1282 full(r0:r0 + nr - 1, lo:hi) = full(r0:r0 + nr - 1, lo:hi) + &
1283 phase(1, iblock)*snapshot%blocks(iblock)%data(:, lo - c0 + 1:hi - c0 + 1)
1284 END IF
1285 IF (route%entry(8, iblock) /= 0) THEN
1286 lo = max(first, r0)
1287 hi = min(last, r0 + nr - 1)
1288 IF (lo <= hi) THEN
1289 full(c0:c0 + nc - 1, lo:hi) = full(c0:c0 + nc - 1, lo:hi) + &
1290 conjg(phase(1, iblock))* &
1291 transpose(snapshot%blocks(iblock)%data(lo - r0 + 1:hi - r0 + 1, :))
1292 END IF
1293 END IF
1294 cycle
1295 END IF
1296 DO part = 1, 2
1297 IF (part == 2 .AND. route%entry(8, iblock) == 0) cycle
1298 lo = max(first, route%tile(3, part, iblock))
1299 hi = min(last, route%tile(4, part, iblock))
1300 IF (lo > hi) cycle
1301 factors(1:SIZE(phase, 1)) = phase(:, iblock)
1302 IF (part == 2) factors(1:SIZE(phase, 1)) = conjg(factors(1:SIZE(phase, 1)))
1303 DO j = lo, hi
1304 col = route%matrix_struct%col_indices(j)
1305 DO row = route%tile(1, part, iblock), route%tile(2, part, iblock)
1306 source_row = route%matrix_struct%row_indices(row)
1307 source_col = col
1308 IF (part == 2) THEN
1309 source_col = source_row
1310 source_row = col
1311 END IF
1312 element = snapshot%blocks(iblock)%data(source_row - r0 + 1, source_col - c0 + 1)
1313 IF (ASSOCIATED(output)) THEN
1314!$OMP SIMD
1315 DO k = 1, SIZE(phase, 1)
1316 output(k, row, j) = output(k, row, j) + factors(k)*element
1317 END DO
1318 ELSE
1319 value = factors(1)*element
1320 fm_re%local_data(row, j) = fm_re%local_data(row, j) + real(value, dp)
1321 IF (PRESENT(fm_im)) fm_im%local_data(row, j) = fm_im%local_data(row, j) + aimag(value)
1322 END IF
1323 END DO
1324 END DO
1325 END DO
1326 END DO
1327 END SUBROUTINE kpoint_operator_rs_accumulate
1328
1329! **************************************************************************************************
1330!> \brief Size per-worker phase scratch for direct assembly into a solver matrix.
1331!> \param ctx prepared context
1332!> \param nentry number of route entries
1333! **************************************************************************************************
1334 SUBROUTINE kpoint_operator_work_size(ctx, nentry)
1335 TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
1336 INTEGER, INTENT(OUT) :: nentry
1337
1338 cpassert(ALLOCATED(ctx%rs_route))
1339 nentry = SIZE(ctx%rs_route%entry, 2)
1340 END SUBROUTINE kpoint_operator_work_size
1341
1342! **************************************************************************************************
1343!> \brief Evaluate an immutable prepared operator with exclusively caller-owned output and scratch.
1344!> Prepare/refresh/release remain collective and serialized. Evaluate contains no MPI,
1345!> DBCSR lookup, window update or nested OpenMP region. Source blocks must stay unchanged.
1346!> \param ctx prepared read-only context
1347!> \param ispin instance spin
1348!> \param slot operator slot
1349!> \param ik global k index owned by this group
1350!> \param phase private phase scratch, sized with work_size
1351!> \param TARGET optional private complex singleton output
1352!> \param fm_re optional private real output instead of TARGET
1353!> \param fm_im optional private imaginary output, accompanying fm_re for a complex operator
1354! **************************************************************************************************
1355 SUBROUTINE kpoint_operator_evaluate(ctx, ispin, slot, ik, phase, TARGET, fm_re, fm_im)
1356 TYPE(kpoint_operator_context_type), INTENT(IN) :: ctx
1357 INTEGER, INTENT(IN) :: ispin, slot, ik
1358 TYPE(cp_cfm_type), INTENT(INOUT), OPTIONAL :: target
1359 COMPLEX(KIND=dp), INTENT(INOUT), TARGET :: phase(:, :)
1360 TYPE(cp_fm_type), INTENT(INOUT), OPTIONAL :: fm_re, fm_im
1361
1362 REAL(kind=dp), POINTER :: xkp(:, :)
1363 INTEGER :: row
1364
1365 cpassert(ik >= ctx%kpoints%kp_range(1))
1366 cpassert(ik <= ctx%kpoints%kp_range(2))
1367 row = ispin
1368 IF (.NOT. ASSOCIATED(ctx%rs(row, slot)%source)) row = 1
1369 cpassert(ASSOCIATED(ctx%rs(row, slot)%source))
1370 xkp => ctx%kpoints%xkp
1371 CALL kpoint_operator_rs_build_batch(ctx%rs(row, slot), ctx%rs_route, xkp(:, ik:ik), &
1372 TARGET=TARGET, fm_re=fm_re, fm_im=fm_im, &
1373 phase_work=phase, parallel_columns=.false.)
1374 IF (ctx%kpoints%use_real_wfn) THEN
1375 IF (PRESENT(TARGET)) target%local_data(:, :) = cmplx(real(target%local_data, dp), 0.0_dp, dp)
1376 IF (PRESENT(fm_im)) fm_im%local_data(:, :) = 0.0_dp
1377 END IF
1378 END SUBROUTINE kpoint_operator_evaluate
1379
1380! **************************************************************************************************
1381!> \brief Start one operator for one global round and spin. The service walks
1382!> all groups with a k point in this round; groups without one still
1383!> participate in the transfers. Every rank must iterate through the
1384!> maximum local k-point count. Finish every started instance with
1385!> kpoint_operator_finish. The first start also creates the slot
1386!> bookkeeping of the context.
1387!> Collective over the environment-wide communicator: every rank must
1388!> call the same sequence of service calls.
1389!> \param ctx context created by kpoint_operator_context_create
1390!> \param ikp global round, up to the maximum k-point count across groups
1391!> \param ispin instance spin: keys the bookkeeping
1392!> \param rsmat real-space image matrices of the operator. A prepared operator
1393!> batch uses the values captured by prepare; otherwise start reads
1394!> the current values. The matching finish never reads rsmat
1395!> \param slot instance slot, optionally backed by a prepared snapshot
1396!> \param grid prepared reciprocal-grid cache of the caller
1397!> \param use_grid extract from grid instead of a direct phase sum
1398!> \param matrix_row row of rsmat to transform. Spin-free operators such
1399!> as S and T must pass kpoint_spin_free: their first matrix
1400!> dimension counts derivatives, not spins. A spin index is out of
1401!> bounds in plain SCF, but with derivatives stored it reads another
1402!> valid row silently
1403! **************************************************************************************************
1404 SUBROUTINE kpoint_operator_start(ctx, ikp, ispin, rsmat, slot, grid, use_grid, matrix_row)
1405
1406 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1407 INTEGER, INTENT(IN) :: ikp, ispin
1408 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1409 INTEGER, INTENT(IN) :: slot
1410 TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
1411 LOGICAL, INTENT(IN) :: use_grid
1412 INTEGER, INTENT(IN) :: matrix_row
1413
1414 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_start'
1415
1416 INTEGER :: handle, igroup, ik, irs, kpmax
1417 INTEGER, DIMENSION(:, :), POINTER :: kp_dist
1418
1419 CALL timeset(routinen, handle)
1420
1421 cpassert(ASSOCIATED(ctx%kpoints))
1422 cpassert(ASSOCIATED(ctx%kpoints%sab_nl))
1423 ! The instance spin is bounded by the spin space declared at creation;
1424 ! the matrix row by this operator's own rows. The two bounds differ for
1425 ! spin-free operators such as S, whose first dimension counts
1426 ! real-space blocks.
1427 cpassert(ispin >= 1 .AND. ispin <= ctx%nspin)
1428 cpassert(matrix_row >= 1 .AND. matrix_row <= SIZE(rsmat, 1))
1429 kpmax = maxval(ctx%kpoints%kp_dist(2, :) - ctx%kpoints%kp_dist(1, :) + 1)
1430 cpassert(ikp >= 1 .AND. ikp <= kpmax)
1431 cpassert(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
1432 ! State is per round; allocate expensive copy descriptors only when
1433 ! an actual distributed transfer is launched.
1434 IF (.NOT. ALLOCATED(ctx%status)) THEN
1435 ALLOCATE (ctx%status(kpmax, ctx%nspin, kpoint_num_slots))
1436 ctx%status(:, :, :) = kpoint_slot_idle
1437 END IF
1438 cpassert(ctx%status(ikp, ispin, slot) == kpoint_slot_idle)
1439 IF (.NOT. use_grid) THEN
1440 irs = kpoint_operator_rs_match(ctx, rsmat, ispin, slot, matrix_row)
1441 IF (irs > 0) THEN
1442 cpassert(ASSOCIATED(ctx%work_sab, ctx%kpoints%sab_nl))
1443 ctx%status(ikp, ispin, slot) = kpoint_slot_local
1444 CALL timestop(handle)
1445 RETURN
1446 END IF
1447 END IF
1448 CALL kpoint_operator_ensure(ctx, rsmat)
1449
1450 ! Only the transfer backend needs copy descriptors. A prepared-only
1451 ! or get-only context never carries them. Size by the declared spin
1452 ! space, not by the possibly spin-free operator's row count.
1453 IF (.NOT. ALLOCATED(ctx%info)) THEN
1454 ! One copy_info per in-flight part: start takes it INTENT(OUT),
1455 ! so two concurrent starts cannot share one instance.
1456 ALLOCATE (ctx%info(SIZE(ctx%kpoints%xkp, 2), ctx%nspin, &
1457 kpoint_num_slots, merge(1, 2, ctx%kpoints%use_real_wfn)))
1458 END IF
1459
1460 kp_dist => ctx%kpoints%kp_dist
1461 DO igroup = 1, ctx%kpoints%nkp_groups
1462 ik = kp_dist(1, igroup) + ikp - 1
1463 IF (ik > kp_dist(2, igroup)) cycle
1464 CALL kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid, use_grid)
1465 CALL kpoint_operator_launch(ctx%kpoints, ik, ctx%fmwork, ctx%fmlocal, &
1466 ctx%info(ik, ispin, slot, :))
1467 END DO
1468 ctx%status(ikp, ispin, slot) = kpoint_slot_in_flight
1469
1470 CALL timestop(handle)
1471
1472 END SUBROUTINE kpoint_operator_start
1473
1474! **************************************************************************************************
1475!> \brief Finish one started instance: build its prepared group-local operator,
1476!> or merge received parts into cmat (real part first, overwrite then add)
1477!> and reclaim the send side on every group.
1478!> The slot returns to idle.
1479!> Collective over the environment-wide communicator: every rank must
1480!> call the same sequence of service calls.
1481!> \param ctx context that started the instance
1482!> \param ikp global round of the start call
1483!> \param ispin spin component of the start call
1484!> \param slot slot of the start call; the slot must be in flight
1485!> \param cmat caller-allocated complex matrix that receives the operator.
1486!> It must share the distribution of the k-point group. Unchanged on
1487!> groups without a k point in this round
1488!> \param fm_re optional real FM output instead of cmat for real wavefunctions, on the group layout
1489! **************************************************************************************************
1490 SUBROUTINE kpoint_operator_finish(ctx, ikp, ispin, slot, cmat, fm_re)
1491
1492 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1493 INTEGER, INTENT(IN) :: ikp, ispin, slot
1494 TYPE(cp_cfm_type), OPTIONAL :: cmat
1495 TYPE(cp_fm_type), OPTIONAL :: fm_re
1496
1497 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_finish'
1498
1499 INTEGER :: handle, igroup, ik, irs, kpmax
1500 INTEGER, DIMENSION(2) :: kp_range
1501 INTEGER, DIMENSION(:, :), POINTER :: kp_dist
1502 LOGICAL :: layout_ok
1503
1504 CALL timeset(routinen, handle)
1505
1506 cpassert(ASSOCIATED(ctx%kpoints))
1507 cpassert(ALLOCATED(ctx%status))
1508 kpmax = SIZE(ctx%status, 1)
1509 cpassert(ikp >= 1 .AND. ikp <= kpmax)
1510 cpassert(ispin >= 1 .AND. ispin <= ctx%nspin)
1511 cpassert(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
1512 cpassert(ctx%status(ikp, ispin, slot) /= kpoint_slot_idle)
1513 ! The merge writes cmat in place. A foreign distribution would fail
1514 ! inside the BLACS calls, far from this argument.
1515
1516 cpassert(PRESENT(cmat) .NEQV. PRESENT(fm_re))
1517 cpassert(.NOT. PRESENT(fm_re) .OR. ctx%kpoints%use_real_wfn)
1518 kp_range = ctx%kpoints%kp_range
1519 IF (ctx%status(ikp, ispin, slot) == kpoint_slot_local) THEN
1520 irs = ispin
1521 IF (.NOT. ASSOCIATED(ctx%rs(irs, slot)%source)) irs = 1
1522 ik = kp_range(1) + ikp - 1
1523 IF (ik <= kp_range(2)) THEN
1524 CALL kpoint_operator_rs_build(ctx, irs, slot, ik, cmat, fm_re=fm_re)
1525 END IF
1526 ctx%status(ikp, ispin, slot) = kpoint_slot_idle
1527 CALL timestop(handle)
1528 RETURN
1529 END IF
1530 cpassert(ctx%status(ikp, ispin, slot) == kpoint_slot_in_flight)
1531 IF (PRESENT(cmat)) THEN
1532 layout_ok = cp_fm_struct_equivalent(cmat%matrix_struct, ctx%fmlocal%matrix_struct)
1533 ELSE
1534 layout_ok = cp_fm_struct_equivalent(fm_re%matrix_struct, ctx%fmlocal%matrix_struct)
1535 END IF
1536 cpassert(layout_ok)
1537 kp_dist => ctx%kpoints%kp_dist
1538 DO igroup = 1, ctx%kpoints%nkp_groups
1539 ik = kp_dist(1, igroup) + ikp - 1
1540 IF (ik > kp_dist(2, igroup)) cycle
1541 CALL kpoint_operator_receive(ctx%kpoints, ctx%fmlocal, ik, ctx%info(ik, ispin, slot, :), &
1542 cmat=cmat, fm_re=fm_re)
1543 END DO
1544 ctx%status(ikp, ispin, slot) = kpoint_slot_idle
1545
1546 CALL timestop(handle)
1547
1548 END SUBROUTINE kpoint_operator_finish
1549
1550! **************************************************************************************************
1551!> \brief Assemble one operator for one global round and spin in one call:
1552!> run the single-group transfer of the service for every group and
1553!> finish it immediately, through transfer states local to this call.
1554!> The persistent slot machinery stays with start and finish.
1555!> Exactly one output flavor must be present: cmat receives the complex
1556!> operator (real part overwrites, imaginary part adds) through the
1557!> context buffer; the fm pair receives the two parts separately,
1558!> transferred directly into fm_re and fm_im on the owning group (any
1559!> distribution: the general copy redistributes); or the dbcsr pair
1560!> receives the two parts as block-cyclic matrices built by the service,
1561!> replacing the caller-side copy_fm_to_dbcsr_bc detour. The dbcsr
1562!> matrices must be freshly initialized (dbcsr_init_p) and the cmat
1563!> flavor must share the distribution of the k-point group.
1564!> Collective over the environment-wide communicator: every rank must
1565!> call the same sequence of service calls.
1566!> \param ctx context created by kpoint_operator_context_create
1567!> \param ikp global round, up to the maximum k-point count across groups
1568!> \param ispin instance spin: bounds-checked against the context spin space
1569!> \param rsmat real-space image matrices of the operator
1570!> \param cmat complex matrix output
1571!> \param fm_re real-part matrix output
1572!> \param fm_im imaginary-part output; may be omitted for real wavefunctions
1573!> \param mat_re real-part block-cyclic dbcsr output
1574!> \param mat_im imaginary-part block-cyclic dbcsr output
1575!> \param matrix_row row of rsmat to transform. Spin-free operators such
1576!> as S and T must pass kpoint_spin_free: their first matrix
1577!> dimension counts derivatives, not spins. A spin index is out of
1578!> bounds in plain SCF, but with derivatives stored it reads another
1579!> valid row silently
1580!> \param slot optional prepared slot: opt in to its immutable snapshot.
1581!> Without slot, get always reads current source values. A missing or
1582!> mismatched snapshot retains the original transfer path. Prepared FM
1583!> targets must match the layout selected at prepare; inactive groups
1584!> do not inspect or modify their output arguments.
1585!> \param grid optional lattice-FFT grid, overriding a prepared snapshot for this request
1586! **************************************************************************************************
1587 SUBROUTINE kpoint_operator_get(ctx, ikp, ispin, rsmat, cmat, fm_re, fm_im, &
1588 mat_re, mat_im, matrix_row, slot, grid)
1589
1590 TYPE(kpoint_operator_context_type), INTENT(INOUT) :: ctx
1591 INTEGER, INTENT(IN) :: ikp, ispin
1592 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1593 TYPE(cp_cfm_type), OPTIONAL :: cmat
1594 TYPE(cp_fm_type), OPTIONAL :: fm_re, fm_im
1595 TYPE(dbcsr_type), INTENT(INOUT), OPTIONAL :: mat_re, mat_im
1596 INTEGER, INTENT(IN) :: matrix_row
1597 INTEGER, INTENT(IN), OPTIONAL :: slot
1598 TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
1599
1600 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_get'
1601
1602 INTEGER :: handle, igroup, ik, irs, kpmax
1603 INTEGER, DIMENSION(:, :), POINTER :: kp_dist
1604 LOGICAL :: layout_ok
1605 TYPE(copy_info_type) :: info(2)
1606
1607 CALL timeset(routinen, handle)
1608
1609 cpassert(count([PRESENT(cmat), PRESENT(fm_re), PRESENT(mat_re)]) == 1)
1610 IF (PRESENT(fm_re) .OR. PRESENT(fm_im)) THEN
1611 cpassert(PRESENT(fm_re))
1612 cpassert(PRESENT(fm_im) .OR. ctx%kpoints%use_real_wfn)
1613 END IF
1614 IF (PRESENT(mat_re) .OR. PRESENT(mat_im)) THEN
1615 cpassert(PRESENT(mat_re) .AND. PRESENT(mat_im))
1616 END IF
1617
1618 cpassert(ASSOCIATED(ctx%kpoints))
1619 cpassert(ASSOCIATED(ctx%kpoints%sab_nl))
1620 cpassert(ispin >= 1 .AND. ispin <= ctx%nspin)
1621 cpassert(matrix_row >= 1 .AND. matrix_row <= SIZE(rsmat, 1))
1622 kpmax = maxval(ctx%kpoints%kp_dist(2, :) - ctx%kpoints%kp_dist(1, :) + 1)
1623 cpassert(ikp >= 1 .AND. ikp <= kpmax)
1624 IF (PRESENT(slot) .AND. .NOT. PRESENT(grid)) THEN
1625 cpassert(slot >= kpoint_slot_ks .AND. slot <= kpoint_slot_t)
1626 irs = kpoint_operator_rs_match(ctx, rsmat, ispin, slot, matrix_row)
1627 IF (irs > 0) THEN
1628 ik = ctx%kpoints%kp_range(1) + ikp - 1
1629 IF (ik <= ctx%kpoints%kp_range(2)) THEN
1630 IF (PRESENT(mat_re)) THEN
1631 IF (.NOT. ALLOCATED(ctx%fmimag)) THEN
1632 ALLOCATE (ctx%fmimag)
1633 CALL cp_fm_create(ctx%fmimag, ctx%fmlocal%matrix_struct)
1634 END IF
1635 CALL kpoint_operator_rs_build(ctx, irs, slot, ik, fm_re=ctx%fmlocal, fm_im=ctx%fmimag)
1636 CALL copy_fm_to_dbcsr_bc(ctx%fmlocal, mat_re)
1637 CALL copy_fm_to_dbcsr_bc(ctx%fmimag, mat_im)
1638 ELSE
1639 CALL kpoint_operator_rs_build(ctx, irs, slot, ik, TARGET=cmat, fm_re=fm_re, fm_im=fm_im)
1640 END IF
1641 END IF
1642 CALL timestop(handle)
1643 RETURN
1644 END IF
1645 END IF
1646 CALL kpoint_operator_ensure(ctx, rsmat)
1647 IF (PRESENT(cmat)) THEN
1648 ! The merge writes cmat in place. A foreign distribution would fail
1649 ! inside the BLACS calls, far from this argument.
1650 layout_ok = cp_fm_struct_equivalent(cmat%matrix_struct, ctx%fmlocal%matrix_struct)
1651 cpassert(layout_ok)
1652 END IF
1653
1654 kp_dist => ctx%kpoints%kp_dist
1655 DO igroup = 1, ctx%kpoints%nkp_groups
1656 ik = kp_dist(1, igroup) + ikp - 1
1657 IF (ik > kp_dist(2, igroup)) cycle
1658 ! Absent fm targets are forwarded as absent optional arguments. Both
1659 ! CFM and DBCSR outputs use the same context-owned receive buffer.
1660 CALL kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid=grid, use_grid=PRESENT(grid))
1661 CALL kpoint_operator_launch(ctx%kpoints, ik, ctx%fmwork, ctx%fmlocal, info, &
1662 target_re=fm_re, target_im=fm_im)
1663 CALL kpoint_operator_receive(ctx%kpoints, ctx%fmlocal, ik, info, cmat=cmat, fm_re=fm_re, fm_im=fm_im, &
1664 mat_re=mat_re, mat_im=mat_im)
1665 END DO
1666
1667 CALL timestop(handle)
1668
1669 END SUBROUTINE kpoint_operator_get
1670
1671! **************************************************************************************************
1672!> \brief Read an operator inside a group-local OT loop, after collective cache preparation.
1673!> One group retains the sparse transform and distribution required by OT's DBCSR kernels.
1674!> Multiple groups read their owned cache; no world collective may occur in this loop.
1675!> \param matrix_rs current real-space operator
1676!> \param kpoints host environment
1677!> \param kp local k-point state
1678!> \param ispin source row (one for spin-free S/T)
1679!> \param cache_re owned real part, used with multiple groups
1680!> \param cache_im owned imaginary part, used with multiple groups
1681!> \param matrix_re newly initialized sparse real output
1682!> \param matrix_im newly initialized sparse imaginary output
1683! **************************************************************************************************
1684 SUBROUTINE kpoint_operator_get_local(matrix_rs, kpoints, kp, ispin, cache_re, cache_im, &
1685 matrix_re, matrix_im)
1686 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: matrix_rs
1687 TYPE(kpoint_type), POINTER :: kpoints
1688 TYPE(kpoint_env_type), POINTER :: kp
1689 INTEGER, INTENT(IN) :: ispin
1690 TYPE(cp_fm_type), INTENT(IN) :: cache_re, cache_im
1691 TYPE(dbcsr_type), POINTER :: matrix_re, matrix_im
1692
1693 INTEGER :: nkp_groups
1694 INTEGER, DIMENSION(:, :, :), POINTER :: cell_to_index
1695 REAL(dp), DIMENSION(:, :), POINTER :: xkp
1696 TYPE(neighbor_list_set_p_type), DIMENSION(:), &
1697 POINTER :: sab_nl
1698
1699 NULLIFY (cell_to_index, sab_nl, xkp)
1700 cpassert(ASSOCIATED(matrix_rs))
1701 cpassert(ispin >= 1 .AND. ispin <= SIZE(matrix_rs, 1))
1702 cpassert(ASSOCIATED(kpoints))
1703 cpassert(ASSOCIATED(kp))
1704
1705 CALL dbcsr_init_p(matrix_re)
1706 CALL dbcsr_init_p(matrix_im)
1707 nkp_groups = kpoints%nkp_groups
1708 xkp => kpoints%xkp
1709 cell_to_index => kpoints%cell_to_index
1710 sab_nl => kpoints%sab_nl
1711 IF (nkp_groups == 1) THEN
1712 cpassert(ASSOCIATED(xkp))
1713 cpassert(ASSOCIATED(cell_to_index))
1714 cpassert(ASSOCIATED(sab_nl))
1715 CALL dbcsr_create(matrix_re, template=matrix_rs(1, 1)%matrix, &
1716 matrix_type=dbcsr_type_symmetric)
1717 CALL dbcsr_create(matrix_im, template=matrix_rs(1, 1)%matrix, &
1718 matrix_type=dbcsr_type_antisymmetric)
1719 CALL cp_dbcsr_alloc_block_from_nbl(matrix_re, sab_nl)
1720 CALL cp_dbcsr_alloc_block_from_nbl(matrix_im, sab_nl)
1721 CALL dbcsr_set(matrix_re, 0.0_dp)
1722 CALL dbcsr_set(matrix_im, 0.0_dp)
1723 CALL rskp_transform(rmatrix=matrix_re, cmatrix=matrix_im, rsmat=matrix_rs, ispin=ispin, &
1724 xkp=xkp(1:3, kp%nkpoint), cell_to_index=cell_to_index, sab_nl=sab_nl)
1725 ELSE
1726 CALL copy_fm_to_dbcsr_bc(cache_re, matrix_re)
1727 CALL copy_fm_to_dbcsr_bc(cache_im, matrix_im)
1728 END IF
1729 END SUBROUTINE kpoint_operator_get_local
1730
1731! **************************************************************************************************
1732!> \brief Fill caller-owned matrices for one operator row, including independent history snapshots.
1733!> Prepare and fallback transfers are collective. Prepared one-rank groups evaluate k points
1734!> concurrently, with no MPI or matrix allocation/refcounts inside the worker loop.
1735!> \param kpoints host environment
1736!> \param ao_ao_fm global-layout template
1737!> \param rsmat current real-space images
1738!> \param slot operator slot
1739!> \param cmat allocated group-local outputs, one per locally owned k point
1740!> \param matrix_row source row; S/T use kpoint_spin_free
1741! **************************************************************************************************
1742 SUBROUTINE kpoint_operator_get_all(kpoints, ao_ao_fm, rsmat, slot, cmat, matrix_row)
1743 TYPE(kpoint_type), POINTER :: kpoints
1744 TYPE(cp_fm_type), INTENT(IN) :: ao_ao_fm
1745 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1746 INTEGER, INTENT(IN) :: slot
1747 TYPE(cp_cfm_type), DIMENSION(:), INTENT(INOUT) :: cmat
1748 INTEGER, INTENT(IN) :: matrix_row
1749
1750 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_get_all'
1751
1752 COMPLEX(KIND=dp), ALLOCATABLE :: phase(:, :)
1753 INTEGER :: handle, ik, ikp, kplocal, kpmax, nentry, &
1754 nworkers
1755 LOGICAL :: layout_ok, prepared
1756 TYPE(kpoint_operator_context_type) :: ctx
1757
1758 CALL timeset(routinen, handle)
1759 kplocal = SIZE(kpoints%kp_env)
1760 cpassert(SIZE(cmat) == kplocal .AND. kplocal > 0)
1761 CALL kpoint_operator_context_create(ctx, kpoints, ao_ao_fm, 1)
1762 CALL kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row=matrix_row)
1763 IF (prepared .AND. kpoints%para_env_kp%num_pe == 1) THEN
1764 DO ikp = 1, kplocal
1765 layout_ok = cp_fm_struct_equivalent(cmat(ikp)%matrix_struct, ctx%rs_route%matrix_struct)
1766 cpassert(layout_ok)
1767 END DO
1768 CALL kpoint_operator_work_size(ctx, nentry)
1769 nworkers = 1
1770!$ nworkers = MIN(omp_get_max_threads(), kplocal)
1771!$OMP PARALLEL DEFAULT(NONE) NUM_THREADS(nworkers) SHARED(ctx,kpoints,cmat,slot,kplocal,nentry) &
1772!$OMP PRIVATE(phase,ikp,ik)
1773 ALLOCATE (phase(1, nentry))
1774!$OMP DO SCHEDULE(DYNAMIC, 1)
1775 DO ikp = 1, kplocal
1776 ik = kpoints%kp_range(1) + ikp - 1
1777 CALL kpoint_operator_evaluate(ctx, 1, slot, ik, phase, cmat(ikp))
1778 END DO
1779!$OMP END DO
1780 DEALLOCATE (phase)
1781!$OMP END PARALLEL
1782 ELSE
1783 kpmax = maxval(kpoints%kp_dist(2, :) - kpoints%kp_dist(1, :) + 1)
1784 DO ikp = 1, kpmax
1785 CALL kpoint_operator_get(ctx, ikp, 1, rsmat, cmat=cmat(min(ikp, kplocal)), &
1786 matrix_row=matrix_row, slot=slot)
1787 END DO
1788 END IF
1790 CALL timestop(handle)
1791 END SUBROUTINE kpoint_operator_get_all
1792
1793! **************************************************************************************************
1794!> \brief Prepare owned split-FM operators for OT without entering an SCF eigensolver.
1795!> H (or the initialization density) retains all source spins; S/T are built only once.
1796!> Outputs survive the transient context. Compatible cache storage is reused, values refreshed.
1797!> \param kpoints host environment and owned OT caches
1798!> \param ao_ao_fm global-layout template
1799!> \param matrix_ks Hamiltonian or initialization density
1800!> \param matrix_s spin-free overlap
1801!> \param matrix_t optional spin-free kinetic operator; absent leaves the existing kinetic cache alone
1802! **************************************************************************************************
1803 SUBROUTINE kpoint_operator_store(kpoints, ao_ao_fm, matrix_ks, matrix_s, matrix_t)
1804 TYPE(kpoint_type), POINTER :: kpoints
1805 TYPE(cp_fm_type), INTENT(IN) :: ao_ao_fm
1806 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: matrix_ks, matrix_s
1807 TYPE(dbcsr_p_type), DIMENSION(:, :), OPTIONAL, &
1808 POINTER :: matrix_t
1809
1810 CHARACTER(LEN=*), PARAMETER :: routinen = 'kpoint_operator_store'
1811 INTEGER :: handle, ic, ikp, ispin, ncomp, nspin
1812 LOGICAL :: kinetic, reuse
1813 TYPE(cp_fm_struct_type), POINTER :: matrix_struct
1814 TYPE(kpoint_env_type), POINTER :: kp
1815 TYPE(kpoint_operator_context_type) :: ctx
1816
1817 CALL timeset(routinen, handle)
1818 nspin = SIZE(matrix_ks, 1)
1819 ncomp = merge(1, 2, kpoints%use_real_wfn)
1820 kinetic = PRESENT(matrix_t)
1821 IF (kinetic) kinetic = ASSOCIATED(matrix_t)
1822 CALL kpoint_operator_context_create(ctx, kpoints, ao_ao_fm, nspin)
1823 CALL kpoint_operator_ensure_group(ctx)
1824 matrix_struct => ctx%fmlocal%matrix_struct
1825 DO ikp = 1, SIZE(kpoints%kp_env)
1826 kp => kpoints%kp_env(ikp)%kpoint_env
1827 reuse = .false.
1828 IF (ASSOCIATED(kp%ot_hmat)) THEN
1829 IF (all(shape(kp%ot_hmat) == [ncomp, nspin])) THEN
1830 reuse = cp_fm_struct_equivalent(kp%ot_hmat(1, 1)%matrix_struct, matrix_struct)
1831 END IF
1832 END IF
1833 IF (.NOT. reuse) THEN
1834 CALL cp_fm_release(kp%ot_hmat)
1835 ALLOCATE (kp%ot_hmat(ncomp, nspin))
1836 DO ispin = 1, nspin
1837 DO ic = 1, ncomp
1838 CALL cp_fm_create(kp%ot_hmat(ic, ispin), matrix_struct)
1839 END DO
1840 END DO
1841 END IF
1842 CALL allocate_spin_free(kp%ot_smat)
1843 IF (kinetic) CALL allocate_spin_free(kp%ot_tmat)
1844 END DO
1845 CALL store_slot(matrix_ks, kpoint_slot_ks, nspin)
1846 CALL store_slot(matrix_s, kpoint_slot_s, 1)
1847 IF (kinetic) CALL store_slot(matrix_t, kpoint_slot_t, 1)
1849 CALL timestop(handle)
1850
1851 CONTAINS
1852
1853! **************************************************************************************************
1854!> \brief ...
1855!> \param cache ...
1856! **************************************************************************************************
1857 SUBROUTINE allocate_spin_free(cache)
1858 TYPE(cp_fm_type), DIMENSION(:), POINTER :: cache
1859
1860 INTEGER :: ic
1861 LOGICAL :: reuse
1862
1863 reuse = .false.
1864 IF (ASSOCIATED(cache)) THEN
1865 IF (SIZE(cache) == ncomp) reuse = cp_fm_struct_equivalent(cache(1)%matrix_struct, matrix_struct)
1866 END IF
1867 IF (reuse) RETURN
1868 CALL cp_fm_release(cache)
1869 ALLOCATE (cache(ncomp))
1870 DO ic = 1, ncomp
1871 CALL cp_fm_create(cache(ic), matrix_struct)
1872 END DO
1873 END SUBROUTINE allocate_spin_free
1874
1875! **************************************************************************************************
1876!> \brief ...
1877!> \param rsmat ...
1878!> \param slot ...
1879!> \param nrows ...
1880! **************************************************************************************************
1881 SUBROUTINE store_slot(rsmat, slot, nrows)
1882 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1883 INTEGER, INTENT(IN) :: slot, nrows
1884
1885 COMPLEX(KIND=dp), ALLOCATABLE :: phase(:, :)
1886 INTEGER :: ikp, ispin, kplocal, kpmax, nentry, &
1887 nworkers
1888 LOGICAL :: prepared, ready, use_grid
1889 TYPE(rskp_grid_type), ALLOCATABLE :: grids(:)
1890
1891 nworkers = 1
1892!$ nworkers = omp_get_max_threads()
1893 use_grid = kpoints%lattice_fft == lattice_fft_on .OR. &
1894 (kpoints%lattice_fft == lattice_fft_auto .AND. SIZE(kpoints%xkp, 2) >= 27)
1895 prepared = .false.
1896 ! Match STANDARD's single-world-rank, single-thread preference for the forward lattice FFT.
1897 IF (.NOT. use_grid .OR. nworkers > 1 .OR. kpoints%blacs_env_all%para_env%num_pe > 1) THEN
1898 IF (slot == kpoint_slot_ks) THEN
1899 CALL kpoint_operator_prepare(ctx, rsmat, slot, prepared)
1900 ELSE
1901 CALL kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row=kpoint_spin_free)
1902 END IF
1903 END IF
1904 kplocal = SIZE(kpoints%kp_env)
1905 IF (prepared .AND. kpoints%para_env_kp%num_pe == 1) THEN
1906 CALL kpoint_operator_work_size(ctx, nentry)
1907 nworkers = min(nworkers, kplocal*nrows)
1908!$OMP PARALLEL DEFAULT(NONE) NUM_THREADS(nworkers) SHARED(ctx,kpoints,rsmat,slot,nrows,kplocal,nentry) &
1909!$OMP PRIVATE(phase,ikp,ispin)
1910 ALLOCATE (phase(1, nentry))
1911!$OMP DO COLLAPSE(2) SCHEDULE(DYNAMIC, 1)
1912 DO ikp = 1, kplocal
1913 DO ispin = 1, nrows
1914 CALL store_round(rsmat, slot, ikp, ispin, phase)
1915 END DO
1916 END DO
1917!$OMP END DO
1918 DEALLOCATE (phase)
1919!$OMP END PARALLEL
1920 ELSE
1921 kpmax = maxval(kpoints%kp_dist(2, :) - kpoints%kp_dist(1, :) + 1)
1922 ! Keep the existing forward lattice-FFT choice when distributed snapshots cannot be prepared.
1923 use_grid = use_grid .AND. .NOT. prepared
1924 IF (use_grid) THEN
1925 CALL kpoint_operator_ensure(ctx, rsmat)
1926 ALLOCATE (grids(nrows))
1927 DO ispin = 1, nrows
1928 CALL rskp_transform_grid_prepare(grids(ispin), ctx%rmatrix, rsmat, ispin, &
1929 kpoints%xkp, kpoints%nkp_grid, kpoints%cell_to_index, &
1930 kpoints%sab_nl, ready, &
1931 max_storage_bytes=512_int_8*1024_int_8**2/int(nspin + 1, int_8))
1932 use_grid = use_grid .AND. ready
1933 END DO
1934 END IF
1935 DO ikp = 1, kpmax
1936 DO ispin = 1, nrows
1937 IF (use_grid) THEN
1938 CALL store_round(rsmat, slot, ikp, ispin, grid=grids(ispin))
1939 ELSE
1940 CALL store_round(rsmat, slot, ikp, ispin)
1941 END IF
1942 END DO
1943 END DO
1944 IF (ALLOCATED(grids)) THEN
1945 DO ispin = 1, nrows
1946 CALL rskp_transform_grid_release(grids(ispin))
1947 END DO
1948 DEALLOCATE (grids)
1949 END IF
1950 END IF
1951 END SUBROUTINE store_slot
1952
1953! **************************************************************************************************
1954!> \brief ...
1955!> \param rsmat ...
1956!> \param slot ...
1957!> \param ikp ...
1958!> \param ispin ...
1959!> \param phase ...
1960!> \param grid ...
1961! **************************************************************************************************
1962 SUBROUTINE store_round(rsmat, slot, ikp, ispin, phase, grid)
1963 TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER :: rsmat
1964 INTEGER, INTENT(IN) :: slot, ikp, ispin
1965 COMPLEX(KIND=dp), INTENT(INOUT), OPTIONAL :: phase(:, :)
1966 TYPE(rskp_grid_type), INTENT(IN), OPTIONAL :: grid
1967
1968 INTEGER :: ik
1969 TYPE(cp_fm_type), DIMENSION(:), POINTER :: cache
1970 TYPE(cp_fm_type), POINTER :: imag
1971 TYPE(kpoint_env_type), POINTER :: kp
1972
1973 ! A valid dummy cache is passed on inactive tail groups; get leaves it untouched.
1974 kp => kpoints%kp_env(min(ikp, SIZE(kpoints%kp_env)))%kpoint_env
1975 SELECT CASE (slot)
1976 CASE (kpoint_slot_ks)
1977 cache => kp%ot_hmat(:, ispin)
1978 CASE (kpoint_slot_s)
1979 cache => kp%ot_smat
1980 CASE (kpoint_slot_t)
1981 cache => kp%ot_tmat
1982 END SELECT
1983 NULLIFY (imag)
1984 IF (SIZE(cache) == 2) imag => cache(2)
1985 IF (PRESENT(phase)) THEN
1986 ik = kpoints%kp_range(1) + ikp - 1
1987 CALL kpoint_operator_evaluate(ctx, ispin, slot, ik, phase, fm_re=cache(1), fm_im=imag)
1988 ELSE
1989 CALL kpoint_operator_get(ctx, ikp, ispin, rsmat, fm_re=cache(1), fm_im=imag, &
1990 matrix_row=ispin, slot=slot, grid=grid)
1991 END IF
1992 END SUBROUTINE store_round
1993 END SUBROUTINE kpoint_operator_store
1994
1995END MODULE qs_kpoint_operators
Basic linear algebra operations for complex full matrices.
subroutine, public cp_cfm_scale_and_add_fm(alpha, matrix_a, beta, matrix_b)
Scale and add two BLACS matrices (a = alpha*a + beta*b). where b is a real matrix (adapted from cp_cf...
Represents a complex full matrix distributed on many processors.
subroutine, public dbcsr_get_readonly_block_p(matrix, row, col, block, found, row_size, col_size)
Like dbcsr_get_block_p() but with matrix being INTENT(IN). When invoking this routine,...
subroutine, public dbcsr_deallocate_matrix(matrix)
...
subroutine, public dbcsr_get_info(matrix, nblkrows_total, nblkcols_total, nfullrows_total, nfullcols_total, nblkrows_local, nblkcols_local, nfullrows_local, nfullcols_local, my_prow, my_pcol, local_rows, local_cols, proc_row_dist, proc_col_dist, row_blk_size, col_blk_size, row_blk_offset, col_blk_offset, distribution, name, matrix_type, group)
...
subroutine, public dbcsr_init_p(matrix)
...
subroutine, public dbcsr_set(matrix, alpha)
...
DBCSR operations in CP2K.
subroutine, public copy_dbcsr_to_fm(matrix, fm, plan, check_plan)
Copy a DBCSR matrix to a BLACS matrix.
subroutine, public copy_fm_to_dbcsr_bc(fm, bc_mat)
Copy a BLACS matrix to a dbcsr matrix with a special block-cyclic distribution, which requires no com...
pool for for elements that are retained and released
subroutine, public fm_pool_create_fm(pool, element, name)
returns an element, allocating it if none is in the pool
subroutine, public fm_pool_give_back_fm(pool, element)
returns the element to the pool
represent the structure of a full matrix
logical function, public cp_fm_struct_equivalent(fmstruct1, fmstruct2)
returns true if the two matrix structures are equivalent, false otherwise.
subroutine, public cp_fm_struct_retain(fmstruct)
retains a full matrix structure
subroutine, public cp_fm_struct_release(fmstruct)
releases a full matrix structure
represent a full matrix distributed on many processors
Definition cp_fm_types.F:15
subroutine, public cp_fm_start_copy_general(source, destination, para_env, info)
Initiates the copy operation: get distribution data, post MPI isend and irecvs.
subroutine, public cp_fm_cleanup_copy_general(info)
Completes the copy operation: wait for comms clean up MPI state.
subroutine, public cp_fm_get_info(matrix, name, nrow_global, ncol_global, nrow_block, ncol_block, nrow_local, ncol_local, row_indices, col_indices, local_data, context, nrow_locals, ncol_locals, matrix_struct, para_env)
returns all kind of information about the full matrix
subroutine, public cp_fm_create(matrix, matrix_struct, name, nrow, ncol, set_zero)
creates a new full matrix with the given structure
subroutine, public cp_fm_finish_copy_general(destination, info)
Completes the copy operation: wait for comms, unpack, clean up MPI state.
function that build the kpoints section of the input
integer, parameter, public lattice_fft_auto
integer, parameter, public lattice_fft_on
Defines the basic variable types.
Definition kinds.F:23
integer, parameter, public int_8
Definition kinds.F:54
integer, parameter, public dp
Definition kinds.F:34
Routines needed for kpoint calculation.
subroutine, public rskp_transform(rmatrix, cmatrix, rsmat, ispin, xkp, cell_to_index, sab_nl, is_complex, rs_sign)
Transformation of real space matrices to a kpoint.
subroutine, public rskp_transform_grid_extract(grid, ikp, rmatrix, cmatrix)
Extract one reciprocal-grid matrix from a prepared local DBCSR block cache.
subroutine, public rskp_transform_grid_release(grid)
Release a reciprocal-grid transformation cache.
subroutine, public kp_transform_plan_create(plan, sab_nl, cell_to_index, nimg, block_template, group_entries)
Build the immutable neighbor-list traversal shared by R-to-K and K-to-R transforms.
subroutine, public rskp_transform_grid_prepare(grid, rmatrix, rsmat, ispin, xkp, nkp_grid, cell_to_index, sab_nl, used_fft, is_complex, rs_sign, max_storage_bytes)
Prepare a batched real-cell to complete reciprocal-grid transform for local DBCSR blocks....
Types and basic routines needed for a kpoint calculation.
Definition of mathematical constants and functions.
complex(kind=dp), parameter, public z_one
complex(kind=dp), parameter, public gaussi
real(kind=dp), parameter, public twopi
complex(kind=dp), parameter, public z_zero
Interface to the message passing library MPI.
Assembly of real and complex k-point operators from real-space DBCSR matrices. Outputs use the distri...
subroutine, public kpoint_operator_context_create(ctx, kpoints, ao_ao_fm, nspin)
Create an empty operator context. The context owns no buffers after this call; they are created on fi...
subroutine, public kpoint_operator_prepare(ctx, rsmat, slot, prepared, matrix_row, target)
Prepare an immutable operator slot for group-local assembly. Refresh after the source values change,...
subroutine, public kpoint_operator_finish(ctx, ikp, ispin, slot, cmat, fm_re)
Finish one started instance: build its prepared group-local operator, or merge received parts into cm...
subroutine, public kpoint_operator_work_size(ctx, nentry)
Size per-worker phase scratch for direct assembly into a solver matrix.
subroutine, public kpoint_operator_get(ctx, ikp, ispin, rsmat, cmat, fm_re, fm_im, mat_re, mat_im, matrix_row, slot, grid)
Assemble one operator for one global round and spin in one call: run the single-group transfer of the...
integer, parameter, public kpoint_spin_free
Row of rsmat that carries the spin-free operator image. S and T carry no spin copies: drivers launch ...
subroutine, public kpoint_operator_evaluate(ctx, ispin, slot, ik, phase, target, fm_re, fm_im)
Evaluate an immutable prepared operator with exclusively caller-owned output and scratch....
integer, parameter, public kpoint_slot_s
integer, parameter rs_batch_size
subroutine kpoint_operator_densify(ctx, rsmat, matrix_row, ik, grid, use_grid)
Transform and densify one operator using the context-owned workspaces and routing plans.
subroutine, public kpoint_operator_store(kpoints, ao_ao_fm, matrix_ks, matrix_s, matrix_t)
Prepare owned split-FM operators for OT without entering an SCF eigensolver. H (or the initialization...
integer, parameter, public kpoint_slot_t
integer, parameter, public kpoint_slot_ks
subroutine, public kpoint_operator_get_local(matrix_rs, kpoints, kp, ispin, cache_re, cache_im, matrix_re, matrix_im)
Read an operator inside a group-local OT loop, after collective cache preparation....
subroutine, public kpoint_operator_get_all(kpoints, ao_ao_fm, rsmat, slot, cmat, matrix_row)
Fill caller-owned matrices for one operator row, including independent history snapshots....
subroutine, public kpoint_operator_start(ctx, ikp, ispin, rsmat, slot, grid, use_grid, matrix_row)
Start one operator for one global round and spin. The service walks all groups with a k point in this...
subroutine, public kpoint_operator_context_release(ctx)
Release the context together with every buffer it created. Call once, at the exit of the routine that...
wrapper for the pools of matrixes
subroutine, public mpools_get(mpools, ao_mo_fm_pools, ao_ao_fm_pools, mo_mo_fm_pools, ao_mosub_fm_pools, mosub_mosub_fm_pools, maxao_maxmo_fm_pool, maxao_maxao_fm_pool, maxmo_maxmo_fm_pool)
returns various attributes of the mpools (notably the pools contained in it)
Define the neighbor list data types and the corresponding functionality.
Represent a complex full matrix.
represent a pool of elements with the same structure
keeps the information about the structure of a full matrix
Stores the state of a copy between cp_fm_start_copy_general and cp_fm_finish_copy_general.
represent a full matrix
Keeps information about a specific k-point.
Contains information about kpoints.
stores all the informations relevant to an mpi environment
Caller-owned assembly context for one driver call. Created and released in the same driver routine an...