add gpu_meta

Author: ilhamv

mcdc/code_factory/code_factory.py

@@ -315,18 +315,6 @@ def generate_numba_objects(simulation):
 
     structures["simulation"] = new_structure + structures["simulation"]
 
-    # GPU interop.
-    structures["simulation"] += [
-        ("gpu_state_pointer", "u8"),
-        ("source_program_pointer", "u8"),
-        ("precursor_program_pointer", "u8"),
-        ("source_seed", "u8"),
-    ]
-    records["simulation"]["gpu_state_pointer"] = 0
-    records["simulation"]["source_program_pointer"] = 0
-    records["simulation"]["precursor_program_pointer"] = 0
-    records["simulation"]["source_seed"] = 0
-
     # Print the fields
     if MPI.COMM_WORLD.Get_rank() == 0:
         with open(f"{Path(mcdc.__file__).parent}/object_/numba_types.py", "w") as f:
@@ -365,7 +353,7 @@ def generate_numba_objects(simulation):
         if len(item) == 3:
             size = item[2][0]
 
-        # Skip particle banks
+        # Skip particular attributes
         if field in bank_names:
             continue
 
@@ -388,6 +376,11 @@ def generate_numba_objects(simulation):
                             singular_field
                         ][i][sub_item[0]]
 
+
+    # Manually set particle bank attributes
+    for name in bank_names:
+        mcdc_simulation[name]['tag'] = getattr(simulation, name).tag
+
     return mcdc_simulation_arr, data["array"]
 
 

mcdc/main.py

@@ -215,16 +215,6 @@ def preparation():
     importlib.reload(mcdc_get)
     importlib.reload(mcdc_set)
 
-    # ==================================================================================
-    # Set some Numba object fields
-    # ==================================================================================
-
-    # Particle bank tags
-    mcdc["bank_active"]["tag"] = "active"
-    mcdc["bank_census"]["tag"] = "census"
-    mcdc["bank_source"]["tag"] = "source"
-    mcdc["bank_future"]["tag"] = "future"
-
     # ==================================================================================
     # Platform Targeting, Adapters, Toggles, etc
     # ==================================================================================
@@ -263,6 +253,8 @@ def preparation():
     #   TODO: Use parallel h5py
     # ==================================================================================
 
+    import h5py
+
     # All ranks, take turn
     for i in range(mcdc["mpi_size"]):
         if mcdc["mpi_rank"] == i:
@@ -285,13 +277,9 @@ def preparation():
         MPI.COMM_WORLD.Barrier()
 
     # ==================================================================================
-    # Finalize data: wrapping into a tuple
+    # Adapt functions
     # ==================================================================================
 
-    from mcdc.transport.simulation import setup_gpu
-
-    setup_gpu(mcdc)
-
     # Pick physics model
     import mcdc.transport.physics as physics
 
@@ -310,7 +298,13 @@ def preparation():
         rng.wrapping_add = rng.wrapping_add_python
         rng.wrapping_mul = rng.wrapping_mul_python
 
-    # TODO: Delete Python objects if running in Numba mode
+    # ==================================================================================
+    # Finalize data: wrapping into a tuple
+    # ==================================================================================
+
+    from mcdc.transport.simulation import setup_gpu
+
+    setup_gpu(mcdc)
 
     return mcdc_arr, data
 

mcdc/mcdc_get/__init__.py

@@ -44,6 +44,8 @@
 
 import mcdc.mcdc_get.tabulated_energy_angle_distribution as tabulated_energy_angle_distribution
 
+import mcdc.mcdc_get.gpu_meta as gpu_meta
+
 import mcdc.mcdc_get.native_material as native_material
 
 import mcdc.mcdc_get.multigroup_material as multigroup_material

mcdc/mcdc_get/gpu_meta.py

@@ -0,0 +1,3 @@
+# The following is automatically generated by code_factory.py
+
+from numba import njit

mcdc/mcdc_set/__init__.py

@@ -44,6 +44,8 @@
 
 import mcdc.mcdc_set.tabulated_energy_angle_distribution as tabulated_energy_angle_distribution
 
+import mcdc.mcdc_set.gpu_meta as gpu_meta
+
 import mcdc.mcdc_set.native_material as native_material
 
 import mcdc.mcdc_set.multigroup_material as multigroup_material

mcdc/mcdc_set/gpu_meta.py

@@ -0,0 +1,3 @@
+# The following is automatically generated by code_factory.py
+
+from numba import njit

mcdc/object_/gpu_tools.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from numpy import uint64
+
+####
+
+from mcdc.object_.base import ObjectSingleton
+
+@dataclass
+class GPUMeta(ObjectSingleton):
+    # Annotations for Numba mode
+    label: str = "gpu_meta"
+    #
+    state_pointer: uint64 = uint64(0)
+    source_program_pointer: uint64 = uint64(0)
+    precursor_program_pointer: uint64 = uint64(0)
+    structure_pointer: uint64 = uint64(0)
+    data_pointer: uint64 = uint64(0)

mcdc/object_/simulation.py

@@ -22,14 +22,15 @@
 import numpy as np
 
 from mpi4py import MPI
-from numpy import float64, int64, uint64, uintp
+from numpy import float64, int64
 from numpy.typing import NDArray
 
 ####
 
 from mcdc.object_.base import ObjectSingleton
 from mcdc.object_.data import DataBase, DataNone
 from mcdc.object_.distribution import DistributionBase, DistributionNone
+from mcdc.object_.gpu_tools import GPUMeta
 from mcdc.object_.mesh import MeshBase
 from mcdc.object_.particle import ParticleBank
 from mcdc.object_.settings import Settings
@@ -124,6 +125,9 @@ class Simulation(ObjectSingleton):
     runtime_output: float
     runtime_bank_management: float
 
+    # GPU metadata
+    gpu_meta: GPUMeta
+
     def __init__(self):
         super().__init__()
 
@@ -218,6 +222,9 @@ def __init__(self):
         self.runtime_output = 0.0
         self.runtime_bank_management = 0.0
 
+        # GPU metadata
+        self.gpu_meta = GPUMeta()
+
     def set_root_universe(self, cells=[]):
         self.universes[0].cells = cells