pyomp 0.5.0__cp314-cp314t-macosx_11_0_arm64.whl

numba/openmp/libs/pass/IntrinsicsOpenMP.h

@@ -0,0 +1,13 @@
+#ifndef LLVM_TRANSFORMS_INTRINSICS_OPENMP_H
+#define LLVM_TRANSFORMS_INTRINSICS_OPENMP_H
+
+#include <llvm/IR/PassManager.h>
+#include <llvm/Pass.h>
+
+namespace llvm {
+
+ModulePass *createIntrinsicsOpenMPPass();
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_INTRINSICS_OPENMP_H

numba/openmp/libs/pass/IntrinsicsOpenMP_CAPI.h

@@ -0,0 +1,23 @@
+#ifndef LLVM_C_TRANSFORMS_INTRINSICS_OPENMP_H
+#define LLVM_C_TRANSFORMS_INTRINSICS_OPENMP_H
+
+#include <llvm-c/ExternC.h>
+#include <llvm-c/Types.h>
+
+LLVM_C_EXTERN_C_BEGIN
+
+/**
+ * @defgroup LLVMCTransformsIntrinsicsOpenMP IntrinsicsOpenMP transformations
+ * @ingroup LLVMCTransforms
+ *
+ * @{
+ */
+
+/** See llvm::createIntrinsicsOpenMPPass function. */
+void LLVMAddIntrinsicsOpenMPPass(LLVMPassManagerRef PM);
+
+/**
+ * @}
+ */
+LLVM_C_EXTERN_C_END
+#endif

numba/openmp/link_utils.py

@@ -0,0 +1,126 @@
+import os
+import tempfile
+from functools import lru_cache
+from pathlib import Path
+
+# Python 3.12+ removed distutils; use the shim in setuptools.
+try:
+    from setuptools._distutils import ccompiler, sysconfig
+except Exception:  # Python <3.12, or older setuptools
+    from distutils import ccompiler, sysconfig  # type: ignore
+
+
+def link_shared_library(obj_path, out_path):
+    # Generate trampolines for numba/NRT symbols. We use trampolines to link the
+    # absolute symbol addresses from numba to the self-contained shared library
+    # for the OpenMP target CPU module.
+    # TODO: ask numba upstream to provide a static library with these symbols.
+    @lru_cache
+    def generate_trampolines():
+        from numba import _helperlib
+        from numba.core.runtime import _nrt_python as _nrt
+
+        # Signature mapping for numba/NRT functions. Add more as needed.
+        SIGNATURES = {
+            # GIL management
+            "numba_gil_ensure": ("void", []),
+            "numba_gil_release": ("void", []),
+            # Memory allocation
+            "NRT_MemInfo_alloc": ("void*", ["size_t"]),
+            "NRT_MemInfo_alloc_safe": ("void*", ["size_t"]),
+            "NRT_MemInfo_alloc_aligned": ("void*", ["size_t", "size_t"]),
+            "NRT_MemInfo_alloc_safe_aligned": ("void*", ["size_t", "size_t"]),
+            "NRT_MemInfo_free": ("void", ["void*"]),
+            # Helperlib
+            "numba_unpickle": ("void*", ["void*", "int", "void*"]),
+        }
+
+        trampoline_c = """#include <stddef.h>"""
+
+        symbols = []
+        # Process _helperlib symbols
+        for py_name in _helperlib.c_helpers:
+            c_name = "numba_" + py_name
+            c_address = _helperlib.c_helpers[py_name]
+
+            if c_name in SIGNATURES:
+                ret_type, params = SIGNATURES[c_name]
+                symbols.append((c_name, c_address, ret_type, params))
+
+        # Process _nrt symbols
+        for py_name in _nrt.c_helpers:
+            if py_name.startswith("_"):
+                c_name = py_name
+            else:
+                c_name = "NRT_" + py_name
+            c_address = _nrt.c_helpers[py_name]
+
+            if c_name in SIGNATURES:
+                ret_type, params = SIGNATURES[c_name]
+                symbols.append((c_name, c_address, ret_type, params))
+
+        # Generate trampolines
+        for c_name, c_address, ret_type, params in sorted(symbols):
+            # Build parameter list
+            if not params:
+                param_list = "void"
+                arg_list = ""
+            else:
+                param_list = ", ".join(
+                    f"{ptype} arg{i}" for i, ptype in enumerate(params)
+                )
+                arg_list = ", ".join(f"arg{i}" for i in range(len(params)))
+
+            # Build function pointer type
+            func_ptr_type = f"{ret_type} (*)({', '.join(params) if params else 'void'})"
+
+            # Generate the trampoline
+            trampoline_c += f"""
+    __attribute__((visibility("default")))
+    {ret_type} {c_name}({param_list}) {{
+        {"" if ret_type == "void" else "return "}(({func_ptr_type})0x{c_address:x})({arg_list});
+    }}
+    """
+
+        return trampoline_c
+
+    """
+    Produce a shared library from a single object file and link numba C symbols.
+    Uses distutils' compiler.
+    """
+    obj_path = str(Path(obj_path))
+    out_path = str(Path(out_path))
+
+    trampoline_code = generate_trampolines()
+    fd, trampoline_c = tempfile.mkstemp(".c")
+    os.close(fd)
+    with open(trampoline_c, "w") as f:
+        f.write(trampoline_code)
+
+    cc = ccompiler.new_compiler()
+    sysconfig.customize_compiler(cc)
+    extra_pre = []
+    extra_post = []
+
+    try:
+        trampoline_o = cc.compile([trampoline_c])
+    except Exception as e:
+        raise RuntimeError(
+            f"Compilation failed for trampolines in {trampoline_c}"
+        ) from e
+    finally:
+        os.remove(trampoline_c)
+
+    objs = [obj_path] + trampoline_o
+    try:
+        cc.link_shared_object(
+            objects=objs,
+            output_filename=out_path,
+            extra_preargs=extra_pre,
+            extra_postargs=extra_post,
+        )
+    except Exception as e:
+        raise RuntimeError(f"Link failed for {out_path}") from e
+    finally:
+        for file_o in trampoline_o:
+            os.remove(file_o)

numba/openmp/llvm_pass.py

@@ -0,0 +1,48 @@
+import ctypes
+import sys
+import llvmlite.binding as ll
+
+
+from .config import libpath, DEBUG_OPENMP_LLVM_PASS
+
+
+def run_intrinsics_openmp_pass(ll_module):
+    libpass = (
+        libpath
+        / "pass"
+        / f"libIntrinsicsOpenMP.{'dylib' if sys.platform == 'darwin' else 'so'}"
+    )
+
+    # Roundtrip the LLVM module through the intrinsics OpenMP pass.
+    WRITE_CB = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_size_t)
+
+    out = bytearray()
+
+    def _writer_cb(ptr, size):
+        out.extend(ctypes.string_at(ptr, size))
+
+    writer_cb = WRITE_CB(_writer_cb)
+
+    lib = ctypes.CDLL(str(libpass))
+    lib.runIntrinsicsOpenMPPass.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_size_t,
+        WRITE_CB,
+    ]
+    lib.runIntrinsicsOpenMPPass.restype = ctypes.c_int
+
+    bc = ll_module.as_bitcode()
+    buf = ctypes.create_string_buffer(bc)
+    ptr = ctypes.cast(buf, ctypes.c_void_p)
+    rc = lib.runIntrinsicsOpenMPPass(ptr, len(bc), writer_cb)
+    if rc != 0:
+        raise RuntimeError(f"Running IntrinsicsOpenMPPass failed with return code {rc}")
+
+    bc_out = bytes(out)
+
+    lowered_module = ll.parse_bitcode(bc_out)
+    if DEBUG_OPENMP_LLVM_PASS >= 1:
+        with open(f"{ll_module.name}-intrinsics-omp.ll", "w") as f:
+            f.write(str(lowered_module))
+
+    return lowered_module

