def main():
global global_var
global_procs = legion.Tunable.select(legion.Tunable.GLOBAL_PYS).get()
with legion.MustEpochLaunch(): # implicit launch domain
for i in range(global_procs):
hi(i, point=i)
legion.execution_fence(block=True)
assert global_var == 4123
global_var = 456
with legion.MustEpochLaunch([global_procs]):
for i in range(global_procs):
hi(i, point=i)
legion.execution_fence(block=True)
assert global_var == 4456
global_var = 789
with legion.MustEpochLaunch([global_procs]):
legion.index_launch([global_procs], hi, ID)
legion.execution_fence(block=True)
assert global_var == 4789
def solve(n_runs):
n_procs = legion.Tunable.select(legion.Tunable.GLOBAL_PYS).get()
print(f"Working with {n_procs} processes\n")
# Allocate data structures.
n_xpp_events_per_node = 1000
xpp_event_raw_shape = (2, 3, 6)
xpp_data = legion.Region.create((n_xpp_events_per_node,) + xpp_event_raw_shape, {'x': legion.uint16})
legion.fill(xpp_data, 'x', 0)
xpp_part = legion.Partition.create_equal(xpp_data, [n_procs])
gen_data_shape = (N_POINTS,) * 3
data = legion.Region.create(gen_data_shape, {
'amplitudes': legion.float32,
'support': legion.uint8, # should be bool
'rho': legion.complex64})
legion.fill(data, 'amplitudes', 0.)
legion.fill(data, 'support', 0)
legion.fill(data, 'rho', 0.)
complete = False
iteration = 0
fences = []
while not complete or iteration < 10:
if not complete:
# Obtain the newest copy of the data.
with legion.MustEpochLaunch([n_procs]):
for idx in range(n_procs): # legion.IndexLaunch([n_procs]): # FIXME: index launch
data_collector.fill_data_region(xpp_part[idx], point=idx)
# Preprocess data.
for idx in range(n_procs): # legion.IndexLaunch([n_procs]): # FIXME: index launch
preprocess(xpp_part, data)
# Generate data on first run
if not iteration:
generate_data(data)
# Run solver.
solve_step(data, 0, iteration)
if not complete:
# Make sure we don't run more than 2 iterations ahead.
fences.append(legion.execution_fence(future=True))
if iteration - 2 >= 0:
fences[iteration - 2].get()
# Check that all runs have been read.
complete = data_collector.get_num_runs_complete() == n_runs
iteration += 1
def load_run_data(run):
det = run.Detector('xppcspad')
# Hack: psana tries to register top-level task when not in script mode
old_is_script = legion.is_script
legion.is_script = True
f = run.analyze(event_fn=load_event_data, det=det)
legion.is_script = old_is_script
n_procs = legion.Tunable.select(legion.Tunable.GLOBAL_PYS).get()
with legion.MustEpochLaunch([n_procs]):
for idx in range(
n_procs
): # legion.IndexLaunch([n_procs]): # FIXME: index launch
mark_completion(f, point=idx)