Refactor local arrays and objects

mcdc/adapt.py

@@ -383,69 +383,6 @@ def add_IC(particle, prog):
     kernel.add_particle(particle, mcdc["technique"]["IC_bank_neutron_local"])
 
 
-@for_cpu()
-def local_translate():
-    return np.zeros(1, dtype=type_.translate)[0]
-
-
-@for_gpu()
-def local_translate():
-    trans = cuda.local.array(1, type_.translate)[0]
-    for i in range(3):
-        trans["values"][i] = 0
-    return trans
-
-
-@for_cpu()
-def local_group_array():
-    return np.zeros(1, dtype=type_.group_array)[0]
-
-
-@for_gpu()
-def local_group_array():
-    return cuda.local.array(1, type_.group_array)[0]
-
-
-@for_cpu()
-def local_j_array():
-    return np.zeros(1, dtype=type_.j_array)[0]
-
-
-@for_gpu()
-def local_j_array():
-    return cuda.local.array(1, type_.j_array)[0]
-
-
-@for_cpu()
-def local_RPN_array():
-    return np.zeros(1, dtype=type_.RPN_array)[0]
-
-
-@for_gpu()
-def local_RPN_array():
-    return cuda.local.array(1, type_.RPN_array)[0]
-
-
-@for_cpu()
-def local_particle():
-    return np.zeros(1, dtype=type_.particle)[0]
-
-
-@for_gpu()
-def local_particle():
-    return cuda.local.array(1, dtype=type_.particle)[0]
-
-
-@for_cpu()
-def local_particle_record():
-    return np.zeros(1, dtype=type_.particle_record)[0]
-
-
-@for_gpu()
-def local_particle_record():
-    return cuda.local.array(1, dtype=type_.particle_record)[0]
-
-
 @for_cpu()
 def global_add(ary, idx, val):
     result = ary[idx]

mcdc/code_factory.py

@@ -0,0 +1,56 @@
+import numpy as np
+
+from numba import cuda, njit
+
+import mcdc.local as local
+import mcdc.type_ as type_
+
+
+# ==============================================================================
+# Local array and object (see local.py)
+# ==============================================================================
+
+
+def make_locals(input_deck):
+    # Hardware target
+    target = input_deck.setting["target"]
+
+    # Problem-dependent sizes
+    G = input_deck.materials[0].G
+    J = input_deck.materials[0].J
+    N_RPN = max([np.sum(np.array(x._region_RPN) >= 0.0) for x in input_deck.cells])
+
+    # Make the locals
+    local.translation = local_array(type_.float64, 3, target)
+    local.energy_group_array = local_array(type_.float64, G, target)
+    local.precursor_group_array = local_array(type_.float64, J, target)
+    local.RPN_array = local_array(type_.bool_, N_RPN, target)
+    local.particle = local_object(type_.particle, target)
+    local.particle_record = local_object(type_.particle_record, target)
+
+
+def local_array(dtype, size, target):
+    struct = type_.into_dtype([("values", dtype, (size,))])
+
+    @njit
+    def cpu():
+        return np.zeros(1, dtype=struct)[0]
+
+    @cuda.jit(device=True)
+    def gpu():
+        return cuda.local.array(1, dtype=struct)[0]
+
+    return cpu if target == "cpu" else gpu
+
+
+def local_object(dtype, target):
+
+    @njit
+    def cpu():
+        return np.zeros(1, dtype=dtype)[0]
+
+    @cuda.jit(device=True)
+    def gpu():
+        return cuda.local.array(1, dtype=dtype)[0]
+
+    return cpu if target == "cpu" else gpu

mcdc/global_.py

@@ -78,6 +78,8 @@ def reset(self):
             # Below are parameters not copied to mcdc.setting
             "bank_active_buff": 100,
             "bank_census_buff": 1.0,
