diff --git a/examples/c5g7/3d/k-eigenvalue/input.py b/examples/c5g7/3d/k-eigenvalue/input.py index 64441f666..88ce2e27c 100644 --- a/examples/c5g7/3d/k-eigenvalue/input.py +++ b/examples/c5g7/3d/k-eigenvalue/input.py @@ -296,9 +296,9 @@ def set_mat(mat): mcdc.tally.mesh_tally(scores=["flux"], x=x_grid, y=y_grid, z=z_grid, g=g_grid) # Setting -mcdc.setting(N_particle=1e3, census_bank_buff=4) +mcdc.setting(N_particle=2e4, census_bank_buff=4) -mcdc.eigenmode(N_inactive=50, N_active=150, gyration_radius="all") +mcdc.eigenmode(N_inactive=5, N_active=15, gyration_radius="all") mcdc.population_control() # Run diff --git a/examples/fixed_source/azurv1_pl_super/input.py b/examples/fixed_source/azurv1_pl_super/input.py index 254e651d6..1b92ec1f6 100644 --- a/examples/fixed_source/azurv1_pl_super/input.py +++ b/examples/fixed_source/azurv1_pl_super/input.py @@ -43,7 +43,7 @@ ) # Setting -mcdc.setting(N_particle=1e2) +mcdc.setting(N_particle=1e5) # Run mcdc.run() diff --git a/examples/fixed_source/kobayashi3-TD/input.py b/examples/fixed_source/kobayashi3-TD/input.py index 98bef5e10..4ea49e35c 100644 --- a/examples/fixed_source/kobayashi3-TD/input.py +++ b/examples/fixed_source/kobayashi3-TD/input.py @@ -60,7 +60,7 @@ scores=["flux"], x=np.linspace(0.0, 60.0, 31), y=np.linspace(0.0, 100.0, 51), - # t=np.linspace(0.0, 200.0, 21), + t=np.linspace(0.0, 200.0, 21), # g=np.array([-0.5, 3.5, 6.5]) # fast (0, 1, 2, 3) and thermal (4, 5, 6) groups ) @@ -79,7 +79,7 @@ # Setting -mcdc.setting(N_particle=1e2) +mcdc.setting(N_particle=1e5) # Run mcdc.run() diff --git a/mcdc/adapt.py b/mcdc/adapt.py index f84fd9dd2..6d853cd03 100644 --- a/mcdc/adapt.py +++ b/mcdc/adapt.py @@ -5,6 +5,7 @@ import numba import mcdc.type_ as type_ import mcdc.kernel as kernel +import mcdc.config as config if importlib.util.find_spec("harmonize") is None: @@ -74,6 +75,15 @@ 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 @@ -95,6 +105,8 @@ def impl(arg): return impl + + # ============================================================================= # Generic GPU/CPU Local Array Variable Constructors # ============================================================================= @@ -401,12 +413,17 @@ def nopython_mode(is_on): # GPU Type / Extern Functions Forward Declarations # ============================================================================= +@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,9 +434,14 @@ 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) @@ -427,16 +449,25 @@ def gpu_forward_declare(args): 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_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_.global_) - 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,8 +477,8 @@ 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) @@ -463,6 +494,46 @@ 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 +# ============================================================================= + +@numba.njit() +def create_tally_array(width,length): + 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((width,length), dtype=type_.float64) + return data_tally, 0 + + +@numba.njit() +def create_mcdc_array(): + 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=type_.global_) + return mcdc_array, 0 # ============================================================================= # Seperate GPU/CPU Functions to Target Different Platforms @@ -514,62 +585,61 @@ def thread(prog): @for_cpu() -def add_active(particle, prog): - kernel.add_particle(particle, prog["bank_active"]) - +def add_active(P_arr, prog): + kernel.add_particle(P_arr, prog["bank_active"]) @for_gpu() -def add_active(P_reclike, prog): - P = local_array(1, type_.particle) - kernel.recordlike_to_particle(P, P_reclike) +def add_active(P_rec_arr, prog): + P_arr = local_array(1, type_.particle) + kernel.recordlike_to_particle(P_arr, P_rec_arr) if SIMPLE_ASYNC: - step_async(prog, P[0]) + step_async(prog, P_arr[0]) else: - find_cell_async(prog, P[0]) + find_cell_async(prog, P_arr[0]) @for_cpu() -def add_source(particle, prog): - kernel.add_particle(particle, prog["bank_source"]) +def add_source(P_arr, prog): + kernel.add_particle(P_arr, prog["bank_source"]) @for_gpu() -def add_source(particle, prog): +def add_source(P_arr, prog): mcdc = mcdc_global(prog) - kernel.add_particle(particle, mcdc["bank_source"]) + kernel.add_particle(P_arr, mcdc["bank_source"]) @for_cpu() -def add_census(particle, prog): - kernel.add_particle(particle, prog["bank_census"]) +def add_census(P_arr, prog): + kernel.add_particle(P_arr, prog["bank_census"]) @for_gpu() -def add_census(particle, prog): +def add_census(P_arr, prog): mcdc = mcdc_global(prog) - kernel.add_particle(particle, mcdc["bank_census"]) + kernel.add_particle(P_arr, mcdc["bank_census"]) @for_cpu() -def add_future(particle, prog): - kernel.add_particle(particle, prog["bank_future"]) +def add_future(P_arr, prog): + kernel.add_particle(P_arr, prog["bank_future"]) @for_gpu() -def add_future(particle, prog): +def add_future(P_arr, prog): mcdc = mcdc_global(prog) - kernel.add_particle(particle, mcdc["bank_future"]) + kernel.add_particle(P_arr, mcdc["bank_future"]) @for_cpu() -def add_IC(particle, prog): - kernel.add_particle(particle, prog["technique"]["IC_bank_neutron_local"]) +def add_IC(P_arr, prog): + kernel.add_particle(P_arr, prog["technique"]["IC_bank_neutron_local"]) @for_gpu() -def add_IC(particle, prog): +def add_IC(P_arr, prog): mcdc = mcdc_global(prog) - kernel.add_particle(particle, mcdc["technique"]["IC_bank_neutron_local"]) + kernel.add_particle(P_arr, mcdc["technique"]["IC_bank_neutron_local"]) @for_cpu() diff --git a/mcdc/config.py b/mcdc/config.py index 3b80f3d04..557200d2e 100644 --- a/mcdc/config.py +++ b/mcdc/config.py @@ -15,6 +15,14 @@ "--target", type=str, help="Target", choices=["cpu", "gpu"], default="cpu" ) +parser.add_argument( + "--gpu_state_storage", + type=str, + help="Strategy used in GPU execution (event or async).", + choices=["separate","managed", "united"], + default="separate", +) + parser.add_argument( "--gpu_strat", type=str, @@ -73,6 +81,7 @@ mode = args.mode target = args.target +gpu_state_storage = args.gpu_state_storage caching = args.caching clear_cache = args.clear_cache diff --git a/mcdc/kernel.py b/mcdc/kernel.py index 436a7d9fa..ad5dc6fa3 100644 --- a/mcdc/kernel.py +++ b/mcdc/kernel.py @@ -21,6 +21,8 @@ from mcdc.print_ import print_error, print_msg from mcdc.src.algorithm import binary_search, binary_search_with_length +import cffi +ffi = cffi.FFI() @njit def round(float_val): @@ -655,13 +657,13 @@ def source_particle_dd(seed, mcdc): # Position if source["box"]: x = sample_uniform( - max(source["box_x"][0], d_x[0]), min(source["box_x"][1], d_x[1]), P + max(source["box_x"][0], d_x[0]), min(source["box_x"][1], d_x[1]), P_arr ) y = sample_uniform( - max(source["box_y"][0], d_y[0]), min(source["box_y"][1], d_y[1]), P + max(source["box_y"][0], d_y[0]), min(source["box_y"][1], d_y[1]), P_arr ) z = sample_uniform( - max(source["box_z"][0], d_z[0]), min(source["box_z"][1], d_z[1]), P + max(source["box_z"][0], d_z[0]), min(source["box_z"][1], d_z[1]), P_arr ) else: @@ -974,7 +976,7 @@ def add_bank_size(bank, value): @for_cpu() def full_bank_print(bank): with objmode(): - print_error("Particle %s bank is full." % bank["tag"]) + print_error("Particle %s bank is full at count %d." % (bank["tag"],bank["size"])) @for_gpu() @@ -982,6 +984,20 @@ def full_bank_print(bank): pass +@njit +def add_full_particle(P_arr, bank): + P = P_arr[0] + + idx = add_bank_size(bank, 1) + + # Check if bank is full + if idx >= bank["particles"].shape[0]: + full_bank_print(bank) + + # Set particle + copy_particle(bank["particles"][idx : idx + 1], P_arr) + + @njit def add_particle(P_arr, bank): P = P_arr[0] @@ -1463,6 +1479,8 @@ def bank_IC(P_arr, prog): # Accumulate fission SigmaF = material["fission"][g] # mcdc["technique"]["IC_fission_score"][0] += v * SigmaF + + # HAZARD adapt.global_add(mcdc["technique"]["IC_fission_score"], 0, round(v * SigmaF)) # ========================================================================= @@ -1484,6 +1502,7 @@ def bank_IC(P_arr, prog): total += nu_d[j] / decay[j] wp = flux * total * SigmaF / mcdc["k_eff"] + # HAZARD (with loop above) # Material has no precursor if total == 0.0: return @@ -2209,57 +2228,79 @@ def score_cs_tally(P_arr, distance, tally, data_tally, mcdc): @njit -def cs_clip(p, q, t0, t1): +def cs_clip(p, q, t): if p < 0: - t = q / p - if t > t1: - return False, t0, t1 - if t > t0: - t0 = t + tc = q / p + if tc > t[1]: + return False + if tc > t[0]: + t[0] = tc elif p > 0: - t = q / p - if t < t0: - return False, t0, t1 - if t < t1: - t1 = t + tc = q / p + if tc < t[0]: + return False + if tc < t[1]: + t[1] = tc elif q < 0: - return False, t0, t1 - return True, t0, t1 + return False + return True @njit -def cs_tracklength_in_box(start, end, x_min, x_max, y_min, y_max): +def cs_tracklength_in_box(base_start, base_end, x_min, x_max, y_min, y_max): + + start = adapt.local_array(2, type_.float64) + end = adapt.local_array(2, type_.float64) + start [0] = base_start[0] + start [1] = base_start[1] + end [0] = base_end[0] + end [1] = base_end[1] + # Uses Liang-Barsky algorithm for finding tracklength in box - t0, t1 = 0.0, 1.0 + t = adapt.local_array(2, type_.float64) + t[0] = 0.0 + t[1] = 1.0 dx = end[0] - start[0] dy = end[1] - start[1] # Perform clipping for each boundary - result, t0, t1 = cs_clip(-dx, start[0] - x_min, t0, t1) - if not result: - return 0.0 - result, t0, t1 = cs_clip(dx, x_max - start[0], t0, t1) + #result = cs_clip(-dx, start[0] - x_min, t) + #if not result: + # return 0.0 + result = cs_clip(dx, x_max - start[1], t) if not result: return 0.0 - result, t0, t1 = cs_clip(-dy, start[1] - y_min, t0, t1) + result = cs_clip(-dy, start[1] - y_min, t) if not result: return 0.0 - result, t0, t1 = cs_clip(dy, y_max - start[1], t0, t1) + result = cs_clip(dy, y_max - start[1], t) if not result: return 0.0 - # Update start and end points based on clipping results - if t1 < 1: - end[0] = start[0] + t1 * dx - end[1] = start[1] + t1 * dy - if t0 > 0: - start[0] = start[0] + t0 * dx - start[1] = start[1] + t0 * dx + ## Update start and end points based on clipping results + #if t[1] < 1: + # end[0] = start[0] + t[1] * dx + # end[1] = start[1] + t[1] * dy + #if t[0] > 0: + # start[0] = start[0] + t[0] * dx + # start[1] = start[1] + t[0] * dx - # Return the norm - X = end[0] - start[0] - Y = end[1] - start[1] - return math.sqrt(X**2 + Y**2) + #if t[0] > 0: + # start[0] = start[0] + t[0] * dx + # start[1] = start[1] + t[0] * dx + + #start[0] = start[0] + t[0] + #end[0] = start[0] + + + ## Return the norm + #X = end[0] - start[0] + #Y = end[1] - start[1] + #return math.sqrt(X**2 + Y**2) + if t[0] == 0: + return 1 + else : + return 0 @njit @@ -2578,6 +2619,7 @@ def eigenvalue_tally(P_arr, distance, mcdc): if not nuclide["fissionable"]: continue for j in range(J): + # HAZARD nu_d = get_nu_group(NU_FISSION_DELAYED, nuclide, E, j) decay = nuclide["ce_decay"][j] total += nu_d / decay @@ -2858,12 +2900,14 @@ def move_to_event(P_arr, data_tally, mcdc): score_cell_tally(P_arr, distance, tally, data_tally, mcdc) # CS tallies + # HAZARD for tally in mcdc["cs_tallies"]: score_cs_tally(P_arr, distance, tally, data_tally, mcdc) if mcdc["setting"]["mode_eigenvalue"]: eigenvalue_tally(P_arr, distance, mcdc) + # Move particle move_particle(P_arr, distance, mcdc) @@ -3402,7 +3446,7 @@ def sample_phasespace_fission_nuclide(P_arr, nuclide, P_new_arr, mcdc): if mcdc["setting"]["mode_MG"]: fission_MG(P_arr, nuclide, P_new_arr) else: - fission_CE(P_arr, nuclide, P_new_arr) + fission_CE(P_arr, nuclide, P_new_arr, mcdc) @njit @@ -3451,7 +3495,7 @@ def fission_MG(P_arr, nuclide, P_new_arr): @njit -def fission_CE(P_arr, nuclide, P_new_arr): +def fission_CE(P_arr, nuclide, P_new_arr, mcdc): P_new = P_new_arr[0] P = P_arr[0] # Get constants @@ -3463,6 +3507,8 @@ def fission_CE(P_arr, nuclide, P_new_arr): for j in range(J): nu_d[j] = get_nu_group(NU_FISSION_DELAYED, nuclide, E, j) + + # Delayed? prompt = True delayed_group = -1 @@ -3886,6 +3932,7 @@ def sample_Eout(P_new_arr, E_grid, NE, chi): idx = binary_search_with_length(xi, chi, NE) # Linear interpolation + # HAZARD? E1 = E_grid[idx] E2 = E_grid[idx + 1] chi1 = chi[idx] diff --git a/mcdc/loop.py b/mcdc/loop.py index d2a8105f2..09d8ecaf5 100644 --- a/mcdc/loop.py +++ b/mcdc/loop.py @@ -1,5 +1,6 @@ from mpi4py import MPI from numba import njit, objmode +from numba.misc.special import literally import shutil @@ -33,11 +34,11 @@ alloc_state, free_state = [None] * 2 src_alloc_program, src_free_program = [None] * 2 -src_load_constant, src_load_constant, src_store_constant, src_store_data = [None] * 4 +src_load_global, src_load_constant, src_store_global, src_store_data, src_store_pointer_data = [None] * 5 src_init_program, src_exec_program, src_complete, src_clear_flags = [None] * 4 pre_alloc_program, pre_free_program = [None] * 2 -pre_load_constant, pre_load_data, pre_store_constant, pre_store_data = [None] * 4 +pre_load_global, pre_load_data, pre_store_global, pre_store_data = [None] * 4 pre_init_program, pre_exec_program, pre_complete, pre_clear_flags = [None] * 4 @@ -45,7 +46,7 @@ # be redefined to overwrite the above symbols and perform initialization/ # finalization of GPU state @njit -def setup_gpu(mcdc): +def setup_gpu(mcdc,data_tally): pass @@ -64,9 +65,10 @@ def loop_fixed_source(data_tally, mcdc_arr): # Ensure `mcdc` exist for the lifetime of the program # by intentionally leaking their memory - adapt.leak(mcdc_arr) + #adapt.leak(mcdc_arr) mcdc = mcdc_arr[0] + # Loop over batches for idx_batch in range(mcdc["setting"]["N_batch"]): if not mcdc["technique"]["domain_decomposition"]: @@ -76,6 +78,7 @@ def loop_fixed_source(data_tally, mcdc_arr): mcdc["idx_batch"] = idx_batch seed_batch = kernel.split_seed(idx_batch, mcdc["setting"]["rng_seed"]) + # Print multi-batch header if mcdc["setting"]["N_batch"] > 1: with objmode(): @@ -84,6 +87,7 @@ def loop_fixed_source(data_tally, mcdc_arr): seed_uq = kernel.split_seed(seed_batch, SEED_SPLIT_UQ) kernel.uq_reset(mcdc, seed_uq) + # Loop over time censuses for idx_census in range(mcdc["setting"]["N_census"]): mcdc["idx_census"] = idx_census @@ -116,6 +120,8 @@ def loop_fixed_source(data_tally, mcdc_arr): break # Loop over source particles seed_source = kernel.split_seed(seed_census, SEED_SPLIT_SOURCE) + + loop_source(seed_source, data_tally, mcdc) # Loop over source precursors @@ -125,10 +131,12 @@ def loop_fixed_source(data_tally, mcdc_arr): ) loop_source_precursor(seed_source_precursor, data_tally, mcdc) + # Manage particle banks: population control and work rebalance seed_bank = kernel.split_seed(seed_census, SEED_SPLIT_BANK) kernel.manage_particle_banks(seed_bank, mcdc) + # Time census-based tally closeout if mcdc["setting"]["census_based_tally"]: kernel.tally_reduce(data_tally, mcdc) @@ -179,7 +187,7 @@ def loop_fixed_source(data_tally, mcdc_arr): def loop_eigenvalue(data_tally, mcdc_arr): # Ensure `mcdc` exist for the lifetime of the program # by intentionally leaking their memory - adapt.leak(mcdc_arr) + #adapt.leak(mcdc_arr) mcdc = mcdc_arr[0] # Loop over power iteration cycles @@ -461,9 +469,12 @@ def finalize(prog: nb.uintp): base_fns = (initialize, finalize, make_work) + shape = eval(f"{adapt.tally_shape_literal}") + # Just do exec/eval def step(prog: nb.uintp, P_input: adapt.particle_gpu): mcdc = adapt.mcdc_global(prog) - data = adapt.mcdc_data(prog) + data_ptr = adapt.mcdc_data(prog) + data = adapt.harm.array_from_ptr(data_ptr,shape,nb.float64) P_arr = adapt.local_array(1, type_.particle) P_arr[0] = P_input P = P_arr[0] @@ -499,7 +510,7 @@ def gpu_loop_source(seed, data, mcdc): # For async execution iter_count = 655360000 # For event-based execution - batch_size = 1 + batch_size = 64 full_work_size = mcdc["mpi_work_size"] if ASYNC_EXECUTION: @@ -515,29 +526,33 @@ def gpu_loop_source(seed, data, mcdc): mcdc["source_seed"] = seed # Store the global state to the GPU - src_store_constant(mcdc["gpu_state_pointer"], mcdc) - src_store_data(mcdc["gpu_state_pointer"], data) + if config.gpu_state_storage == "separate": + adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["state_pointer"], mcdc) + adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["state_pointer"], data) # Execute the program, and continue to do so until it is done if ASYNC_EXECUTION: - src_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, iter_count) - while not src_complete(mcdc["source_program_pointer"]): + src_exec_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, iter_count) + while not src_complete(mcdc["gpu_meta"]["source_program_pointer"]): kernel.dd_particle_send(mcdc) src_exec_program( - mcdc["source_program_pointer"], BLOCK_COUNT, iter_count + mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, iter_count ) else: - src_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, batch_size) - while not src_complete(mcdc["source_program_pointer"]): + src_exec_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, batch_size) + while not src_complete(mcdc["gpu_meta"]["source_program_pointer"]): kernel.dd_particle_send(mcdc) src_exec_program( - mcdc["source_program_pointer"], BLOCK_COUNT, batch_size + mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, batch_size ) - + src_clear_flags(mcdc["gpu_meta"]["source_program_pointer"]) # Recover the original program state - src_load_constant(mcdc, mcdc["gpu_state_pointer"]) - src_load_data(data, mcdc["gpu_state_pointer"]) - src_clear_flags(mcdc["source_program_pointer"]) + + if config.gpu_state_storage == "separate": + adapt.harm.memcpy_device_to_host(mcdc,mcdc["gpu_meta"]["state_pointer"]) + adapt.harm.memcpy_device_to_host(data,mcdc["gpu_meta"]["state_pointer"]) + + src_clear_flags(mcdc["gpu_meta"]["source_program_pointer"]) mcdc["mpi_work_size"] = full_work_size @@ -772,6 +787,13 @@ def loop_source_precursor(seed, data, mcdc): source_precursor_closeout(mcdc, idx_work, N_prog, data) + + +# ========================================================================= +# GPU Runtime Specifications and Function Re-definitions +# ========================================================================= + + def gpu_precursor_spec(): def make_work(prog: nb.uintp) -> nb.boolean: mcdc = adapt.mcdc_global(prog) @@ -813,9 +835,11 @@ def finalize(prog: nb.uintp): base_fns = (initialize, finalize, make_work) + shape = eval(f"{adapt.tally_shape_literal}") def step(prog: nb.uintp, P_input: adapt.particle_gpu): mcdc = adapt.mcdc_global(prog) - data = adapt.mcdc_data(prog) + data_ptr = adapt.mcdc_data(prog) + data = adapt.harm.array_from_ptr(data_ptr,shape,nb.float64) P_arr = adapt.local_array(1, type_.particle) P_arr[0] = P_input P = P_arr[0] @@ -862,26 +886,31 @@ def gpu_loop_source_precursor(seed, data, mcdc): mcdc["source_seed"] = seed # Store the global state to the GPU - pre_store_constant(mcdc["gpu_state_pointer"], mcdc) - pre_store_data(mcdc["gpu_state_pointer"], data) + if config.gpu_state_storage == "separate": + adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["state_pointer"], mcdc) + adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["state_pointer"], data) # Execute the program, and continue to do so until it is done # Execute the program, and continue to do so until it is done if ASYNC_EXECUTION: - pre_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, iter_count) - while not pre_complete(mcdc["source_program_pointer"]): + pre_exec_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, iter_count) + while not pre_complete(mcdc["gpu_meta"]["source_program_pointer"]): kernel.dd_particle_send(mcdc) - pre_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, iter_count) + pre_exec_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, iter_count) else: - pre_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, batch_size) - while not pre_complete(mcdc["source_program_pointer"]): + pre_exec_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, batch_size) + while not pre_complete(mcdc["gpu_meta"]["source_program_pointer"]): kernel.dd_particle_send(mcdc) - pre_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, batch_size) + pre_exec_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT, batch_size) + # Recover the original program state - pre_load_constant(mcdc, mcdc["gpu_state_pointer"]) - pre_load_data(data, mcdc["gpu_state_pointer"]) - pre_clear_flags(mcdc["source_program_pointer"]) + if config.gpu_state_storage == "separate": + adapt.harm.memcpy_device_to_host(mcdc,mcdc["gpu_meta"]["state_pointer"]) + adapt.harm.memcpy_device_to_host(data,mcdc["gpu_meta"]["state_pointer"]) + + pre_clear_flags(mcdc["gpu_meta"]["source_program_pointer"]) + kernel.set_bank_size(mcdc["bank_active"], 0) @@ -925,14 +954,16 @@ def build_gpu_progs(input_deck, args): free_state = src_fns["free_state"] global src_alloc_program, src_free_program - global src_load_constant, src_store_constant, src_load_data, src_store_data + global src_load_global, src_store_global, src_load_data, src_store_data, src_store_pointer_data global src_init_program, src_exec_program, src_complete, src_clear_flags src_alloc_program = src_fns["alloc_program"] src_free_program = src_fns["free_program"] - src_load_constant = src_fns["load_state_device_global"] - src_store_constant = src_fns["store_state_device_global"] + src_load_global = src_fns["load_state_device_global"] + src_store_global = src_fns["store_state_device_global"] + src_store_pointer_global = src_fns["store_pointer_state_device_global"] src_load_data = src_fns["load_state_device_data"] src_store_data = src_fns["store_state_device_data"] + src_store_pointer_data = src_fns["store_pointer_state_device_data"] src_init_program = src_fns["init_program"] src_exec_program = src_fns["exec_program"] src_complete = src_fns["complete"] @@ -940,14 +971,14 @@ def build_gpu_progs(input_deck, args): src_set_device = src_fns["set_device"] global pre_alloc_program, pre_free_program - global pre_load_constant, pre_store_constant, pre_load_data, pre_store_data + global pre_load_global, pre_store_global, pre_load_data, pre_store_data global pre_init_program, pre_exec_program, pre_complete, pre_clear_flags pre_alloc_state = pre_fns["alloc_state"] pre_free_state = pre_fns["free_state"] pre_alloc_program = pre_fns["alloc_program"] pre_free_program = pre_fns["free_program"] - pre_load_constant = pre_fns["load_state_device_global"] - pre_store_constant = pre_fns["store_state_device_global"] + pre_load_global = pre_fns["load_state_device_global"] + pre_store_global = pre_fns["store_state_device_global"] pre_load_data = pre_fns["load_state_device_data"] pre_store_data = pre_fns["store_state_device_data"] pre_init_program = pre_fns["init_program"] @@ -956,25 +987,34 @@ def build_gpu_progs(input_deck, args): pre_clear_flags = pre_fns["clear_flags"] @njit - def real_setup_gpu(mcdc): + def real_setup_gpu(mcdc_array,data_tally): + mcdc = mcdc_array[0] src_set_device(device_id) arena_size = ARENA_SIZE - mcdc["gpu_state_pointer"] = adapt.cast_voidptr_to_uintp(alloc_state()) - mcdc["source_program_pointer"] = adapt.cast_voidptr_to_uintp( - src_alloc_program(mcdc["gpu_state_pointer"], ARENA_SIZE) + mcdc["gpu_meta"]["state_pointer"] = adapt.cast_voidptr_to_uintp(alloc_state()) + #src_store_global(mcdc["gpu_meta"]["state_pointer"], mcdc_array[0]) + if config.gpu_state_storage == "separate": + src_store_pointer_global(mcdc["gpu_meta"]["state_pointer"], mcdc["gpu_meta"]["global_pointer"]) + src_store_pointer_data(mcdc["gpu_meta"]["state_pointer"], mcdc["gpu_meta"]["tally_pointer"]) + else: + src_store_pointer_global(mcdc["gpu_meta"]["state_pointer"], mcdc_array) + src_store_pointer_data(mcdc["gpu_meta"]["state_pointer"], data_tally) + + mcdc["gpu_meta"]["source_program_pointer"] = adapt.cast_voidptr_to_uintp( + src_alloc_program(mcdc["gpu_meta"]["state_pointer"], ARENA_SIZE) ) - src_init_program(mcdc["source_program_pointer"], BLOCK_COUNT) - mcdc["precursor_program_pointer"] = adapt.cast_voidptr_to_uintp( - pre_alloc_program(mcdc["gpu_state_pointer"], ARENA_SIZE) + src_init_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT) + mcdc["gpu_meta"]["precursor_program_pointer"] = adapt.cast_voidptr_to_uintp( + pre_alloc_program(mcdc["gpu_meta"]["state_pointer"], ARENA_SIZE) ) - pre_init_program(mcdc["precursor_program_pointer"], BLOCK_COUNT) + pre_init_program(mcdc["gpu_meta"]["precursor_program_pointer"], BLOCK_COUNT) return @njit def real_teardown_gpu(mcdc): - src_free_program(adapt.cast_uintp_to_voidptr(mcdc["source_program_pointer"])) - pre_free_program(adapt.cast_uintp_to_voidptr(mcdc["precursor_program_pointer"])) - free_state(adapt.cast_uintp_to_voidptr(mcdc["gpu_state_pointer"])) + src_free_program(adapt.cast_uintp_to_voidptr(mcdc["gpu_meta"]["source_program_pointer"])) + pre_free_program(adapt.cast_uintp_to_voidptr(mcdc["gpu_meta"]["precursor_program_pointer"])) + free_state(adapt.cast_uintp_to_voidptr(mcdc["gpu_meta"]["state_pointer"])) global setup_gpu, teardown_gpu setup_gpu = real_setup_gpu diff --git a/mcdc/main.py b/mcdc/main.py index 63558206e..74ff48fb5 100644 --- a/mcdc/main.py +++ b/mcdc/main.py @@ -82,6 +82,7 @@ def run(): elif mcdc["setting"]["mode_eigenvalue"]: loop_eigenvalue(data_tally, mcdc_arr) else: + print_msg("Starting fixed source") loop_fixed_source(data_tally, mcdc_arr) mcdc["runtime_simulation"] = MPI.Wtime() - simulation_start @@ -99,6 +100,9 @@ def run(): MPI.COMM_WORLD.Barrier() mcdc["runtime_total"] = MPI.Wtime() - total_start + #for i in range(mcdc["bank_log"]["size"][0]): + # print(mcdc["bank_log"]["particles"][i]) + # Closout closeout(mcdc) @@ -200,7 +204,7 @@ def calculate_cs_sparse_solution(data, mcdc, A, b): def cs_reconstruct(data, mcdc): - tally_bin = data[TALLY] + tally_bin = data tally = mcdc["cs_tallies"][0] stride = tally["stride"] bin_idx = stride["tally"] @@ -526,6 +530,7 @@ def prepare(): type_.make_type_domain_decomp(input_deck) type_.make_type_dd_turnstile_event(input_deck) type_.make_type_technique(input_deck) + type_.make_type_gpu_meta() type_.make_type_global(input_deck) type_.make_size_rpn(input_deck) kernel.adapt_rng(nb.config.DISABLE_JIT) @@ -1137,7 +1142,8 @@ def prepare(): tally_bin_N_copies = 3 else: tally_bin_N_copies = 5 - data_tally = np.zeros((tally_bin_N_copies, tally_bin_size), dtype=type_.float64) + + tally_shape = (tally_bin_N_copies, tally_bin_size) # ========================================================================= # Platform Targeting, Adapters, Toggles, etc @@ -1152,7 +1158,7 @@ def prepare(): print_error( "No module named 'harmonize' - GPU functionality not available. " ) - adapt.gpu_forward_declare(config.args) + adapt.gpu_forward_declare(config.args,tally_shape) adapt.set_toggle("iQMC", input_deck.technique["iQMC"]) adapt.set_toggle("domain_decomp", input_deck.technique["domain_decomposition"]) @@ -1162,6 +1168,17 @@ def prepare(): build_gpu_progs(input_deck, config.args) adapt.nopython_mode((config.mode == "numba") or (config.mode == "numba_debug")) + + # ========================================================================= + # Allocate Tally Storage + # ========================================================================= + + data_tally, data_tally_uint = adapt.create_tally_array(tally_shape[0],tally_shape[1]) + + mcdc_arr, mcdc_uint = adapt.create_mcdc_array() + mcdc_arr[0] = mcdc + mcdc = mcdc_arr[0] + # ========================================================================= # Setting # ========================================================================= @@ -1563,7 +1580,7 @@ def prepare(): "w" ] - loop.setup_gpu(mcdc) + loop.setup_gpu(mcdc_arr,data_tally) # ========================================================================= # Finalize data: wrapping into a tuple diff --git a/mcdc/src/geometry.py b/mcdc/src/geometry.py index b8b208996..80d371547 100644 --- a/mcdc/src/geometry.py +++ b/mcdc/src/geometry.py @@ -91,7 +91,7 @@ def inspect_geometry(particle_container, mcdc): # Apply rotation if cell["fill_rotated"]: - _rotate_particle(particle, cell["rotation"]) + _rotate_particle(particle_container, cell["rotation"]) # Universe cell? if cell["fill_type"] == FILL_UNIVERSE: @@ -212,7 +212,7 @@ def locate_particle(particle_container, mcdc): # Apply rotation if cell["fill_rotated"]: - _rotate_particle(particle, cell["rotation"]) + _rotate_particle(particle_container, cell["rotation"]) # Universe cell? if cell["fill_type"] == FILL_UNIVERSE: @@ -262,8 +262,9 @@ def locate_particle(particle_container, mcdc): @nb.njit -def _rotate_particle(particle, rotation): +def _rotate_particle(particle_container, rotation): # Particle initial coordinate + particle = particle_container[0] x = particle["x"] y = particle["y"] z = particle["z"] diff --git a/mcdc/src/surface/common.py b/mcdc/src/surface/common.py index 4d8168941..3ef1df39c 100644 --- a/mcdc/src/surface/common.py +++ b/mcdc/src/surface/common.py @@ -209,8 +209,10 @@ def _get_distance_static(particle_container, surface): return plane_y.get_distance(particle_container, surface) elif surface["type"] & SURFACE_PLANE_Z: return plane_z.get_distance(particle_container, surface) - elif surface["type"] & SURFACE_PLANE: + elif surface["type"] & SURFACE_PLANE: # SHOULD BE REVIEWED return plane.get_distance(particle_container, surface) + else: + return INF else: if surface["type"] & SURFACE_CYLINDER_X: return cylinder_x.get_distance(particle_container, surface) diff --git a/mcdc/type_.py b/mcdc/type_.py index b86bc1696..d75d421e7 100644 --- a/mcdc/type_.py +++ b/mcdc/type_.py @@ -48,8 +48,10 @@ cs_tally = None technique = None -global_ = None +gpu_meta = None +global_ = None +global_size = None # ============================================================================== # MC/DC Member Array Sizes @@ -120,6 +122,10 @@ def align(field_list): offset = 0 pad_id = 0 for field in field_list: + + if isinstance(field[1],list): + raise ValueError("Given type of subfield is a list, but should be a dtype.") + if len(field) > 3: print_error( "Unexpected struct field specification. Specifications \ @@ -235,7 +241,7 @@ def make_type_particle(input_deck): # iQMC vector of weights if iQMC: G = input_deck.materials[0].G - iqmc_struct = [("w", float64, (G,))] + iqmc_struct = into_dtype([("w", float64, (G,))]) struct += [("iqmc", iqmc_struct)] # Save type @@ -273,7 +279,7 @@ def make_type_particle_record(input_deck): # iQMC vector of weights if iQMC: G = input_deck.materials[0].G - iqmc_struct = [("w", float64, (G,))] + iqmc_struct = into_dtype([("w", float64, (G,))]) struct += [("iqmc", iqmc_struct)] # Save type @@ -299,6 +305,14 @@ def make_type_particle_record(input_deck): # Particle bank # ============================================================================== +def full_particle_bank(max_size): + return into_dtype( + [ + ("particles", particle, (max_size,)), + ("size", int64, (1,)), + ("tag", str_), + ] + ) def particle_bank(max_size): return into_dtype( @@ -778,7 +792,7 @@ def make_type_mesh_tally(input_deck): Nmax_x, Nmax_y, Nmax_z = dd_meshtally(input_deck) # Set the filter - filter_ = [ + filter_ = into_dtype([ ("x", float64, (Nmax_x,)), ("y", float64, (Nmax_y,)), ("z", float64, (Nmax_z,)), @@ -793,11 +807,11 @@ def make_type_mesh_tally(input_deck): ("Nmu", int64), ("N_azi", int64), ("Ng", int64), - ] + ]) struct += [("filter", filter_)] # Tally strides - stride = [ + stride = into_dtype([ ("tally", int64), ("sensitivity", int64), ("mu", int64), @@ -807,7 +821,7 @@ def make_type_mesh_tally(input_deck): ("x", int64), ("y", int64), ("z", int64), - ] + ]) struct += [("stride", stride)] # Total number of bins @@ -840,24 +854,24 @@ def make_type_surface_tally(input_deck): Nmax_score = max(Nmax_score, len(card.scores)) # Set the filter - filter_ = [ + filter_ = into_dtype([ ("surface_ID", int64), ("t", float64, (Nmax_t,)), ("mu", float64, (Nmax_mu,)), ("azi", float64, (Nmax_azi,)), ("g", float64, (Nmax_g,)), - ] + ]) struct = [("filter", filter_)] # Tally strides - stride = [ + stride = into_dtype([ ("tally", int64), ("sensitivity", int64), ("mu", int64), ("azi", int64), ("g", int64), ("t", int64), - ] + ]) struct += [("stride", stride)] # Total number of bins @@ -889,7 +903,7 @@ def make_type_cell_tally(input_deck): Nmax_score = max(Nmax_score, len(card.scores)) # Set the filter - filter_ = [ + filter_ = into_dtype([ ("cell_ID", int64), ("t", float64, (Nmax_t,)), ("mu", float64, (Nmax_mu,)), @@ -897,18 +911,18 @@ def make_type_cell_tally(input_deck): ("g", float64, (Nmax_g,)), ("Nt", int64), ("Ng", int64), - ] + ]) struct = [("filter", filter_)] # Tally strides - stride = [ + stride = into_dtype([ ("tally", int64), ("sensitivity", int64), ("mu", int64), ("azi", int64), ("g", int64), ("t", int64), - ] + ]) struct += [("stride", stride)] # Total number of bins @@ -951,7 +965,7 @@ def make_type_cs_tally(input_deck): # Nmax_x, Nmax_y, Nmax_z = dd_meshtally(input_deck) # Set the filter - filter_ = [ + filter_ = into_dtype([ ("N_cs_bins", int), ("cs_bin_size", float64, (2,)), ( @@ -971,12 +985,12 @@ def make_type_cs_tally(input_deck): ("mu", float64, (Nmax_mu,)), ("azi", float64, (Nmax_azi,)), ("g", float64, (Nmax_g,)), - ] + ]) struct += [("filter", filter_)] # Tally strides - stride = [ + stride = into_dtype([ ("tally", int64), ("sensitivity", int64), ("mu", int64), @@ -987,7 +1001,7 @@ def make_type_cs_tally(input_deck): ("y", int64), ("z", int64), # ("N_cs_bins", int64), # TODO: get rid of this line? - ] + ]) struct += [("stride", stride)] # Total number of bins (will be used for the reconstruction) @@ -1423,13 +1437,29 @@ def make_type_domain_decomp(input_deck): param_names = ["tag", "ID", "key", "mean", "delta", "distribution", "rng_seed"] +# ============================================================================== +# GPU Metadata +# ============================================================================== + +def make_type_gpu_meta(): + global gpu_meta + + gpu_meta = into_dtype([ + ("state_pointer", uintp), + ("source_program_pointer", uintp), + ("precursor_program_pointer", uintp), + ("global_pointer",uintp), + ("tally_pointer",uintp), + ]) + + # ============================================================================== # Global # ============================================================================== def make_type_global(input_deck): - global global_ + global global_, global_size # Get modes mode_CE = input_deck.setting["mode_CE"] @@ -1508,7 +1538,7 @@ def make_type_global(input_deck): ) or input_deck.technique["iQMC"]: bank_source = particle_bank(N_work) - # GLobal type + global_ = into_dtype( [ ("nuclides", nuclide, (N_nuclide,)), @@ -1577,13 +1607,15 @@ def make_type_global(input_deck): ("runtime_bank_management", float64), ("precursor_strength", float64), ("mpi_work_iter", int64, (1,)), - ("gpu_state_pointer", uintp), - ("source_program_pointer", uintp), - ("precursor_program_pointer", uintp), + ("gpu_meta",gpu_meta), ("source_seed", uint64), ] ) + # GLobal type + + global_size = global_.itemsize + # ============================================================================== # Util