Merge pull request #212 from mit-crpg/work-distribution

Improve how particles are assigned to processes

src/cmfd_execute.F90

@@ -281,8 +281,7 @@ contains
     use constants,   only: ZERO, ONE
     use error,       only: warning, fatal_error
     use global,      only: n_particles, meshes, source_bank, work,             &
-                           n_user_meshes, message, cmfd, master, mpi_err,      &
-                           bank_first, bank_last
+                           n_user_meshes, message, cmfd, master, mpi_err
     use mesh_header, only: StructuredMesh
     use mesh,        only: count_bank_sites, get_mesh_indices
     use search,      only: binary_search
@@ -296,7 +295,6 @@ contains
     integer :: ijk(3) ! spatial bin location
     integer :: e_bin ! energy bin of source particle
     integer :: n_groups ! number of energy groups
-    integer(8) :: size_bank ! size of source bank
     logical :: outside ! any source sites outside mesh
     logical :: in_mesh ! source site is inside mesh
     logical :: new_weights ! calcualte new weights
@@ -312,9 +310,6 @@ contains
     nz = cmfd%indices(3)
     ng = cmfd%indices(4)
 
-    ! compute size of source bank
-    size_bank = bank_last - bank_first + 1_8 
-
     ! allocate arrays in cmfd object (can take out later extend to multigroup)
     if (.not.allocated(cmfd%sourcecounts)) then 
       allocate(cmfd%sourcecounts(ng,nx,ny,nz))
@@ -337,7 +332,7 @@ contains
 
       ! count bank sites in mesh
       call count_bank_sites(m, source_bank, cmfd%sourcecounts, egrid, &
-           sites_outside=outside, size_bank = size_bank)
+           sites_outside=outside, size_bank=work)
 
       ! check for sites outside of the mesh
       if (master .and. outside) then
@@ -360,7 +355,7 @@ contains
    end if
 
     ! begin loop over source bank
-    do i = 1, int(size_bank, 4) 
+    do i = 1, int(work,4)
 
       ! determine spatial bin
       call get_mesh_indices(m, source_bank(i)%xyz, ijk, in_mesh)

src/eigenvalue.F90

@@ -378,7 +378,7 @@ contains
       end if
 
       ! the last processor should not be sending sites to right
-      finish = bank_last
+      finish = work_index(rank + 1)
     end if
 
     call time_bank_sample % stop()
@@ -394,11 +394,11 @@ contains
     if (start < n_particles) then
       ! Determine the index of the processor which has the first part of the
       ! source_bank for the local processor
-      neighbor = start / maxwork
+      neighbor = binary_search(work_index, n_procs + 1, start) - 1
 
       SEND_SITES: do while (start < finish)
         ! Determine the number of sites to send
-        n = min((neighbor + 1)*maxwork, finish) - start
+        n = min(work_index(neighbor + 1), finish) - start
 
         ! Initiate an asynchronous send of source sites to the neighboring
         ! process
@@ -423,7 +423,7 @@ contains
     ! ==========================================================================
     ! RECEIVE BANK SITES FROM NEIGHBORS OR TEMPORARY BANK
 
-    start = bank_first - 1
+    start = work_index(rank)
     index_local = 1
 
     ! Determine what process has the source sites that will need to be stored at
@@ -435,13 +435,12 @@ contains
       neighbor = binary_search(bank_position, n_procs, start) - 1
     end if
 
-    RECV_SITES: do while (start < bank_last)
+    RECV_SITES: do while (start < work_index(rank + 1))
       ! Determine how many sites need to be received
       if (neighbor == n_procs - 1) then
-        n = min(n_particles, (rank+1)*maxwork) - start
+        n = work_index(rank + 1) - start
       else
-        n = min(bank_position(neighbor+2), min(n_particles, &
-             (rank+1)*maxwork)) - start
+        n = min(bank_position(neighbor + 2), work_index(rank + 1)) - start
       end if
 
       if (neighbor /= rank) then

src/fixed_source.F90

@@ -50,12 +50,12 @@ contains
       PARTICLE_LOOP: do i = 1, work
 
         ! Set unique particle ID
-        p % id = (current_batch - 1)*n_particles + bank_first + i - 1
+        p % id = (current_batch - 1)*n_particles + work_index(rank) + i
 
         ! set particle trace
         trace = .false.
         if (current_batch == trace_batch .and. current_gen == trace_gen .and. &
-             bank_first + i - 1 == trace_particle) trace = .true.
+             work_index(rank) + i == trace_particle) trace = .true.
 
         ! set random number seed
         call set_particle_seed(p % id)

src/global.F90

@@ -160,10 +160,8 @@ module global
   type(Bank), allocatable, target :: source_bank(:)
   type(Bank), allocatable, target :: fission_bank(:)
   integer(8) :: n_bank       ! # of sites in fission bank
-  integer(8) :: bank_first   ! index of first particle in bank
-  integer(8) :: bank_last    ! index of last particle in bank
   integer(8) :: work         ! number of particles per processor
