Line data    Source code

       1              : !--------------------------------------------------------------------------------------------------!
       2              : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3              : !   Copyright 2000-2026 CP2K developers group <https://cp2k.org>                                   !
       4              : !                                                                                                  !
       5              : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6              : !--------------------------------------------------------------------------------------------------!
       7              : 
       8            2 : PROGRAM cp_fm_diag_utils_unittest
       9            2 :    USE cp_fm_diag_utils,                ONLY: cp_fm_max_ncpu_non_zero_column
      10              :    USE cp_fm_struct,                    ONLY: cp_fm_struct_type
      11              :    USE message_passing,                 ONLY: mp_para_env_type
      12              : 
      13              :    IMPLICIT NONE
      14              : 
      15              :    INTEGER, PARAMETER :: orders(2) = [509, 536]
      16              :    INTEGER :: i, itest, n
      17              :    INTEGER, EXTERNAL :: numroc
      18              :    TYPE(cp_fm_struct_type) :: matrix_struct
      19              :    TYPE(mp_para_env_type), TARGET :: para_env
      20              : 
      21              :    ! Layout metadata only: exercise large process grids even in two-rank CI.
      22            2 :    para_env%num_pe = 32
      23            2 :    matrix_struct%para_env => para_env
      24            2 :    matrix_struct%nrow_block = 64
      25            2 :    matrix_struct%ncol_block = 64
      26            2 :    ALLOCATE (matrix_struct%nrow_locals(0:3), matrix_struct%ncol_locals(0:7))
      27              : 
      28            6 :    DO itest = 1, SIZE(orders)
      29            4 :       n = orders(itest)
      30            4 :       matrix_struct%nrow_global = n
      31            4 :       matrix_struct%ncol_global = n
      32           20 :       DO i = 0, 3
      33           20 :          matrix_struct%nrow_locals(i) = numroc(n, 64, i, 0, 4)
      34              :       END DO
      35           36 :       DO i = 0, 7
      36           36 :          matrix_struct%ncol_locals(i) = numroc(n, 64, i, 0, 8)
      37              :       END DO
      38              : 
      39            4 :       IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 32, 64) /= 32) THEN
      40            0 :          ERROR STOP "Valid original redistribution layout changed"
      41              :       END IF
      42              :       ! 20 ranks use a 2x10 grid with empty columns; 18 ranks use a safe 3x6 grid.
      43            4 :       IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 20, 64) /= 18) THEN
      44            0 :          ERROR STOP "Unsafe forced redistribution layout accepted"
      45              :       END IF
      46            4 :       IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 12, 64) /= 12) THEN
      47            0 :          ERROR STOP "Safe redistribution heuristic target changed"
      48              :       END IF
      49            4 :       IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 8, 64) /= 8) THEN
      50            0 :          ERROR STOP "Safe smaller redistribution target changed"
      51              :       END IF
      52            6 :       IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 20, 32) /= 20) THEN
      53            0 :          ERROR STOP "Redistribution check ignored the final block size"
      54              :       END IF
      55              :    END DO
      56              : 
      57              :    ! Mandatory redistribution must still work when forced redistribution is off.
      58            2 :    matrix_struct%nrow_global = 1
      59            2 :    matrix_struct%ncol_global = 1
      60           10 :    matrix_struct%nrow_locals = 0
      61           18 :    matrix_struct%ncol_locals = 0
      62            2 :    matrix_struct%nrow_locals(0) = 1
      63            2 :    matrix_struct%ncol_locals(0) = 1
      64            2 :    IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 32, 64) /= 1) THEN
      65            0 :       ERROR STOP "Empty original redistribution layout was retained"
      66              :    END IF
      67              : 
      68            2 :    DEALLOCATE (matrix_struct%nrow_locals, matrix_struct%ncol_locals)
      69              : 
      70              :    ! Changing blocks also recreates the grid, even when the rank count is unchanged.
      71            2 :    para_env%num_pe = 20
      72            2 :    matrix_struct%nrow_global = 509
      73            2 :    matrix_struct%ncol_global = 509
      74            2 :    matrix_struct%nrow_block = 96
      75            2 :    matrix_struct%ncol_block = 96
      76            2 :    ALLOCATE (matrix_struct%nrow_locals(0:3), matrix_struct%ncol_locals(0:4))
      77           10 :    DO i = 0, 3
      78           10 :       matrix_struct%nrow_locals(i) = numroc(509, 96, i, 0, 4)
      79              :    END DO
      80           12 :    DO i = 0, 4
      81           12 :       matrix_struct%ncol_locals(i) = numroc(509, 96, i, 0, 5)
      82              :    END DO
      83            2 :    IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 20, 64) /= 18) THEN
      84            0 :       ERROR STOP "Block-size change skipped the target grid check"
      85              :    END IF
      86            2 :    DEALLOCATE (matrix_struct%nrow_locals, matrix_struct%ncol_locals)
      87              : 
      88              :    ! A safe default grid must not hide an invalid original custom grid.
      89            2 :    para_env%num_pe = 32
      90            2 :    matrix_struct%nrow_block = 64
      91            2 :    matrix_struct%ncol_block = 64
      92            2 :    ALLOCATE (matrix_struct%nrow_locals(0:1), matrix_struct%ncol_locals(0:15))
      93            6 :    DO i = 0, 1
      94            6 :       matrix_struct%nrow_locals(i) = numroc(509, 64, i, 0, 2)
      95              :    END DO
      96           34 :    DO i = 0, 15
      97           34 :       matrix_struct%ncol_locals(i) = numroc(509, 64, i, 0, 16)
      98              :    END DO
      99            2 :    IF (cp_fm_max_ncpu_non_zero_column(matrix_struct, 32, 64) /= 24) THEN
     100            0 :       ERROR STOP "Invalid original grid skipped mandatory redistribution"
     101              :    END IF
     102            2 :    DEALLOCATE (matrix_struct%nrow_locals, matrix_struct%ncol_locals)
     103            2 : END PROGRAM cp_fm_diag_utils_unittest