pyomp 0.5.0__cp314-cp314-macosx_11_0_arm64.whl

numba/openmp/__init__.py

@@ -0,0 +1,106 @@
+import llvmlite.binding as ll
+import sys
+from ._version import version as __version__  # noqa: F401
+
+from .config import (
+    libpath,
+    DEBUG_OPENMP,
+)
+from .omp_runtime import (  # noqa F401
+    omp_set_num_threads,
+    omp_get_thread_num,
+    omp_get_num_threads,
+    omp_get_wtime,
+    omp_set_dynamic,
+    omp_set_nested,
+    omp_set_max_active_levels,
+    omp_get_max_active_levels,
+    omp_get_max_threads,
+    omp_get_num_procs,
+    omp_in_parallel,
+    omp_get_thread_limit,
+    omp_get_supported_active_levels,
+    omp_get_level,
+    omp_get_active_level,
+    omp_get_ancestor_thread_num,
+    omp_get_team_size,
+    omp_in_final,
+    omp_get_proc_bind,
+    omp_get_num_places,
+    omp_get_place_num_procs,
+    omp_get_place_num,
+    omp_set_default_device,
+    omp_get_default_device,
+    omp_get_num_devices,
+    omp_get_device_num,
+    omp_get_team_num,
+    omp_get_num_teams,
+    omp_is_initial_device,
+    omp_get_initial_device,
+)
+
+from .compiler import (
+    CustomCompiler,
+    CustomFunctionCompiler,
+)
+
+from .exceptions import (  # noqa: F401
+    UnspecifiedVarInDefaultNone,
+    ParallelForExtraCode,
+    ParallelForWrongLoopCount,
+    ParallelForInvalidCollapseCount,
+    NonconstantOpenmpSpecification,
+    NonStringOpenmpSpecification,
+    MultipleNumThreadsClauses,
+)
+from .overloads import omp_shared_array  # noqa: F401
+from .omp_context import _OpenmpContextType
+from .decorators import jit, njit  # noqa: F401
+
+
+def _init():
+    sys_platform = sys.platform
+    from ctypes.util import find_library
+
+    omplib = (
+        libpath
+        / "openmp"
+        / "lib"
+        / f"libomp{'.dylib' if sys_platform == 'darwin' else '.so'}"
+    )
+
+    # Prefer bundled libomp if it exists.
+    if omplib.exists():
+        if DEBUG_OPENMP >= 1:
+            print("Found bundled OpenMP runtime library at", omplib)
+        ll.load_library_permanently(str(omplib))
+    else:
+        # There is no bundled libomp, try to find it in standard library paths.
+        system_omplib = find_library("omp")
+        if system_omplib:
+            if DEBUG_OPENMP >= 1:
+                print(f"Found system OpenMP runtime library: {system_omplib}")
+            ll.load_library_permanently(system_omplib)
+        else:
+            raise RuntimeError(
+                f"OpenMP runtime not found. Bundled library missing at {omplib} "
+                "and no system libomp found via ctypes.util.find_library('omp'). "
+                "Ensure libomp is available in library paths."
+            )
+
+    # libomptarget is unavailable on apple, windows, so return.
+    if sys_platform.startswith("darwin") or sys_platform.startswith("win32"):
+        return
+
+    omptargetlib = libpath / "openmp" / "lib" / "libomptarget.so"
+    if omptargetlib.exists():
+        if DEBUG_OPENMP >= 1:
+            print("Found OpenMP target runtime library at", omptargetlib)
+        ll.load_library_permanently(str(omptargetlib))
+    else:
+        raise RuntimeError(f"OpenMP target runtime not found at {omptargetlib}")
+
+
+_init()
+
+openmp_context = _OpenmpContextType()

numba/openmp/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.5.0'
+__version_tuple__ = version_tuple = (0, 5, 0)
+
+__commit_id__ = commit_id = 'g6d955c9a6'

numba/openmp/analysis.py

