Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2025 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief General methods for testing DBT tensors.
10 : !> \author Patrick Seewald
11 : ! **************************************************************************************************
12 : MODULE dbt_test
13 : #:include "dbt_macros.fypp"
14 : #:set maxdim = maxrank
15 : #:set ndims = range(2,maxdim+1)
16 :
17 : USE dbt_tas_base, ONLY: dbt_tas_info
18 : USE dbm_tests, ONLY: generate_larnv_seed
19 : USE dbt_methods, ONLY: &
20 : dbt_copy, dbt_get_block, dbt_iterator_type, dbt_iterator_blocks_left, &
21 : dbt_iterator_next_block, dbt_iterator_start, dbt_iterator_stop, &
22 : dbt_reserve_blocks, dbt_get_stored_coordinates, dbt_put_block, &
23 : dbt_contract, dbt_inverse_order
24 : USE dbt_block, ONLY: block_nd
25 : USE dbt_types, ONLY: &
26 : dbt_create, dbt_destroy, dbt_type, dbt_distribution_type, &
27 : dbt_distribution_destroy, dims_tensor, ndims_tensor, dbt_distribution_new, &
28 : mp_environ_pgrid, dbt_pgrid_type, dbt_pgrid_create, dbt_pgrid_destroy, dbt_get_info, &
29 : dbt_default_distvec
30 : USE dbt_io, ONLY: &
31 : dbt_write_blocks, dbt_write_block_indices
32 : USE kinds, ONLY: dp, default_string_length, int_8, dp
33 : USE dbt_allocate_wrap, ONLY: allocate_any
34 : USE dbt_index, ONLY: &
35 : combine_tensor_index, get_2d_indices_tensor, dbt_get_mapping_info
36 : USE dbt_tas_test, ONLY: dbt_tas_checksum
37 : USE message_passing, ONLY: mp_comm_type
38 :
39 : #include "../base/base_uses.f90"
40 :
41 : IMPLICIT NONE
42 : PRIVATE
43 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_test'
44 :
45 : PUBLIC :: &
46 : dbt_setup_test_tensor, &
47 : dbt_contract_test, &
48 : dbt_test_formats, &
49 : dbt_checksum, &
50 : dbt_reset_randmat_seed
51 :
52 : INTERFACE dist_sparse_tensor_to_repl_dense_array
53 : #:for ndim in ndims
54 : MODULE PROCEDURE dist_sparse_tensor_to_repl_dense_${ndim}$d_array
55 : #:endfor
56 : END INTERFACE
57 :
58 : INTEGER, SAVE :: randmat_counter = 0
59 : INTEGER, PARAMETER, PRIVATE :: rand_seed_init = 12341313
60 :
61 : CONTAINS
62 :
63 : ! **************************************************************************************************
64 : !> \brief check if two (arbitrarily mapped and distributed) tensors are equal.
65 : !> \author Patrick Seewald
66 : ! **************************************************************************************************
67 172 : FUNCTION dbt_equal(tensor1, tensor2)
68 : TYPE(dbt_type), INTENT(INOUT) :: tensor1, tensor2
69 : LOGICAL :: dbt_equal
70 :
71 1204 : TYPE(dbt_type) :: tensor2_tmp
72 : TYPE(dbt_iterator_type) :: iter
73 172 : TYPE(block_nd) :: blk_data1, blk_data2
74 172 : INTEGER, DIMENSION(ndims_tensor(tensor1)) :: blk_size, ind_nd
75 : LOGICAL :: found
76 :
77 : ! create a copy of tensor2 that has exact same data format as tensor1
78 172 : CALL dbt_create(tensor1, tensor2_tmp)
79 :
80 172 : CALL dbt_reserve_blocks(tensor1, tensor2_tmp)
81 172 : CALL dbt_copy(tensor2, tensor2_tmp)
82 :
83 172 : dbt_equal = .TRUE.
84 :
85 : !$OMP PARALLEL DEFAULT(NONE) SHARED(tensor1,tensor2_tmp,dbt_equal) &
86 172 : !$OMP PRIVATE(iter,ind_nd, blk_size,blk_data1,blk_data2,found)
87 : CALL dbt_iterator_start(iter, tensor1)
88 :
89 : DO WHILE (dbt_iterator_blocks_left(iter))
90 : CALL dbt_iterator_next_block(iter, ind_nd, blk_size=blk_size)
91 : CALL dbt_get_block(tensor1, ind_nd, blk_data1, found)
92 : IF (.NOT. found) CPABORT("Tensor block 1 not found")
93 : CALL dbt_get_block(tensor2_tmp, ind_nd, blk_data2, found)
94 : IF (.NOT. found) CPABORT("Tensor block 2 not found")
95 :
96 : IF (.NOT. blocks_equal(blk_data1, blk_data2)) THEN
97 : !$OMP CRITICAL
98 : dbt_equal = .FALSE.
99 : !$OMP END CRITICAL
100 : END IF
101 : END DO
102 :
103 : CALL dbt_iterator_stop(iter)
104 : !$OMP END PARALLEL
105 :
106 172 : CALL dbt_destroy(tensor2_tmp)
107 344 : END FUNCTION
108 :
109 : ! **************************************************************************************************
110 : !> \brief check if two blocks are equal
111 : !> \author Patrick Seewald
112 : ! **************************************************************************************************
113 1464 : PURE FUNCTION blocks_equal(block1, block2)
114 : TYPE(block_nd), INTENT(IN) :: block1, block2
115 : LOGICAL :: blocks_equal
116 :
117 1809044 : blocks_equal = MAXVAL(ABS(block1%blk - block2%blk)) .LT. 1.0E-12_dp
118 :
119 1464 : END FUNCTION
120 :
121 : ! **************************************************************************************************
122 : !> \brief Compute factorial
123 : !> \author Patrick Seewald
124 : ! **************************************************************************************************
125 12 : PURE FUNCTION factorial(n)
126 : INTEGER, INTENT(IN) :: n
127 : INTEGER :: k
128 : INTEGER :: factorial
129 96 : factorial = PRODUCT((/(k, k=1, n)/))
130 12 : END FUNCTION
131 :
132 : ! **************************************************************************************************
133 : !> \brief Compute all permutations p of (1, 2, ..., n)
134 : !> \author Patrick Seewald
135 : ! **************************************************************************************************
136 6 : SUBROUTINE permute(n, p)
137 : INTEGER, INTENT(IN) :: n
138 : INTEGER :: i, c
139 12 : INTEGER, DIMENSION(n) :: pp
140 : INTEGER, DIMENSION(n, factorial(n)), INTENT(OUT) :: p
141 :
142 48 : pp = [(i, i=1, n)]
143 6 : c = 1
144 6 : CALL perm(1)
145 : CONTAINS
146 108 : RECURSIVE SUBROUTINE perm(i)
147 : INTEGER, INTENT(IN) :: i
148 : INTEGER :: j, t
149 108 : IF (i == n) THEN
150 300 : p(:, c) = pp(:)
151 64 : c = c + 1
152 : ELSE
153 146 : DO j = i, n
154 102 : t = pp(i)
155 102 : pp(i) = pp(j)
156 102 : pp(j) = t
157 102 : call perm(i + 1)
158 102 : t = pp(i)
159 102 : pp(i) = pp(j)
160 146 : pp(j) = t
161 : END DO
162 : END IF
163 108 : END SUBROUTINE
164 : END SUBROUTINE
165 :
166 : ! **************************************************************************************************
167 : !> \brief Test equivalence of all tensor formats, using a random distribution.
168 : !> \param blk_size_i block sizes along respective dimension
169 : !> \param blk_ind_i index along respective dimension of non-zero blocks
170 : !> \param ndims tensor rank
171 : !> \param unit_nr output unit, needs to be a valid unit number on all mpi ranks
172 : !> \param verbose if .TRUE., print all tensor blocks
173 : !> \author Patrick Seewald
174 : ! **************************************************************************************************
175 6 : SUBROUTINE dbt_test_formats(ndims, mp_comm, unit_nr, verbose, &
176 6 : ${varlist("blk_size")}$, &
177 6 : ${varlist("blk_ind")}$)
178 : INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: ${varlist("blk_size")}$
179 : INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: ${varlist("blk_ind")}$
180 : INTEGER, INTENT(IN) :: ndims
181 : INTEGER, INTENT(IN) :: unit_nr
182 : LOGICAL, INTENT(IN) :: verbose
183 : TYPE(mp_comm_type), INTENT(IN) :: mp_comm
184 84 : TYPE(dbt_distribution_type) :: dist1, dist2
185 78 : TYPE(dbt_type) :: tensor1, tensor2
186 : INTEGER :: isep, iblk
187 6 : INTEGER, DIMENSION(:), ALLOCATABLE :: ${varlist("dist1")}$, &
188 6 : ${varlist("dist2")}$
189 : INTEGER :: nblks, imap
190 12 : INTEGER, DIMENSION(ndims) :: pdims, myploc
191 : LOGICAL :: eql
192 : INTEGER :: iperm, idist, icount
193 6 : INTEGER, DIMENSION(:), ALLOCATABLE :: map1, map2, map1_ref, map2_ref
194 6 : INTEGER, DIMENSION(ndims, factorial(ndims)) :: perm
195 : INTEGER :: io_unit
196 : INTEGER :: mynode
197 18 : TYPE(dbt_pgrid_type) :: comm_nd
198 : CHARACTER(LEN=default_string_length) :: tensor_name
199 :
200 : ! Process grid
201 24 : pdims(:) = 0
202 6 : CALL dbt_pgrid_create(mp_comm, pdims, comm_nd)
203 6 : mynode = mp_comm%mepos
204 :
205 6 : io_unit = 0
206 6 : IF (mynode .EQ. 0) io_unit = unit_nr
207 :
208 6 : CALL permute(ndims, perm)
209 26 : ALLOCATE (map1_ref, source=perm(1:ndims/2, 1))
210 28 : ALLOCATE (map2_ref, source=perm(ndims/2 + 1:ndims, 1))
211 :
212 6 : IF (io_unit > 0) THEN
213 3 : WRITE (io_unit, *)
214 3 : WRITE (io_unit, '(A)') repeat("-", 80)
215 3 : WRITE (io_unit, '(A,1X,I1)') "Testing matrix representations of tensor rank", ndims
216 3 : WRITE (io_unit, '(A)') repeat("-", 80)
217 3 : WRITE (io_unit, '(A)') "Block sizes:"
218 :
219 : #:for dim in range(1, maxdim+1)
220 11 : IF (ndims >= ${dim}$) THEN
221 9 : WRITE (io_unit, '(T4,A,1X,I1,A,1X)', advance='no') 'Dim', ${dim}$, ':'
222 82 : DO iblk = 1, SIZE(blk_size_${dim}$)
223 82 : WRITE (io_unit, '(I2,1X)', advance='no') blk_size_${dim}$ (iblk)
224 : END DO
225 9 : WRITE (io_unit, *)
226 : END IF
227 : #:endfor
228 :
229 3 : WRITE (io_unit, '(A)') "Non-zero blocks:"
230 40 : DO iblk = 1, SIZE(blk_ind_1)
231 : #:for ndim in ndims
232 58 : IF (ndims == ${ndim}$) THEN
233 : WRITE (io_unit, '(T4,A, I3, A, ${ndim}$I3, 1X, A)') &
234 37 : 'Block', iblk, ': (', ${varlist("blk_ind", nmax=ndim, suffix='(iblk)')}$, ')'
235 : END IF
236 : #:endfor
237 : END DO
238 :
239 3 : WRITE (io_unit, *)
240 3 : WRITE (io_unit, '(A,1X)', advance='no') "Reference map:"
241 3 : WRITE (io_unit, '(A1,1X)', advance='no') "("
242 7 : DO imap = 1, SIZE(map1_ref)
243 7 : WRITE (io_unit, '(I1,1X)', advance='no') map1_ref(imap)
244 : END DO
245 3 : WRITE (io_unit, '(A1,1X)', advance='no') "|"
246 8 : DO imap = 1, SIZE(map2_ref)
247 8 : WRITE (io_unit, '(I1,1X)', advance='no') map2_ref(imap)
248 : END DO
249 3 : WRITE (io_unit, '(A1)') ")"
250 :
251 : END IF
252 :
253 6 : icount = 0
254 70 : DO iperm = 1, factorial(ndims)
255 242 : DO isep = 1, ndims - 1
256 172 : icount = icount + 1
257 :
258 844 : ALLOCATE (map1, source=perm(1:isep, iperm))
259 844 : ALLOCATE (map2, source=perm(isep + 1:ndims, iperm))
260 :
261 172 : mynode = mp_comm%mepos
262 172 : CALL mp_environ_pgrid(comm_nd, pdims, myploc)
263 :
264 : #:for dim in range(1, maxdim+1)
265 684 : IF (${dim}$ <= ndims) THEN
266 656 : nblks = SIZE(blk_size_${dim}$)
267 1968 : ALLOCATE (dist1_${dim}$ (nblks))
268 1312 : ALLOCATE (dist2_${dim}$ (nblks))
269 656 : CALL dbt_default_distvec(nblks, pdims(${dim}$), blk_size_${dim}$, dist1_${dim}$)
270 656 : CALL dbt_default_distvec(nblks, pdims(${dim}$), blk_size_${dim}$, dist2_${dim}$)
271 : END IF
272 : #:endfor
273 :
274 172 : WRITE (tensor_name, '(A,1X,I3,1X)') "Test", icount
275 :
276 172 : IF (io_unit > 0) THEN
277 86 : WRITE (io_unit, *)
278 86 : WRITE (io_unit, '(A,A,1X)', advance='no') TRIM(tensor_name), ':'
279 86 : WRITE (io_unit, '(A1,1X)', advance='no') "("
280 250 : DO imap = 1, SIZE(map1)
281 250 : WRITE (io_unit, '(I1,1X)', advance='no') map1(imap)
282 : END DO
283 86 : WRITE (io_unit, '(A1,1X)', advance='no') "|"
284 250 : DO imap = 1, SIZE(map2)
285 250 : WRITE (io_unit, '(I1,1X)', advance='no') map2(imap)
286 : END DO
287 86 : WRITE (io_unit, '(A1)') ")"
288 :
289 86 : WRITE (io_unit, '(T4,A)') "Reference distribution:"
290 : #:for dim in range(1, maxdim+1)
291 342 : IF (${dim}$ <= ndims) THEN
292 328 : WRITE (io_unit, '(T7,A,1X)', advance='no') "Dist vec ${dim}$:"
293 2354 : DO idist = 1, SIZE(dist2_${dim}$)
294 2354 : WRITE (io_unit, '(I2,1X)', advance='no') dist2_${dim}$ (idist)
295 : END DO
296 328 : WRITE (io_unit, *)
297 : END IF
298 : #:endfor
299 :
300 86 : WRITE (io_unit, '(T4,A)') "Test distribution:"
301 : #:for dim in range(1, maxdim+1)
302 342 : IF (${dim}$ <= ndims) THEN
303 328 : WRITE (io_unit, '(T7,A,1X)', advance='no') "Dist vec ${dim}$:"
304 2354 : DO idist = 1, SIZE(dist2_${dim}$)
305 2354 : WRITE (io_unit, '(I2,1X)', advance='no') dist1_${dim}$ (idist)
306 : END DO
307 328 : WRITE (io_unit, *)
308 : END IF
309 : #:endfor
310 : END IF
311 :
312 : #:for ndim in ndims
313 200 : IF (ndims == ${ndim}$) THEN
314 172 : CALL dbt_distribution_new(dist2, comm_nd, ${varlist("dist2", nmax=ndim)}$)
315 : CALL dbt_create(tensor2, "Ref", dist2, map1_ref, map2_ref, &
316 172 : ${varlist("blk_size", nmax=ndim)}$)
317 172 : CALL dbt_setup_test_tensor(tensor2, comm_nd%mp_comm_2d, .TRUE., ${varlist("blk_ind", nmax=ndim)}$)
318 : END IF
319 : #:endfor
320 :
321 172 : IF (verbose) CALL dbt_write_blocks(tensor2, io_unit, unit_nr)
322 :
323 : #:for ndim in ndims
324 200 : IF (ndims == ${ndim}$) THEN
325 172 : CALL dbt_distribution_new(dist1, comm_nd, ${varlist("dist1", nmax=ndim)}$)
326 : CALL dbt_create(tensor1, tensor_name, dist1, map1, map2, &
327 172 : ${varlist("blk_size", nmax=ndim)}$)
328 172 : CALL dbt_setup_test_tensor(tensor1, comm_nd%mp_comm_2d, .TRUE., ${varlist("blk_ind", nmax=ndim)}$)
329 : END IF
330 : #:endfor
331 :
332 172 : IF (verbose) CALL dbt_write_blocks(tensor1, io_unit, unit_nr)
333 :
334 172 : eql = dbt_equal(tensor1, tensor2)
335 :
336 172 : IF (.NOT. eql) THEN
337 0 : IF (io_unit > 0) WRITE (io_unit, '(A,1X,A)') TRIM(tensor_name), 'Test failed!'
338 0 : CPABORT('')
339 : ELSE
340 172 : IF (io_unit > 0) WRITE (io_unit, '(A,1X,A)') TRIM(tensor_name), 'Test passed!'
341 : END IF
342 172 : DEALLOCATE (map1, map2)
343 :
344 172 : CALL dbt_destroy(tensor1)
345 172 : CALL dbt_distribution_destroy(dist1)
346 :
347 172 : CALL dbt_destroy(tensor2)
348 172 : CALL dbt_distribution_destroy(dist2)
349 :
350 : #:for dim in range(1, maxdim+1)
351 748 : IF (${dim}$ <= ndims) THEN
352 656 : DEALLOCATE (dist1_${dim}$, dist2_${dim}$)
353 : END IF
354 : #:endfor
355 :
356 : END DO
357 : END DO
358 6 : CALL dbt_pgrid_destroy(comm_nd)
359 6 : END SUBROUTINE
360 :
361 : ! **************************************************************************************************
362 : !> \brief Allocate and fill test tensor - entries are enumerated by their index s.t. they only depend
363 : !> on global properties of the tensor but not on distribution, matrix representation, etc.
364 : !> \param mp_comm communicator
365 : !> \param blk_ind_i index along respective dimension of non-zero blocks
366 : !> \author Patrick Seewald
367 : ! **************************************************************************************************
368 404 : SUBROUTINE dbt_setup_test_tensor(tensor, mp_comm, enumerate, ${varlist("blk_ind")}$)
369 : TYPE(dbt_type), INTENT(INOUT) :: tensor
370 : CLASS(mp_comm_type), INTENT(IN) :: mp_comm
371 : LOGICAL, INTENT(IN) :: enumerate
372 : INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: ${varlist("blk_ind")}$
373 : INTEGER :: mynode
374 :
375 : INTEGER :: i, ib, my_nblks_alloc, nblks_alloc, proc, nze
376 404 : INTEGER, ALLOCATABLE, DIMENSION(:) :: ${varlist("my_blk_ind")}$
377 404 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: blk_index, blk_offset, blk_size, &
378 404 : tensor_dims
379 404 : INTEGER, DIMENSION(:, :), ALLOCATABLE :: ind_nd
380 : #:for ndim in ndims
381 : REAL(KIND=dp), ALLOCATABLE, &
382 404 : DIMENSION(${shape_colon(ndim)}$) :: blk_values_${ndim}$
383 : #:endfor
384 : TYPE(dbt_iterator_type) :: iterator
385 : INTEGER, DIMENSION(4) :: iseed
386 : INTEGER, DIMENSION(2) :: blk_index_2d, nblks_2d
387 :
388 404 : nblks_alloc = SIZE(blk_ind_1)
389 404 : mynode = mp_comm%mepos
390 :
391 404 : IF (.NOT. enumerate) THEN
392 60 : CPASSERT(randmat_counter .NE. 0)
393 :
394 60 : randmat_counter = randmat_counter + 1
395 : END IF
396 :
397 1616 : ALLOCATE (ind_nd(nblks_alloc, ndims_tensor(tensor)))
398 :
399 : !$OMP PARALLEL DEFAULT(NONE) SHARED(ind_nd,nblks_alloc,tensor,mynode,enumerate,randmat_counter) &
400 : !$OMP SHARED(${varlist("blk_ind")}$) &
401 : !$OMP PRIVATE(my_nblks_alloc,ib,proc,i,iterator,blk_offset,blk_size,blk_index) &
402 : !$OMP PRIVATE(blk_index_2d,nblks_2d,iseed,nze,tensor_dims) &
403 : !$OMP PRIVATE(${varlist("blk_values", nmin=2)}$) &
404 404 : !$OMP PRIVATE(${varlist("my_blk_ind")}$)
405 :
406 : my_nblks_alloc = 0
407 : !$OMP DO
408 : DO ib = 1, nblks_alloc
409 : #:for ndim in ndims
410 : IF (ndims_tensor(tensor) == ${ndim}$) THEN
411 : ind_nd(ib, :) = [${varlist("blk_ind", nmax=ndim, suffix="(ib)")}$]
412 : END IF
413 : #:endfor
414 : CALL dbt_get_stored_coordinates(tensor, ind_nd(ib, :), proc)
415 : IF (proc == mynode) THEN
416 : my_nblks_alloc = my_nblks_alloc + 1
417 : END IF
418 : END DO
419 : !$OMP END DO
420 :
421 : #:for dim in range(1, maxdim+1)
422 : IF (ndims_tensor(tensor) >= ${dim}$) THEN
423 : ALLOCATE (my_blk_ind_${dim}$ (my_nblks_alloc))
424 : END IF
425 : #:endfor
426 :
427 : i = 0
428 : !$OMP DO
429 : DO ib = 1, nblks_alloc
430 : CALL dbt_get_stored_coordinates(tensor, ind_nd(ib, :), proc)
431 : IF (proc == mynode) THEN
432 : i = i + 1
433 : #:for dim in range(1, maxdim+1)
434 : IF (ndims_tensor(tensor) >= ${dim}$) THEN
435 : my_blk_ind_${dim}$ (i) = blk_ind_${dim}$ (ib)
436 : END IF
437 : #:endfor
438 : END IF
439 : END DO
440 : !$OMP END DO
441 :
442 : #:for ndim in ndims
443 : IF (ndims_tensor(tensor) == ${ndim}$) THEN
444 : CALL dbt_reserve_blocks(tensor, ${varlist("my_blk_ind", nmax=ndim)}$)
445 : END IF
446 : #:endfor
447 :
448 : CALL dbt_iterator_start(iterator, tensor)
449 : DO WHILE (dbt_iterator_blocks_left(iterator))
450 : CALL dbt_iterator_next_block(iterator, blk_index, blk_size=blk_size, blk_offset=blk_offset)
451 :
452 : IF (.NOT. enumerate) THEN
453 : blk_index_2d = INT(get_2d_indices_tensor(tensor%nd_index_blk, blk_index))
454 : CALL dbt_get_mapping_info(tensor%nd_index_blk, dims_2d=nblks_2d)
455 : iseed = generate_larnv_seed(blk_index_2d(1), nblks_2d(1), blk_index_2d(2), nblks_2d(2), randmat_counter)
456 : nze = PRODUCT(blk_size)
457 : END IF
458 :
459 : #:for ndim in ndims
460 : IF (ndims_tensor(tensor) == ${ndim}$) THEN
461 : CALL allocate_any(blk_values_${ndim}$, shape_spec=blk_size)
462 : CALL dims_tensor(tensor, tensor_dims)
463 : IF (enumerate) THEN
464 : CALL enumerate_block_elements(blk_size, blk_offset, tensor_dims, blk_${ndim}$=blk_values_${ndim}$)
465 : ELSE
466 : CALL dlarnv(1, iseed, nze, blk_values_${ndim}$)
467 : END IF
468 : CALL dbt_put_block(tensor, blk_index, blk_size, blk_values_${ndim}$)
469 : DEALLOCATE (blk_values_${ndim}$)
470 : END IF
471 : #:endfor
472 : END DO
473 : CALL dbt_iterator_stop(iterator)
474 : !$OMP END PARALLEL
475 :
476 808 : END SUBROUTINE
477 :
478 : ! **************************************************************************************************
479 : !> \brief Enumerate tensor entries in block
480 : !> \param blk_2 block values for 2 dimensions
481 : !> \param blk_3 block values for 3 dimensions
482 : !> \param blk_size size of block
483 : !> \param blk_offset block offset (indices of first element)
484 : !> \param tensor_size global tensor sizes
485 : !> \author Patrick Seewald
486 : ! **************************************************************************************************
487 2928 : SUBROUTINE enumerate_block_elements(blk_size, blk_offset, tensor_size, ${varlist("blk", nmin=2)}$)
488 : INTEGER, DIMENSION(:), INTENT(IN) :: blk_size, blk_offset, tensor_size
489 : #:for ndim in ndims
490 : REAL(KIND=dp), DIMENSION(${shape_colon(ndim)}$), &
491 : OPTIONAL, INTENT(OUT) :: blk_${ndim}$
492 : #:endfor
493 : INTEGER :: ndim
494 5856 : INTEGER, DIMENSION(SIZE(blk_size)) :: arr_ind, tens_ind
495 : INTEGER :: ${varlist("i")}$
496 :
497 2928 : ndim = SIZE(tensor_size)
498 :
499 : #:for ndim in ndims
500 3120 : IF (ndim == ${ndim}$) THEN
501 : #:for idim in range(ndim,0,-1)
502 4323620 : DO i_${idim}$ = 1, blk_size(${idim}$)
503 : #:endfor
504 18053064 : arr_ind(:) = [${varlist("i", nmax=ndim)}$]
505 18053064 : tens_ind(:) = arr_ind(:) + blk_offset(:) - 1
506 4199108 : blk_${ndim}$ (${arrlist("arr_ind", nmax=ndim)}$) = combine_tensor_index(tens_ind, tensor_size)
507 : #:for idim in range(ndim,0,-1)
508 : END DO
509 : #:endfor
510 : END IF
511 : #:endfor
512 :
513 2928 : END SUBROUTINE
514 :
515 : #:for ndim in ndims
516 : ! **************************************************************************************************
517 : !> \brief Transform a distributed sparse tensor to a replicated dense array. This is only useful for
518 : !> testing tensor contraction by matrix multiplication of dense arrays.
519 : !> \author Patrick Seewald
520 : ! **************************************************************************************************
521 80 : SUBROUTINE dist_sparse_tensor_to_repl_dense_${ndim}$d_array(tensor, array)
522 :
523 : TYPE(dbt_type), INTENT(INOUT) :: tensor
524 : REAL(dp), ALLOCATABLE, DIMENSION(${shape_colon(ndim)}$), &
525 : INTENT(OUT) :: array
526 80 : REAL(dp), ALLOCATABLE, DIMENSION(${shape_colon(ndim)}$) :: block
527 160 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: dims_nd, ind_nd, blk_size, blk_offset
528 : TYPE(dbt_iterator_type) :: iterator
529 : INTEGER :: idim
530 80 : INTEGER, DIMENSION(ndims_tensor(tensor)) :: blk_start, blk_end
531 : LOGICAL :: found
532 :
533 0 : CPASSERT(ndims_tensor(tensor) .EQ. ${ndim}$)
534 80 : CALL dbt_get_info(tensor, nfull_total=dims_nd)
535 80 : CALL allocate_any(array, shape_spec=dims_nd)
536 33063268 : array(${shape_colon(ndim)}$) = 0.0_dp
537 :
538 : !$OMP PARALLEL DEFAULT(NONE) SHARED(tensor,array) &
539 80 : !$OMP PRIVATE(iterator,ind_nd,blk_size,blk_offset,block,found,idim,blk_start,blk_end)
540 : CALL dbt_iterator_start(iterator, tensor)
541 : DO WHILE (dbt_iterator_blocks_left(iterator))
542 : CALL dbt_iterator_next_block(iterator, ind_nd, blk_size=blk_size, blk_offset=blk_offset)
543 : CALL dbt_get_block(tensor, ind_nd, block, found)
544 : CPASSERT(found)
545 :
546 : DO idim = 1, ndims_tensor(tensor)
547 : blk_start(idim) = blk_offset(idim)
548 : blk_end(idim) = blk_offset(idim) + blk_size(idim) - 1
549 : END DO
550 : array(${", ".join(["blk_start("+str(idim)+"):blk_end("+str(idim)+")" for idim in range(1, ndim + 1)])}$) = &
551 : block(${shape_colon(ndim)}$)
552 :
553 : DEALLOCATE (block)
554 : END DO
555 : CALL dbt_iterator_stop(iterator)
556 : !$OMP END PARALLEL
557 80 : CALL tensor%pgrid%mp_comm_2d%sum(array)
558 :
559 160 : END SUBROUTINE
560 : #:endfor
561 :
562 : ! **************************************************************************************************
563 : !> \brief test tensor contraction
564 : !> \note for testing/debugging, simply replace a call to dbt_contract with a call to this routine
565 : !> \author Patrick Seewald
566 : ! **************************************************************************************************
567 20 : SUBROUTINE dbt_contract_test(alpha, tensor_1, tensor_2, beta, tensor_3, &
568 20 : contract_1, notcontract_1, &
569 20 : contract_2, notcontract_2, &
570 20 : map_1, map_2, &
571 : unit_nr, &
572 8 : bounds_1, bounds_2, bounds_3, &
573 : log_verbose, write_int)
574 :
575 : REAL(dp), INTENT(IN) :: alpha
576 : TYPE(dbt_type), INTENT(INOUT) :: tensor_1, tensor_2, tensor_3
577 : REAL(dp), INTENT(IN) :: beta
578 : INTEGER, DIMENSION(:), INTENT(IN) :: contract_1, contract_2, &
579 : notcontract_1, notcontract_2, &
580 : map_1, map_2
581 : INTEGER, INTENT(IN) :: unit_nr
582 : INTEGER, DIMENSION(2, SIZE(contract_1)), &
583 : OPTIONAL :: bounds_1
584 : INTEGER, DIMENSION(2, SIZE(notcontract_1)), &
585 : OPTIONAL :: bounds_2
586 : INTEGER, DIMENSION(2, SIZE(notcontract_2)), &
587 : OPTIONAL :: bounds_3
588 : LOGICAL, INTENT(IN), OPTIONAL :: log_verbose
589 : LOGICAL, INTENT(IN), OPTIONAL :: write_int
590 : INTEGER :: io_unit, mynode
591 : TYPE(mp_comm_type) :: mp_comm
592 20 : INTEGER, DIMENSION(:), ALLOCATABLE :: size_1, size_2, size_3, &
593 20 : order_t1, order_t2, order_t3
594 40 : INTEGER, DIMENSION(2, ndims_tensor(tensor_1)) :: bounds_t1
595 40 : INTEGER, DIMENSION(2, ndims_tensor(tensor_2)) :: bounds_t2
596 :
597 : #:for ndim in ndims
598 : REAL(KIND=dp), ALLOCATABLE, &
599 20 : DIMENSION(${shape_colon(ndim)}$) :: array_1_${ndim}$d, &
600 20 : array_2_${ndim}$d, &
601 20 : array_3_${ndim}$d, &
602 20 : array_1_${ndim}$d_full, &
603 20 : array_2_${ndim}$d_full, &
604 20 : array_3_0_${ndim}$d, &
605 20 : array_1_rs${ndim}$d, &
606 20 : array_2_rs${ndim}$d, &
607 20 : array_3_rs${ndim}$d, &
608 20 : array_3_0_rs${ndim}$d
609 : #:endfor
610 : REAL(KIND=dp), ALLOCATABLE, &
611 20 : DIMENSION(:, :) :: array_1_mm, &
612 20 : array_2_mm, &
613 20 : array_3_mm, &
614 20 : array_3_test_mm
615 : LOGICAL :: eql, notzero
616 : LOGICAL, PARAMETER :: debug = .FALSE.
617 : REAL(KIND=dp) :: cs_1, cs_2, cs_3, eql_diff
618 : LOGICAL :: do_crop_1, do_crop_2
619 :
620 20 : mp_comm = tensor_1%pgrid%mp_comm_2d
621 20 : mynode = mp_comm%mepos
622 20 : io_unit = -1
623 20 : IF (mynode .EQ. 0) io_unit = unit_nr
624 :
625 20 : cs_1 = dbt_checksum(tensor_1)
626 20 : cs_2 = dbt_checksum(tensor_2)
627 20 : cs_3 = dbt_checksum(tensor_3)
628 :
629 20 : IF (io_unit > 0) THEN
630 10 : WRITE (io_unit, *)
631 10 : WRITE (io_unit, '(A)') repeat("-", 80)
632 10 : WRITE (io_unit, '(A,1X,A,1X,A,1X,A,1X,A,1X,A)') "Testing tensor contraction", &
633 20 : TRIM(tensor_1%name), "x", TRIM(tensor_2%name), "=", TRIM(tensor_3%name)
634 10 : WRITE (io_unit, '(A)') repeat("-", 80)
635 : END IF
636 :
637 : IF (debug) THEN
638 : IF (io_unit > 0) THEN
639 : WRITE (io_unit, "(A, E9.2)") "checksum ", TRIM(tensor_1%name), cs_1
640 : WRITE (io_unit, "(A, E9.2)") "checksum ", TRIM(tensor_2%name), cs_2
641 : WRITE (io_unit, "(A, E9.2)") "checksum ", TRIM(tensor_3%name), cs_3
642 : END IF
643 : END IF
644 :
645 : IF (debug) THEN
646 : CALL dbt_write_block_indices(tensor_1, io_unit, unit_nr)
647 : CALL dbt_write_blocks(tensor_1, io_unit, unit_nr, write_int)
648 : END IF
649 :
650 : SELECT CASE (ndims_tensor(tensor_3))
651 : #:for ndim in ndims
652 : CASE (${ndim}$)
653 20 : CALL dist_sparse_tensor_to_repl_dense_array(tensor_3, array_3_0_${ndim}$d)
654 : #:endfor
655 : END SELECT
656 :
657 : CALL dbt_contract(alpha, tensor_1, tensor_2, beta, tensor_3, &
658 : contract_1, notcontract_1, &
659 : contract_2, notcontract_2, &
660 : map_1, map_2, &
661 : bounds_1=bounds_1, bounds_2=bounds_2, bounds_3=bounds_3, &
662 : filter_eps=1.0E-12_dp, &
663 20 : unit_nr=io_unit, log_verbose=log_verbose)
664 :
665 20 : cs_3 = dbt_checksum(tensor_3)
666 :
667 : IF (debug) THEN
668 : IF (io_unit > 0) THEN
669 : WRITE (io_unit, "(A, E9.2)") "checksum ", TRIM(tensor_3%name), cs_3
670 : END IF
671 : END IF
672 :
673 20 : do_crop_1 = .FALSE.; do_crop_2 = .FALSE.!; do_crop_3 = .FALSE.
674 :
675 : ! crop tensor as first step
676 82 : bounds_t1(1, :) = 1
677 82 : CALL dbt_get_info(tensor_1, nfull_total=bounds_t1(2, :))
678 :
679 80 : bounds_t2(1, :) = 1
680 80 : CALL dbt_get_info(tensor_2, nfull_total=bounds_t2(2, :))
681 :
682 20 : IF (PRESENT(bounds_1)) THEN
683 22 : bounds_t1(:, contract_1) = bounds_1
684 10 : do_crop_1 = .TRUE.
685 22 : bounds_t2(:, contract_2) = bounds_1
686 20 : do_crop_2 = .TRUE.
687 : END IF
688 :
689 20 : IF (PRESENT(bounds_2)) THEN
690 14 : bounds_t1(:, notcontract_1) = bounds_2
691 20 : do_crop_1 = .TRUE.
692 : END IF
693 :
694 20 : IF (PRESENT(bounds_3)) THEN
695 14 : bounds_t2(:, notcontract_2) = bounds_3
696 20 : do_crop_2 = .TRUE.
697 : END IF
698 :
699 : ! Convert tensors to simple multidimensional arrays
700 : #:for i in range(1,4)
701 : SELECT CASE (ndims_tensor(tensor_${i}$))
702 : #:for ndim in ndims
703 : CASE (${ndim}$)
704 : #:if i < 3
705 : CALL dist_sparse_tensor_to_repl_dense_array(tensor_${i}$, array_${i}$_${ndim}$d_full)
706 162 : CALL allocate_any(array_${i}$_${ndim}$d, shape_spec=SHAPE(array_${i}$_${ndim}$d_full))
707 22512580 : array_${i}$_${ndim}$d = 0.0_dp
708 : array_${i}$_${ndim}$d(${", ".join(["bounds_t" + str(i) + "(1, " + str(idim) + "):bounds_t" + str(i) + "(2, " + str(idim) + ")" for idim in range(1, ndim+1)])}$) = &
709 19082870 : array_${i}$_${ndim}$d_full(${", ".join(["bounds_t" + str(i) + "(1, " + str(idim) + "):bounds_t" + str(i) + "(2, " + str(idim) + ")" for idim in range(1, ndim+1)])}$)
710 : #:else
711 20 : CALL dist_sparse_tensor_to_repl_dense_array(tensor_${i}$, array_${i}$_${ndim}$d)
712 : #:endif
713 :
714 : #:endfor
715 : END SELECT
716 : #:endfor
717 :
718 : ! Get array sizes
719 :
720 : #:for i in range(1,4)
721 : SELECT CASE (ndims_tensor(tensor_${i}$))
722 : #:for ndim in ndims
723 : CASE (${ndim}$)
724 392 : ALLOCATE (size_${i}$, source=SHAPE(array_${i}$_${ndim}$d))
725 :
726 : #:endfor
727 : END SELECT
728 : #:endfor
729 :
730 : #:for i in range(1,4)
731 140 : ALLOCATE (order_t${i}$ (ndims_tensor(tensor_${i}$)))
732 : #:endfor
733 :
734 : ASSOCIATE (map_t1_1 => notcontract_1, map_t1_2 => contract_1, &
735 : map_t2_1 => notcontract_2, map_t2_2 => contract_2, &
736 : map_t3_1 => map_1, map_t3_2 => map_2)
737 :
738 : #:for i in range(1,4)
739 612 : order_t${i}$ (:) = dbt_inverse_order([map_t${i}$_1, map_t${i}$_2])
740 :
741 6 : SELECT CASE (ndims_tensor(tensor_${i}$))
742 : #:for ndim in ndims
743 : CASE (${ndim}$)
744 6 : CALL allocate_any(array_${i}$_rs${ndim}$d, source=array_${i}$_${ndim}$d, order=order_t${i}$)
745 60 : CALL allocate_any(array_${i}$_mm, sizes_2d(size_${i}$, map_t${i}$_1, map_t${i}$_2))
746 246 : array_${i}$_mm(:, :) = RESHAPE(array_${i}$_rs${ndim}$d, SHAPE(array_${i}$_mm))
747 : #:endfor
748 : END SELECT
749 : #:endfor
750 :
751 0 : SELECT CASE (ndims_tensor(tensor_3))
752 : #:for ndim in ndims
753 : CASE (${ndim}$)
754 0 : CALL allocate_any(array_3_0_rs${ndim}$d, source=array_3_0_${ndim}$d, order=order_t3)
755 20 : CALL allocate_any(array_3_test_mm, sizes_2d(size_3, map_t3_1, map_t3_2))
756 80 : array_3_test_mm(:, :) = RESHAPE(array_3_0_rs${ndim}$d, SHAPE(array_3_mm))
757 : #:endfor
758 : END SELECT
759 :
760 5101228 : array_3_test_mm(:, :) = beta*array_3_test_mm(:, :) + alpha*MATMUL(array_1_mm, transpose(array_2_mm))
761 :
762 : END ASSOCIATE
763 :
764 5101208 : eql_diff = MAXVAL(ABS(array_3_test_mm(:, :) - array_3_mm(:, :)))
765 5101208 : notzero = MAXVAL(ABS(array_3_test_mm(:, :))) .GT. 1.0E-12_dp
766 :
767 20 : eql = eql_diff .LT. 1.0E-11_dp
768 :
769 20 : IF (.NOT. eql .OR. .NOT. notzero) THEN
770 0 : IF (io_unit > 0) WRITE (io_unit, *) 'Test failed!', eql_diff
771 0 : CPABORT('')
772 : ELSE
773 20 : IF (io_unit > 0) WRITE (io_unit, *) 'Test passed!', eql_diff
774 : END IF
775 :
776 48 : END SUBROUTINE
777 :
778 : ! **************************************************************************************************
779 : !> \brief mapped sizes in 2d
780 : !> \author Patrick Seewald
781 : ! **************************************************************************************************
782 80 : FUNCTION sizes_2d(nd_sizes, map1, map2)
783 : INTEGER, DIMENSION(:), INTENT(IN) :: nd_sizes, map1, map2
784 : INTEGER, DIMENSION(2) :: sizes_2d
785 206 : sizes_2d(1) = PRODUCT(nd_sizes(map1))
786 200 : sizes_2d(2) = PRODUCT(nd_sizes(map2))
787 : END FUNCTION
788 :
789 : ! **************************************************************************************************
790 : !> \brief checksum of a tensor consistent with block_checksum
791 : !> \author Patrick Seewald
792 : ! **************************************************************************************************
793 80 : FUNCTION dbt_checksum(tensor)
794 : TYPE(dbt_type), INTENT(IN) :: tensor
795 : REAL(KIND=dp) :: dbt_checksum
796 80 : dbt_checksum = dbt_tas_checksum(tensor%matrix_rep)
797 80 : END FUNCTION
798 :
799 : ! **************************************************************************************************
800 : !> \brief Reset the seed used for generating random matrices to default value
801 : !> \author Patrick Seewald
802 : ! **************************************************************************************************
803 2 : SUBROUTINE dbt_reset_randmat_seed()
804 2 : randmat_counter = rand_seed_init
805 2 : END SUBROUTINE
806 :
807 20 : END MODULE
|