numba/openmp/llvmlite_extensions.py

@@ -0,0 +1,75 @@
+import llvmlite.ir as lir
+
+
+def get_decl(alloca):
+    if not isinstance(alloca, lir.instructions.AllocaInstr):
+        raise TypeError("Expected AllocaInstr, got %s" % type(alloca))
+    return '{0} %"{1}"'.format(alloca.type, alloca._get_name())
+
+
+# TODO: Upstream to llvmlite, it's part of the langref.
+class TokenType(lir.Type):
+    """
+    The type for tokens.  From the LLVM Language Reference.
+
+      'The token type is used when a value is associated with an
+       instruction but all uses of the value must not attempt to
+       introspect or obscure it. As such, it is not appropriate
+       to have a phi or select of type token.'
+    """
+
+    def _to_string(self):
+        return "token"
+
+    def __eq__(self, other):
+        return isinstance(other, TokenType)
+
+    def __hash__(self):
+        return hash(TokenType)
+
+
+class CallInstrWithOperandBundle(lir.instructions.CallInstr):
+    def set_tags(self, tags):
+        self.tags = tags
+
+    # TODO: This is ugly duplication, we should upstream to llvmlite.
+    def descr(self, buf, add_metadata=True):
+        def descr_arg(i, a):
+            if i in self.arg_attributes:
+                attrs = " ".join(self.arg_attributes[i]._to_list()) + " "
+            else:
+                attrs = ""
+            return "{0} {1}{2}".format(a.type, attrs, a.get_reference())
+
+        args = ", ".join([descr_arg(i, a) for i, a in enumerate(self.args)])
+
+        fnty = self.callee.function_type
+        # Only print function type if variable-argument
+        if fnty.var_arg:
+            ty = fnty
+        # Otherwise, just print the return type.
+        else:
+            # Fastmath flag work only in this case
+            ty = fnty.return_type
+        callee_ref = "{0} {1}".format(ty, self.callee.get_reference())
+        if self.cconv:
+            callee_ref = "{0} {1}".format(self.cconv, callee_ref)
+
+        tail_marker = ""
+        if self.tail:
+            tail_marker = "{0} ".format(self.tail)
+
+        buf.append(
+            "{tail}{op}{fastmath} {callee}({args}){attr}{tags}{meta}\n".format(
+                tail=tail_marker,
+                op=self.opname,
+                fastmath="".join([" " + attr for attr in self.fastmath]),
+                callee=callee_ref,
+                args=args,
+                attr="".join([" " + attr for attr in self.attributes]),
+                tags=(" " + self.tags if self.tags is not None else ""),
+                meta=(
+                    self._stringify_metadata(leading_comma=True) if add_metadata else ""
+                ),
+            )
+        )

