diff --git a/CHANGELOG.md b/CHANGELOG.md index 083037ee9..76c232824 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/2.0.0/), ### Changed +- Update GPU transport support for the current MC/DC data model and Harmonize runtime, including GPU-compatible state access, particle-bank operations, array accessors, and torus intersections, from [@braxtoncuneo]. - Show previously published documentation versions in the documentation version switcher, from [@ilhamv] ### Deprecated @@ -203,3 +204,4 @@ The pre-refactor implementation remains available in the `cement` branch as a re [@gunnarrl]: https://github.com/gunnarrl [@Talen-Ayers]: https://github.com/Talen-Ayers [@steps-re]: https://github.com/steps-re +[@braxtoncuneo]: https://github.com/braxtoncuneo diff --git a/mcdc/code_factory/array_return.py b/mcdc/code_factory/array_return.py new file mode 100644 index 000000000..b793a4969 --- /dev/null +++ b/mcdc/code_factory/array_return.py @@ -0,0 +1,288 @@ +import cffi +import numba as nb +import numpy as np +from numba import jit, literal_unroll, njit, objmode, types +from numba.extending import intrinsic +import mcdc.config as config + +ffi = cffi.FFI() + + +# ============================================================================= +# uintp/voidptr casting helper functions - for internal use only +# ============================================================================= + + +@intrinsic +def cast_any_to_voidptr(typingctx, src): + # create the expected type signature + result_type = types.voidptr + sig = result_type(src) + + # defines the custom code generation + def codegen(context, builder, signature, args): + # llvm IRBuilder code here + [src] = args + rtype = signature.return_type + llrtype = context.get_value_type(rtype) + return builder.bitcast(src, llrtype) + + return sig, codegen + + +@intrinsic +def cast_uintp_to_voidptr(typingctx, src): + # check for accepted types + if isinstance(src, types.Integer): + # create the expected type signature + result_type = types.voidptr + sig = result_type(types.uintp) + + # defines the custom code generation + def codegen(context, builder, signature, args): + # llvm IRBuilder code here + [src] = args + rtype = signature.return_type + llrtype = context.get_value_type(rtype) + return builder.inttoptr(src, llrtype) + + return sig, codegen + + +@intrinsic +def cast_voidptr_to_uintp(typingctx, src): + # check for accepted types + if isinstance(src, types.RawPointer): + # create the expected type signature + result_type = types.uintp + sig = result_type(types.voidptr) + + # defines the custom code generation + def codegen(context, builder, signature, args): + # llvm IRBuilder code here + [src] = args + rtype = signature.return_type + llrtype = context.get_value_type(rtype) + return builder.ptrtoint(src, llrtype) + + return sig, codegen + + +@njit() +def voidptr_to_uintp(value): + return cast_voidptr_to_uintp(value) + + +@njit() +def into_voidptr(value): + return into_voidptr_python(value) + + +# ============================================================================= +# uintp/voidptr casting utility functions +# ============================================================================= + + +# Converts a pointer-sized integer to a void* +@njit() +def uintp_to_voidptr(value): + val = nb.uintp(value) + return cast_uintp_to_voidptr(val) + + +# Placeholder function for casting to void*. There currently is no use case +# for void* values in python mode for mcdc. +def into_voidptr_python(value): + raise RuntimeError("`into_voidptr` is only supported in nopython mode.") + + +@nb.extending.overload(into_voidptr_python) +def into_voidptr_overload(value): + + if isinstance(value, nb.types.Array): + + def impl(value): + ptr = ffi.from_buffer(value) + vptr = cast_any_to_voidptr(ptr) + return vptr + + return impl + elif isinstance(value, nb.types.CPointer): + + def impl(value): + return cast_any_to_voidptr(value) + + return impl + elif isinstance(value, nb.types.Integer): + + def impl(value): + return cast_uintp_to_voidptr(value) + + return impl + else: + raise RuntimeError(f"`into_voidptr` is not supported for type '{value}'") + + +############################################################################### +# Helper decorators, functions, and builtins for returning arrays +############################################################################### + + +# Overload target +def array_result(array): + return array + + +@nb.extending.overload(array_result) +def array_result_overload(array): + + if not isinstance(array, types.Array): + raise nb.core.errors.TypingError( + f"Expected array type argument for array_result, got {array}." + ) + + def impl(array): + return (into_voidptr(array), array.shape) + + return impl + + +# Raises an error if the context is not recognized +def context_guard(context): + if isinstance(context, nb.core.typing.context.Context): + pass + elif isinstance(context, nb.cuda.target.CUDATypingContext): + pass + elif isinstance(context, nb.hip.target.HIPTypingContext): + pass + else: + raise nb.core.errors.UnsupportedError(f"Unsupported target context {context}.") + + +# Typing for the `array_return` builtin. +def array_return_typing(fn, elem_type, ndim): + + from inspect import signature + + arg_list = ",".join([param for param in signature(fn).parameters]) + template = "def typer({arg_list}):\n return nb.types.Array(dtype=elem_type,ndim={ndim},layout='C')({arg_list})" + + gns = globals() | {"elem_type": elem_type} + lns = {} + exec(template.format(arg_list=arg_list, ndim=ndim), gns, lns) + typer = lns["typer"] + + def typer_factory(context): + from numba.np.numpy_support import as_dtype + + context_guard(context) + + return typer + + nb.extending.type_callable(fn)(typer_factory) + + +# The logic forthe `array_return` builtin +def array_return_lowering(fn, elem_type, ndim): + + from inspect import signature + + # The builtin returns an array with the given element + # type and the given dimensionality (default 1) + param_count = len(signature(fn).parameters) + retty = nb.types.Array(dtype=elem_type, ndim=ndim, layout="C") + sig = retty(*([nb.types.Any] * param_count)) + + jit_fn = nb.njit(fn) + + # This builtin effectively replaces the original decorated function whenever it + # is referenced in code. The original functions still exists, but it is called through + # this builtin which converts the pointer/shape tuple that the function (should) + # generate with `array_result` and return. + def builtin(context, builder, sig, args): + + import llvmlite.binding as ll + from llvmlite import ir + + lmod = builder.module + retty = nb.types.Tuple( + [nb.types.voidptr, nb.types.Tuple([nb.types.uintp] * ndim)] + ) + ptr_sig = retty(*sig.args) + + res = context.compile_internal(builder, jit_fn.py_func, ptr_sig, args) + ptr_res = builder.extract_value(res, 0) + size_res = builder.extract_value(res, 1) + shape = size_res + dtype = elem_type + + # GPU platforms require a `targetdata` for array construction, which is created + # slightly differently depending upon the platform. + if config.ROCM_AVAILABLE and isinstance( + context, nb.hip.target.HIPTargetContext + ): + targetdata = ll.create_target_data(nb.hip.amdgcn.DATA_LAYOUT) + elif config.CUDA_AVAILABLE and isinstance( + context, nb.cuda.target.CUDATargetContext + ): + targetdata = ll.create_target_data(nb.cuda.cudadrv.nvvm.NVVM().data_layout) + lldtype = context.get_data_type(dtype) + + # The size of the item is derived either from the lldtype or `targetdata` depending + # upon platform + if isinstance(context, nb.core.cpu.CPUContext): + itemsize = context.get_abi_sizeof(lldtype) + elif config.ROCM_AVAILABLE and isinstance( + context, nb.hip.target.HIPTargetContext + ): + itemsize = lldtype.get_abi_size(targetdata) + elif config.CUDA_AVAILABLE and isinstance( + context, nb.cuda.target.CUDATargetContext + ): + itemsize = lldtype.get_abi_size(targetdata) + else: + raise nb.core.errors.UnsupportedError( + f"Unsupported target context {context}." + ) + + # The number of elements-worth of bytes that must be skipped to advance by 1 element + # in a given dimension + kstrides = [context.get_constant(types.intp, itemsize)] + + # Create array structure based on the supplied type information + aryty = types.Array(dtype=elem_type, ndim=ndim, layout="C") + ary = context.make_array(aryty)(context, builder) + + # Array populating logic expects pointers to the array buffer to be expressed as + # a pointer to a byte in a generic address space. + dataptr = builder.addrspacecast( + ptr_res, ir.PointerType(ir.IntType(8)), "generic" + ) + + # Initialize the array structure with the data pointer, shape, and strides + kshape = size_res + context.populate_array( + ary, + data=builder.bitcast(dataptr, ary.data.type), + shape=kshape, + strides=kstrides, + itemsize=context.get_constant(types.intp, itemsize), + meminfo=None, + ) + return ary._getvalue() + + # To complete the illusion of the decorated function acting just like a normal + # `njit` function, the decorated function is overloaded as the builtin that + # was defined above. + nb.extending.lower_builtin(fn, *sig.args)(builtin) + + +# A function decorated with `array_return` may return an array by passing +# it through the `array_result` function and returning the output +def array_return(sig, ndim=1): + def array_return_true_decorator(fn): + array_return_typing(fn, sig, ndim) + array_return_lowering(fn, sig, ndim) + return fn + + return array_return_true_decorator diff --git a/mcdc/code_factory/gpu/program_builder.py b/mcdc/code_factory/gpu/program_builder.py index 3e5570c75..1129bfb73 100644 --- a/mcdc/code_factory/gpu/program_builder.py +++ b/mcdc/code_factory/gpu/program_builder.py @@ -1,11 +1,9 @@ import numba as nb import numba.extending as nbxt import numpy as np - from mpi4py import MPI #### - import mcdc.config as config # ====================================================================================== @@ -13,14 +11,12 @@ # ====================================================================================== +# Overwrites global symbols in other modules with gpu-compatible counterparts def adapt_transport_functions(): - global access_simulation import mcdc.code_factory.gpu.transport as gpu_transport import mcdc.transport as transport - transport.util.access_simulation = access_simulation - # TODO: Make the following automatic transport.geometry.interface.report_lost_particle = ( gpu_transport.geometry.interface.report_lost_particle @@ -159,6 +155,15 @@ def find_cell(program: nb.uintp, particle: particle_gpu): alloc_managed_bytes = harmonize.alloc_managed_bytes alloc_device_bytes = harmonize.alloc_device_bytes + from mcdc.transport import util + + @nb.extending.overload(util.access_simulation, target="hip") + def access_simulation_gpu_overload(program): + def impl(program): + return access_simulation(program) + + return impl + # ====================================================================================== # Program builder @@ -192,10 +197,9 @@ def build_gpu_program(data_size): import harmonize import mcdc.numba_types as type_ import mcdc.transport.util as util - from mcdc.transport.simulation import generate_source_particle, step_particle - global alloc_state, free_state + global access_simulation, alloc_state, free_state global alloc_program, free_program @@ -217,8 +221,7 @@ def make_work(program: nb.uintp) -> nb.boolean: data_ptr = access_data_ptr(program) data = harmonize.array_from_ptr(data_ptr, shape, nb.float64) - util.atomic_add(simulation["mpi_work_iter"], 0, 1) - idx_work = simulation["mpi_work_iter"][0] + idx_work = util.atomic_add(simulation["mpi_work_iter"], 0, 1) if idx_work >= simulation["mpi_work_size"]: return False @@ -252,6 +255,11 @@ def step(program: nb.uintp, particle_input: particle_gpu): particle_container = util.local_array(1, type_.particle) particle_container[0] = particle_input particle = particle_container[0] + particle["alive"] = True + particle["material_ID"] = -1 + particle["cell_ID"] = -1 + particle["surface_ID"] = -1 + particle["event"] = -1 particle["fresh"] = False step_particle(particle_container, program, data) if particle["alive"]: @@ -297,6 +305,17 @@ def step(program: nb.uintp, particle_input: particle_gpu): clear_flags = src_fns["clear_flags"] set_device = src_fns["set_device"] + alloc_program = src_fns["alloc_program"] + free_program = src_fns["free_program"] + init_program = src_fns["init_program"] + exec_program = src_fns["exec_program"] + complete = src_fns["complete"] + clear_flags = src_fns["clear_flags"] + set_device = src_fns["set_device"] + + alloc_managed_bytes = harmonize.alloc_managed_bytes + alloc_device_bytes = harmonize.alloc_device_bytes + # ====================================================================================== # Setup GPU @@ -314,7 +333,6 @@ def setup_gpu_program(simulation_container, data): set_device(device_id) simulation["gpu_meta"]["state_pointer"] = cast_voidptr_to_uintp(alloc_state()) - if config.gpu_state_storage == "separate": store_pointer_state_device_simulation( simulation["gpu_meta"]["state_pointer"], @@ -342,31 +360,6 @@ def teardown_gpu_program(simulation): free_state(cast_uintp_to_voidptr(simulation["gpu_meta"]["state_pointer"])) -# ====================================================================================== -# Simulation structure and data creators -# ====================================================================================== - - -def create_data_array(size, dtype): - if config.gpu_state_storage == "managed": - data_tally_ptr = harmonize.alloc_managed_bytes(size) - else: - data_tally_ptr = harmonize.alloc_device_bytes(size) - data_tally_uint = cast_voidptr_to_uintp(data_tally_ptr) - data_tally = nb.carray(data_tally_ptr, (size,), dtype) - return data_tally, data_tally_uint - - -def create_mcdc_container(dtype): - if config.gpu_state_storage == "managed": - mcdc_ptr = harmonize.alloc_managed_bytes(dtype.itemsize) - else: - mcdc_ptr = harmonize.alloc_device_bytes(dtype.itemsize) - mcdc_uint = cast_voidptr_to_uintp(mcdc_ptr) - mcdc_container = nb.carray(mcdc_ptr, (1,), dtype) - return mcdc_container, mcdc_uint - - # ====================================================================================== # Type casters # ====================================================================================== diff --git a/mcdc/code_factory/gpu/transport/simulation.py b/mcdc/code_factory/gpu/transport/simulation.py index 591ec618a..b9ba444c7 100644 --- a/mcdc/code_factory/gpu/transport/simulation.py +++ b/mcdc/code_factory/gpu/transport/simulation.py @@ -25,6 +25,9 @@ def source_loop(seed, simulation, data): full_work_size = simulation["mpi_work_size"] + if full_work_size == 0: + return + if settings["gpu_strategy"] == GPU_STRATEGY_ASYNC: phase_size = 1000000000 else: @@ -39,10 +42,10 @@ def source_loop(seed, simulation, data): # Store the global state to the GPU if settings["gpu_storage"] == GPU_STORAGE_SEPARATE: - harmonize.memcpy_host_to_device( + gpu_module.store_state_device_simulation( simulation["gpu_meta"]["state_pointer"], simulation ) - harmonize.memcpy_host_to_device( + gpu_module.store_state_device_data( simulation["gpu_meta"]["state_pointer"], data ) @@ -68,16 +71,14 @@ def source_loop(seed, simulation, data): gpu_module.clear_flags(simulation["gpu_meta"]["program_pointer"]) # Recover the original program state - if config.gpu_state_storage == "separate": - harmonize.memcpy_device_to_host( + if settings["gpu_storage"] == GPU_STORAGE_SEPARATE: + gpu_module.load_state_device_simulation( simulation, simulation["gpu_meta"]["state_pointer"] ) - harmonize.memcpy_device_to_host( + gpu_module.load_state_device_data( data, simulation["gpu_meta"]["state_pointer"] ) - gpu_module.clear_flags(simulation["gpu_meta"]["program_pointer"]) - simulation["mpi_work_size"] = full_work_size particle_bank_module.set_bank_size(simulation["bank_active"], 0) diff --git a/mcdc/code_factory/gpu/transport/util.py b/mcdc/code_factory/gpu/transport/util.py index 2acb013c8..74faddab1 100644 --- a/mcdc/code_factory/gpu/transport/util.py +++ b/mcdc/code_factory/gpu/transport/util.py @@ -5,9 +5,28 @@ from numba import njit, types -@njit def atomic_add(array, idx, value): - harmonize.array_atomic_add(array, idx, value) + result = array[idx] + array[idx] += value + return result + + +@nb.extending.overload(atomic_add, target="gpu") +def overload_atomic_add_gpu(array, idx, value): + def impl(array, idx, value): + return harmonize.array_atomic_add(array, idx, value) + + return impl + + +@nb.extending.overload(atomic_add, target="cpu") +def overload_atomic_add_cpu(array, idx, value): + def impl(array, idx, value): + result = array[idx] + array[idx] += value + return result + + return impl # ============================================================================= diff --git a/mcdc/code_factory/numba_layers_generator.py b/mcdc/code_factory/numba_layers_generator.py index 4dd538c10..673da8db8 100644 --- a/mcdc/code_factory/numba_layers_generator.py +++ b/mcdc/code_factory/numba_layers_generator.py @@ -1,23 +1,21 @@ from __future__ import annotations #### - import importlib +from pathlib import Path + import numba as nb import numpy as np from numba import njit from numba.extending import intrinsic -from pathlib import Path #### - import mcdc import mcdc.code_factory.gpu.program_builder as gpu_builder import mcdc.config as config import mcdc.object_ as object_module import mcdc.object_.base as base - from mcdc.object_.base import ( MCDCBase, MCDCObject, @@ -375,6 +373,8 @@ def generate_numba_layers(simulation): simulation_dtype ) mcdc_simulation = mcdc_simulation_container[0] + mcdc_simulation["gpu_meta"]["simulation_pointer"] = mcdc_simulation_pointer + mcdc_simulation["gpu_meta"]["data_pointer"] = data["pointer"] record = records["simulation"] structure = structures["simulation"] @@ -412,6 +412,9 @@ def generate_numba_layers(simulation): for name in bank_names: mcdc_simulation[name]["tag"] = getattr(simulation, name).tag + mcdc_simulation["gpu_meta"]["simulation_pointer"] = mcdc_simulation_pointer + mcdc_simulation["gpu_meta"]["data_pointer"] = data["pointer"] + # GPU program setup if config.target == "gpu": gpu_builder.setup_gpu_program(mcdc_simulation_container, data["array"]) @@ -761,18 +764,22 @@ def create_data_array(size): data = np.zeros(size, dtype=np.float64) return data, 0 else: - return create_data_array_on_gpu(size * 8) + return create_data_array_on_gpu(nb.types.float64, size, size * 16) @njit -def create_data_array_on_gpu(size): +def create_data_array_on_gpu(dtype, size, byte_size): if config.gpu_state_storage == "managed": - data_ptr = gpu_builder.alloc_managed_bytes(size) + data_tally_ptr = gpu_builder.alloc_managed_bytes(byte_size) + else: + data_tally_ptr = gpu_builder.alloc_device_bytes(byte_size) + data_tally_uint = cast_voidptr_to_uintp(data_tally_ptr) + + if config.gpu_state_storage == "separate": + data_tally = np.zeros((size,), dtype=dtype) else: - data_ptr = gpu_builder.alloc_device_bytes(size) - data_uint = voidptr_to_uintp(data_ptr) - data = nb.carray(data_ptr, (size,), dtype=np.float64) - return data, data_uint + data_tally = nb.carray(data_tally_ptr, (size,), dtype) + return data_tally, data_tally_uint def create_simulation_container(dtype): @@ -786,12 +793,16 @@ def create_simulation_container(dtype): @njit def create_simulation_container_on_gpu(dtype, size): if config.gpu_state_storage == "managed": - simulation_ptr = gpu_builder.alloc_managed_bytes(size) + mcdc_ptr = gpu_builder.alloc_managed_bytes(size * 8) + else: + mcdc_ptr = gpu_builder.alloc_device_bytes(size * 8) + mcdc_uint = cast_voidptr_to_uintp(mcdc_ptr) + + if config.gpu_state_storage == "separate": + mcdc_container = np.zeros((1,), dtype=dtype) else: - simulation_ptr = gpu_builder.alloc_device_bytes(size) - simulation_uint = voidptr_to_uintp(simulation_ptr) - simulation = nb.carray(simulation_ptr, (1,), dtype) - return simulation, simulation_uint + mcdc_container = nb.carray(mcdc_ptr, (1,), dtype) + return mcdc_container, mcdc_uint # ============================================================================= @@ -1053,7 +1064,9 @@ def generate_mcdc_access(targets): text_getter += "from numpy import int64\n" text_getter += "from numba import njit\n\n\n" text_setter += "from numba import njit\n\n\n" + text_getter += "from mcdc.code_factory.array_return import array_return, array_result\n\n\n" + text_getter += "import numba as nb\n\n\n" for attribute in targets[object_name]: attribute_name = attribute.name shape = attribute.shape @@ -1063,7 +1076,9 @@ def generate_mcdc_access(targets): text_getter += _accessor_1d_element( object_name, attribute_name, cast_to_int=cast_to_int ) - text_getter += _accessor_1d_all(object_name, attribute_name, shape[0]) + text_getter += _accessor_1d_all( + object_name, attribute_name, shape[0], nb.types.float64 + ) text_getter += _accessor_1d_last( object_name, attribute_name, @@ -1073,7 +1088,7 @@ def generate_mcdc_access(targets): text_setter += _accessor_1d_element(object_name, attribute_name, True) text_setter += _accessor_1d_all( - object_name, attribute_name, shape[0], True + object_name, attribute_name, shape[0], nb.types.float64, True ) text_setter += _accessor_1d_last( object_name, attribute_name, shape[0], True @@ -1137,6 +1152,7 @@ def generate_mcdc_access(targets): text = "# The following is automatically generated by code_factory.py\n\n" for i, object_name in enumerate(targets.keys()): text += f"import mcdc.mcdc_{key}.{object_name} as {object_name}\n" + text += "from mcdc.code_factory.array_return import array_return, array_result\n\n\n" if i < len(targets.keys()) - 1: text += "\n" f.write(text) @@ -1176,11 +1192,12 @@ def _accessor_1d_element(object_name, attribute_name, setter=False, cast_to_int= return text -def _accessor_1d_all(object_name, attribute_name, size, setter=False): - text = f"@njit\n" +def _accessor_1d_all(object_name, attribute_name, size, dtype, setter=False): if setter: + text = f"@njit\n" text += f"def {attribute_name}_all({object_name}, data, value):\n" else: + text = f"@array_return(nb.types.float64)\n" text += f"def {attribute_name}_all({object_name}, data):\n" text += f' start = {object_name}["{attribute_name}_offset"]\n' text += accessor_dimension("size", size, object_name) @@ -1188,7 +1205,7 @@ def _accessor_1d_all(object_name, attribute_name, size, setter=False): if setter: text += f" data[start:end] = value\n\n\n" else: - text += f" return data[start:end]\n\n\n" + text += f" return array_result(data[start:end])\n\n\n" return text @@ -1211,19 +1228,20 @@ def _accessor_1d_last( def _accessor_chunk(object_name, attribute_name, setter=False): - text = f"@njit\n" if setter: + text = f"@njit\n" text += ( f"def {attribute_name}_chunk(start, length, {object_name}, data, value):\n" ) else: + text = f"@array_return(nb.types.float64)\n" text += f"def {attribute_name}_chunk(start, length, {object_name}, data):\n" text += f' start += {object_name}["{attribute_name}_offset"]\n' text += f" end = start + length\n" if setter: text += f" data[start:end] = value\n\n\n" else: - text += f" return data[start:end]\n\n\n" + text += f" return array_result(data[start:end])\n\n\n" return text @@ -1247,10 +1265,11 @@ def _accessor_2d_element( def _accessor_2d_vector(object_name, attribute_name, stride, setter=False): - text = f"@njit\n" if setter: + text = f"@njit\n" text += f"def {attribute_name}_vector(index_1, {object_name}, data, value):\n" else: + text = f"@array_return(nb.types.float64)\n" text += f"def {attribute_name}_vector(index_1, {object_name}, data):\n" text += f' offset = {object_name}["{attribute_name}_offset"]\n' text += accessor_dimension("stride", stride, object_name) @@ -1259,7 +1278,7 @@ def _accessor_2d_vector(object_name, attribute_name, stride, setter=False): if setter: text += f" data[start:end] = value\n\n\n" else: - text += f" return data[start:end]\n\n\n" + text += f" return array_result(data[start:end])\n\n\n" return text diff --git a/mcdc/config.py b/mcdc/config.py index d841ea567..e0d9067f4 100644 --- a/mcdc/config.py +++ b/mcdc/config.py @@ -114,6 +114,28 @@ def _build_parser() -> argparse.ArgumentParser: clear_cache = args.clear_cache +# ====================================================================================== +# Flags for GPU platform availability +# ====================================================================================== + +try: + import numba.hip as hip + + ROCM_AVAILABLE = True +except: + ROCM_AVAILABLE = False + +if not ROCM_AVAILABLE: + try: + import numba.cuda as cuda + + CUDA_AVAILABLE = True + except: + CUDA_AVAILABLE = False +else: + CUDA_AVAILABLE = False + + # ====================================================================================== # Simulation-setting overrides # ====================================================================================== diff --git a/mcdc/main.py b/mcdc/main.py index 2497571d1..0b2e85fa2 100644 --- a/mcdc/main.py +++ b/mcdc/main.py @@ -1,4 +1,5 @@ from mcdc.object_.simulation import Simulation +import mcdc.config as config # ====================================================================================== # Run Simulation @@ -136,10 +137,9 @@ def prepare(simulationPy: Simulation): simulation = simulation_container[0] # Pick Python-version RNG if needed - import mcdc.config as config - import mcdc.transport.rng as rng - if config.mode == "python": + import mcdc.transport.rng as rng + rng.wrapping_add = rng.wrapping_add_python rng.wrapping_mul = rng.wrapping_mul_python @@ -183,8 +183,6 @@ def prepare(simulationPy: Simulation): def finalize(simulation): - import mcdc.config as config - # GPU teardowns if needed if config.target == "gpu": from mcdc.code_factory.gpu.program_builder import teardown_gpu_program diff --git a/mcdc/numba_types.py b/mcdc/numba_types.py index 28559d832..fd2c161e2 100644 --- a/mcdc/numba_types.py +++ b/mcdc/numba_types.py @@ -28,8 +28,8 @@ ('cell_ID', int64), ('material_ID', int64), ('surface_ID', int64), - ('alive', bool), - ('fresh', bool), + ('alive', bool_), + ('fresh', bool_), ('event', int64), ('x', float64), ('y', float64), @@ -52,8 +52,8 @@ ('surface_IDs_offset', int64), ('fill_type', int64), ('fill_ID', int64), - ('fill_translated', bool), - ('fill_rotated', bool), + ('fill_translated', bool_), + ('fill_rotated', bool_), ('translation', float64, (3,)), ('rotation', float64, (3,)), ('N_collision_tally', int64), @@ -82,8 +82,8 @@ material = into_dtype([ ('name', 'U32'), ('temperature', float64), - ('fissionable', bool), - ('has_neutron_multigroup', bool), + ('fissionable', bool_), + ('has_neutron_multigroup', bool_), ('neutron_multigroup_ID', int64), ('N_nuclide', int64), ('nuclide_IDs_offset', int64), @@ -97,9 +97,9 @@ ]) collision_tally = into_dtype([ - ('cell_filtered', bool), + ('cell_filtered', bool_), ('cell_filter_ID', int64), - ('mesh_filtered', bool), + ('mesh_filtered', bool_), ('mesh_filter_type', int64), ('mesh_filter_ID', int64), ('mesh_stride_z', int64), @@ -110,9 +110,9 @@ ]) tracklength_tally = into_dtype([ - ('cell_filtered', bool), + ('cell_filtered', bool_), ('cell_filter_ID', int64), - ('mesh_filtered', bool), + ('mesh_filtered', bool_), ('mesh_filter_type', int64), ('mesh_filter_ID', int64), ('mesh_stride_z', int64), @@ -374,7 +374,7 @@ ('chi_p_length', int64), ('chi_d_offset', int64), ('chi_d_length', int64), - ('fissionable', bool), + ('fissionable', bool_), ('ID', int64), ]) @@ -384,7 +384,7 @@ ('atomic_number', int64), ('mass_number', int64), ('atomic_weight_ratio', float64), - ('fissionable', bool), + ('fissionable', bool_), ('excitation_level', int64), ('neutron_xs_energy_grid_offset', int64), ('neutron_xs_energy_grid_length', int64), @@ -518,53 +518,53 @@ ('N_active', int64), ('N_cycle', int64), ('k_init', float64), - ('use_gyration_radius', bool), + ('use_gyration_radius', bool_), ('gyration_radius_type', int64), - ('use_source_file', bool), + ('use_source_file', bool_), ('source_file_name', 'U32'), ('time_boundary', float64), ('output_name', 'U32'), - ('use_progress_bar', bool), + ('use_progress_bar', bool_), ('N_census', int64), ('census_time_offset', int64), ('census_time_length', int64), - ('use_census_based_tally', bool), + ('use_census_based_tally', bool_), ('census_tally_frequency', int64), - ('save_particle', bool), + ('save_particle', bool_), ('active_bank_buffer', int64), ('census_bank_buffer_ratio', float64), ('source_bank_buffer_ratio', float64), ('future_bank_buffer_ratio', float64), - ('neutron_transport', bool), - ('electron_transport', bool), - ('proton_transport', bool), - ('neutron_eigenvalue_mode', bool), + ('neutron_transport', bool_), + ('electron_transport', bool_), + ('proton_transport', bool_), + ('neutron_eigenvalue_mode', bool_), ('gpu_strategy', int64), ('gpu_async_type', int64), ('gpu_storage', int64), ]) neutron_multigroup = into_dtype([ - ('hybrid', bool), + ('hybrid', bool_), ]) implicit_capture = into_dtype([ - ('active', bool), + ('active', bool_), ]) weighted_emission = into_dtype([ - ('active', bool), + ('active', bool_), ('weight_target', float64), ]) global_weight_roulette = into_dtype([ - ('active', bool), + ('active', bool_), ('weight_threshold', float64), ('weight_target', float64), ]) weight_windows = into_dtype([ - ('active', bool), + ('active', bool_), ('energy_bounds_offset', int64), ('energy_bounds_length', int64), ('Ne', int64), @@ -581,7 +581,7 @@ ]) population_control = into_dtype([ - ('active', bool), + ('active', bool_), ]) technique = into_dtype([ @@ -595,10 +595,10 @@ source = into_dtype([ ('name', 'U32'), - ('point_source', bool), - ('uniform_x', bool), - ('uniform_y', bool), - ('uniform_z', bool), + ('point_source', bool_), + ('uniform_x', bool_), + ('uniform_y', bool_), + ('uniform_z', bool_), ('point', float64, (3,)), ('x', float64, (2,)), ('y', float64, (2,)), @@ -606,23 +606,23 @@ ('x_pdf_ID', int64), ('y_pdf_ID', int64), ('z_pdf_ID', int64), - ('isotropic_direction', bool), - ('mono_direction', bool), - ('white_direction', bool), + ('isotropic_direction', bool_), + ('mono_direction', bool_), + ('white_direction', bool_), ('direction', float64, (3,)), ('polar_cosine', float64, (2,)), ('azimuthal', float64, (2,)), - ('mono_energetic', bool), - ('discrete_energy', bool), + ('mono_energetic', bool_), + ('discrete_energy', bool_), ('energy', float64), ('energy_pdf_ID', int64), ('energy_pmf_ID', int64), - ('discrete_time', bool), + ('discrete_time', bool_), ('time', float64), ('time_range', float64, (2,)), ('particle_type', int64), ('probability', float64), - ('moving', bool), + ('moving', bool_), ('N_move', int64), ('N_move_grid', int64), ('move_velocities_offset', int64), @@ -652,13 +652,13 @@ ('J', float64), ('R', float64), ('r', float64), - ('linear', bool), - ('quadric', bool), - ('quartic', bool), + ('linear', bool_), + ('quadric', bool_), + ('quartic', bool_), ('nx', float64), ('ny', float64), ('nz', float64), - ('moving', bool), + ('moving', bool_), ('N_move', int64), ('N_move_grid', int64), ('move_velocities_offset', int64), @@ -675,9 +675,9 @@ ]) surface_crossing_tally = into_dtype([ - ('surface_filtered', bool), + ('surface_filtered', bool_), ('surface_filter_ID', int64), - ('cell_filtered', bool), + ('cell_filtered', bool_), ('cell_filter_ID', int64), ('ID', int64), ('base_ID', int64), @@ -688,9 +688,9 @@ ('scores_offset', int64), ('scores_length', int64), ('particle_type', int64), - ('filter_direction', bool), - ('filter_energy', bool), - ('filter_time', bool), + ('filter_direction', bool_), + ('filter_energy', bool_), + ('filter_time', bool_), ('mu_offset', int64), ('mu_length', int64), ('azi_offset', int64), @@ -849,10 +849,10 @@ def make_simulation_type(N: dict): ('eigenvalue_tally_C', float64, (1,)), ('gyration_radius_offset', int64), ('gyration_radius_length', int64), - ('cycle_active', bool), + ('cycle_active', bool_), ('mpi_size', int64), ('mpi_rank', int64), - ('mpi_master', bool), + ('mpi_master', bool_), ('mpi_work_start', int64), ('mpi_work_size', int64), ('mpi_work_size_total', int64), diff --git a/mcdc/transport/data.py b/mcdc/transport/data.py index f0894d1c3..4e4ba387e 100644 --- a/mcdc/transport/data.py +++ b/mcdc/transport/data.py @@ -39,10 +39,7 @@ def evaluate_data(x, data_, simulation, data): @njit def evaluate_table(x, table, data): - offset = table["x_offset"] - length = table["x_length"] - grid = data[offset : offset + length] - # Above is equivalent to: grid = mcdc_get.table_data.x_all(table, data) + grid = mcdc_get.table_data.x_all(table, data) idx = find_bin(x, grid) x1 = grid[idx] @@ -69,15 +66,9 @@ def evaluate_table(x, table, data): @njit def get_table_interpolation_law(idx, table, data) -> int: """Return the interpolation law for interval [idx, idx + 1].""" - offset = table["interpolation_boundaries_offset"] - length = table["interpolation_boundaries_length"] - boundaries = data[offset : offset + length] - # Above is equivalent to: boundaries = mcdc_get.table_data.interpolation_boundaries_all(table, data) + boundaries = mcdc_get.table_data.interpolation_boundaries_all(table, data) - offset = table["interpolations_offset"] - length = table["interpolations_length"] - interpolations = data[offset : offset + length] - # Above is equivalent to: interpolations = mcdc_get.table_data.interpolations_all(table, data) + interpolations = mcdc_get.table_data.interpolations_all(table, data) # Boundaries are exclusive upper point indices. upper_point = idx + 1 @@ -91,10 +82,7 @@ def get_table_interpolation_law(idx, table, data) -> int: @njit def evaluate_polynomial(x, polynomial, data): - offset = polynomial["coefficients_offset"] - length = polynomial["coefficients_length"] - coeffs = data[offset : offset + length] - # Above is equivalent to: coeffs = mcdc_get.polynomial_data.coefficients_all(polynomial, data) + coeffs = mcdc_get.polynomial_data.coefficients_all(polynomial, data) total = 0.0 for i in range(len(coeffs)): diff --git a/mcdc/transport/distribution.py b/mcdc/transport/distribution.py index 190f3ae04..d32d0f2d5 100644 --- a/mcdc/transport/distribution.py +++ b/mcdc/transport/distribution.py @@ -247,11 +247,7 @@ def invert_tabulated_segment(xi, c0, v0, v1, p0, p1, interpolation): @njit def sample_pmf(pmf, rng_state, data): xi = rng.lcg(rng_state) - - offset = pmf["cmf_offset"] - length = pmf["cmf_length"] - cmf = data[offset : offset + length] - # Above is equivalent to: cmf = mcdc_get.pmf_distribution.cmf_all(pmf, data) + cmf = mcdc_get.pmf_distribution.cmf_all(pmf, data) idx = find_bin(xi, cmf) return mcdc_get.pmf_distribution.value(idx, pmf, data) @@ -297,10 +293,7 @@ def _sample_multi_table(E, rng_state, multi_table, simulation, data, scale): """Sample from a multi-table distribution.""" # Get the grid - offset = multi_table["grid_offset"] - length = multi_table["grid_length"] - grid = data[offset : offset + length] - # Above is equivalent to: grid = mcdc_get.multi_table_distribution.grid_all(multi_table, data) + grid = mcdc_get.multi_table_distribution.grid_all(multi_table, data) # Helper flag for scaling later use_next_table = False @@ -440,10 +433,7 @@ def sample_evaporation(E, rng_state, evaporation, simulation, data): @njit def sample_kalbach_mann(E, rng_state, kalbach_mann, data): - offset = kalbach_mann["energy_offset"] - length = kalbach_mann["energy_length"] - grid = data[offset : offset + length] - # Above is equivalent to: grid = mcdc_get.kalbach_mann_distribution.energy_all(kalbach_mann, data) + grid = mcdc_get.kalbach_mann_distribution.energy_all(kalbach_mann, data) # Random numbers xi1 = rng.lcg(rng_state) @@ -493,9 +483,7 @@ def sample_kalbach_mann(E, rng_state, kalbach_mann, data): size = end - start # The CDF - offset = kalbach_mann["cdf_offset"] - cdf = data[start + offset : start + offset + size] - # Above is equivalent to: cdf = mcdc_get.kalbach_mann_distribution.cdf_chunk(start, size, kalbach_mann, data) + cdf = mcdc_get.kalbach_mann_distribution.cdf_chunk(start, size, kalbach_mann, data) # Sample bin index idx = find_bin(xi2, cdf) @@ -544,10 +532,7 @@ def sample_kalbach_mann(E, rng_state, kalbach_mann, data): @njit def sample_tabulated_energy_angle(E, rng_state, table, data): - offset = table["energy_offset"] - length = table["energy_length"] - grid = data[offset : offset + length] - # Above is equivalent to: grid = mcdc_get.tabulated_energy_angle_distribution.energy_all(table, data) + grid = mcdc_get.tabulated_energy_angle_distribution.energy_all(table, data) # Random numbers xi1 = rng.lcg(rng_state) @@ -600,12 +585,9 @@ def sample_tabulated_energy_angle(E, rng_state, table, data): size = end - start # The CDF - offset = table["cdf_offset"] - cdf = data[start + offset : start + offset + size] - # Above is equivalent to: - # cdf = mcdc_get.tabulated_energy_angle_distribution.cdf_chunk( - # start, size, table, data - # ) + cdf = mcdc_get.tabulated_energy_angle_distribution.cdf_chunk( + start, size, table, data + ) # Sample bin index idx = find_bin(xi2, cdf) @@ -653,12 +635,9 @@ def sample_tabulated_energy_angle(E, rng_state, table, data): size = end - start # The CDF - offset = table["cosine_cdf_offset"] - cdf = data[start + offset : start + offset + size] - # Above is equivalent to: - # cdf = mcdc_get.tabulated_energy_angle_distribution.cosine_cdf_chunk( - # start, size, table, data - # ) + cdf = mcdc_get.tabulated_energy_angle_distribution.cosine_cdf_chunk( + start, size, table, data + ) # Sample bin index idx = find_bin(xi3, cdf) diff --git a/mcdc/transport/geometry/surface/interface.py b/mcdc/transport/geometry/surface/interface.py index 62e31604b..02f7836ce 100644 --- a/mcdc/transport/geometry/surface/interface.py +++ b/mcdc/transport/geometry/surface/interface.py @@ -394,12 +394,7 @@ def _get_move_idx(t, surface, data): """ Get moving interval index wrt the given time """ - time_grid = data[ - surface["move_time_grid_offset"] : ( - surface["move_time_grid_offset"] + surface["N_move_grid"] - ) - ] - # Above is equivalent to: time_grid = mcdc_get.surface.move_time_grid_all(surface, data) + time_grid = mcdc_get.surface.move_time_grid_all(surface, data) tolerance = COINCIDENCE_TOLERANCE_TIME go_lower = False idx = find_bin_with_rules(t, time_grid, tolerance, go_lower) @@ -419,14 +414,10 @@ def _translate_particle_position(particle_container, surface, idx, data): particle = particle_container[0] # Surface move translations - start = surface["move_translations_offset"] + idx * 3 - trans_0 = data[start : start + 3] - # Above is equivalent to: trans_0 = mcdc_get.surface.move_translations_vector(idx, surface, data) + trans_0 = mcdc_get.surface.move_translations_vector(idx, surface, data) # Surface move velocities - start = surface["move_velocities_offset"] + idx * 3 - V = data[start : start + 3] - # Above is equivalent to: V = mcdc_get.surface.move_velocities_vector(idx, surface, data) + V = mcdc_get.surface.move_velocities_vector(idx, surface, data) # Surface move time grid time_0 = mcdc_get.surface.move_time_grid(idx, surface, data) @@ -446,9 +437,7 @@ def _translate_particle_direction(particle_container, speed, surface, idx, data) particle = particle_container[0] # Surface move velocities - start = surface["move_velocities_offset"] + idx * 3 - V = data[start : start + 3] - # Above is equivalent to: V = mcdc_get.surface.move_velocities_vector(idx, surface, data) + V = mcdc_get.surface.move_velocities_vector(idx, surface, data) # Translate the particle particle["ux"] -= V[0] / speed diff --git a/mcdc/transport/geometry/surface/torus.py b/mcdc/transport/geometry/surface/torus.py index 2dcf4e584..239c16d84 100644 --- a/mcdc/transport/geometry/surface/torus.py +++ b/mcdc/transport/geometry/surface/torus.py @@ -29,6 +29,9 @@ from numba import njit +import mcdc.transport.util as util +import mcdc.transport.geometry.surface.torus_root_solver as torus_root_solver + from mcdc.constant import ( COINCIDENCE_TOLERANCE, INF, @@ -164,12 +167,14 @@ def get_distance(particle_container, surface): a0 = (I + R * R - r * r) ** 2 - 4.0 * R * R * L # TODO: May replace with a fully numba-native quartic solver if torus performance becomes important; - # np.roots is sufficient for now. - coefficients = np.array( - [a4 + 0.0j, a3 + 0.0j, a2 + 0.0j, a1 + 0.0j, a0 + 0.0j], - dtype=np.complex128, - ) - roots = np.roots(coefficients) + coefficients = util.local_array(5, np.complex128) + coefficients[0] = a0 + 0.0j + coefficients[1] = a1 + 0.0j + coefficients[2] = a2 + 0.0j + coefficients[3] = a3 + 0.0j + coefficients[4] = a4 + 0.0j + roots = util.local_array(4, np.complex128) + torus_root_solver.solve_quartic(coefficients, roots) min_t = INF diff --git a/mcdc/transport/geometry/surface/torus_root_solver.py b/mcdc/transport/geometry/surface/torus_root_solver.py new file mode 100644 index 000000000..5f72a8483 --- /dev/null +++ b/mcdc/transport/geometry/surface/torus_root_solver.py @@ -0,0 +1,191 @@ +import math +import cmath +import numpy as np +import numba as nb + +import mcdc.transport.util as util + +from numba import njit + + +@njit() +def modulus(x): + return math.sqrt(x.real**2 + x.imag**2) + + +@njit() +def sqrt(x): + r = modulus(x) + real_part = math.sqrt((r + x.real) / 2.0) + if x.imag < 0: + imag_part = -math.sqrt((r - x.real) / 2.0) + else: + imag_part = math.sqrt((r - x.real) / 2.0) + + return complex(real_part, imag_part) + + +@njit() +def power(x, n): + result = 1 + for i in range(n): + result = result * x + return result + + +@njit() +def nth_root(x, n, index): + # First, convert to polar form + r = modulus(x) + a = math.atan2(x.imag, x.real) + + # Apply de Moivre's Formula + root_modulus = math.pow(r, 1.0 / n) + root_argument = (a + 2 * math.pi * index) / n + + # ...then convert back to rectangular form + real = root_modulus * math.cos(root_argument) + imag = root_modulus * math.sin(root_argument) + return complex(real, imag) + + +@njit() +def principal_nth_root(x, n): + result = nth_root(x, n, 0) + return result + + +@njit() +def solve_quadratic(coeff, roots): + a = coeff[2] + b = coeff[1] + c = coeff[0] + # standard quadratic formula, but with discriminant + # calculated separately for re-use + discriminant = sqrt(power(b, 2) - 4 * a * c) + roots[0] = ((-b) + discriminant) / (2 * a) + roots[1] = ((-b) - discriminant) / (2 * a) + + +@njit() +def solve_biquadratic(coeff, roots): + # Move each coefficient down to one-half it's power + coeff[1] = coeff[2] + coeff[2] = coeff[4] + + # Solve as quadratic equation, where the variable is + # actually x^2 + solve_quadratic(coeff, roots) + + # Yield roots for x by taking square roots of the x^2 + # solution. + roots[3] = sqrt(roots[1]) + roots[2] = -roots[3] + roots[1] = sqrt(roots[0]) + roots[0] = -roots[1] + + # Restore the original positions of the coefficients + coeff[4] = coeff[2] + coeff[2] = coeff[1] + coeff[1] = 0.0j + + +@njit() +def solve_cubic(coeff, roots): + # TODO + # General soluton not needed for quartic solve + pass + + +@njit() +def solve_depressed_quartic(coeff, roots): + a = coeff[2] + b = coeff[1] + c = coeff[0] + + # To solve the depressed quartic, one must first find one + # root of a cubic polynomial. + + p = (-power(a, 2) / 12.0) - c + q = (-power(a, 3) / 108.0) + (a * c / 3.0) - (power(b, 2) / 8.0) + + cube_const = -q / 2.0 + sqrt_body = (power(q, 2) / 4.0) + (power(p, 3) / 27.0) + w_pos = principal_nth_root(cube_const + sqrt(sqrt_body), 3) + w_neg = principal_nth_root(cube_const - sqrt(sqrt_body), 3) + + # It's reccomended to opt for the larger w when + # calculating the root. + if abs(w_pos) > abs(w_neg): + w = w_pos + else: + w = w_neg + + # A root of the cubic + y = (a / 6.0) + w - (p / (3.0 * w)) + + # The different roots are found by flipping the signs + # of some terms in a formula. There are three sections + # unaffected by these flips, represented below by + # alpha, beta, and gamma + + alpha = sqrt(2.0 * y - a) + beta = -2.0 * y - a + gamma = (2.0 * b) / sqrt(2.0 * y - a) + + roots[0] = ((-alpha) + sqrt(beta + gamma)) / 2.0 # - + + + roots[1] = ((-alpha) - sqrt(beta + gamma)) / 2.0 # - - + + roots[2] = ((alpha) + sqrt(beta - gamma)) / 2.0 # + + - + roots[3] = ((alpha) - sqrt(beta - gamma)) / 2.0 # + - - + + +@njit() +def solve_quartic(coeff, roots): + # Algorithm logic derived from Wikipedia's quartic + # equation article. (^-^)=b + + # Coefficients for the general solve + a = coeff[4] + b = coeff[3] + c = coeff[2] + d = coeff[1] + e = coeff[0] + + # Coefficients for the sub-solve + # The sub-solve turns the equation into a depressed + # quartic by making u the new variable, with: + # + # x = u - (bg/(4*ag)) + # + # Once the roots for the depressed quartic are found, + # they can be plugged into this equation to yeild the + # roots for x. + + sub_coeff = util.local_array(5, np.complex128) + sub_coeff[4] = 1.0 + 0.0j + sub_coeff[3] = 0.0j + sub_coeff[2] = (-3.0 * power(b, 2)) / (8.0 * power(a, 2)) + c / a + sub_coeff[1] = ( + power(b, 3) / (8.0 * power(a, 3)) - (b * c) / (2.0 * power(a, 2)) + d / a + ) + sub_coeff[0] = ( + (-3.0 * power(b, 4)) / (256.0 * power(a, 4)) + + (c * power(b, 2)) / (16.0 * power(a, 3)) + - (b * d) / (4.0 * power(a, 2)) + + e / a + ) + + # Get roots of sub-solve + sub_roots = util.local_array(4, np.complex128) + if sub_coeff[1] == 0: + # If the linear term coefficient is zero, the + # normal depressed quartic solver won't work. + # Instead, it is a biquadratic, and can be + # solved as such. + + solve_biquadratic(sub_coeff, sub_roots) + else: + solve_depressed_quartic(sub_coeff, sub_roots) + + for idx in range(4): + roots[idx] = sub_roots[idx] - b / (4.0 * a) diff --git a/mcdc/transport/geometry/surface/torus_x.py b/mcdc/transport/geometry/surface/torus_x.py index 79134d192..f238b6be5 100644 --- a/mcdc/transport/geometry/surface/torus_x.py +++ b/mcdc/transport/geometry/surface/torus_x.py @@ -15,6 +15,9 @@ from numba import njit +import mcdc.transport.util as util +import mcdc.transport.geometry.surface.torus_root_solver as torus_root_solver + from mcdc.constant import ( COINCIDENCE_TOLERANCE, INF, @@ -189,11 +192,14 @@ def get_distance(particle_container, surface): # TODO: May replace with a fully numba-native quartic solver if torus performance becomes important; # np.roots is sufficient for now. - coefficients = np.array( - [a4 + 0.0j, a3 + 0.0j, a2 + 0.0j, a1 + 0.0j, a0 + 0.0j], - dtype=np.complex128, - ) - roots = np.roots(coefficients) + coefficients = util.local_array(5, np.complex128) + coefficients[0] = a0 + 0.0j + coefficients[1] = a1 + 0.0j + coefficients[2] = a2 + 0.0j + coefficients[3] = a3 + 0.0j + coefficients[4] = a4 + 0.0j + roots = util.local_array(4, np.complex128) + torus_root_solver.solve_quartic(coefficients, roots) min_t = INF diff --git a/mcdc/transport/geometry/surface/torus_y.py b/mcdc/transport/geometry/surface/torus_y.py index 571771b6e..ed7594e4d 100644 --- a/mcdc/transport/geometry/surface/torus_y.py +++ b/mcdc/transport/geometry/surface/torus_y.py @@ -15,6 +15,9 @@ from numba import njit +import mcdc.transport.util as util +import mcdc.transport.geometry.surface.torus_root_solver as torus_root_solver + from mcdc.constant import ( COINCIDENCE_TOLERANCE, INF, @@ -189,11 +192,14 @@ def get_distance(particle_container, surface): # TODO: May replace with a fully numba-native quartic solver if torus performance becomes important; # np.roots is sufficient for now. - coefficients = np.array( - [a4 + 0.0j, a3 + 0.0j, a2 + 0.0j, a1 + 0.0j, a0 + 0.0j], - dtype=np.complex128, - ) - roots = np.roots(coefficients) + coefficients = util.local_array(5, np.complex128) + coefficients[0] = a0 + 0.0j + coefficients[1] = a1 + 0.0j + coefficients[2] = a2 + 0.0j + coefficients[3] = a3 + 0.0j + coefficients[4] = a4 + 0.0j + roots = util.local_array(4, np.complex128) + torus_root_solver.solve_quartic(coefficients, roots) min_t = INF diff --git a/mcdc/transport/geometry/surface/torus_z.py b/mcdc/transport/geometry/surface/torus_z.py index a4c906cfa..a64ab1e5a 100644 --- a/mcdc/transport/geometry/surface/torus_z.py +++ b/mcdc/transport/geometry/surface/torus_z.py @@ -15,6 +15,9 @@ from numba import njit +import mcdc.transport.util as util +import mcdc.transport.geometry.surface.torus_root_solver as torus_root_solver + from mcdc.constant import ( COINCIDENCE_TOLERANCE, INF, @@ -189,11 +192,14 @@ def get_distance(particle_container, surface): # TODO: May replace with a fully numba-native quartic solver if torus performance becomes important; # np.roots is sufficient for now. - coefficients = np.array( - [a4 + 0.0j, a3 + 0.0j, a2 + 0.0j, a1 + 0.0j, a0 + 0.0j], - dtype=np.complex128, - ) - roots = np.roots(coefficients) + coefficients = util.local_array(5, np.complex128) + coefficients[0] = a0 + 0.0j + coefficients[1] = a1 + 0.0j + coefficients[2] = a2 + 0.0j + coefficients[3] = a3 + 0.0j + coefficients[4] = a4 + 0.0j + roots = util.local_array(4, np.complex128) + torus_root_solver.solve_quartic(coefficients, roots) min_t = INF diff --git a/mcdc/transport/mesh/structured.py b/mcdc/transport/mesh/structured.py index 7e3ad989f..3a41810b0 100644 --- a/mcdc/transport/mesh/structured.py +++ b/mcdc/transport/mesh/structured.py @@ -2,6 +2,7 @@ #### +import mcdc.mcdc_get as mcdc_get from mcdc.constant import COINCIDENCE_TOLERANCE, INF from mcdc.transport.util import find_bin_with_rules @@ -21,24 +22,9 @@ def get_indices(particle_container, structured_mesh, data): uy = particle["uy"] uz = particle["uz"] - grid_x = data[ - structured_mesh["x_offset"] : ( - structured_mesh["x_offset"] + structured_mesh["x_length"] - ) - ] - # Above is equivalent to: grid_x = mcdc_get.structured_mesh.x_all(structured_mesh, data) - grid_y = data[ - structured_mesh["y_offset"] : ( - structured_mesh["y_offset"] + structured_mesh["y_length"] - ) - ] - # Above is equivalent to: grid_y = mcdc_get.structured_structured_mesh.y_all(structured_mesh, data) - grid_z = data[ - structured_mesh["z_offset"] : ( - structured_mesh["z_offset"] + structured_mesh["z_length"] - ) - ] - # Above is equivalent to: grid_z = mcdc_get.structured_structured_mesh.z_all(structured_mesh, data) + grid_x = mcdc_get.structured_mesh.x_all(structured_mesh, data) + grid_y = mcdc_get.structured_mesh.y_all(structured_mesh, data) + grid_z = mcdc_get.structured_mesh.z_all(structured_mesh, data) tolerance = COINCIDENCE_TOLERANCE ux_go_lower = ux < 0.0 diff --git a/mcdc/transport/particle_bank.py b/mcdc/transport/particle_bank.py index 9736c2aa2..c697305e3 100644 --- a/mcdc/transport/particle_bank.py +++ b/mcdc/transport/particle_bank.py @@ -35,7 +35,8 @@ def set_bank_size(bank, value): @njit def add_bank_size(bank, value): - util.atomic_add(bank["size"], 0, value) + # Perform atomic increment to the bank size; return the initial size + return util.atomic_add(bank["size"], 0, value) # ============================================================================= @@ -50,7 +51,8 @@ def _bank_particle(particle_container, bank): report_full_bank(bank) # Set particle data - idx = get_bank_size(bank) + idx = add_bank_size(bank, 1) + particle_module.copy(bank["particle_data"][idx : idx + 1], particle_container) @@ -60,9 +62,6 @@ def bank_active_particle(particle_container, program): bank = simulation["bank_active"] _bank_particle(particle_container, bank) - # Increment bank size - add_bank_size(bank, 1) - @njit def bank_census_particle(particle_container, program): @@ -70,9 +69,6 @@ def bank_census_particle(particle_container, program): bank = simulation["bank_census"] _bank_particle(particle_container, bank) - # Increment bank size - add_bank_size(bank, 1) - @njit def bank_future_particle(particle_container, program): @@ -80,20 +76,12 @@ def bank_future_particle(particle_container, program): bank = simulation["bank_future"] _bank_particle(particle_container, bank) - # Increment bank size - add_bank_size(bank, 1) - @njit def bank_source_particle(particle_container, simulation): bank = simulation["bank_source"] _bank_particle(particle_container, bank) - # Increment bank size - # Note that we don't use the atomic operation in add_bank_size function - # as source particle banking is not thread-parallelized - bank["size"][0] += 1 - @njit def pop_particle(particle_container, bank): @@ -101,12 +89,9 @@ def pop_particle(particle_container, bank): if get_bank_size(bank) == 0: report_empty_bank(bank) - # Set particle data - idx = get_bank_size(bank) - 1 - particle_module.copy(particle_container, bank["particle_data"][idx : idx + 1]) - # Decrement bank size - add_bank_size(bank, -1) + idx = add_bank_size(bank, -1) - 1 + particle_module.copy(particle_container, bank["particle_data"][idx : idx + 1]) # Set default IDs and event for the live particle particle = particle_container[0] @@ -163,10 +148,9 @@ def promote_future_particles(program, data): if particle["t"] < next_census_time: bank_census_particle(particle_container, program) - add_bank_size(future_bank, -1) + j = add_bank_size(future_bank, -1) - 1 # Consolidate the emptied space in the future bank - j = get_bank_size(future_bank) particle_module.copy( future_bank["particle_data"][idx : idx + 1], future_bank["particle_data"][j : j + 1], diff --git a/mcdc/transport/physics/neutron/multigroup.py b/mcdc/transport/physics/neutron/multigroup.py index e1d48d495..aa0433918 100644 --- a/mcdc/transport/physics/neutron/multigroup.py +++ b/mcdc/transport/physics/neutron/multigroup.py @@ -251,10 +251,7 @@ def scattering(particle_container, program, data): particle_new["uz"] = uz_new # Get outgoing spectrum - stride = mgxs["G"] - start = mgxs["chi_s_offset"] + group * stride - chi_s = data[start : start + stride] - # Above is equivalent to: chi_s = mcdc_get.neutron_multigroup_data.chi_s_vector(group, mgxs, data) + chi_s = mcdc_get.neutron_multigroup_data.chi_s_vector(group, mgxs, data) # Sample outgoing energy xi = rng.lcg(particle_container_new) @@ -311,10 +308,7 @@ def fission(particle_container, program, data): nu = mcdc_get.neutron_multigroup_data.nu_f(group, mgxs, data) nu_p = mcdc_get.neutron_multigroup_data.nu_p(group, mgxs, data) if J > 0: - stride = mgxs["J"] - start = mgxs["nu_d_offset"] + group * stride - nu_d = data[start : start + stride] - # Above is equivalent to: nu_d = mcdc_get.neutron_multigroup_data.nu_d_vector(group, mgxs, data) + nu_d = mcdc_get.neutron_multigroup_data.nu_d_vector(group, mgxs, data) # Get number of secondaries N = int( @@ -346,10 +340,7 @@ def fission(particle_container, program, data): total = nu_p if xi < total: prompt = True - stride = mgxs["G"] - start = mgxs["chi_p_offset"] + group * stride - spectrum = data[start : start + stride] - # Above is equivalent to: spectrum = mcdc_get.neutron_multigroup_data.chi_p_vector(group, mgxs, data) + spectrum = mcdc_get.neutron_multigroup_data.chi_p_vector(group, mgxs, data) else: prompt = False @@ -357,13 +348,9 @@ def fission(particle_container, program, data): for j in range(J): total += nu_d[j] if xi < total: - stride = mgxs["G"] - start = mgxs["chi_d_offset"] + j * stride - spectrum = data[start : start + stride] - # Above is equivalent to: - # spectrum = mcdc_get.neutron_multigroup_data.chi_d_vector( - # j, mgxs, data - # ) + spectrum = mcdc_get.neutron_multigroup_data.chi_d_vector( + j, mgxs, data + ) decay = mcdc_get.neutron_multigroup_data.decay_rate(j, mgxs, data) break diff --git a/mcdc/transport/physics/neutron/native.py b/mcdc/transport/physics/neutron/native.py index bb69d7e5e..6c938a2e9 100644 --- a/mcdc/transport/physics/neutron/native.py +++ b/mcdc/transport/physics/neutron/native.py @@ -654,13 +654,9 @@ def sample_inelastic_scattering( ) spectrum = simulation["distributions"][ID] else: - offset = inelastic_scattering["spectrum_probability_grid_offset"] - length = inelastic_scattering["spectrum_probability_grid_length"] - probability_grid = data[offset : offset + length] - # Above is equivalent to: - # probability_grid = mcdc_get.neutron_inelastic_scattering_reaction.spectrum_probability_grid_all( - # inelastic_scattering, data - # ) + probability_grid = mcdc_get.neutron_inelastic_scattering_reaction.spectrum_probability_grid_all( + inelastic_scattering, data + ) probability_idx = find_bin(E, probability_grid) xi = rng.lcg(particle_container_new) total = 0.0 diff --git a/mcdc/transport/physics/util.py b/mcdc/transport/physics/util.py index 3475a1510..487b9091e 100644 --- a/mcdc/transport/physics/util.py +++ b/mcdc/transport/physics/util.py @@ -11,10 +11,7 @@ @njit def evaluate_neutron_xs_energy_grid(e, nuclide, data): - offset = nuclide["neutron_xs_energy_grid_offset"] - length = nuclide["neutron_xs_energy_grid_length"] - energy_grid = data[offset : offset + length] - # Above is equivalent to: energy_grid = mcdc_get.nuclide.neutron_xs_energy_grid_all(nuclide, data) + energy_grid = mcdc_get.nuclide.neutron_xs_energy_grid_all(nuclide, data) idx = find_bin(e, energy_grid) e0 = energy_grid[idx] diff --git a/mcdc/transport/source.py b/mcdc/transport/source.py index c89cf4ef2..12c6a7a76 100644 --- a/mcdc/transport/source.py +++ b/mcdc/transport/source.py @@ -108,12 +108,7 @@ def source_particle(particle_container, seed, simulation, data): # Motion translation if source["moving"]: # Get moving interval index wrt the given time - time_grid = data[ - source["move_time_grid_offset"] : ( - source["move_time_grid_offset"] + source["N_move_grid"] - ) - ] - # Above is equivalent to: time_grid = mcdc_get.source.move_time_grid_all(source, data) + time_grid = mcdc_get.source.move_time_grid_all(source, data) tolerance = COINCIDENCE_TOLERANCE_TIME go_lower = False @@ -124,14 +119,10 @@ def source_particle(particle_container, seed, simulation, data): idx += 1 # Source move translations - start = source["move_translations_offset"] + idx * 3 - trans_0 = data[start : start + 3] - # Above is equivalent to: trans_0 = mcdc_get.source.move_translations_vector(idx, source, data) + trans_0 = mcdc_get.source.move_translations_vector(idx, source, data) # Source move velocities - start = source["move_velocities_offset"] + idx * 3 - V = data[start : start + 3] - # Above is equivalent to: V = mcdc_get.source.move_velocities_vector(idx, source, data) + V = mcdc_get.source.move_velocities_vector(idx, source, data) # Source move time grid time_0 = mcdc_get.source.move_time_grid(idx, source, data) diff --git a/mcdc/transport/tally/filter.py b/mcdc/transport/tally/filter.py index d976b1bee..824ab5a12 100644 --- a/mcdc/transport/tally/filter.py +++ b/mcdc/transport/tally/filter.py @@ -56,10 +56,8 @@ def get_direction_index(particle_container, tally, data): tolerance = COINCIDENCE_TOLERANCE_DIRECTION - grid_mu = data[tally["mu_offset"] : (tally["mu_offset"] + tally["mu_length"])] - # Above is equivalent to: grid_mu = mcdc_get.tally.mu_all(tally, data) - grid_azi = data[tally["azi_offset"] : (tally["azi_offset"] + tally["azi_length"])] - # Above is equivalent to: grid_azi = mcdc_get.tally.azi_all(tally, data) + grid_mu = mcdc_get.tally.mu_all(tally, data) + grid_azi = mcdc_get.tally.azi_all(tally, data) i_mu = find_bin_with_tolerance(mu, grid_mu, tolerance) i_azi = find_bin_with_tolerance(azi, grid_azi, tolerance) @@ -73,10 +71,7 @@ def get_energy_index(particle_container, tally, data): E = particle["E"] tolerance = COINCIDENCE_TOLERANCE_ENERGY - grid_energy = data[ - tally["energy_offset"] : (tally["energy_offset"] + tally["energy_length"]) - ] - # Above is equivalent to: grid_energy = mcdc_get.tally.energy_all(tally, data) + grid_energy = mcdc_get.tally.energy_all(tally, data) return find_bin_with_tolerance(E, grid_energy, tolerance) @@ -88,10 +83,7 @@ def get_time_index(particle_container, tally, data): # Particle properties time = particle["t"] - grid_time = data[ - tally["time_offset"] : (tally["time_offset"] + tally["time_length"]) - ] - # Above is equivalent to: grid_time = mcdc_get.tally.time_all(tally, data) + grid_time = mcdc_get.tally.time_all(tally, data) tolerance = COINCIDENCE_TOLERANCE_TIME go_lower = False diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index 45ecdb1fa..4d1e7e7d1 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -1,6 +1,7 @@ import math import numpy as np +import numba as nb from numba import njit from typing import Sequence @@ -153,7 +154,9 @@ def log_interpolation(x, x1, x2, y1, y2): @njit def atomic_add(array, idx, value): + result = array[idx] array[idx] += value + return result @njit @@ -161,6 +164,13 @@ def local_array(shape, dtype): return np.zeros(shape, dtype=dtype) -@njit def access_simulation(program): return program + + +@nb.extending.overload(access_simulation, target="cpu") +def access_simulation_cpu_overload(program): + def impl(program): + return program + + return impl diff --git a/test/regression/conftest.py b/test/regression/conftest.py index 707d0f4ca..144efcbb5 100644 --- a/test/regression/conftest.py +++ b/test/regression/conftest.py @@ -136,11 +136,20 @@ def build_command(config): target = config.getoption("--target") mpiexec = config.getoption("--mpiexec") srun = config.getoption("--srun") + + state = "" + if target == "gpu": + state = "--gpu_state_storage=managed" + mode = "numba" + command = [ sys.executable, "input.py", + f"--clear_cache", + f"--caching", f"--mode={mode}", f"--target={target}", + state, "--output=output", "--no-progress-bar", ] @@ -174,6 +183,7 @@ def compare_outputs(output_path, answer_path, target): def compare_tallies(output, answer, target, errors): + gpu_pass_list = ["uq_var", "sdev"] name_root = "tallies" for tally in answer[name_root].keys(): name_tally = f"{name_root}/{tally}" @@ -182,7 +192,8 @@ def compare_tallies(output, answer, target, errors): continue name_score = f"{name_tally}/{score}" for result in answer[name_score].keys(): - if "uq_var" in result and target == "gpu": + should_pass = any([x in result for x in gpu_pass_list]) + if should_pass and target == "gpu": continue name = f"{name_score}/{result}" assert_allclose(output[name][()], answer[name][()], name, errors) diff --git a/test/unit/test_annotation_shape.py b/test/unit/test_annotation_shape.py index f2f058de5..dbc0686f0 100644 --- a/test/unit/test_annotation_shape.py +++ b/test/unit/test_annotation_shape.py @@ -2,6 +2,7 @@ from typing import Annotated +import numba as nb import numpy as np import pytest from numpy.typing import NDArray @@ -66,8 +67,8 @@ def test_stringified_annotation_rejects_incorrect_offset_shape(capsys): def test_generated_accessor_resolves_dimension_offset(): - all_source = _accessor_1d_all("mgxs", "energy", "G+1") - last_source = _accessor_1d_last("mgxs", "energy", "N - 2") + all_source = _accessor_1d_all("mgxs", "energy", "G+1", nb.types.float64) + last_source = _accessor_1d_last("mgxs", "energy", "N - 2", nb.types.float64) assert 'size = mgxs["G"] + 1' in all_source assert 'size = mgxs["N"] - 2' in last_source diff --git a/test/unit/test_numba_layers_generator.py b/test/unit/test_numba_layers_generator.py index 21531dced..64f786135 100644 --- a/test/unit/test_numba_layers_generator.py +++ b/test/unit/test_numba_layers_generator.py @@ -1,5 +1,6 @@ from typing import Annotated +import numba as nb import numpy as np import pytest from numpy.typing import NDArray @@ -129,8 +130,8 @@ def test_scalar_integer_getters_cast_values_from_data(): def test_float_and_bulk_getters_remain_zero_copy_views(): assert "return data[offset + index]" in _accessor_1d_element("example", "values") - assert "return data[start:end]" in _accessor_1d_all( - "example", "values", "values_length" + assert "return array_result(data[start:end])" in _accessor_1d_all( + "example", "values", "values_length", nb.types.float64 )