Repository: libAtoms/QUIP (language: python)

Issue:
Feat: decouple blocksizes
Cheat `pdtrtrs` to think blocksizes of A are identical asserting that there is only a single process column. Change `mpi_blocksize` to `mpi_blocksize_rows` and `mpi_blocksize_cols`.

:warning: This is a breaking change wrt. `mpi_blocksize`. Should not be much of a problem as it's an optional argument, only relevant for people who have tried out MPI and set it manually.

Depends on https://github.com/libAtoms/GAP/pull/42


Discussion / hints:

Fix patch:
diff --git a/src/GAP b/src/GAP
index 17468dca62..ade605f83d 160000
--- a/src/GAP
+++ b/src/GAP
@@ -1,1 +1,1 @@
-Subproject commit 17468dca621ff08a01cceb95bf66650e877331ff
+Subproject commit ade605f83dbea5a4a2c192a7daa4147b488cb17b
diff --git a/src/libAtoms/ScaLAPACK.f95 b/src/libAtoms/ScaLAPACK.f95
index 3d18f670e2..ee82f34187 100644
--- a/src/libAtoms/ScaLAPACK.f95
+++ b/src/libAtoms/ScaLAPACK.f95
@@ -48,7 +48,7 @@ module ScaLAPACK_module
 private
 
 #ifdef SCALAPACK
-integer, external :: indxl2g, numroc
+integer, external :: ilcm, indxg2p, indxl2g, numroc
 #endif
 
 integer, parameter :: dlen_ = 50
@@ -154,6 +154,20 @@ module ScaLAPACK_module
   module procedure ScaLAPACK_diag_spinorZ, ScaLAPACK_diag_spinorD
 end interface diag_spinor
 
+public :: get_lwork_pdgeqrf
+interface get_lwork_pdgeqrf
+  module procedure get_lwork_pdgeqrf_i32o64
+  module procedure ScaLAPACK_get_lwork_pdgeqrf
+  module procedure ScaLAPACK_matrix_get_lwork_pdgeqrf
+end interface
+
+public :: get_lwork_pdormqr
+interface get_lwork_pdormqr
+  module procedure get_lwork_pdormqr_i32o64
+  module procedure ScaLAPACK_get_lwork_pdormqr
+  module procedure ScaLAPACK_matrix_get_lwork_pdormqr
+end interface
+
 public :: ScaLAPACK_pdgeqrf_wrapper, ScaLAPACK_pdtrtrs_wrapper, ScaLAPACK_pdormqr_wrapper
 public :: ScaLAPACK_matrix_QR_solve, ScaLAPACK_to_array1d, ScaLAPACK_to_array2d
 
@@ -1373,7 +1387,6 @@ subroutine ScaLAPACK_pdormqr_wrapper(A_info, A_data, C_info, C_data, tau, work)
   n = C_info%N_C
   k = size(tau)
 
-  call reallocate(tau, k) ! @fixme circular logic
   call reallocate(work, 1)
   call pdormqr('L', 'T', m, n, k, A_data, 1, 1, A_info%desc, &
     tau, C_data, 1, 1, C_info%desc, work, -1, info)
@@ -1461,4 +1474,134 @@ subroutine ScaLAPACK_to_array2d(A_info, A_data, array)
 #endif
 end subroutine ScaLAPACK_to_array2d
 