numba/openmp/omp_context.py

@@ -0,0 +1,242 @@
+from numba.core import ir
+from numba.core.ir_utils import (
+    build_definitions,
+    get_definition,
+    dump_blocks,
+    dprint_func_ir,
+    compute_cfg_from_blocks,
+    compute_use_defs,
+    compute_live_map,
+)
+from numba.core.withcontexts import WithContext
+import sys
+import os
+
+
+from .parser import var_collector_parser
+from .analysis import get_name_var_table
+from .config import DEBUG_OPENMP, OPENMP_DISABLED
+from .compiler import LowerNoSROA
+from .omp_ir import (
+    openmp_region_start,
+    openmp_region_end,
+    _lower_openmp_region_start,
+    _lower_openmp_region_end,
+)
+from .analysis import in_openmp_region
+from .omp_lower import VarCollector, remove_ssa_from_func_ir, _add_openmp_ir_nodes
+
+
+class PythonOpenmp:
+    def __init__(self, *args):
+        self.args = args
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, typ, val, tb):
+        pass
+
+
+def iscall(x):
+    if isinstance(x, ir.Assign):
+        return isinstance(x.value, ir.Expr) and x.value.op == "call"
+    elif isinstance(x, ir.Expr):
+        return x.op == "call"
+    else:
+        return False
+
+
+def extract_args_from_openmp(func_ir):
+    """Find all the openmp context calls in the function and then
+    use the VarCollector transformer to find all the Python variables
+    referenced in the openmp clauses.  We then add those variables as
+    regular arguments to the openmp context call just so Numba's
+    usedef analysis is able to keep variables alive that are only
+    referenced in openmp clauses.
+    """
+    func_ir._definitions = build_definitions(func_ir.blocks)
+    var_table = get_name_var_table(func_ir.blocks)
+    for block in func_ir.blocks.values():
+        for inst in block.body:
+            if iscall(inst):
+                func_def = get_definition(func_ir, inst.value.func)
+                if isinstance(func_def, ir.Global) and isinstance(
+                    func_def.value, _OpenmpContextType
+                ):
+                    str_def = get_definition(func_ir, inst.value.args[0])
+                    if not isinstance(str_def, ir.Const) or not isinstance(
+                        str_def.value, str
+                    ):
+                        # The non-const openmp string error is handled later.
+                        continue
+                    assert isinstance(str_def, ir.Const) and isinstance(
+                        str_def.value, str
+                    )
+                    parse_res = var_collector_parser.parse(str_def.value)
+                    visitor = VarCollector()
+                    try:
+                        visit_res = visitor.transform(parse_res)
+                        inst.value.args.extend([var_table[x] for x in visit_res])
+                    except Exception as e:
+                        print(f"generic transform exception: {e}")
+                        exc_type, exc_obj, exc_tb = sys.exc_info()
+                        fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+                        print(exc_type, fname, exc_tb.tb_lineno)
+                        sys.exit(-2)
+
+
+def remove_empty_blocks(blocks):
+    found = True
+    while found:
+        found = False
+        empty_block = None
+        for label, block in blocks.items():
+            if len(block.body) == 1:
+                assert isinstance(block.body[-1], ir.Jump)
+                empty_block = label
+                next_block = block.body[-1].target
+                break
+
+        if empty_block is not None:
+            del blocks[empty_block]
+
+            found = True
+            for block in blocks.values():
+                last_stmt = block.body[-1]
+                if isinstance(last_stmt, ir.Jump):
+                    if last_stmt.target == empty_block:
+                        block.body[-1] = ir.Jump(next_block, last_stmt.loc)
+                elif isinstance(last_stmt, ir.Branch):
+                    if last_stmt.truebr == empty_block:
+                        block.body[-1] = ir.Branch(
+                            last_stmt.cond, next_block, last_stmt.falsebr, last_stmt.loc
+                        )
+                    elif block.body[-1].falsebr == empty_block:
+                        block.body[-1] = ir.Branch(
+                            last_stmt.cond, last_stmt.truebr, next_block, last_stmt.loc
+                        )
+                elif isinstance(last_stmt, ir.Return):
+                    # Intentionally do nothing.
+                    pass
+                else:
+                    print(type(last_stmt))
+                    assert False
+
+
+class _OpenmpContextType(WithContext):
+    is_callable = True
+    first_time = True
+    blk_end_live_map = set()
+
+    def do_numba_fixups(self):
+        from numba import core
+
+        orig_lower_inst = core.lowering.Lower.lower_inst
+        core.lowering.Lower.orig_lower_inst = orig_lower_inst
+
+        orig_lower = core.lowering.Lower.lower
+        core.lowering.Lower.orig_lower = orig_lower
+
+        # Use method to retrieve the outside region live map, which is updated
+        # during the with-context mutation.
+        def get_blk_end_live_map():
+            return self.blk_end_live_map
+
+        def new_lower(self, inst):
+            if not isinstance(self, LowerNoSROA):
+                self.__class__ = LowerNoSROA
+            if isinstance(inst, openmp_region_start):
+                return _lower_openmp_region_start(self, inst)
+            elif isinstance(inst, openmp_region_end):
+                return _lower_openmp_region_end(self, inst)
+            # TODO: instead of monkey patching for Del instructions outside the
+            # openmp region do: (1) either outline to create a function scope
+            # that will decouple the lifetime of variables inside the OpenMP
+            # region, (2) or subclass the PostProcessor to extend use-def
+            # analysis with OpenMP lifetime information.
+            elif isinstance(inst, ir.Del):
+                # Lower Del normally in the openmp region.
+                if in_openmp_region(self.builder):
+                    return self.orig_lower_inst(inst)
+
+                # Lower the Del instruction ONLY if the variable is not live
+                # after the openmp region.
+                if inst.value not in get_blk_end_live_map():
+                    return self.orig_lower_inst(inst)
+            elif isinstance(inst, ir.Assign):
+                return self.lower_assign_inst(orig_lower_inst, inst)
+            elif isinstance(inst, ir.Return):
+                return self.lower_return_inst(orig_lower_inst, inst)
+            else:
+                return self.orig_lower_inst(inst)
+
+        core.lowering.Lower.lower_inst = new_lower
+
+    def mutate_with_body(
+        self,
+        func_ir,
+        blocks,
+        blk_start,
+        blk_end,
+        body_blocks,
+        dispatcher_factory,
+        extra,
+    ):
+        if _OpenmpContextType.first_time:
+            _OpenmpContextType.first_time = False
+            self.do_numba_fixups()
+
+        if DEBUG_OPENMP >= 1:
+            print("pre-dead-code")
+            dump_blocks(blocks)
+        if not OPENMP_DISABLED and not hasattr(func_ir, "has_openmp_region"):
+            # We can't do dead code elimination at this point because if an argument
+            # is used only in an openmp clause then it is detected as dead and is
+            # eliminated.  We'd have to run through the IR and find all the
+            # openmp regions and extract the vars used there and then modify the
+            # IR with something fake just to take the var alive.  The other approach
+            # would be to modify dead code elimination to find the vars referenced
+            # in openmp context strings.
+            extract_args_from_openmp(func_ir)
+            # dead_code_elimination(func_ir)
+            remove_ssa_from_func_ir(func_ir)
+            # remove_empty_blocks(blocks)
+            func_ir.has_openmp_region = True
+        if DEBUG_OPENMP >= 1:
+            print("pre-with-removal")
+            dump_blocks(blocks)
+        if OPENMP_DISABLED:
+            # If OpenMP disabled, do nothing except remove the enter_with marker.
+            sblk = blocks[blk_start]
+            sblk.body = sblk.body[1:]
+        else:
+            if DEBUG_OPENMP >= 1:
+                print("openmp:mutate_with_body")
+                dprint_func_ir(func_ir, "func_ir")
+                print("blocks:", blocks, type(blocks))
+                print("blk_start:", blk_start, type(blk_start))
+                print("blk_end:", blk_end, type(blk_end))
+                print("body_blocks:", body_blocks, type(body_blocks))
+                print("extra:", extra, type(extra))
+            assert extra is not None
+            _add_openmp_ir_nodes(
+                func_ir, blocks, blk_start, blk_end, body_blocks, extra
+            )
+            func_ir._definitions = build_definitions(blocks)
+            if DEBUG_OPENMP >= 1:
+                print("post-with-removal")
+                dump_blocks(blocks)
+            dispatcher = dispatcher_factory(func_ir)
+            dispatcher.can_cache = True
+
+            # Find live variables after the region to make sure we don't Del
+            # them if they are defined in the openmp region.
+            cfg = compute_cfg_from_blocks(blocks)
+            usedefs = compute_use_defs(blocks)
+            live_map = compute_live_map(cfg, blocks, usedefs.usemap, usedefs.defmap)
+            self.blk_end_live_map = live_map[blk_end]
+            return dispatcher
+
+    def __call__(self, args):
+        return PythonOpenmp(args)