start merging adapt

Author: ilhamv

mcdc/code_factory/adapt.py

@@ -1,11 +1,10 @@
 import importlib
-import numpy as np
-from numba import njit, jit, objmode, literal_unroll, types
-from numba.extending import intrinsic
+import inspect
 import numba
-import mcdc.object_.numba_types as type_
-import mcdc.transport.particle_bank as particle_bank_module
+import numpy as np
 
+from numba import njit, jit, types
+from numba.extending import intrinsic
 
 if importlib.util.find_spec("harmonize") is None:
     HAS_HARMONIZE = False
@@ -14,13 +13,12 @@
 
     HAS_HARMONIZE = True
 
+####
 
-import math
-import inspect
 
-from mcdc.print_ import print_error
+import mcdc.config as config
 
-import mcdc.code_factory.adapt as adapt
+from mcdc.print_ import print_error
 
 
 # =============================================================================
@@ -75,6 +73,17 @@ def codegen(context, builder, signature, args):
         return sig, codegen
 
 
+@njit()
+def uintp_to_voidptr(value):
+    val = numba.uintp(value)
+    return cast_uintp_to_voidptr(val)
+
+
+@njit()
+def voidptr_to_uintp(value):
+    return cast_voidptr_to_uintp(value)
+
+
 def leak(arg):
     pass
 
@@ -100,10 +109,10 @@ def impl(arg):
 # =============================================================================
 
 
-"""
 def local_array(shape, dtype):
     return np.zeros(shape, dtype=dtype)
 
+
 @numba.extending.type_callable(local_array)
 def type_local_array(context):
 
@@ -252,7 +261,7 @@ def builtin_local_array(context, builder, sig, args):
         raise numba.core.errors.UnsupportedError(
             f"Unsupported target context {context}."
         )
-"""
+
 
 # =============================================================================
 # Decorators
@@ -402,11 +411,17 @@ def nopython_mode(is_on):
 # =============================================================================
 
 
+@numba.njit()
+def alloc_bytes_placeholder(size):
+    return uintp_to_voidptr(0)
+
+
 SIMPLE_ASYNC = True
 
 none_type = None
 mcdc_global_type = None
 mcdc_data_type = None
+mcdc_shared_type = None
 state_spec = None
 mcdc_global_gpu = None
 mcdc_data_gpu = None
@@ -417,26 +432,41 @@ def nopython_mode(is_on):
 step_async = None
 halt_early = None
 find_cell_async = None
+tally_width = None
+tally_length = None
+tally_size = None
+alloc_managed_bytes = alloc_bytes_placeholder
+alloc_device_bytes = alloc_bytes_placeholder
+tally_shape_literal = None
 
 
-def gpu_forward_declare(args):
+def gpu_forward_declare(args, tally_shape):
 
     if args.gpu_rocm_path != None:
         harm.config.set_rocm_path(args.gpu_rocm_path)
 
     if args.gpu_cuda_path != None:
         harm.config.set_cuda_path(args.gpu_cuda_path)
 
-    global none_type, mcdc_global_type, mcdc_data_type
+    global none_type, mcdc_global_type, mcdc_data_type, mcdc_shared_type
     global state_spec
     global mcdc_global_gpu, mcdc_data_gpu
     global group_gpu, thread_gpu
-    global particle_gpu, particle_data_gpu
+    global particle_gpu, particle_record_gpu
     global step_async, find_cell_async, halt_early
+    global tally_width, tally_length, tally_size
+
+    tally_size = tally_shape[0] * tally_shape[1] * 8
+
+    global tally_shape_literal
+    tally_shape_literal = tally_shape
 
     none_type = numba.from_dtype(np.dtype([]))
