distribute particles based off of how many procs are in the domain

spasmann · Jan 25, 2024 · 0547954 · 0547954
1 parent 47dbe86
commit 0547954
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 14 deletions.
diff --git a/mcdc/kernel.py b/mcdc/kernel.py
@@ -570,11 +570,20 @@ def particle_in_domain(P, mcdc):
 
 @njit
 def distribute_work_dd(N, mcdc, precursor=False):
+    domain_size = len(mcdc["technique"]["work_ratio"])    
+
     # Total # of work
     work_size_total = N
+
+    # Evenly distribute work per domain
+    work_size = math.floor(N / domain_size)
+    d_idx = mcdc["d_idx"]
+    # Evenly distribute work per processor in domain
+    work_size = math.floor(work_size / mcdc["technique"]["work_ratio"][d_idx])
+
     work_start = 0
-    work_size = work_size_total
 
+    # print("\n work_size = ", work_size)
     if not precursor:
         mcdc["mpi_work_start"] = work_start
         mcdc["mpi_work_size"] = work_size
@@ -3185,8 +3194,8 @@ def iqmc_prepare_particles(mcdc):
     Q = iqmc["source"]
     mesh = iqmc["mesh"]
 
-    if mcdc["technique"]["domain_decomp"]:
-        distribute_work(N_particle, mcdc)
+    # if mcdc["technique"]["domain_decomp"]:
+        # distribute_work_dd(N_particle, mcdc)
 
     # total number of spatial cells (of all domains)
     N_total = iqmc["Nc_total"]

diff --git a/mcdc/main.py b/mcdc/main.py
@@ -670,7 +670,10 @@ def prepare():
         N_dim = mcdc["technique"]["iqmc"]["N_dim"]
         mcdc["mpi_size"] = MPI.COMM_WORLD.Get_size()
         mcdc["mpi_rank"] = MPI.COMM_WORLD.Get_rank()
-        kernel.distribute_work(N_particle, mcdc)
+        if input_deck.technique["domain_decomp"]:
+            kernel.distribute_work_dd(N_particle, mcdc)
+        else:
+            kernel.distribute_work(N_particle, mcdc)
         N_work = int(mcdc["mpi_work_size"])
         N_start = int(mcdc["mpi_work_start"])
 

diff --git a/mcdc/type_.py b/mcdc/type_.py
@@ -531,16 +531,27 @@ def make_type_technique(N_particle, G, card):
         Nx = Ny = Nz = Nt = N_particle = Ng = N_dim = 0
 
     iqmc_list += [("mesh", mesh)]
-    # Low-discprenecy sequence
-    size = MPI.COMM_WORLD.Get_size()
-    rank = MPI.COMM_WORLD.Get_rank()
-    # Evenly distribute work
-    work_size = math.floor(N_particle / size)
-    # Count reminder
-    rem = N_particle % size
-    # Assign reminder and update starting index
-    if rank < rem:
-        work_size += 1
+
+    # size of LDS
+    if card['domain_decomp']:
+        domain_size = len(card["work_ratio"])   
+        print("\n work_ratio ", card["work_ratio"])
+        print("\n d_idx ", card["d_idx"])
+        # Evenly distribute work per domain
+        work_size = math.floor(N_particle / domain_size)
+        d_idx = card["d_idx"]
+        # Evenly distribute work per processor in domain
+        work_size = math.floor(work_size / card["work_ratio"][d_idx])
+    else:
+        size = MPI.COMM_WORLD.Get_size()
+        rank = MPI.COMM_WORLD.Get_rank()
+        # Evenly distribute work by # of processors
+        work_size = math.floor(N_particle / size)
+        # Count reminder
+        rem = N_particle % size
+        # Assign reminder and update starting index
+        if rank < rem:
+            work_size += 1
 
     iqmc_list += [("lds", float64, (work_size, N_dim))]
     iqmc_list += [("fixed_source", float64, (Ng, Nt, Nx, Ny, Nz))]