kernels

def jit(
    function: Callable | None = None, *, parallel: bool = False, contract: bool = False
) -> Callable:
    """Compile ``function`` with numba, cached on disk where numba can write.

    The compiled function releases the GIL, so Python threads can run it at once.
    Numba refuses to define a cached function when no cache directory is
    writable (e.g. a read-only install without ``NUMBA_CACHE_DIR``), so such
    installs compile in every process instead. With ``contract``, LLVM may fuse
    the function's multiplies and adds.
    """
    if function is None:
        return functools.partial(jit, parallel=parallel, contract=contract)
    # NOTE: numba names its cache files after the function and ignores nogil when
    # it reads them, so the name carries the flags: an entry cached without nogil
    # would keep holding the GIL.
    function = _renamed(function, "_nogil_contract" if contract else "_nogil")
    options = _prange_only() if parallel else False
    fastmath = {"contract"} if contract else False
    try:
        return numba.njit(cache=True, parallel=options, nogil=True, fastmath=fastmath)(
            function
        )
    except RuntimeError:
        return numba.njit(parallel=options, nogil=True, fastmath=fastmath)(function)

def kernel(function: Callable) -> Callable:
    """Compile ``function`` serially and in parallel; ``run(work, *args)`` picks one.

    ``work`` estimates the call's cost in simple operations; calls below
    ``PARALLEL_WORK`` run serially, and so do calls on the worker threads of
    ``run_tasks`` and calls while numba has one thread. In the serial
    compilation every ``numba.prange`` is a plain range, so both compute the
    same result.
    """
    serial = jit(function)
    # NOTE: numba names its cache files after the function, so the parallel
    # compilation needs another name to keep its own cache entry.
    parallel = jit(_renamed(function, "_prange"), parallel=True)

    def run(work: int, *args):
        if (
            work < PARALLEL_WORK
            or getattr(_TASK_THREAD, "active", False)
            or numba.get_num_threads() == 1
        ):
            return serial(*args)
        with _LOCK:
            return parallel(*args)

    run.__doc__ = function.__doc__
    return run

def run_tasks(
    function: Callable[[_Task], _Result], tasks: Iterable[_Task], threads: int
) -> list[_Result]:
    """Return ``[function(task) for task in tasks]``, on up to ``threads`` threads.

    With more than one thread the tasks run on worker threads, and every kernel
    in them runs serially: the tasks keep the cores busy, and independent tasks
    scale far better than numba's threads inside each kernel. Otherwise, and in
    a task, the tasks run in the calling thread. Every task finishes before the
    first error, if any, is raised.
    """
    tasks = list(tasks)
    if min(threads, len(tasks)) <= 1 or getattr(_TASK_THREAD, "active", False):
        return [function(task) for task in tasks]
    slots = threading.Semaphore(threads)

    def run(task: _Task) -> _Result:
        with slots:
            return function(task)

    futures = [_pool().submit(run, task) for task in tasks]
    wait(futures)
    return [future.result() for future in futures]