Repository: libAtoms/QUIP (language: python)

Issue:
Test: gap_fit update



Discussion / hints:

Fix patch:
diff --git a/src/GAP b/src/GAP
index bf7b1ff694..3874340304 160000
--- a/src/GAP
+++ b/src/GAP
@@ -1,1 +1,1 @@
-Subproject commit bf7b1ff6943f60765151fa9d0a656f9f9220542d
+Subproject commit 38743403045febf51cba1140c3334196d9b52997
diff --git a/src/Programs/Makefile b/src/Programs/Makefile
index 39e2e5f88c..3b29ffe3e2 100644
--- a/src/Programs/Makefile
+++ b/src/Programs/Makefile
@@ -43,7 +43,7 @@ PROGRAMS = quip md \
   quip_wrapper_example descriptors_wrapper_example \
   slice_sample order_atoms_as_molecules md_gid \
   quip_wrapper_simple_example quip_wrapper_simple_example_C \
-  get_qw
+  get_qw test_task_manager
 
 ifeq (${HAVE_TB},1)
    PROGRAMS += NRL_TB_to_xml DFTB_to_xml calc_n_poles 
diff --git a/src/libAtoms/Matrix.f95 b/src/libAtoms/Matrix.f95
index a0cf1fd109..05bdb69dc7 100644
--- a/src/libAtoms/Matrix.f95
+++ b/src/libAtoms/Matrix.f95
@@ -283,24 +283,32 @@ subroutine MatrixD_QR_Solve(A_SP, B_SP)
   endif
 end subroutine MatrixD_QR_Solve
 
-subroutine SP_Matrix_QR_Solve(A, B, X, procs, ScaLAPACK_obj)
+subroutine SP_Matrix_QR_Solve(A, B, X, ScaLAPACK_obj, blocksize)
   real(dp), intent(in), dimension(:,:), target :: A
   real(dp), intent(in), dimension(:), target :: B
   real(dp), intent(out), dimension(:) :: X
-  integer, intent(in) :: procs
   type(ScaLAPACK), intent(in) :: ScaLAPACK_obj
+  integer, intent(in) :: blocksize
 
-  integer :: m, n, mb, nb
+  integer :: mb, nb, ml, nl, mg, ng
   type(MatrixD) :: A_SP, B_SP
 
-  mb = size(A, 1)
-  nb = size(A, 2)
-  m = mb * procs
-  n = nb
+  ml = size(A, 1)
+  nl = size(A, 2)
+  mg = ml * ScaLAPACK_obj%n_proc_rows
+  ng = nl * ScaLAPACK_obj%n_proc_cols
 
