diff --git a/mcdc/adapt.py b/mcdc/adapt.py index efe36bcf0..d4f1f5394 100644 --- a/mcdc/adapt.py +++ b/mcdc/adapt.py @@ -421,11 +421,9 @@ def alloc_bytes_placeholder(size): 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 group_gpu = None thread_gpu = None particle_gpu = None @@ -435,7 +433,6 @@ def alloc_bytes_placeholder(size): 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 @@ -449,36 +446,25 @@ def gpu_forward_declare(args, tally_shape): if args.gpu_cuda_path != None: harm.config.set_cuda_path(args.gpu_cuda_path) - global none_type, mcdc_global_type, mcdc_data_type, mcdc_shared_type + global none_type, mcdc_global_type, mcdc_shared_type global state_spec - global mcdc_global_gpu, mcdc_data_gpu + global mcdc_global_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.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, - "data": mcdc_data_type, }, none_type, none_type, ) access_fns = harm.RuntimeSpec.access_fns(state_spec) 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) @@ -505,27 +491,6 @@ def find_cell(prog: numba.uintp, P: particle_gpu): # ============================================================================= -@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) - if config.gpu_state_storage == "separate": - data_tally = np.empty((width, length), type_.float64) - else: - data_tally = numba.carray(data_tally_ptr, (width, length), type_.float64) - for i in range(data_tally.shape[0]): - for j in range(data_tally.shape[1]): - data_tally[i, j] = 0 - 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": @@ -538,6 +503,9 @@ def create_mcdc_array(): mcdc_array = np.empty((1,), type_.global_) else: mcdc_array = numba.carray(mcdc_ptr, (1,), type_.global_) + for i in range(len(mcdc_array[0]["data"])): + for j in range(len(mcdc_array[0]["data"][0])): + mcdc_array[0]["data"][i][j] = 0 return mcdc_array, mcdc_uint else: mcdc_array = np.zeros((1,), dtype=type_.global_) @@ -570,7 +538,7 @@ def mcdc_data(prog): @for_gpu() def mcdc_data(prog): - return mcdc_data_gpu(prog) + return mcdc_global_gpu(prog).data @for_cpu() diff --git a/mcdc/loop.py b/mcdc/loop.py index 8b8a76906..943ff2780 100644 --- a/mcdc/loop.py +++ b/mcdc/loop.py @@ -34,17 +34,11 @@ alloc_state, free_state = [None] * 2 src_alloc_program, src_free_program = [None] * 2 -( - src_load_global, - src_load_constant, - src_store_global, - src_store_data, - src_store_pointer_data, -) = [None] * 5 +src_load_global, src_store_global = [None] * 2 src_init_program, src_exec_program, src_complete, src_clear_flags = [None] * 4 pre_alloc_program, pre_free_program = [None] * 2 -pre_load_global, pre_load_data, pre_store_global, pre_store_data = [None] * 4 +pre_load_global, pre_store_global = [None] * 2 pre_init_program, pre_exec_program, pre_complete, pre_clear_flags = [None] * 4 @@ -52,7 +46,7 @@ # be redefined to overwrite the above symbols and perform initialization/ # finalization of GPU state @njit -def setup_gpu(mcdc, data_tally): +def setup_gpu(mcdc): pass @@ -67,12 +61,13 @@ def teardown_gpu(mcdc): @njit -def loop_fixed_source(data_tally, mcdc_arr): +def loop_fixed_source(mcdc_arr): # Ensure `mcdc` exist for the lifetime of the program # by intentionally leaking their memory # adapt.leak(mcdc_arr) mcdc = mcdc_arr[0] + data_tally = mcdc["data"] # Loop over batches for idx_batch in range(mcdc["setting"]["N_batch"]): @@ -184,11 +179,12 @@ def loop_fixed_source(data_tally, mcdc_arr): @njit -def loop_eigenvalue(data_tally, mcdc_arr): +def loop_eigenvalue(mcdc_arr): # Ensure `mcdc` exist for the lifetime of the program # by intentionally leaking their memory # adapt.leak(mcdc_arr) mcdc = mcdc_arr[0] + data_tally = mcdc["data"] # Loop over power iteration cycles for idx_cycle in range(mcdc["setting"]["N_cycle"]): @@ -469,19 +465,10 @@ def finalize(prog: nb.uintp): base_fns = (initialize, finalize, make_work) - lns = {} - exec( - f"shape = ({adapt.tally_shape_literal[0]},{adapt.tally_shape_literal[1]})", - globals(), - lns, - ) - shape = lns["shape"] - # Just do exec/eval def step(prog: nb.uintp, P_input: adapt.particle_gpu): mcdc = adapt.mcdc_global(prog) - data_ptr = adapt.mcdc_data(prog) - data = adapt.harm.array_from_ptr(data_ptr, shape, nb.float64) + data = mcdc["data"] P_arr = adapt.local_array(1, type_.particle) P_arr[0] = P_input P = P_arr[0] @@ -535,7 +522,6 @@ def gpu_loop_source(seed, data, mcdc): # Store the global state to the GPU if config.gpu_state_storage == "separate": adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["global_pointer"], mcdc) - adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["tally_pointer"], data) # Execute the program, and continue to do so until it is done if ASYNC_EXECUTION: @@ -561,7 +547,6 @@ def gpu_loop_source(seed, data, mcdc): if config.gpu_state_storage == "separate": adapt.harm.memcpy_device_to_host(mcdc, mcdc["gpu_meta"]["global_pointer"]) - adapt.harm.memcpy_device_to_host(data, mcdc["gpu_meta"]["tally_pointer"]) src_clear_flags(mcdc["gpu_meta"]["source_program_pointer"]) @@ -844,18 +829,10 @@ def finalize(prog: nb.uintp): base_fns = (initialize, finalize, make_work) - lns = {} - exec( - f"shape = ({adapt.tally_shape_literal[0]},{adapt.tally_shape_literal[1]})", - globals(), - lns, - ) - shape = lns["shape"] def step(prog: nb.uintp, P_input: adapt.particle_gpu): mcdc = adapt.mcdc_global(prog) - data_ptr = adapt.mcdc_data(prog) - data = adapt.harm.array_from_ptr(data_ptr, shape, nb.float64) + data = mcdc["data"] P_arr = adapt.local_array(1, type_.particle) P_arr[0] = P_input P = P_arr[0] @@ -904,7 +881,6 @@ def gpu_loop_source_precursor(seed, data, mcdc): # Store the global state to the GPU if config.gpu_state_storage == "separate": adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["global_pointer"], mcdc) - adapt.harm.memcpy_host_to_device(mcdc["gpu_meta"]["tally_pointer"], data) # Execute the program, and continue to do so until it is done @@ -931,7 +907,6 @@ def gpu_loop_source_precursor(seed, data, mcdc): # Recover the original program state if config.gpu_state_storage == "separate": adapt.harm.memcpy_device_to_host(mcdc, mcdc["gpu_meta"]["global_pointer"]) - adapt.harm.memcpy_device_to_host(data, mcdc["gpu_meta"]["tally_pointer"]) pre_clear_flags(mcdc["gpu_meta"]["source_program_pointer"]) @@ -977,16 +952,13 @@ def build_gpu_progs(input_deck, args): free_state = src_fns["free_state"] global src_alloc_program, src_free_program - global src_load_global, src_store_global, src_load_data, src_store_data, src_store_pointer_data + global src_load_global, src_store_global 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_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"] @@ -994,7 +966,7 @@ def build_gpu_progs(input_deck, args): src_set_device = src_fns["set_device"] global pre_alloc_program, pre_free_program - global pre_load_global, pre_store_global, pre_load_data, pre_store_data + global pre_load_global, pre_store_global 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"] @@ -1002,30 +974,24 @@ def build_gpu_progs(input_deck, args): pre_free_program = pre_fns["free_program"] 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"] pre_exec_program = pre_fns["exec_program"] pre_complete = pre_fns["complete"] pre_clear_flags = pre_fns["clear_flags"] @njit - def real_setup_gpu(mcdc_array, data_tally): + def real_setup_gpu(mcdc_array): mcdc = mcdc_array[0] src_set_device(device_id) arena_size = 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) @@ -1054,3 +1020,5 @@ def real_teardown_gpu(mcdc): global loop_source, loop_source_precursor loop_source = gpu_loop_source loop_source_precursor = gpu_loop_source_precursor + + diff --git a/mcdc/main.py b/mcdc/main.py index c90f44a60..a3b8c29c3 100644 --- a/mcdc/main.py +++ b/mcdc/main.py @@ -64,7 +64,7 @@ def run(): if input_deck.technique["iQMC"]: iqmc_validate_inputs(input_deck) - data_tally, mcdc_arr = prepare() + mcdc_arr = prepare() mcdc = mcdc_arr[0] mcdc["runtime_preparation"] = MPI.Wtime() - preparation_start @@ -80,20 +80,20 @@ def run(): if mcdc["technique"]["iQMC"]: iqmc_simulation(mcdc_arr) elif mcdc["setting"]["mode_eigenvalue"]: - loop_eigenvalue(data_tally, mcdc_arr) + loop_eigenvalue(mcdc_arr) else: print_msg("Starting fixed source") - loop_fixed_source(data_tally, mcdc_arr) + loop_fixed_source(mcdc_arr) mcdc["runtime_simulation"] = MPI.Wtime() - simulation_start # Compressed sensing reconstruction N_cs_bins = mcdc["cs_tallies"]["filter"]["N_cs_bins"][0] if N_cs_bins != 0: - cs_reconstruct(data_tally, mcdc) + cs_reconstruct(mcdc) # Output: generate hdf5 output files output_start = MPI.Wtime() - generate_hdf5(data_tally, mcdc) + generate_hdf5(mcdc) mcdc["runtime_output"] = MPI.Wtime() - output_start # Stop timer @@ -203,8 +203,8 @@ def calculate_cs_sparse_solution(data, mcdc, A, b): return sparse_solution -def cs_reconstruct(data, mcdc): - tally_bin = data +def cs_reconstruct(mcdc): + tally_bin = mcdc["data"] tally = mcdc["cs_tallies"][0] stride = tally["stride"] bin_idx = stride["tally"] @@ -214,8 +214,8 @@ def cs_reconstruct(data, mcdc): b = tally_bin[TALLY_SUM, bin_idx : bin_idx + N_cs_bins] - A, T_inv = calculate_cs_A(data, mcdc) - x = calculate_cs_sparse_solution(data, mcdc, A, b) + A, T_inv = calculate_cs_A(tally_bin, mcdc) + x = calculate_cs_sparse_solution(tally_bin, mcdc, A, b) recon = T_inv @ x recon_reshaped = recon.reshape(Ny, Nx) @@ -531,8 +531,8 @@ def prepare(): 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) + type_.make_type_global_prep(input_deck) kernel.adapt_rng(nb.config.DISABLE_JIT) input_deck.setting["target"] = config.target @@ -542,7 +542,7 @@ def prepare(): # TODO: Better alternative? # ========================================================================= - mcdc_arr = np.zeros(1, dtype=type_.global_) + mcdc_arr = np.zeros(1, dtype=type_.global_prep) mcdc = mcdc_arr[0] # Now, set up the global variable container @@ -1144,6 +1144,7 @@ def prepare(): tally_bin_N_copies = 5 tally_shape = (tally_bin_N_copies, tally_bin_size) + type_.make_type_global(input_deck,tally_shape) # ========================================================================= # Platform Targeting, Adapters, Toggles, etc @@ -1172,16 +1173,17 @@ def prepare(): # 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 + + for name in type_.global_.names: + if name == "data": + continue + mcdc_arr[0][name] = mcdc[name] + mcdc = mcdc_arr[0] mcdc["gpu_meta"]["global_pointer"] = mcdc_uint - mcdc["gpu_meta"]["tally_pointer"] = data_tally_uint + # ========================================================================= # Setting @@ -1584,13 +1586,13 @@ def prepare(): "w" ] - loop.setup_gpu(mcdc_arr, data_tally) + loop.setup_gpu(mcdc_arr) # ========================================================================= # Finalize data: wrapping into a tuple # ========================================================================= - return data_tally, mcdc_arr + return mcdc_arr def cardlist_to_h5group(dictlist, input_group, name): @@ -1642,13 +1644,15 @@ def dict_to_h5group(dict_, group): group[k] = v -def dd_mergetally(mcdc, data_tally): +def dd_mergetally(mcdc): """ Performs tally recombination on domain-decomposed mesh tallies. Gathers and re-organizes tally data into a single array as it would appear in a non-decomposed simulation. """ + data_tally = mcdc["data"] + # create bin for recomposed tallies d_Nx = input_deck.technique["dd_mesh"]["x"].size - 1 d_Ny = input_deck.technique["dd_mesh"]["y"].size - 1 @@ -1751,12 +1755,15 @@ def dd_mergetally(mcdc, data_tally): return dd_tally -def dd_mergemesh(mcdc, data_tally): +def dd_mergemesh(mcdc): """ Performs mesh recombination on domain-decomposed mesh tallies. Gathers and re-organizes mesh data into a single array as it would appear in a non-decomposed simulation. """ + + data_tally = mcdc["data"] + d_Nx = input_deck.technique["dd_mesh"]["x"].size - 1 d_Ny = input_deck.technique["dd_mesh"]["y"].size - 1 d_Nz = input_deck.technique["dd_mesh"]["z"].size - 1 @@ -1840,7 +1847,9 @@ def dd_mergemesh(mcdc, data_tally): return dd_mesh -def generate_hdf5(data_tally, mcdc): +def generate_hdf5(mcdc): + + data_tally = mcdc["data"] if mcdc["technique"]["domain_decomposition"]: dd_tally = dd_mergetally(mcdc, data_tally) @@ -2458,7 +2467,7 @@ def visualize( """ # TODO: add input error checkers - _, mcdc_container = prepare() + mcdc_container = prepare() mcdc = mcdc_container[0] # Color assignment for materials (by material ID) diff --git a/mcdc/type_.py b/mcdc/type_.py index 5a185e8cf..9c0910597 100644 --- a/mcdc/type_.py +++ b/mcdc/type_.py @@ -51,8 +51,10 @@ gpu_meta = None global_ = None +global_prep = None global_size = None + # ============================================================================== # MC/DC Member Array Sizes # ============================================================================== @@ -1467,7 +1469,6 @@ def make_type_gpu_meta(): ("source_program_pointer", uintp), ("precursor_program_pointer", uintp), ("global_pointer", uintp), - ("tally_pointer", uintp), ] ) @@ -1477,8 +1478,8 @@ def make_type_gpu_meta(): # ============================================================================== -def make_type_global(input_deck): - global global_, global_size +def make_global_fields(input_deck,tally_shape): + global global_, global_prep, global_size # Get modes mode_CE = input_deck.setting["mode_CE"] @@ -1557,8 +1558,7 @@ def make_type_global(input_deck): ) or input_deck.technique["iQMC"]: bank_source = particle_bank(N_work) - global_ = into_dtype( - [ + global_fields = [ ("nuclides", nuclide, (N_nuclide,)), ("materials", material, (N_material,)), ("surfaces", surface, (N_surface,)), @@ -1627,12 +1627,30 @@ def make_type_global(input_deck): ("mpi_work_iter", int64, (1,)), ("gpu_meta", gpu_meta), ("source_seed", uint64), + ("data",float64,tally_shape), ] - ) + + return global_fields - # GLobal type + +def make_type_global(input_deck,tally_shape): + global global_, global_size + + global_fields = make_global_fields(input_deck,tally_shape) + global_ = into_dtype(global_fields) + + # GLobal type global_size = global_.itemsize + + + +def make_type_global_prep(input_deck): + global global_prep + + global_fields = make_global_fields(input_deck,(1,)) + global_prep = into_dtype(global_fields[:-1]) + # ==============================================================================