codon/stdlib/openmp.codon

673 lines
21 KiB
Python

# OpenMP interface
# Ref: https://github.com/llvm/llvm-project/tree/main/openmp
Routine = Function[[i32, cobj], i32]
_KMP_IDENT_IMB = 0x01
_KMP_IDENT_KMPC = 0x02
_KMP_IDENT_AUTOPAR = 0x08
_KMP_IDENT_ATOMIC_REDUCE = 0x10
_KMP_IDENT_BARRIER_EXPL = 0x20
_KMP_IDENT_BARRIER_IMPL = 0x0040
_KMP_IDENT_BARRIER_IMPL_MASK = 0x01C0
_KMP_IDENT_BARRIER_IMPL_FOR = 0x0040
_KMP_IDENT_BARRIER_IMPL_SECTIONS = 0x00C0
_KMP_IDENT_BARRIER_IMPL_SINGLE = 0x0140
_KMP_IDENT_BARRIER_IMPL_WORKSHARE = 0x01C0
_KMP_IDENT_WORK_LOOP = 0x200
_KMP_IDENT_WORK_SECTIONS = 0x400
_KMP_IDENT_WORK_DISTRIBUTE = 0x800
_KMP_IDENT_ATOMIC_HINT_MASK = 0xFF0000
_KMP_IDENT_ATOMIC_HINT_UNCONTENDED = 0x010000
_KMP_IDENT_ATOMIC_HINT_CONTENDED = 0x020000
_KMP_IDENT_ATOMIC_HINT_NONSPECULATIVE = 0x040000
_KMP_IDENT_ATOMIC_HINT_SPECULATIVE = 0x080000
_KMP_IDENT_OPENMP_SPEC_VERSION_MASK = 0xFF000000
@tuple
class Lock:
a1: i32
a2: i32
a3: i32
a4: i32
a5: i32
a6: i32
a7: i32
a8: i32
def __new__():
z = i32(0)
return Lock(z, z, z, z, z, z, z, z)
@tuple
class Ident:
reserved_1: i32
flags: i32
reserved_2: i32
reserved_3: i32
psource: cobj
def __new__(flags: int = 0, source: str = ";unknown;unknown;0;0;;"):
return Ident(i32(0), i32(flags | _KMP_IDENT_KMPC), i32(0), i32(0), source.ptr)
@tuple
class LRData:
routine: Routine
@tuple
class Task:
shareds: cobj
routine: Routine
flags: i32
x: LRData
y: LRData
@tuple
class TaskWithPrivates[T]:
task: Task
data: T
_DEFAULT_IDENT = Ident()
_STATIC_LOOP_IDENT = Ident(_KMP_IDENT_WORK_LOOP)
_REDUCTION_IDENT = Ident(_KMP_IDENT_ATOMIC_REDUCE)
def _default_loc():
return __ptr__(_DEFAULT_IDENT)
_default_loc()
def _static_loop_loc():
return __ptr__(_STATIC_LOOP_IDENT)
_static_loop_loc()
def _reduction_loc():
return __ptr__(_REDUCTION_IDENT)
_reduction_loc()
def _critical_begin(loc_ref: Ptr[Ident], gtid: int, lck: cobj):
from C import __kmpc_critical(Ptr[Ident], i32, cobj)
__kmpc_critical(loc_ref, i32(gtid), lck)
def _critical_end(loc_ref: Ptr[Ident], gtid: int, lck: cobj):
from C import __kmpc_end_critical(Ptr[Ident], i32, cobj)
__kmpc_end_critical(loc_ref, i32(gtid), lck)
def _single_begin(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_single(Ptr[Ident], i32) -> i32
return int(__kmpc_single(loc_ref, i32(gtid)))
def _single_end(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_end_single(Ptr[Ident], i32)
__kmpc_end_single(loc_ref, i32(gtid))
def _master_begin(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_master(Ptr[Ident], i32) -> i32
return int(__kmpc_master(loc_ref, i32(gtid)))
def _master_end(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_end_master(Ptr[Ident], i32)
__kmpc_end_master(loc_ref, i32(gtid))
def _ordered_begin(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_ordered(Ptr[Ident], i32) -> i32
return int(__kmpc_ordered(loc_ref, i32(gtid)))
def _ordered_end(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_end_ordered(Ptr[Ident], i32)
__kmpc_end_ordered(loc_ref, i32(gtid))
def _taskwait(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_omp_taskwait(Ptr[Ident], i32)
__kmpc_omp_taskwait(loc_ref, i32(gtid))
def _taskgroup_begin(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_taskgroup(Ptr[Ident], i32)
__kmpc_taskgroup(loc_ref, i32(gtid))
def _taskgroup_end(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_end_taskgroup(Ptr[Ident], i32)
__kmpc_end_taskgroup(loc_ref, i32(gtid))
def _task_alloc(loc_ref: Ptr[Ident], gtid: int, flags: int, size_of_task: int, size_of_shareds: int, task_entry: Routine):
from C import __kmpc_omp_task_alloc(Ptr[Ident], i32, i32, int, int, Routine) -> cobj
return __kmpc_omp_task_alloc(loc_ref, i32(gtid), i32(flags), size_of_task, size_of_shareds, task_entry)
def _task_run(loc_ref: Ptr[Ident], gtid: int, new_task: cobj):
from C import __kmpc_omp_task(Ptr[Ident], i32, cobj) -> i32
return int(__kmpc_omp_task(loc_ref, i32(gtid), new_task))
def _barrier(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_barrier(Ptr[Ident], i32)
__kmpc_barrier(loc_ref, i32(gtid))
def _flush(loc_ref: Ptr[Ident]):
from C import __kmpc_flush(Ptr[Ident])
__kmpc_flush(loc_ref)
def flush():
_flush(_default_loc())
def _static_init(loc_ref: Ptr[Ident], gtid: int, schedtype: int, loop: range, incr: int, chunk: int):
from C import __kmpc_for_static_init_8(Ptr[Ident], i32, i32, Ptr[i32], Ptr[int], Ptr[int], Ptr[int], int, int)
last = i32(0)
lower = 0
upper = len(loop) - 1
stride = 1
__kmpc_for_static_init_8(loc_ref, i32(gtid), i32(schedtype), __ptr__(last), __ptr__(lower), __ptr__(upper), __ptr__(stride), incr, chunk)
return bool(last), range(loop._get(lower), loop._get(upper + 1), loop.step), stride
def _static_fini(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_for_static_fini(Ptr[Ident], i32)
__kmpc_for_static_fini(loc_ref, i32(gtid))
def _dynamic_init(loc_ref: Ptr[Ident], gtid: int, schedtype: int, loop: range, chunk: int):
from C import __kmpc_dispatch_init_8(Ptr[Ident], i32, i32, int, int, int, int)
lower = 0
upper = len(loop) - 1
stride = 1
__kmpc_dispatch_init_8(loc_ref, i32(gtid), i32(schedtype), lower, upper, stride, chunk)
def _dynamic_next(loc_ref: Ptr[Ident], gtid: int, loop: range):
from C import __kmpc_dispatch_next_8(Ptr[Ident], i32, Ptr[i32], Ptr[int], Ptr[int], Ptr[int]) -> i32
last = i32(0)
lower = 0
upper = 0
stride = 0
more = __kmpc_dispatch_next_8(loc_ref, i32(gtid), __ptr__(last), __ptr__(lower), __ptr__(upper), __ptr__(stride))
return bool(more), bool(last), range(loop._get(lower), loop._get(upper + 1), loop.step)
def _dynamic_fini(loc_ref: Ptr[Ident], gtid: int):
from C import __kmpc_dispatch_fini_8(Ptr[Ident], i32)
__kmpc_dispatch_fini_8(loc_ref, i32(gtid))
def _reduce[T](loc_ref: Ptr[Ident], gtid: int, reduce_data: T, reduce_func: cobj, lck: cobj):
from internal.gc import sizeof
from C import __kmpc_reduce(Ptr[Ident], i32, i32, int, cobj, cobj, cobj) -> i32
num_vars = staticlen(reduce_data)
reduce_size = sizeof(T)
return int(__kmpc_reduce(loc_ref, i32(gtid), i32(num_vars), reduce_size, __ptr__(reduce_data).as_byte(), reduce_func, lck))
def _end_reduce(loc_ref: Ptr[Ident], gtid: int, lck: cobj):
from C import __kmpc_end_reduce(Ptr[Ident], i32, cobj)
__kmpc_end_reduce(loc_ref, i32(gtid), lck)
def _reduce_nowait[T](loc_ref: Ptr[Ident], gtid: int, reduce_data: T, reduce_func: cobj, lck: cobj):
from internal.gc import sizeof
from C import __kmpc_reduce_nowait(Ptr[Ident], i32, i32, int, cobj, cobj, cobj) -> i32
num_vars = staticlen(reduce_data)
reduce_size = sizeof(T)
return int(__kmpc_reduce_nowait(loc_ref, i32(gtid), i32(num_vars), reduce_size, __ptr__(reduce_data).as_byte(), reduce_func, lck))
def _end_reduce_nowait(loc_ref: Ptr[Ident], gtid: int, lck: cobj):
from C import __kmpc_end_reduce_nowait(Ptr[Ident], i32, cobj)
__kmpc_end_reduce_nowait(loc_ref, i32(gtid), lck)
def _fork_call(microtask: cobj, args):
from C import __kmpc_fork_call(Ptr[Ident], i32, cobj, ...)
loc_ref = _default_loc() # TODO: pass real loc?
__kmpc_fork_call(loc_ref, i32(1), microtask, __ptr__(args))
def _static_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
def _loop_step():
return 1
def _loop_loc_and_gtid(loc_ref: Ptr[Ident], reduction_loc_ref: Ptr[Ident], gtid: int):
pass
def _loop_body_stub(i, args):
pass
def _loop_schedule():
return (1 << 30) | 35 # nonmonotonic, dynamic chunked
def _loop_shared_updates(args):
pass
def _loop_reductions(args):
pass
chunk, start, stop, extra = args[0]
step = _loop_step()
gtid = int(gtid_ptr[0])
loc_ref = _default_loc()
static_loop_loc_ref = _static_loop_loc()
reduction_loc_ref = _reduction_loc()
_loop_loc_and_gtid(loc_ref, reduction_loc_ref, gtid)
loop = range(start, stop, step)
schedule = _loop_schedule()
last, subloop, stride = _static_init(static_loop_loc_ref, gtid, schedtype=schedule, loop=loop, incr=1, chunk=1)
i = subloop.start
stop = min(subloop.stop, loop.stop) if step >= 0 else max(subloop.stop, loop.stop)
while (step >= 0 and i < stop) or (step < 0 and i > stop):
_loop_body_stub(i, extra)
i += step
_static_fini(static_loop_loc_ref, gtid)
if last:
_loop_shared_updates(extra)
_loop_reductions(extra)
def _static_chunked_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
def _loop_step():
return 1
def _loop_loc_and_gtid(loc_ref: Ptr[Ident], reduction_loc_ref: Ptr[Ident], gtid: int):
pass
def _loop_body_stub(i, args):
pass
def _loop_schedule():
return (1 << 30) | 35 # nonmonotonic, dynamic chunked
def _loop_shared_updates(args):
pass
def _loop_reductions(args):
pass
chunk, start, stop, extra = args[0]
step = _loop_step()
gtid = int(gtid_ptr[0])
loc_ref = _default_loc()
static_loop_loc_ref = _static_loop_loc()
reduction_loc_ref = _reduction_loc()
_loop_loc_and_gtid(loc_ref, reduction_loc_ref, gtid)
loop = range(start, stop, step)
schedule = _loop_schedule()
last, subloop, stride = _static_init(static_loop_loc_ref, gtid, schedtype=schedule, loop=loop, incr=1, chunk=chunk)
start = subloop.start
stop = min(subloop.stop, loop.stop) if step >= 0 else max(subloop.stop, loop.stop)
while (step >= 0 and start < loop.stop) or (step < 0 and start > loop.stop):
i = start
while (step >= 0 and i < stop) or (step < 0 and i > stop):
_loop_body_stub(i, extra)
i += step
start += stride * step
stop += stride * step
stop = min(stop, loop.stop) if step >= 0 else max(stop, loop.stop)
_static_fini(static_loop_loc_ref, gtid)
if last:
_loop_shared_updates(extra)
_loop_reductions(extra)
def _dynamic_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
def _loop_step():
return 1
def _loop_loc_and_gtid(loc_ref: Ptr[Ident], reduction_loc_ref: Ptr[Ident], gtid: int):
pass
def _loop_body_stub(i, args):
pass
def _loop_schedule():
return (1 << 30) | 35 # nonmonotonic, dynamic chunked
def _loop_shared_updates(args):
pass
def _loop_reductions(args):
pass
def _loop_ordered():
return False
chunk, start, stop, extra = args[0]
step = _loop_step()
gtid = int(gtid_ptr[0])
loc_ref = _default_loc()
reduction_loc_ref = _reduction_loc()
_loop_loc_and_gtid(loc_ref, reduction_loc_ref, gtid)
loop = range(start, stop, step)
schedule = _loop_schedule()
ordered = _loop_ordered()
_dynamic_init(loc_ref, gtid, schedtype=schedule, loop=loop, chunk=chunk)
while True:
more, last, subloop = _dynamic_next(loc_ref, gtid, loop)
if not more:
break
i = subloop.start
while (step >= 0 and i < subloop.stop) or (step < 0 and i > subloop.stop):
_loop_body_stub(i, extra)
i += step
if ordered:
_dynamic_fini(loc_ref, gtid)
if last:
_loop_shared_updates(extra)
_loop_reductions(extra)
# P = privates; tuple of types
# S = shareds; tuple of pointers
def _spawn_and_run_task[P,S](loc_ref: Ptr[Ident], gtid: int, routine: cobj, priv: P, shared: S):
from internal.gc import sizeof, add_roots
TaskThunk = TaskWithPrivates[P]
flags = 1
size_of_kmp_task_t = sizeof(TaskThunk)
size_of_privs = sizeof(P)
size_of_shareds = sizeof(S)
loc_ref = _default_loc()
task = Ptr[TaskThunk](_task_alloc(loc_ref, gtid, flags, size_of_kmp_task_t, size_of_shareds, Routine(routine)))
if staticlen(shared) != 0:
shared_ptr = task[0].task.shareds
str.memcpy(shared_ptr, __ptr__(shared).as_byte(), size_of_shareds)
add_roots(shared_ptr, shared_ptr + size_of_shareds)
if staticlen(priv) != 0:
priv_ptr = task.as_byte() + sizeof(Task)
str.memcpy(priv_ptr, __ptr__(priv).as_byte(), size_of_privs)
add_roots(priv_ptr, priv_ptr + size_of_privs)
_task_run(loc_ref, gtid, task.as_byte())
# Note: this is different than OpenMP's "taskloop" -- this template simply
# spawns a new task for each loop iteration.
def _task_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
def _routine_stub[P,S](gtid: i32, data: cobj):
def _task_loop_body_stub(priv, shared):
pass
task = Ptr[TaskWithPrivates[P]](data)[0]
priv = task.data
if staticlen(S()) != 0:
shared = Ptr[S](task.task.shareds)[0]
_task_loop_body_stub(priv, shared)
else:
shared = ()
_task_loop_body_stub(priv, shared)
return i32(0)
def _insert_new_loop_var(i, priv, shared):
return priv, shared
iterable, priv, shared = args[0]
P = type(priv)
S = type(shared)
gtid = int(gtid_ptr[0])
loc_ref = _default_loc()
if _single_begin(loc_ref, gtid) != 0:
_taskgroup_begin(loc_ref, gtid)
try:
for i in iterable:
priv, shared = _insert_new_loop_var(i, priv, shared)
_spawn_and_run_task(loc_ref, gtid, _routine_stub(P=P,S=S,...).__raw__(), priv, shared)
finally:
_taskgroup_end(loc_ref, gtid)
_single_end(loc_ref, gtid)
@pure
def get_num_threads():
from C import omp_get_num_threads() -> i32
return int(omp_get_num_threads())
@pure
def get_thread_num():
from C import omp_get_thread_num() -> i32
return int(omp_get_thread_num())
@pure
def get_max_threads():
from C import omp_get_max_threads() -> i32
return int(omp_get_max_threads())
@pure
def get_num_procs():
from C import omp_get_num_procs() -> i32
return int(omp_get_num_procs())
def set_num_threads(num_threads: int):
from C import omp_set_num_threads(i32)
omp_set_num_threads(i32(num_threads))
@pure
def in_parallel():
from C import omp_in_parallel() -> i32
return bool(omp_in_parallel())
def set_dynamic(dynamic_threads: bool = True):
from C import omp_set_dynamic(i32)
omp_set_dynamic(i32(1 if dynamic_threads else 0))
@pure
def get_dynamic():
from C import omp_get_dynamic() -> i32
return bool(omp_get_dynamic())
@pure
def get_cancellation():
from C import omp_get_cancellation() -> i32
return bool(omp_get_cancellation())
def set_schedule(kind: str, chunk_size: int = 0):
from C import omp_set_schedule(i32, i32)
if kind == 'static':
omp_set_schedule(i32(1), i32(chunk_size))
elif kind == 'dynamic':
omp_set_schedule(i32(2), i32(chunk_size))
elif kind == 'guided':
omp_set_schedule(i32(3), i32(chunk_size))
elif kind == 'auto':
if chunk_size != 0:
raise ValueError('cannot specify chunk size for auto schedule')
omp_set_schedule(i32(4), i32(chunk_size))
else:
raise ValueError("invalid schedule kind; valid ones are: 'static', 'dynamic', 'guided', 'auto'")
@pure
def get_schedule():
from C import omp_get_schedule(Ptr[i32], Ptr[i32])
kind_code = i32(0)
chunk_size = i32(0)
omp_get_schedule(__ptr__(kind_code), __ptr__(chunk_size))
idx = int(kind_code)
kind = ('static', 'dynamic', 'guided', 'auto')[idx - 1] if 1 < idx <= 4 else 'unknown'
return kind, int(chunk_size)
@pure
def get_thread_limit():
from C import omp_get_thread_limit() -> i32
return int(omp_get_thread_limit())
def set_max_active_levels(max_levels: int):
from C import omp_set_max_active_levels(i32)
omp_set_max_active_levels(i32(max_levels))
@pure
def get_max_active_levels():
from C import omp_get_max_active_levels() -> i32
return int(omp_get_max_active_levels())
@pure
def get_level():
from C import omp_get_level() -> i32
return int(omp_get_level())
@pure
def get_ancestor_thread_num(level: int):
from C import omp_get_ancestor_thread_num(i32) -> i32
return int(omp_get_ancestor_thread_num(i32(level)))
@pure
def get_team_size(level: int):
from C import omp_get_team_size(i32) -> i32
return int(omp_get_team_size(i32(level)))
@pure
def get_active_level():
from C import omp_get_active_level() -> i32
return int(omp_get_active_level())
@pure
def in_final():
from C import omp_in_final() -> i32
return bool(omp_in_final())
@pure
def get_proc_bind():
from C import omp_get_proc_bind() -> i32
result = int(omp_get_proc_bind())
if result < 0 or result > 4:
return 'unknown'
return ('false', 'true', 'master', 'close', 'spread')[result]
def set_default_device(device_num: int):
from C import omp_set_default_device(i32)
omp_set_default_device(i32(device_num))
@pure
def get_default_device():
from C import omp_get_default_device() -> i32
return int(omp_get_default_device())
@pure
def get_num_devices():
from C import omp_get_num_devices() -> i32
return int(omp_get_num_devices())
@pure
def get_num_teams():
from C import omp_get_num_teams() -> i32
return int(omp_get_num_teams())
@pure
def get_team_num():
from C import omp_get_team_num() -> i32
return int(omp_get_team_num())
@pure
def is_initial_device():
from C import omp_is_initial_device() -> i32
return bool(omp_is_initial_device())
@pure
def get_wtime():
from C import omp_get_wtime() -> float
return omp_get_wtime()
@pure
def get_wtick():
from C import omp_get_wtick() -> float
return omp_get_wtick()
def single(func):
def _wrapper(*args, **kwargs):
gtid = get_thread_num()
loc = _default_loc()
if _single_begin(loc, gtid) != 0:
try:
func(*args, **kwargs)
finally:
_single_end(loc, gtid)
return _wrapper
def master(func):
def _wrapper(*args, **kwargs):
gtid = get_thread_num()
loc = _default_loc()
if _master_begin(loc, gtid) != 0:
try:
func(*args, **kwargs)
finally:
_master_end(loc, gtid)
return _wrapper
def ordered(func):
def _wrapper(*args, **kwargs):
gtid = get_thread_num()
loc = _default_loc()
if _ordered_begin(loc, gtid) != 0:
try:
func(*args, **kwargs)
finally:
_ordered_end(loc, gtid)
return _wrapper
_default_lock = Lock()
def critical(func):
def _wrapper(*args, **kwargs):
gtid = get_thread_num()
loc = _default_loc()
_critical_begin(loc, gtid, __ptr__(_default_lock).as_byte())
try:
func(*args, **kwargs)
finally:
_critical_end(loc, gtid, __ptr__(_default_lock).as_byte())
return _wrapper
def _push_num_threads(num_threads: int):
from C import __kmpc_push_num_threads(Ptr[Ident], i32, i32)
gtid = get_thread_num()
loc = _default_loc()
__kmpc_push_num_threads(loc, i32(gtid), i32(num_threads))
@llvm
def _atomic_int_add(a: Ptr[int], b: int) -> void:
%old = atomicrmw add i64* %a, i64 %b monotonic
ret void
def _atomic_int_mul(a: Ptr[int], b: int):
from C import __kmpc_atomic_fixed8_mul(Ptr[Ident], i32, Ptr[int], int)
__kmpc_atomic_fixed8_mul(_default_loc(), i32(0), a, b)
@llvm
def _atomic_int_and(a: Ptr[int], b: int) -> void:
%old = atomicrmw and i64* %a, i64 %b monotonic
ret void
@llvm
def _atomic_int_or(a: Ptr[int], b: int) -> void:
%old = atomicrmw or i64* %a, i64 %b monotonic
ret void
@llvm
def _atomic_int_xor(a: Ptr[int], b: int) -> void:
%old = atomicrmw xor i64* %a, i64 %b monotonic
ret void
@llvm
def _atomic_int_min(a: Ptr[int], b: int) -> void:
%old = atomicrmw min i64* %a, i64 %b monotonic
ret void
@llvm
def _atomic_int_max(a: Ptr[int], b: int) -> void:
%old = atomicrmw max i64* %a, i64 %b monotonic
ret void
@llvm
def _atomic_float_add(a: Ptr[float], b: float) -> void:
%old = atomicrmw fadd double* %a, double %b monotonic
ret void
def _atomic_float_mul(a: Ptr[float], b: float):
from C import __kmpc_atomic_float8_mul(Ptr[Ident], i32, Ptr[float], float)
__kmpc_atomic_float8_mul(_default_loc(), i32(get_thread_num()), a, b)
def _atomic_float_min(a: Ptr[float], b: float):
from C import __kmpc_atomic_float8_min(Ptr[Ident], i32, Ptr[float], float)
__kmpc_atomic_float8_min(_default_loc(), i32(get_thread_num()), a, b)
def _atomic_float_max(a: Ptr[float], b: float) -> void:
from C import __kmpc_atomic_float8_max(Ptr[Ident], i32, Ptr[float], float)
__kmpc_atomic_float8_max(_default_loc(), i32(get_thread_num()), a, b)
def for_par(
num_threads: int = -1,
chunk_size: int = -1,
schedule: Static[str] = "static",
ordered: Static[int] = False
):
pass