+            # Portability
+            "target": "cpu",
         }
 
         self.technique = {

mcdc/iqmc/iqmc_kernel.py

@@ -1,23 +1,24 @@
-from mpi4py import MPI
-
-import mcdc.adapt as adapt
 import numpy as np
 
+from mpi4py import MPI
 from numpy import ascontiguousarray as cga
 from numba import njit, objmode, literal_unroll
-from mcdc.loop import caching
-from mcdc.type_ import iqmc_score_list
-from mcdc.kernel import distance_to_boundary, distance_to_mesh
 
-# from mcdc.iqmc.iqmc_loop import iqmc_loop_source
+import mcdc.adapt as adapt
+import mcdc.local as local
+
 from mcdc.adapt import toggle, for_cpu, for_gpu
 from mcdc.constant import *
+from mcdc.loop import caching
+from mcdc.type_ import iqmc_score_list
 from mcdc.kernel import (
-    move_particle,
-    set_bank_size,
+    distance_to_boundary,
+    distance_to_mesh,
     get_particle_cell,
     get_particle_material,
     mesh_get_index,
+    move_particle,
+    set_bank_size,
 )
 
 # =========================================================================
@@ -285,7 +286,7 @@ def iqmc_prepare_particles(mcdc):
 
     for n in range(N_work):
         # Create new particle
-        P_new = adapt.local_particle_record()
+        P_new = local.particle_record()
         # assign initial group, time, and rng_seed (not used)
         P_new["g"] = 0
         P_new["t"] = 0
@@ -456,10 +457,10 @@ def iqmc_generate_material_idx(mcdc):
     Nz = len(mesh["z"]) - 1
     dx = dy = dz = 1
     # variables for cell finding functions
-    trans_struct = adapt.local_translate()
+    trans_struct = local.translation()
     trans = trans_struct["values"]
     # create particle to utilize cell finding functions
-    P_temp = adapt.local_particle()
+    P_temp = local.particle()
     # set default attributes
     P_temp["alive"] = True
     P_temp["material_ID"] = -1

mcdc/iqmc/iqmc_loop.py

@@ -1,19 +1,21 @@
-import mcdc.kernel as kernel
-import mcdc.iqmc.iqmc_kernel as iqmc_kernel
-import mcdc.adapt as adapt
 import numpy as np
 
 from numpy import ascontiguousarray as cga
 from numba import njit, objmode
-from mcdc.loop import caching
-from mcdc.constant import *
 
+import mcdc.adapt as adapt
+import mcdc.iqmc.iqmc_kernel as iqmc_kernel
+import mcdc.kernel as kernel
+import mcdc.local as local
+
+from mcdc.constant import *
+from mcdc.loop import caching
 from mcdc.print_ import (
-    print_progress,
-    print_progress_iqmc,
-    print_iqmc_eigenvalue_progress,
     print_iqmc_eigenvalue_exit_code,
+    print_iqmc_eigenvalue_progress,
     print_msg,
+    print_progress,
+    print_progress_iqmc,
 )
 
 
@@ -48,7 +50,7 @@ def iqmc_step_particle(P, prog):
 
     # Find cell from root universe if unknown
     if P["cell_ID"] == -1:
-        trans_struct = adapt.local_translate()
+        trans_struct = local.translation()
         trans = trans_struct["values"]
         P["cell_ID"] = kernel.get_particle_cell(P, 0, trans, mcdc)
 
@@ -100,7 +102,7 @@ def iqmc_loop_source(mcdc):
 
         # Loop until active bank is exhausted
         while mcdc["bank_active"]["size"] > 0:
-            P = adapt.local_particle()
+            P = local.particle()
             # Get particle from active bank
             kernel.get_particle(P, mcdc["bank_active"], mcdc)
             # Particle loop

mcdc/kernel.py

@@ -4,6 +4,7 @@
 from numba import njit, objmode, literal_unroll
 import numba
 
+import mcdc.local as local
 import mcdc.type_ as type_
 
 from mcdc.constant import *
@@ -852,7 +853,7 @@ def rng_array(seed, shape, size):
 
 @njit
 def source_particle(seed, mcdc):
-    P: type_.particle_record = adapt.local_particle_record()
+    P = local.particle_record()
     P["rng_seed"] = seed
 
     # Sample source
@@ -1581,7 +1582,7 @@ def get_particle_cell(P, universe_ID, trans, mcdc):
 @njit
 def get_particle_material(P, mcdc):
     # Translation accumulator
-    trans_struct = adapt.local_translate()
+    trans_struct = local.translation()
     trans = trans_struct["values"]
 
     # Top level cell
@@ -1645,14 +1646,14 @@ def copy_recordlike(P_new, P):
 
 @njit
 def copy_record(P):
-    P_new = adapt.local_particle_record()
+    P_new = local.particle_record()
     copy_recordlike(P_new, P)
     return P_new
 
 
 @njit
 def recordlike_to_particle(P_rec):
-    P_new = adapt.local_particle()
+    P_new = local.particle()
     copy_recordlike(P_new, P_rec)
     P_new["fresh"] = True
     P_new["alive"] = True
@@ -1705,7 +1706,7 @@ def cell_check(P, cell, trans, mcdc):
     N_token = mcdc["cell_region_data"][idx]
 
     # Create local value array
-    value_struct = adapt.local_RPN_array()
+    value_struct = local.RPN_array()
     value = value_struct["values"]
     N_value = 0
 
@@ -2645,7 +2646,7 @@ def distance_to_boundary(P, mcdc):
     event = 0
 
     # Translation accumulator
-    trans_struct = adapt.local_translate()
+    trans_struct = local.translation()
     trans = trans_struct["values"]
 
     # Top level cell
@@ -2779,7 +2780,7 @@ def surface_crossing(P, prog):
 
     mcdc = adapt.device(prog)
 
-    trans_struct = adapt.local_translate()
+    trans_struct = local.translation()
     trans = trans_struct["values"]
     trans = P["translation"]
 
@@ -2801,7 +2802,7 @@ def surface_crossing(P, prog):
     if P["alive"] and not surface["BC"] == BC_REFLECTIVE:
         cell = mcdc["cells"][P["cell_ID"]]
         if not cell_check(P, cell, trans, mcdc):
-            trans_struct = adapt.local_translate()
+            trans_struct = local.translation()
             trans = trans_struct["values"]
             P["cell_ID"] = get_particle_cell(P, UNIVERSE_ROOT, trans, mcdc)
 
@@ -3324,7 +3325,7 @@ def fission_CE(P, nuclide, P_new):
     J = 6
     nu = get_nu(NU_FISSION, nuclide, E)
     nu_p = get_nu(NU_FISSION_PROMPT, nuclide, E)
-    nu_d_struct = adapt.local_j_array()
+    nu_d_struct = local.precursor_group_array()
     nu_d = nu_d_struct["values"]
     for j in range(J):
         nu_d[j] = get_nu_group(NU_FISSION_DELAYED, nuclide, E, j)

mcdc/local.py

@@ -0,0 +1,13 @@
+"""
+Collection of functions to make local array and object
+
+All of the functions are target-dependent (CPU or GPU mode)
+and defined in code_factory.py
+"""
+
+translation = None
+energy_group_array = None
+precursor_group_array = None
+RPN_array = None
+particle = None
+particle_record = None