+! returns 64bit minimal work array length for pdgeqrf from 32bit sources
+! adapted from documentation of pdgeqrf
+function get_lwork_pdgeqrf_i32o64(m, n, ia, ja, mb_a, nb_a, &
+    myrow, mycol, rsrc_a, csrc_a, nprow, npcol) result(lwork)
+  integer, intent(in) :: m, n, ia, ja, mb_a, nb_a
+  integer, intent(in) :: myrow, mycol, rsrc_a, csrc_a, nprow, npcol
+  integer(idp) :: lwork
+
+  integer :: iarow, iacol, iroff, icoff
+  integer(idp) :: mp0, nq0, nb64
+
+  lwork = 0
+
+#ifdef SCALAPACK
+  iroff = mod(ia-1, mb_a)
+  icoff = mod(ja-1, nb_a)
+  iarow = indxg2p(ia, mb_a, myrow, rsrc_a, nprow)
+  iacol = indxg2p(ja, nb_a, mycol, csrc_a, npcol)
+  mp0 = numroc(m+iroff, mb_a, myrow, iarow, nprow)
+  nq0 = numroc(n+icoff, nb_a, mycol, iacol, npcol)
+
+  nb64 = int(nb_a, idp)
+  lwork = nb64 * (mp0 + nq0 + nb64)
+#endif
+end function get_lwork_pdgeqrf_i32o64
+
+function ScaLAPACK_get_lwork_pdgeqrf(this, m, n, mb_a, nb_a) result(lwork)
+  type(ScaLAPACK), intent(in) :: this
+  integer, intent(in) :: m, n, mb_a, nb_a
+  integer(idp) :: lwork
+
+  integer, parameter :: ia = 1, ja = 1, rsrc_a = 0, csrc_a = 0
+
+  lwork = 0
+
+#ifdef SCALAPACK
+  lwork = get_lwork_pdgeqrf(m, n, ia, ja, mb_a, nb_a, &
+    this%my_proc_row, this%my_proc_col, rsrc_a, csrc_a, &
+    this%n_proc_rows, this%n_proc_cols)
+#endif
+end function ScaLAPACK_get_lwork_pdgeqrf
+
+function ScaLAPACK_matrix_get_lwork_pdgeqrf(this) result(lwork)
+  type(Matrix_ScaLAPACK_Info), intent(in) :: this
+  integer(idp) :: lwork
+
+  lwork = 0
+
+#ifdef SCALAPACK
+  lwork = get_lwork_pdgeqrf(this%ScaLAPACK_obj, this%N_R, this%N_C, this%NB_R, this%NB_C)
+#endif
+end function ScaLAPACK_matrix_get_lwork_pdgeqrf
+
+! returns 64bit minimal work array length for pdormqr from 32bit sources
+! adapted from documentation of pdormqr
+function get_lwork_pdormqr_i32o64(side, m, n, ia, ja, mb_a, nb_a, ic, jc, &
+    mb_c, nb_c, myrow, mycol, rsrc_a, csrc_a, rsrc_c, csrc_c, nprow, npcol) result(lwork)
+  character :: side
+  integer, intent(in) :: m, n, ia, ja, mb_a, nb_a
+  integer, intent(in) :: ic, jc, mb_c, nb_c
+  integer, intent(in) :: rsrc_a, csrc_a, rsrc_c, csrc_c
+  integer, intent(in) :: myrow, mycol, nprow, npcol
+  integer(idp) :: lwork
+
+  integer :: lcm, lcmq
+  integer :: iarow, iroffa, icoffa, iroffc, icoffc, icrow, iccol
+  integer(idp) :: npa0, mpc0, nqc0, nr, nb64, lwork1, lwork2
+
+  lwork = 0
+
+#ifdef SCALAPACK
+  iroffc = mod(ic-1, mb_c)
+  icoffc = mod(jc-1, nb_c)
+  icrow = indxg2p(ic, mb_c, myrow, rsrc_c, nprow)
+  iccol = indxg2p(jc, nb_c, mycol, csrc_c, npcol)
+  mpc0 = numroc(m+iroffc, mb_c, myrow, icrow, nprow)
+  nqc0 = numroc(n+icoffc, nb_c, mycol, iccol, npcol)
+
+  nb64 = int(nb_a, idp)
+  lwork1 = (nb64 * (nb64 - 1)) / 2
+  if (side == 'L') then
+    lwork2 = (nqc0 + mpc0) * nb64
+  else if (side == 'R') then
+    iroffa = mod(ia-1, mb_a)
+    icoffa = mod(ja-1, nb_a)
+    iarow = indxg2p(ia, mb_a, myrow, rsrc_a, nprow)
+    npa0 = numroc(n+iroffa, mb_a, myrow, iarow, nprow)
+
+    lcm = ilcm(nprow, npcol)
+    lcmq = lcm / npcol
+    nr = numroc(n+icoffc, nb_a, 0, 0, npcol)
+    nr = numroc(int(nr, isp), nb_a, 0, 0, lcmq)
+    nr = max(npa0 + nr, mpc0)
+    lwork2 = (nqc0 + nr) * nb64
+  end if
+  lwork = max(lwork1, lwork2) + nb64 * nb64
+#endif
+end function get_lwork_pdormqr_i32o64
+
+function ScaLAPACK_get_lwork_pdormqr(this, side, m, n, mb_a, nb_a, mb_c, nb_c) result(lwork)
+  type(ScaLAPACK), intent(in) :: this
+  character :: side
+  integer, intent(in) :: m, n, mb_a, nb_a, mb_c, nb_c
+  integer(idp) :: lwork
+
+  integer, parameter :: ia = 1, ja = 1, rsrc_a = 0, csrc_a = 0
+  integer, parameter :: ic = 1, jc = 1, rsrc_c = 0, csrc_c = 0
+
+  lwork = 0
+
+#ifdef SCALAPACK
+  lwork = get_lwork_pdormqr(side, m, n, ia, ja, mb_a, nb_a, ic, jc, mb_c, nb_c, &
+    this%my_proc_row, this%my_proc_col, rsrc_a, csrc_a, rsrc_c, csrc_c, &
+    this%n_proc_rows, this%n_proc_cols)
+#endif
+end function ScaLAPACK_get_lwork_pdormqr
+
+function ScaLAPACK_matrix_get_lwork_pdormqr(A_info, C_info, side) result(lwork)
+  type(Matrix_ScaLAPACK_Info), intent(in) :: A_info, C_info
+  character :: side
+  integer(idp) :: lwork
+
+  lwork = 0
+
+#ifdef SCALAPACK
+  lwork = get_lwork_pdormqr(A_info%ScaLAPACK_obj, side, C_info%N_R, &
+    C_info%N_C, A_info%NB_R, A_info%NB_C, A_info%NB_R, A_info%NB_C)
+#endif
+end function ScaLAPACK_matrix_get_lwork_pdormqr
+
 end module ScaLAPACK_module


Test patch:
diff --git a/tests/test_gapfit.py b/tests/test_gapfit.py
index ea9a849f28..98078b7157 100644
--- a/tests/test_gapfit.py
+++ b/tests/test_gapfit.py
@@ -142,5 +142,13 @@ def test_gap_fit_silicon_scalapack(self):
         self.check_gap_fit()
         self.check_latest_sparsex_file_hash()
 
+    @unittest.skipIf(os.environ.get('HAVE_SCALAPACK') != '1', 'ScaLAPACK support not enabled')
+    def test_gap_fit_silicon_scalapack_big_blocksize(self):
+        command_line = self.cl_template.safe_substitute(SPARSE_METHOD='sparse_method=FILE sparse_file=si_gap_fit_sparseX.inp')
+        command_line += ' mpi_blocksize_cols=37838'  # too large for 32bit integer
+        self.run_gap_fit(command_line, prefix='mpirun -np 2')
+        with open(self.log_name) as f:
+            self.assertTrue("too large for 32bit work array in ScaLAPACK" in f.read())
+
 if __name__ == '__main__':
     unittest.main()