Skip to content

Proof of concept for combined state #419

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
braxtoncuneo wants to merge 1 commit into
base: cement
Choose a base branch
from
Draft

Proof of concept for combined state #419

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 6 additions & 38 deletions mcdc/adapt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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":
Expand All @@ -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_)
Expand Down Expand Up @@ -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()
Expand Down
62 changes: 15 additions & 47 deletions mcdc/loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,25 +34,19 @@
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


# If GPU execution is supported and selected, the functions shown below will
# 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


Expand All @@ -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"]):
Expand Down Expand Up @@ -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"]):
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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:
Expand All @@ -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"])

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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

Expand All @@ -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"])

Expand Down Expand Up @@ -977,55 +952,46 @@ 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"]
src_clear_flags = src_fns["clear_flags"]
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"]
pre_alloc_program = pre_fns["alloc_program"]
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)
Expand Down Expand Up @@ -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


Loading
Loading