@@ -0,0 +1,251 @@
+from numba.core import ir, types, typing
+from numba.core.analysis import _fix_loop_exit, compute_cfg_from_blocks
+from numba.core.ir_utils import visit_vars
+from numba.extending import intrinsic
+
+
+def remove_ssa(var_name, scope, loc):
+    # Get the base name of a variable, removing the SSA extension.
+    var = ir.Var(scope, var_name, loc)
+    return var.unversioned_name
+
+
+def user_defined_var(var):
+    if not isinstance(var, str):
+        return False
+    return not var.startswith("$")
+
+
+def is_dsa(name):
+    return (
+        name
+        in [
+            "QUAL.OMP.FIRSTPRIVATE",
+            "QUAL.OMP.PRIVATE",
+            "QUAL.OMP.SHARED",
+            "QUAL.OMP.LASTPRIVATE",
+            "QUAL.OMP.TARGET.IMPLICIT",
+        ]
+        or name.startswith("QUAL.OMP.REDUCTION")
+        or name.startswith("QUAL.OMP.MAP")
+    )
+
+
+def is_private(x):
+    return x in [
+        "QUAL.OMP.PRIVATE",
+        "QUAL.OMP.FIRSTPRIVATE",
+        "QUAL.OMP.LASTPRIVATE",
+        "QUAL.OMP.TARGET.IMPLICIT",
+    ]
+
+
+def is_internal_var(var):
+    # Determine if a var is a Python var or an internal Numba var.
+    if var.is_temp:
+        return True
+    return var.unversioned_name != var.name
+
+
+def has_user_defined_var(the_set):
+    for x in the_set:
+        if user_defined_var(x):
+            return True
+    return False
+
+
+def get_user_defined_var(the_set):
+    ret = set()
+    for x in the_set:
+        if user_defined_var(x):
+            ret.add(x)
+    return ret
+
+
+def get_enclosing_region(func_ir, cur_block):
+    if not hasattr(func_ir, "openmp_enclosing"):
+        func_ir.openmp_enclosing = {}
+    if cur_block in func_ir.openmp_enclosing:
+        return func_ir.openmp_enclosing[cur_block]
+    else:
+        return None
+
+
+def add_enclosing_region(func_ir, blocks, openmp_node):
+    if not hasattr(func_ir, "openmp_enclosing"):
+        func_ir.openmp_enclosing = {}
+    if not hasattr(func_ir, "openmp_regions"):
+        func_ir.openmp_regions = {}
+    func_ir.openmp_regions[openmp_node] = sorted(blocks)
+    for b in blocks:
+        if b not in func_ir.openmp_enclosing:
+            func_ir.openmp_enclosing[b] = []
+        func_ir.openmp_enclosing[b].append(openmp_node)
+
+
+def filter_nested_loops(cfg, loops):
+    blocks_in_loop = set()
+    # get loop bodies
+    for loop in loops.values():
+        insiders = set(loop.body) | set(loop.entries) | set(loop.exits)
+        insiders.discard(loop.header)
+        blocks_in_loop |= insiders
+    # find loop that is not part of other loops
+    for loop in loops.values():
+        if loop.header not in blocks_in_loop:
+            yield _fix_loop_exit(cfg, loop)
+
+
+def remove_privatized(x):
+    if isinstance(x, ir.Var):
+        x = x.name
+
+    if isinstance(x, str) and x.endswith("%privatized"):
+        return x[: len(x) - len("%privatized")]
+    else:
+        return x
+
+
+def get_var_from_enclosing(enclosing_regions, var):
+    if not enclosing_regions:
+        return None
+    if len(enclosing_regions) == 0:
+        return None
+    return enclosing_regions[-1].get_var_dsa(var)
+
+
+def remove_indirections(clause):
+    if not isinstance(clause, list):
+        return clause
+    while len(clause) == 1 and isinstance(clause[0], list):
+        clause = clause[0]
+    return clause
+
+
+def add_tags_to_enclosing(func_ir, cur_block, tags):
+    enclosing_region = get_enclosing_region(func_ir, cur_block)
+    if enclosing_region:
+        for region in enclosing_region:
+            for tag in tags:
+                region.add_tag(tag)
+
+
+def get_blocks_between_start_end(blocks, start_block, end_block):
+    cfg = compute_cfg_from_blocks(blocks)
+    blocks_in_region = [start_block]
+
+    def add_in_region(cfg, blk, blocks_in_region, end_block):
+        """For each successor in the CFG of the block we're currently
+        adding to blocks_in_region, add that successor to
+        blocks_in_region if it isn't the end_block.  Then,
+        recursively call this routine for the added block to add
+        its successors.
+        """
+        for out_blk, _ in cfg.successors(blk):
+            if out_blk != end_block and out_blk not in blocks_in_region:
+                blocks_in_region.append(out_blk)
+                add_in_region(cfg, out_blk, blocks_in_region, end_block)
+
+    # Calculate all the Numba IR blocks in the target region.
+    add_in_region(cfg, start_block, blocks_in_region, end_block)
+    return blocks_in_region
+
+
+@intrinsic
+def get_itercount(typingctx, it):
+    if isinstance(it, types.RangeIteratorType):
+        sig = typing.signature(it.yield_type, it)
+
+        def codegen(context, builder, signature, args):
+            assert len(args) == 1
+            val = args[0]
+            pair = context.make_helper(builder, it, val)
+            return builder.load(pair.count)
+
+        return sig, codegen
+
+
+def remove_all_privatized(x):
+    new_x = None
+    while new_x != x:
+        new_x = x
+        x = remove_privatized(new_x)
+
+    return new_x
+
+
+def typemap_lookup(typemap, x):
+    orig_x = x
+    if isinstance(x, ir.Var):
+        x = x.name
+
+    while True:
+        if x in typemap:
+            return typemap[x]
+        new_x = remove_privatized(x)
+        if new_x == x:
+            break
+        else:
+            x = new_x
+
+    tkeys = typemap.keys()
+
+    # Get basename (without privatized)
+    x = remove_all_privatized(x)
+
+    potential_keys = list(filter(lambda y: y.startswith(x), tkeys))
+
+    for pkey in potential_keys:
+        pkey_base = remove_all_privatized(pkey)
+        if pkey_base == x:
+            return typemap[pkey]
+
+    raise KeyError(f"{orig_x} and all of its non-privatized names not found in typemap")
+
+
+def is_target_tag(x):
+    ret = x.startswith("DIR.OMP.TARGET") and x not in [
+        "DIR.OMP.TARGET.DATA",
+        "DIR.OMP.TARGET.ENTER.DATA",
+        "DIR.OMP.TARGET.EXIT.DATA",
+    ]
+    return ret
+
+
+def is_target_arg(name):
+    return (
+        name in ["QUAL.OMP.FIRSTPRIVATE", "QUAL.OMP.TARGET.IMPLICIT"]
+        or name.startswith("QUAL.OMP.MAP")
+        or name.startswith("QUAL.OMP.REDUCTION")
+    )
+
+
+def in_openmp_region(builder):
+    if hasattr(builder, "_lowerer_push_alloca_callbacks"):
+        return builder._lowerer_push_alloca_callbacks > 0
+    else:
+        return False
+
+
+def get_name_var_table(blocks):
+    """create a mapping from variable names to their ir.Var objects"""
+
+    def get_name_var_visit(var, namevar):
+        namevar[var.name] = var
+        return var
+
+    namevar = {}
+    visit_vars(blocks, get_name_var_visit, namevar)
+    return namevar
+
+
+def is_pointer_target_arg(name, typ):
+    if name.startswith("QUAL.OMP.REDUCTION"):
+        return True
+    if name.startswith("QUAL.OMP.MAP"):
+        return True
+    if name in ["QUAL.OMP.TARGET.IMPLICIT"]:
+        if isinstance(typ, types.npytypes.Array):
+            return True
+
+    return False