Interface to OpenMP Runtime Library Routines.

omp.py

import ctypes


class OpenMPRuntime(object):

    """
    Interface to OpenMP Runtime Library Routines.

    >>> gnu = OpenMPRuntime('libgomp.so')
    >>> gnu.omp_set_num_threads(3)
    >>> gnu.omp_get_max_threads()
    3
    >>> gnu.omp_set_schedule('static', 3)
    >>> gnu.omp_get_schedule()
    ('static', 3)
    >>> gnu.omp_set_schedule('auto')
    >>> gnu.omp_get_schedule()
    ('auto', 3)

    See:
    https://gcc.gnu.org/onlinedocs/libgomp/Runtime-Library-Routines.html
    """

    def __init__(self, libname):
        self.libname = libname

    @cached_property
    def lib(self):
        return ctypes.cdll.LoadLibrary(self.libname)

    @cached_property
    def omp_set_num_threads(self):
        fun = self.lib.omp_set_num_threads
        fun.argtypes = [ctypes.c_int]
        fun.restype = None
        return fun

    @cached_property
    def omp_get_num_threads(self):
        fun = self.lib.omp_get_num_threads
        fun.restype = ctypes.c_int
        return fun

    @cached_property
    def omp_get_max_threads(self):
        fun = self.lib.omp_get_max_threads
        fun.restype = ctypes.c_int
        return fun

    @cached_property
    def omp_get_thread_limit(self):
        fun = self.lib.omp_get_thread_limit
        fun.restype = ctypes.c_int
        return fun

    omp_sched = dict(
        static=1,
        dynamic=2,
        guided=3,
        auto=4,
    )
# https://github.com/gcc-mirror/gcc/blob/gcc-6_3_0-release/libgomp/omp.h.in#L48
# TODO: check that this is the same for Intel and LLVM.

    omp_sched_r = {v: k for k, v in omp_sched.items()}

    _omp_sched_t = ctypes.c_uint
    # http://stackoverflow.com/a/1546467
    # http://stackoverflow.com/q/3100704

    @cached_property
    def _omp_set_schedule(self):
        fun = self.lib.omp_set_schedule
        fun.argtypes = [self._omp_sched_t, ctypes.c_int]
        fun.restype = None
        return fun

    @cached_property
    def _omp_get_schedule(self):
        fun = self.lib.omp_get_schedule
        fun.argtypes = [ctypes.POINTER(self._omp_sched_t),
                        ctypes.POINTER(ctypes.c_int)]
        fun.restype = None
        return fun

    def omp_set_schedule(self, kind, chunk_size=None):
        if kind != 'auto' and chunk_size is None:
            raise ValueError("chunk_size is required for non-auto kind.")
        if chunk_size is None:
            chunk_size = 0
        self._omp_set_schedule(self.omp_sched[kind], chunk_size)

    def omp_get_schedule(self):
        kind = self._omp_sched_t()
        chunk_size = ctypes.c_int()
        self._omp_get_schedule(ctypes.pointer(kind),
                               ctypes.pointer(chunk_size))
        return self.omp_sched_r[kind.value], chunk_size.value


gnu = OpenMPRuntime('libgomp.so')
intel = OpenMPRuntime('libiomp5.so')
llvm = OpenMPRuntime('libomp.so')