-    mcdc_global_type = numba.from_dtype(type_.simulation)
-    mcdc_data_type = numba.from_dtype(type_.tally)
+    mcdc_global_type = numba.types.Array(numba.from_dtype(type_.global_), (1,), "C")
+    # mcdc_global_type = numba.from_dtype(type_.global_)
+
+    tally_dims = len(tally_shape)
+    mcdc_data_type = numba.types.Array(numba.float64, tally_dims, "C")
     state_spec = (
         {
             "global": mcdc_global_type,
@@ -446,12 +476,12 @@ def gpu_forward_declare(args):
         none_type,
     )
     access_fns = harm.RuntimeSpec.access_fns(state_spec)
-    mcdc_global_gpu = access_fns["device"]["global"]
-    mcdc_data_gpu = access_fns["device"]["data"]
+    mcdc_global_gpu = access_fns["device"]["global"]["indirect"]
+    mcdc_data_gpu = access_fns["device"]["data"]["direct"]
     group_gpu = access_fns["group"]
     thread_gpu = access_fns["thread"]
     particle_gpu = numba.from_dtype(type_.particle)
-    particle_data_gpu = numba.from_dtype(type_.particle_data)
+    particle_record_gpu = numba.from_dtype(type_.particle_record)
 
     def step(prog: numba.uintp, P: particle_gpu):
         pass
@@ -463,6 +493,44 @@ def find_cell(prog: numba.uintp, P: particle_gpu):
     interface = adapt.harm.RuntimeSpec.program_interface()
     halt_early = interface["halt_early"]
 
+    global alloc_managed_bytes
+    global alloc_device_bytes
+    alloc_managed_bytes = harm.alloc_managed_bytes
+    alloc_device_bytes = harm.alloc_device_bytes
+
+
+# =============================================================================
+# Global GPU/CPU Arry Variable Constructors
+# =============================================================================
+
+
+def create_data_array(size, dtype):
+    if config.target == "gpu":
+        if config.gpu_state_storage == "managed":
+            data_tally_ptr = alloc_managed_bytes(tally_size)
+        else:
+            data_tally_ptr = alloc_device_bytes(tally_size)
+        data_tally_uint = voidptr_to_uintp(data_tally_ptr)
+        data_tally = numba.carray(data_tally_ptr, (width, length), type_.float64)
+        return data_tally, data_tally_uint
+    else:
+        data_tally = np.zeros(size, dtype=dtype)
+        return data_tally, 0
+
+
+def create_mcdc_array(dtype):
+    if config.target == "gpu":
+        if config.gpu_state_storage == "managed":
+            mcdc_ptr = alloc_managed_bytes(type_.global_size)
+        else:
+            mcdc_ptr = alloc_device_bytes(type_.global_size)
+        mcdc_uint = voidptr_to_uintp(mcdc_ptr)
+        mcdc_array = numba.carray(mcdc_ptr, (1,), type_.global_)
+        return mcdc_array, mcdc_uint
+    else:
+        mcdc_array = np.zeros((1,), dtype=dtype)
+        return mcdc_array, 0
+
 
 # =============================================================================
 # Seperate GPU/CPU Functions to Target Different Platforms
@@ -514,51 +582,14 @@ def thread(prog):
 
 
 @for_cpu()
-def add_active(particle, prog):
-    particle_bank_module.add_particle(particle, prog["bank_active"])
-
-
-@for_gpu()
-def add_active(P_reclike, prog):
-    P = local_array(1, type_.particle)
-    particle_bank_module.recordlike_to_particle(P, P_reclike)
-    if SIMPLE_ASYNC:
-        step_async(prog, P[0])
-    else:
-        find_cell_async(prog, P[0])
-
-
-@for_cpu()
-def add_source(particle, prog):
-    particle_bank_module.add_particle(particle, prog["bank_source"])
-
-
-@for_gpu()
-def add_source(particle, prog):
-    mcdc = mcdc_global(prog)
-    particle_bank_module.add_particle(particle, mcdc["bank_source"])
-
-
-@for_cpu()
-def add_census(particle, prog):
-    particle_bank_module.add_particle(particle, prog["bank_census"])
-
-
-@for_gpu()
-def add_census(particle, prog):
-    mcdc = mcdc_global(prog)
-    particle_bank_module.add_particle(particle, mcdc["bank_census"])
-
-
-@for_cpu()
-def add_future(particle, prog):
-    particle_bank_module.add_particle(particle, prog["bank_future"])
+def add_IC(P_arr, prog):
+    particle_bank.add_particle(P_arr, prog["technique"]["IC_bank_neutron_local"])
 
 
 @for_gpu()
-def add_future(particle, prog):
+def add_IC(P_arr, prog):
     mcdc = mcdc_global(prog)
-    particle_bank_module.add_particle(particle, mcdc["bank_future"])
+    particle_bank.add_particle(P_arr, mcdc["technique"]["IC_bank_neutron_local"])
 
 
 @for_cpu()

mcdc/code_factory/code_factory.py

@@ -10,6 +10,7 @@
 ####
 
 import mcdc
+import mcdc.code_factory.adapt as adapt
 import mcdc.object_ as object_module
 import mcdc.object_.base as base
 
@@ -253,7 +254,10 @@ def generate_numba_objects(simulation):
     set_object(simulation, annotations, structures, records, data)
 
     # Allocate the flattened data and re-set the objects
-    data["array"] = np.zeros(data["size"], dtype=type_map[float])
+    #data["array"] = np.zeros(data["size"], dtype=type_map[float])
+    data["array"], data["pointer"] = adapt.create_data_array(
+        data['size'], type_map[float]
+    )
 
     data["size"] = 0
     records = {}
@@ -341,7 +345,8 @@ def generate_numba_objects(simulation):
     # ==================================================================================
 
     # The global structure/variable container
-    mcdc_simulation_arr = np.zeros(1, dtype=into_dtype(structures["simulation"]))
+    #mcdc_simulation_arr = np.zeros(1, dtype=into_dtype(structures["simulation"]))
+    mcdc_simulation_arr, mcdc_simulation_pointer = adapt.create_mcdc_array(into_dtype(structures["simulation"]))
     mcdc_simulation = mcdc_simulation_arr[0]
 
     record = records["simulation"]

mcdc/main.py

@@ -221,6 +221,7 @@ def preparation():
 
     # Adapt kernels
     import numba as nb
+    import mcdc.code_factory.adapt as adapt
     import mcdc.config as config
     import mcdc.transport.mpi as mpi
 
@@ -238,15 +239,14 @@ def preparation():
             print_error(
                 "No module named 'harmonize' - GPU functionality not available. "
             )
-        adapt.gpu_forward_declare(config.args)
-
-    from mcdc.code_factory.adapt import eval_toggle, target_for, nopython_mode
+        adapt.gpu_forward_declare(config.args, tally_shape)
 
-    eval_toggle()
-    target_for(config.target)
+    adapt.eval_toggle()
+    adapt.target_for(config.target)
     if config.target == "gpu":
         build_gpu_progs(input_deck, config.args)
-    nopython_mode((config.mode == "numba") or (config.mode == "numba_debug"))
+    adapt.nopython_mode((config.mode == "numba") or (config.mode == "numba_debug"))
+
 
     # ==================================================================================
     # Source file