test_complete_redistribute Subroutine

subroutine test_complete_redistribute()

Arguments

None

Source Code

   SUBROUTINE test_complete_redistribute()
      INTEGER, POINTER, CONTIGUOUS :: rbs(:), cbs(:), rd(:), cd(:)
      INTEGER, POINTER, CONTIGUOUS :: new_rbs(:), new_cbs(:), new_rd(:), new_cd(:)
      INTEGER :: itype, icase, r, c, i, j, k, n, p, ro, co, offset
      INTEGER :: rows(16), cols(16)
      INTEGER, ALLOCATABLE :: index_before(:)
      TYPE(dbcsr_distribution_obj) :: dist, new_dist
      TYPE(dbcsr_type) :: source, target, converted
      TYPE(dbcsr_data_obj) :: values, actual
      COMPLEX(dp) :: reference(10, 10), expected(10, 10)
      LOGICAL :: keep

      ALLOCATE (rbs(4), cbs(4), rd(4), cd(4), new_rbs(3), new_cbs(3), new_rd(3), new_cd(3))
      rbs = [2, 3, 1, 4]; cbs = [3, 1, 2, 4]
      new_rbs = [3, 4, 3]; new_cbs = [2, 5, 3]
      rd = [(MOD(i - 1, npdims(1)), i=1, 4)]
      cd = [(MOD(i - 1, npdims(2)), i=1, 4)]
      new_rd = [(MOD(i, npdims(1)), i=1, 3)]
      new_cd = [(MOD(i, npdims(2)), i=1, 3)]
      CALL dbcsr_distribution_new(dist, mp_env, rd, cd)
      CALL dbcsr_distribution_new(new_dist, mp_env, new_rd, new_cd)
      CALL dbcsr_data_init(values)
      CALL dbcsr_data_init(actual)
      CALL dbcsr_data_new(values, dbcsr_type_complex_8, data_size=100)
      CALL dbcsr_data_new(actual, dbcsr_type_complex_8_2d, 10, 10)
      DO itype = 1, SIZE(data_types)
         reference = CMPLX(0.0_dp, 0.0_dp, dp)
         CALL dbcsr_create(source, "reblock source", dist, dbcsr_type_no_symmetry, &
                           row_blk_size=rbs, col_blk_size=cbs, data_type=data_types(itype))
         n = 0
         DO r = 1, SIZE(rbs)
            ro = SUM(rbs(:r - 1))
            DO c = 1, SIZE(cbs)
               IF (MOD(r + c, 3) == 0) CYCLE
               co = SUM(cbs(:c - 1))
               DO j = 1, cbs(c)
                  DO i = 1, rbs(r)
                     reference(ro + i, co + j) = CMPLX(100*(ro + i) + co + j, 0, dp)
                  END DO
               END DO
               CALL mp_cart_rank(group, [rd(r), cd(c)], p)
               IF (p /= mynode) CYCLE
               n = n + 1; rows(n) = r; cols(n) = c
            END DO
         END DO
         IF (itype > 2) reference = reference*CMPLX(1.0_dp, -1.0_dp, dp)
         CALL dbcsr_reserve_blocks(source, rows(:n), cols(:n))
         DO k = 1, n
            r = rows(k); c = cols(k)
            ro = SUM(rbs(:r - 1)); co = SUM(cbs(:c - 1))
            offset = source%blk_p(k)
            values%d%c_dp(offset:offset + rbs(r)*cbs(c) - 1) = &
               RESHAPE(reference(ro + 1:ro + rbs(r), co + 1:co + cbs(c)), [rbs(r)*cbs(c)])
         END DO
         CALL dbcsr_data_set_size_referenced(values, source%nze)
         CALL dbcsr_data_convert(source%data_area, values)
         DO icase = 1, 4
            ! Empty target, populated target, retained sparsity, and empty source.
            keep = icase == 3
            CALL dbcsr_create(target, "reblock target", new_dist, dbcsr_type_no_symmetry, &
                              row_blk_size=new_rbs, col_blk_size=new_cbs, data_type=data_types(itype))
            IF (icase > 1) THEN
               n = 0
               DO r = 1, SIZE(new_rbs)
                  DO c = 1, SIZE(new_cbs)
                     IF (MOD(r + c, 2) == 0) CYCLE
                     CALL mp_cart_rank(group, [new_rd(r), new_cd(c)], p)
                     IF (p /= mynode) CYCLE
                     n = n + 1; rows(n) = r; cols(n) = c
                  END DO
               END DO
               CALL dbcsr_reserve_blocks(target, rows(:n), cols(:n))
               values%d%c_dp(:) = CMPLX(-7.0_dp, 0.0_dp, dp)
               CALL dbcsr_data_set_size_referenced(values, target%nze)
               CALL dbcsr_data_convert(target%data_area, values)
            END IF
            IF (icase == 4) THEN
               CALL dbcsr_clear(source)
               reference = CMPLX(0.0_dp, 0.0_dp, dp)
            END IF
            ALLOCATE (index_before, SOURCE=target%index)
            CALL dbcsr_complete_redistribute(source, target, keep_sparsity=keep)
            CALL dbcsr_verify_matrix(target)
            IF (.NOT. ASSOCIATED(target%dist%d, new_dist%d)) &
               DBCSR_ABORT("Reblocking changed distribution")
            IF (keep .OR. icase == 4) THEN
               IF (ANY(target%index /= index_before)) &
                  DBCSR_ABORT("Reblocking rebuilt an unchanged index")
            END IF
            DEALLOCATE (index_before)
            IF (icase == 1) THEN
               ! CSR conversion requires newly built matrices to have packed row-wise data.
               offset = 1
               DO r = 1, SIZE(new_rbs)
                  DO k = target%row_p(r) + 1, target%row_p(r + 1)
                     IF (target%blk_p(k) /= offset) &
                        DBCSR_ABORT("Reblocking changed data order")
                     offset = offset + new_rbs(r)*new_cbs(target%col_i(k))
                  END DO
               END DO
            END IF
            CALL dbcsr_create(converted, template=target, data_type=dbcsr_type_complex_8)
            CALL dbcsr_copy(converted, target)
            CALL dbcsr_to_dense_local(converted, actual)
            CALL dbcsr_release(converted)
            expected = reference
            ! Compare locally, including zeros outside this rank's ownership.
            DO r = 1, SIZE(new_rbs)
               ro = SUM(new_rbs(:r - 1))
               DO c = 1, SIZE(new_cbs)
                  co = SUM(new_cbs(:c - 1))
                  CALL mp_cart_rank(group, [new_rd(r), new_cd(c)], p)
                  IF (p /= mynode .OR. (keep .AND. MOD(r + c, 2) == 0)) &
                     expected(ro + 1:ro + new_rbs(r), co + 1:co + new_cbs(c)) = 0.0_dp
               END DO
            END DO
            IF (ANY(ABS(actual%d%c2_dp - expected) > 1.0E-6_dp)) &
               DBCSR_ABORT("Reblocking changed matrix values")
            CALL dbcsr_release(target)
         END DO
         CALL dbcsr_release(source)
      END DO
      CALL dbcsr_data_release(values)
      CALL dbcsr_data_release(actual)
      CALL dbcsr_distribution_release(dist)
      CALL dbcsr_distribution_release(new_dist)
      DEALLOCATE (rbs, cbs, rd, cd, new_rbs, new_cbs, new_rd, new_cd)
   END SUBROUTINE test_complete_redistribute