-  integer(8) :: maxwork      ! maximum number of particles per processor
+  integer(8), allocatable :: work_index(:) ! starting index in source bank for each process
   integer(8) :: current_work ! index in source bank of current history simulated
 
   ! Temporary k-effective values
@@ -442,6 +440,9 @@ contains
     if (allocated(source_bank)) deallocate(source_bank)
     if (allocated(entropy_p)) deallocate(entropy_p)
 
+    ! Deallocate array of work indices
+    if (allocated(work_index)) deallocate(work_index)
+
     ! Deallocate cmfd
     call deallocate_cmfd(cmfd)
 

src/initialize.F90

@@ -754,15 +754,37 @@ contains
 
   subroutine calculate_work()
 
-    ! Determine maximum amount of particles to simulate on each processor
-    maxwork = ceiling(real(n_particles)/n_procs,8)
+    integer    :: i         ! loop index
+    integer    :: remainder ! Number of processors with one extra particle
+    integer(8) :: i_bank    ! Running count of number of particles
+    integer(8) :: min_work  ! Minimum number of particles on each proc
+    integer(8) :: work_i    ! Number of particles on rank i
 
-    ! ID's of first and last source particles
-    bank_first = rank*maxwork + 1
-    bank_last  = min((rank+1)*maxwork, n_particles)
+    allocate(work_index(0:n_procs))
 
-    ! number of particles for this processor
-    work = bank_last - bank_first + 1
+    ! Determine minimum amount of particles to simulate on each processor
+    min_work = n_particles/n_procs
+
+    ! Determine number of processors that have one extra particle
+    remainder = int(mod(n_particles, int(n_procs,8)), 4)
+
+    i_bank = 0
+    work_index(0) = 0
+    do i = 0, n_procs - 1
+      ! Number of particles for rank i
+      if (i < remainder) then
+        work_i = min_work + 1
+      else
+        work_i = min_work
+      end if
+
+      ! Set number of particles
+      if (rank == i) work = work_i
+
+      ! Set index into source bank for rank i
+      i_bank = i_bank + work_i
+      work_index(i+1) = i_bank
+    end do
 
   end subroutine calculate_work
 
@@ -775,7 +797,7 @@ contains
     integer    :: alloc_err  ! allocation error code
 
     ! Allocate source bank
-    allocate(source_bank(maxwork), STAT=alloc_err)
+    allocate(source_bank(work), STAT=alloc_err)
 
     ! Check for allocation errors 
     if (alloc_err /= 0) then
@@ -784,7 +806,7 @@ contains
     end if
 
     ! Allocate fission bank
-    allocate(fission_bank(3*maxwork), STAT=alloc_err)
+    allocate(fission_bank(3*work), STAT=alloc_err)
 
     ! Check for allocation errors 
     if (alloc_err /= 0) then

src/output_interface.F90

@@ -947,7 +947,7 @@ contains
     call h5dget_space_f(dset, dspace, hdf5_err)
 
     ! Select hyperslab for this dataspace
-    offset(1) = bank_first - 1_8
+    offset(1) = work_index(rank)
     call h5sselect_hyperslab_f(dspace, H5S_SELECT_SET_F, offset, dims1, hdf5_err)
 
     ! Set up the property list for parallel writing
@@ -1009,7 +1009,7 @@ contains
 
     ! Set the proper offset for source data on this processor
     call MPI_TYPE_SIZE(MPI_BANK, size_bank, mpi_err)
-    offset = offset + size_bank*maxwork*rank
+    offset = offset + size_bank*work_index(rank)
 
     ! Write all source sites
     call MPI_FILE_WRITE_AT(mpi_fh, offset, source_bank(1), work, MPI_BANK, &
@@ -1056,7 +1056,7 @@ contains
     call h5dget_space_f(dset, dspace, hdf5_err)
 
     ! Select hyperslab for this dataspace
-    offset(1) = bank_first - 1_8
+    offset(1) = work_index(rank)
     call h5sselect_hyperslab_f(dspace, H5S_SELECT_SET_F, offset, dims1, hdf5_err)
 
     ! Set up the property list for parallel writing
@@ -1109,7 +1109,7 @@ contains
 
     ! Set the proper offset for source data on this processor
     call MPI_TYPE_SIZE(MPI_BANK, size_bank, mpi_err)
-    offset = offset + size_bank*maxwork*rank
+    offset = offset + size_bank*work_index(rank)
 
     ! Write all source sites
     call MPI_FILE_READ_AT(mpi_fh, offset, source_bank(1), work, MPI_BANK, &

src/source.F90

@@ -47,7 +47,7 @@ contains
         src => source_bank(i)
 
         ! initialize random number seed
-        id = bank_first + i - 1
+        id = work_index(rank) + i
         call set_particle_seed(id)
 
         ! sample external source distribution
@@ -165,7 +165,7 @@ contains
     call copy_source_attributes(p, src)
 
     ! set identifier for particle
-    p % id = bank_first + index_source - 1
+    p % id = work_index(rank) + index_source
 
     ! set random number seed
     particle_seed = (overall_gen - 1)*n_particles + p % id