-  ! Scalapack needs mb == nb for p?trtrs; choose nb for smaller work arrays
-  call initialise(A_SP, m, n, nb, nb, scalapack_obj=ScaLAPACK_obj, use_allocate=.false.)
-  call initialise(B_SP, m, 1, nb, 1, scalapack_obj=ScaLAPACK_obj, use_allocate=.false.)
+  mb = blocksize
+  nb = blocksize
+
+  if (ScaLAPACK_obj%my_proc_row < ScaLAPACK_obj%n_proc_rows .and. mod(ml, mb) /= 0) &
+    call system_abort("SP_Matrix_QR_Solve: nrows is not a multiple of blocksize: "//ml//" "//mb)
+  if (ScaLAPACK_obj%my_proc_col < ScaLAPACK_obj%n_proc_cols .and. mod(nl, nb) /= 0) &
+    call system_abort("SP_Matrix_QR_Solve: ncols is not a multiple of blocksize: "//nl//" "//nb)
+
+  ! Scalapack needs mb == nb for p?trtrs
+  call initialise(A_SP, mg, ng, mb, nb, scalapack_obj=ScaLAPACK_obj, use_allocate=.false.)
+  call initialise(B_SP, mg, 1, nb, 1, scalapack_obj=ScaLAPACK_obj, use_allocate=.false.)
 
   A_SP%data => A
   B_SP%data(lbound(B,1):ubound(B,1),1:1) => B(:)
diff --git a/src/libAtoms/System.f95 b/src/libAtoms/System.f95
index 61fc00a635..cceb391b2f 100644
--- a/src/libAtoms/System.f95
+++ b/src/libAtoms/System.f95
@@ -403,6 +403,7 @@ subroutine c_mem_info(total_mem,free_mem)
   public :: link_run_directory
   public :: wait_for_file_to_exist
   public :: is_open
+  public :: increase_to_multiple
 contains
 
 #ifdef NO_FORTRAN_ISNAN
@@ -3639,5 +3640,11 @@ subroutine progress_timer(total, current, name, elapsed_seconds)
 
    end subroutine progress_timer
 
+   function increase_to_multiple(a, m) result(res)
+     integer, intent(in) :: a, m
+     integer :: res
+     res = (a / m) * m
+     if (res < a) res = res + m
+   end function increase_to_multiple
 
 end module system_module
diff --git a/src/libAtoms/task_manager.f95 b/src/libAtoms/task_manager.f95
index 76bc313de1..e61f990d28 100644
--- a/src/libAtoms/task_manager.f95
+++ b/src/libAtoms/task_manager.f95
@@ -1,9 +1,16 @@
 module task_manager_module
   !> A task manager determines the distribution of tasks to workers with
   !> respect to the tasks' demands.
+  !>
+  !> Each task stores integer data in %idata.
+  !> %idata(1) is used for distribution.
+  !>
+  !> The special %worker_id=SHARED is used for shared tasks. Each worker
+  !> has a local entry with local %idata. %idata(1) contains the local
+  !> size, %idata(2) contains the global offset.
 
   use libatoms_module, only : initialise, finalise, print_title, print, &
-      operator(//), system_abort, inoutput, optional_default, OUTPUT
+      operator(//), system_abort, inoutput, optional_default, OUTPUT, PRINT_VERBOSE
   use linearalgebra_module, only : heap_sort
   use MPI_context_module, only : MPI_Context, print
   use ScaLAPACK_module, only : Scalapack, print
@@ -11,17 +18,19 @@ module task_manager_module
   implicit none
   private
 
-  public :: task_manager_type, task_type, worker_type
+  public :: task_manager_type, task_type, worker_type, SHARED
   public :: task_manager_init_tasks, task_manager_init_workers, task_manager_deallocate
+  public :: task_manager_init_idata
   public :: task_manager_add_task, task_manager_add_worker
-  public :: task_manager_distribute_tasks
+  public :: task_manager_distribute_tasks, task_manager_check_distribution
   public :: task_manager_print, task_manager_export_distribution
 
   integer, parameter :: UNDEFINED = -1
+  integer, parameter :: SHARED = 0
 
   type task_type
     integer :: index = UNDEFINED !> initial task index
-    integer :: memory = 0 !> memory needed for this task
+    integer, dimension(:), allocatable :: idata !> integer data for this task
     integer :: worker_id = UNDEFINED !> ID of worker assigned to this task
   end type task_type
 
@@ -29,7 +38,7 @@ module task_manager_module
     integer :: index = UNDEFINED !> initial worker index
     integer :: n_tasks = 0 !> number of tasks assigned to this worker
     integer :: n_padding = 0 !> number of padding rows to reach unified workload
-    integer, dimension(:), allocatable :: task_ids !> IDs of tasks assigned to this worker
+    type(task_type), dimension(:), allocatable :: tasks !> tasks assigned to this worker
   end type worker_type
 
   type task_manager_type
@@ -40,6 +49,9 @@ module task_manager_module
 
     integer :: n_tasks = 0 !> number of tasks
     integer :: n_workers = 0 !> number of workers
+    integer :: n_shared = 0 !> number of shared tasks
+
+    integer, dimension(:), allocatable :: idata !< custom shared integer data
 
     type(MPI_Context) :: MPI_obj
     type(ScaLAPACK) :: ScaLAPACK_obj
@@ -51,16 +63,24 @@ module task_manager_module
   contains
 
   !> Add task to list
-  subroutine task_manager_add_task(this, memory, worker_id)
+  subroutine task_manager_add_task(this, idata1, n_idata, worker_id)
     type(task_manager_type), intent(inout) :: this
-    integer, intent(in) :: memory !> memory needed for new task
+    integer, intent(in), optional :: idata1 !> first idata entry for new task
+    integer, intent(in), optional :: n_idata !> number of idata entries for new task
     integer, intent(in), optional :: worker_id !> preset worker for this task
 
+    integer :: my_n_idata
+
+    if (.not. this%active) return
+    my_n_idata = optional_default(1, n_idata)
+
     this%n_tasks = this%n_tasks + 1
     if (this%n_tasks > size(this%tasks)) call system_abort("More tasks added than allocated.")
     associate(task => this%tasks(this%n_tasks))
+        allocate(task%idata(my_n_idata))
+        task%idata = UNDEFINED
         task%index = this%n_tasks
-        task%memory = memory
+        if (present(idata1)) task%idata(1) = idata1
         if (present(worker_id)) task%worker_id = worker_id
     end associate
   end subroutine task_manager_add_task
@@ -69,6 +89,8 @@ end subroutine task_manager_add_task
   subroutine task_manager_add_worker(this)
     type(task_manager_type), intent(inout) :: this
 
+    if (.not. this%active) return
+
     this%n_workers = this%n_workers + 1
     if (this%n_workers > size(this%workers)) call system_abort("More workers added than allocated.")
     associate(worker => this%workers(this%n_workers))
@@ -77,84 +99,149 @@ subroutine task_manager_add_worker(this)
     end associate
   end subroutine task_manager_add_worker
 
-  !> @brief Distribute tasks among workers for minimal memory usage
-  !> Assertions: %tasks(:)%memory is set, %n_workers > 0
+  !> @brief Distribute tasks among workers for minimal idata(1) usage
   subroutine task_manager_distribute_tasks(this)
     type(task_manager_type), intent(inout) :: this
 
     integer :: i, t, w ! indices
     integer :: n
     integer, dimension(:), allocatable :: workloads
-    integer, dimension(:), allocatable :: memory_list
+    integer, dimension(:), allocatable :: idata1_list
     integer, dimension(:), allocatable :: index_list
 
+    if (.not. this%active) return
+    if (this%n_tasks < 1) return
+    if (this%n_workers < 1) call system_abort("task_manager_distribute_tasks: No workers initialised.")
+
     allocate(workloads(this%n_workers))
-    allocate(memory_list(this%n_tasks))
+    allocate(idata1_list(this%n_tasks))
     allocate(index_list(this%n_tasks))
 
-    ! copy to temp arrays, sort both by memory
-    memory_list = this%tasks(1:this%n_tasks)%memory
-    index_list = this%tasks(1:this%n_tasks)%index
-    call heap_sort(memory_list, i_data=index_list) ! index_list sorted via memory_list
+    ! copy to temp arrays, sort both by idata1
+    do i = 1, this%n_tasks
+      if (.not. allocated(this%tasks(i)%idata)) call system_abort("task_manager_distribute_tasks: idata of task "//i//" is unallocated.")
+      idata1_list(i) = this%tasks(i)%idata(1)
+      index_list(i) = this%tasks(i)%index
+    end do
+    call heap_sort(idata1_list, i_data=index_list) ! index_list sorted via idata1_list
 
     ! count workload of preset tasks
     workloads = 0
     do t = 1, this%n_tasks
         i = index_list(t)
         w = this%tasks(i)%worker_id
-        if (w /= UNDEFINED) then
+        if (w >= 1) then
           this%workers(w)%n_tasks = this%workers(w)%n_tasks + 1
-          workloads(w) = workloads(w) + memory_list(t)
+          workloads(w) = workloads(w) + idata1_list(t)
         end if
     end do
 
-    ! distribute tasks to workers to minimize maximum memory usage
+    ! distribute tasks to workers to minimize maximum idata1 sum
     do t = this%n_tasks, 1, -1
         i = index_list(t)
-        if (this%tasks(i)%worker_id /= UNDEFINED) cycle ! this preset tasks
+        if (this%tasks(i)%worker_id /= UNDEFINED) cycle ! skip preset tasks
         w = minloc(workloads, 1)
         this%tasks(i)%worker_id = w
         this%workers(w)%n_tasks = this%workers(w)%n_tasks + 1
-        workloads(w) = workloads(w) + memory_list(t)
+        workloads(w) = workloads(w) + idata1_list(t)
     end do
 
     this%unified_workload = maxval(workloads)
 
     ! fill workers with task id lists (replication for simpler retrieval)
+    this%n_shared = count(this%tasks(:)%worker_id == SHARED)
     do w = 1, this%n_workers
-        n = this%workers(w)%n_tasks
-        allocate(this%workers(w)%task_ids(n))
+        n = this%workers(w)%n_tasks + this%n_shared
+        allocate(this%workers(w)%tasks(n))
         this%workers(w)%n_tasks = 0
     end do
     do t = 1, this%n_tasks
         w = this%tasks(t)%worker_id
+        if (w < 1) cycle
         n = this%workers(w)%n_tasks + 1
-        this%workers(w)%task_ids(n) = t
+        this%workers(w)%tasks(n) = this%tasks(t)
         this%workers(w)%n_tasks = n
     end do
 
     do w = 1, this%n_workers
-        this%workers(w)%n_padding = this%unified_workload - workloads(w)
+      this%workers(w)%n_padding = this%unified_workload - workloads(w)
     end do
 
+    call task_manager_distribute_shared_tasks(this)
+
     this%distributed = .true.
   end subroutine task_manager_distribute_tasks
 
+  subroutine task_manager_distribute_shared_tasks(this)
+    ! fill padding first, then distribute rows evenly
+    type(task_manager_type), intent(inout) :: this
+  
+    integer :: i, n, t, w
+    integer :: rest, n_split, n_extra, n_plus, offset
+
+    if (.not. this%active) return
+
+    do t = 1, this%n_tasks
+      if (this%tasks(t)%worker_id /= SHARED) cycle
+
+      rest = max(this%tasks(t)%idata(1) - sum(this%workers(:)%n_padding), 0)
+      n_split = rest / this%n_workers
+      n_extra = mod(rest, this%n_workers)
+      n_plus = n_split
+      if (n_extra > 0) n_plus = n_split + 1
+
+      rest = this%tasks(t)%idata(1)
+      offset = 0
+      do w = 1, this%n_workers
+        n = this%workers(w)%n_tasks + 1
+        associate(worker => this%workers(w), w_task => this%workers(w)%tasks(n))
+          call task_deepcopy(w_task, this%tasks(t))
+
+          w_task%idata(1) = worker%n_padding + n_split
+          if (w <= n_extra) w_task%idata(1) =  w_task%idata(1) + 1
+          w_task%idata(1) = min(w_task%idata(1), rest)
+          w_task%idata(2) = offset
+          worker%n_padding = worker%n_padding + n_plus - w_task%idata(1)
+
+          rest = rest - w_task%idata(1)
+          offset =  offset + w_task%idata(1)
+          worker%n_tasks = n
+        end associate
+      end do
+      this%unified_workload = this%unified_workload + n_plus
+    end do
+  end subroutine task_manager_distribute_shared_tasks
+
+  subroutine task_manager_check_distribution(this)
+    type(task_manager_type), intent(in) :: this
+
+    integer :: t, w, n, tmp
+
+    if (.not. this%active) return
+    if (.not. this%distributed) return
+
+    tmp = task_manager_get_max_idata1_sum(this)
+    if (tmp /= this%unified_workload) call system_abort("task_manager_check_distribution: unified_workload is inconsistent: "//tmp//" "//this%unified_workload)
+    if (any(this%workers(:)%n_padding < 0)) call system_abort("task_manager_check_distribution: negative padding: "//this%workers(:)%n_padding)
+  end subroutine task_manager_check_distribution
+
   subroutine task_manager_show_distribution(this)
     type(task_manager_type), intent(in) :: this
 
-    integer :: i, t, w, load, n_tasks
+    integer :: i, t, w, idata1_sum, n_tasks
+
+    if (.not. this%active) return
 
     do w = 1, this%n_workers
         call print("Tasks of worker "//w)
-        load = 0
+        idata1_sum = 0
         n_tasks = 0
         do i = 1, this%workers(w)%n_tasks
-          t = this%workers(w)%task_ids(i)
-          call print(i // ": " // t // ": " // this%tasks(t)%memory)
-          load = load + this%tasks(t)%memory
+          t = this%workers(w)%tasks(i)%index
+          call print(i // ": " // t // ": " // this%workers(w)%tasks(i)%idata)
+          idata1_sum = idata1_sum + this%tasks(t)%idata(1)
         end do
-        call print("Summary: "//this%workers(w)%n_tasks//" tasks, "//load//" memory, "//this%workers(w)%n_padding//" padding")
+        call print("Summary: "//this%workers(w)%n_tasks//" tasks, "//idata1_sum//" idata1_sum, "//this%workers(w)%n_padding//" padding")
     end do
   end subroutine task_manager_show_distribution
 
@@ -166,15 +253,16 @@ subroutine task_manager_export_distribution(this, only_id)
     integer :: only_id_opt
     type(Inoutput) :: file
 
+    if (.not. this%active) return
     only_id_opt = optional_default(UNDEFINED, only_id)
     if (only_id_opt >= 0 .and. only_id_opt /= this%my_worker_id) return
 
     call initialise(file, "tm.dist", OUTPUT)
-    call print("worker_id task_id_local task_id_global memory", file=file)
+    call print("worker_id task_id_local task_id_global idata", file=file)
     do w = 1, this%n_workers
         do i = 1, this%workers(w)%n_tasks
-          t = this%workers(w)%task_ids(i)
-          call print(w // " " // i // " " // t // " " // this%tasks(t)%memory, file=file)
+          t = this%workers(w)%tasks(i)%index
+          call print(w // " " // i // " " // t // " " // this%workers(w)%tasks(i)%idata, file=file)
         end do
     end do
     call finalise(file)
@@ -184,7 +272,9 @@ subroutine task_manager_init_tasks(this, n_tasks)
     type(task_manager_type), intent(inout) :: this
     integer, intent(in) :: n_tasks !> number of tasks to allocate
 
-    call print("Allocate %tasks for "//n_tasks//" tasks.")
+    if (.not. this%active) return
+
+    call print("Allocate task_manager%tasks for "//n_tasks//" tasks.", PRINT_VERBOSE)
     allocate(this%tasks(n_tasks))
   end subroutine task_manager_init_tasks
 
@@ -192,16 +282,32 @@ subroutine task_manager_init_workers(this, n_workers)
     type(task_manager_type), intent(inout) :: this
     integer, intent(in) :: n_workers !> number of workers to allocate and set
 
-    this%n_workers = n_workers
-    allocate(this%workers(this%n_workers))
+    integer :: w
+
+    if (.not. this%active) return
+
+    allocate(this%workers(n_workers))
+    do w = 1, n_workers
+      call task_manager_add_worker(this)
+    end do
   end subroutine task_manager_init_workers
 
+  subroutine task_manager_init_idata(this, n_idata)
+    type(task_manager_type), intent(inout) :: this
+    integer, intent(in) :: n_idata
+
+    if (.not. this%active) return
+
+    allocate(this%idata(n_idata))
+  end subroutine task_manager_init_idata
+
   subroutine task_manager_deallocate(this)
     type(task_manager_type), intent(inout) :: this
 
     this%distributed = .false.
     this%n_tasks = 0
     this%n_workers = 0
+    if (allocated(this%idata)) deallocate(this%idata)
     if (allocated(this%tasks)) deallocate(this%tasks)
     if (allocated(this%workers)) deallocate(this%workers)
   end subroutine task_manager_deallocate
@@ -225,4 +331,28 @@ subroutine task_manager_print(this)
     call print_title('')
   end subroutine task_manager_print
 
+  function task_manager_get_max_idata1_sum(this) result(res)
+    type(task_manager_type), intent(in) :: this
+
+    integer :: res
+    integer :: n, t, w
+
+    if (.not. this%active) return
+
+    res = 0
+    do w = 1, this%n_workers
+      n = sum([(this%workers(w)%tasks(t)%idata(1), t = 1, this%workers(w)%n_tasks)])
+      res = max(n, res)
+    end do
+  end function task_manager_get_max_idata1_sum
+
+  subroutine task_deepcopy(this, from)
+    type(task_type), intent(inout) :: this
+    type(task_type), intent(in) :: from
+
+    this%index = from%index
+    allocate(this%idata, source=from%idata)
+    this%worker_id = from%worker_id
+  end subroutine task_deepcopy
+
 end module task_manager_module


Test patch:
diff --git a/src/Programs/test_task_manager.f95 b/src/Programs/test_task_manager.f95
new file mode 100644
index 0000000000..9e2fa9b207
--- /dev/null
+++ b/src/Programs/test_task_manager.f95
@@ -0,0 +1,56 @@
+program test_task_manager
+  use system_module
+  use task_manager_module
+
+  call system_initialise()
+  call test_1to9_for_2()
+  call print("test_task_manager OK")
+  call system_finalise()
+
+  contains
+
+  subroutine test_1to9_for_2()
+    integer, parameter :: n_workers = 2
+    integer, parameter :: idata1_normal(*) = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+    integer, parameter :: idata1_shared(*) = [10]
+
+    integer, parameter :: workers_n_tasks(*) = [6, 5]
+    integer, parameter :: worker1_task_ids(*) = [1, 2, 5, 6, 9, 10]
+    integer, parameter :: worker2_task_ids(*) = [3, 4, 7, 8, 10]
+    integer, parameter :: workers_paddings(*) = [0, 1]
+    integer, parameter :: unified_workload = 28
+
+    type(task_manager_type) :: task_manager
+
+    call test_core(task_manager, n_workers, idata1_normal, idata1_shared)
+    if (any(task_manager%workers(:)%n_tasks /= workers_n_tasks)) call system_abort("test_1to9_for_2:workers_n_tasks FAIL")
+    if (any(task_manager%workers(1)%tasks(:)%index /= worker1_task_ids)) call system_abort("test_1to9_for_2:worker1_task_ids FAIL")
+    if (any(task_manager%workers(2)%tasks(:)%index /= worker2_task_ids)) call system_abort("test_1to9_for_2:worker2_task_ids FAIL")
+    if (task_manager%unified_workload /= unified_workload) call system_abort("test_1to9_for_2:unified_workload FAIL")
+    if (any(task_manager%workers(:)%n_padding /= workers_paddings)) call system_abort("test_1to9_for_2:workers_paddings FAIL")
+  end subroutine test_1to9_for_2
+
+  subroutine test_core(task_manager, n_workers, idata1_normal, idata1_shared)
+    integer, intent(in) :: n_workers
+    integer, intent(in) :: idata1_normal(:), idata1_shared(:)
+    
+    type(task_manager_type), intent(inout) :: task_manager
+
+    integer :: t
+
+    task_manager%active = .true.
+    call task_manager_init_workers(task_manager, n_workers)
+    call task_manager_init_tasks(task_manager, size(idata1_normal) + size(idata1_shared))
+
+    do t = 1, size(idata1_normal)
+      call task_manager_add_task(task_manager, idata1_normal(t))
+    end do
+
+    do t = 1, size(idata1_shared)
+      call task_manager_add_task(task_manager, idata1_shared(t), n_idata=2, worker_id=SHARED)
+    end do
+
+    call task_manager_distribute_tasks(task_manager)
+  end subroutine test_core
+
+end program test_task_manager
diff --git a/tests/test_task_manager.sh b/tests/test_task_manager.sh
new file mode 100755
index 0000000000..078d1f8523
--- /dev/null
+++ b/tests/test_task_manager.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+set -e
+
+if [ -z $QUIP_ARCH ]; then
+   echo "$0: Need QUIP_ARCH defined"
+   exit 1
+fi
+
+mydir=`dirname $0`
+bindir=$mydir/../build/$QUIP_ARCH
+QUIP_ROOT=${mydir}/..
+
+if [ ! -x $bindir/test_task_manager ]; then
+    (cd $QUIP_ROOT && make Programs) || exit 2
+fi
+
+$bindir/test_task_manager &> test_task_manager.log
+grep "test_task_manager OK" test_task_manager.log &> /dev/null
+error=$?
+
+rm -f test_task_manager.log
+exit $error