numba-vulkan 0.1.0__py3-none-any.whl

numba_vulkan/__init__.py

@@ -0,0 +1,79 @@
+"""Proof of concept: a Vulkan compute target for Numba."""
+
+from numba_vulkan.autotune import autotune
+from numba_vulkan.codegen import CompiledKernel
+from numba_vulkan.dispatcher import VulkanDispatcher, jit
+from numba_vulkan.errors import (
+    SpirvCodegenError,
+    VulkanPerformanceWarning,
+    VulkanPrecisionWarning,
+    VulkanSupportError,
+    VulkanUnsupportedError,
+    VulkanValidationWarning,
+)
+from numba_vulkan.runtime import (
+    DeviceArray,
+    device_array,
+    device_array_like,
+    event,
+    event_elapsed_time,
+    get_device,
+    list_devices,
+    pinned_array,
+    select_device,
+    stream,
+    synchronize,
+    to_device,
+)
+from numba_vulkan.stubs import (
+    atomic,
+    barrier,
+    global_id,
+    group_id,
+    local,
+    local_id,
+    local_size,
+    num_groups,
+    shared,
+    subgroup,
+    syncthreads,
+)
+
+__all__ = [
+    "CompiledKernel",
+    "DeviceArray",
+    "SpirvCodegenError",
+    "VulkanDispatcher",
+    "VulkanPerformanceWarning",
+    "VulkanPrecisionWarning",
+    "VulkanSupportError",
+    "VulkanUnsupportedError",
+    "VulkanValidationWarning",
+    "atomic",
+    "autotune",
+    "barrier",
+    "device_array",
+    "device_array_like",
+    "event",
+    "event_elapsed_time",
+    "get_device",
+    "global_id",
+    "group_id",
+    "jit",
+    "list_devices",
+    "local",
+    "local_id",
+    "local_size",
+    "num_groups",
+    "pinned_array",
+    "select_device",
+    "shared",
+    "stream",
+    "subgroup",
+    "synchronize",
+    "syncthreads",
+    "to_device",
+]
+
+# Registers target="vulkan" with numba.vectorize and numba.guvectorize.
+from numba_vulkan import vectorizers  # noqa: F401

numba_vulkan/_emit.py

@@ -0,0 +1,98 @@
+"""Child process that turns LLVM IR into SPIR-V.
+
+LLVM's SPIR-V backend aborts the whole process on input it cannot handle,
+and it keeps state between modules that corrupts the second module
+translated by one process. It therefore runs outside the user's
+interpreter, and in a process of its own for every module.
+
+Starting Python and importing llvmlite takes much longer than the
+translation, so this process does that once and then forks for every
+request, where the platform allows it. It serves requests until its
+standard input is closed. A request is the length of the IR as an unsigned
+64-bit little-endian integer followed by the IR; the reply is framed in the
+same way, with the length `FAILED` and no data if the backend failed. Its
+message is then on standard error. See `numba_vulkan.codegen.Emitter`.
+"""
+
+import os
+import signal
+import struct
+import sys
+
+import llvmlite.binding as llvm
+
+HEADER = struct.Struct("<Q")
+FAILED = 2**64 - 1
+CAN_FORK = hasattr(os, "fork")
+# Lets the backend use float atomics; they are emitted only for devices that
+# support them (see numba_vulkan.narrowing.Mode).
+EXTENSIONS = "+SPV_EXT_shader_atomic_float_add"
+
+
+def _read(stream, count):
+    """Read exactly `count` bytes, or return ``None`` at end of input."""
+    data = b""
+    while len(data) < count:
+        chunk = stream.read(count - len(data))
+        if not chunk:
+            return None
+        data += chunk
+    return data
+
+
+def _translate(target, text, stdout):
+    """Translate one module and write the reply."""
+    module = llvm.parse_assembly(text.decode())
+    # codemodel="default": with llvmlite's "jitdefault", Windows gets "-elf"
+    # appended to the triple, and ELF output makes the backend abort.
+    machine = target.create_target_machine(opt=0, codemodel="default")
+    spirv = machine.emit_object(module)
+    stdout.write(HEADER.pack(len(spirv)) + spirv)
+    stdout.flush()
+
+
+def main():
+    """Serve translation requests on standard input and output.
+
+    The target triple is taken from the first command line argument.
+    Without ``fork``, the process translates a single module itself and
+    exits.
+    """
+    triple = sys.argv[1]
+    llvm.set_option("", f"--spirv-ext={EXTENSIONS}")
+    llvm.initialize_all_targets()
+    llvm.initialize_all_asmprinters()
+    target = llvm.Target.from_triple(triple)
+    stdin, stdout = sys.stdin.buffer, sys.stdout.buffer
+    while True:
+        header = _read(stdin, HEADER.size)
+        if header is None:
+            return
+        text = _read(stdin, HEADER.unpack(header)[0])
+        if text is None:
+            return
+        if not CAN_FORK:
+            _translate(target, text, stdout)
+            return
+        pid = os.fork()
+        if pid == 0:
+            status = 1
+            try:
+                _translate(target, text, stdout)
+                status = 0
+            except BaseException as exc:  # noqa: BLE001 - reported to the parent
+                print(f"{type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
+            finally:
+                os._exit(status)
+        status = os.waitpid(pid, 0)[1]
+        if os.WIFSIGNALED(status):
+            # A crash leaves no message of its own.
+            name = signal.Signals(os.WTERMSIG(status)).name
+            print(f"the backend was killed by {name}", file=sys.stderr, flush=True)
+        if status != 0:
+            stdout.write(HEADER.pack(FAILED))
+            stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

numba_vulkan/arrayfuncs.py

@@ -0,0 +1,467 @@
+"""Reductions over arrays, written in Python on top of element access.
+
+Numba's own ``sum``, ``min``... walk the array through its data pointer,
+which arrays on this target do not have. The versions here are overloads
+for the ``vulkan`` target that read the elements by position instead (see
+`numba_vulkan.vkimpl.flat_item`), so they work for arrays and views of any
+dimensionality. They reduce over all elements, or, with a constant
+``axis``, along one axis: kernels cannot allocate the resulting array, so
+``a.sum(axis=0)`` is an array expression (see `AxisReduction`), whose
+elements are computed where they are read.
+
+Each is available as a method (``a.sum()``) and as the NumPy function
+(``np.sum(a)``).
+"""
+
+import functools
+from dataclasses import dataclass
+
+import numpy as np
+from numba.core import cgutils, errors, types
+from numba.core.extending import intrinsic, overload, overload_method
+
+from numba_vulkan.vkimpl import flat_item
+from numba_vulkan.vktypes import VulkanArray, VulkanExpr
+
+TARGET = "vulkan"
+
+
+def _accumulator(dtype):
+    """Type in which sums and products of ``dtype`` elements are formed.
+
+    As in NumPy, integers and booleans widen to the platform integer.
+
+    Parameters
+    ----------
+    dtype : numba.types.Type
+        Element type.
+
+    Returns
+    -------
+    numba.types.Type
+    """
+    if isinstance(dtype, types.Boolean):
+        return types.intp
+    if isinstance(dtype, types.Integer):
+        return types.intp if dtype.signed else types.uintp
+    return dtype
+
+
+def _check_whole_axis(a, axis):
+    """Raise unless `axis` is None, or a valid axis of a 1-d array."""
+
+
+@overload(_check_whole_axis, target=TARGET)
+def _check_whole_axis_impl(a, axis):
+    """Build `_check_whole_axis`: reductions over all elements take it."""
+    if axis is None or isinstance(axis, (types.NoneType, types.Omitted)):
+        return lambda a, axis: None
+
+    def impl(a, axis):
+        if axis != 0 and axis != -1:
+            raise ValueError("axis is out of bounds for an array of dimension 1")
+
+    return impl
+
+
+def _sum(a):
+    """Build ``sum``: the sum of all elements."""
+    acc = _accumulator(a.dtype)
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        total = acc(0)
+        for k in range(a.size):
+            total += flat_item(a, k)
+        return total
+
+    return impl
+
+
+def _prod(a):
+    """Build ``prod``: the product of all elements."""
+    acc = _accumulator(a.dtype)
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        total = acc(1)
+        for k in range(a.size):
+            total *= flat_item(a, k)
+        return total
+
+    return impl
+
+
+def _mean(a):
+    """Build ``mean``: float64 for integers, the element type otherwise."""
+    acc = a.dtype if isinstance(a.dtype, types.Float) else types.float64
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        total = acc(0)
+        for k in range(a.size):
+            total += flat_item(a, k)
+        return total / acc(a.size)
+
+    return impl
+
+
+def _min(a):
+    """Build ``min``; raises ``ValueError`` for an empty array."""
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        if a.size == 0:
+            raise ValueError("zero-size array has no minimum")
+        best = flat_item(a, 0)
+        for k in range(1, a.size):
+            value = flat_item(a, k)
+            # written so that a NaN, once found, is kept, as in NumPy
+            if value < best or value != value:  # noqa: PLR0124 - NaN, for any type
+                best = value
+        return best
+
+    return impl
+
+
+def _max(a):
+    """Build ``max``; raises ``ValueError`` for an empty array."""
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        if a.size == 0:
+            raise ValueError("zero-size array has no maximum")
+        best = flat_item(a, 0)
+        for k in range(1, a.size):
+            value = flat_item(a, k)
+            if value > best or value != value:  # noqa: PLR0124 - NaN, for any type
+                best = value
+        return best
+
+    return impl
+
+
+def _argmin(a):
+    """Build ``argmin``: row-major position of the first minimum."""
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        if a.size == 0:
+            raise ValueError("attempt to get argmin of an empty sequence")
+        best, where = flat_item(a, 0), 0
+        for k in range(1, a.size):
+            value = flat_item(a, k)
+            if value < best:
+                best, where = value, k
+        return where
+
+    return impl
+
+
+def _argmax(a):
+    """Build ``argmax``: row-major position of the first maximum."""
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        if a.size == 0:
+            raise ValueError("attempt to get argmax of an empty sequence")
+        best, where = flat_item(a, 0), 0
+        for k in range(1, a.size):
+            value = flat_item(a, k)
+            if value > best:
+                best, where = value, k
+        return where
+
+    return impl
+
+
+def _any(a):
+    """Build ``any``: whether some element is non-zero."""
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        for k in range(a.size):
+            if flat_item(a, k):
+                return True
+        return False
+
+    return impl
+
+
+def _all(a):
+    """Build ``all``: whether every element is non-zero."""
+
+    def impl(a, axis=None):
+        _check_whole_axis(a, axis)
+        for k in range(a.size):
+            if not flat_item(a, k):
+                return False
+        return True
+
+    return impl
+
+
+@dataclass(frozen=True)
+class AxisReduction:
+    """The operation of an array expression that reduces along an axis.
+
+    ``a.sum(axis=k)`` is a `numba_vulkan.vktypes.VulkanExpr` with this as
+    its operation, and ``a`` and the axis as its operands; the axis is a
+    run-time value, as in Numba's ``sum(axis=...)`` on the CPU, since the
+    number of dimensions of the result does not depend on it. Its shape is
+    that of ``a`` without the axis; each element is a loop along the axis,
+    written in Python and compiled where the element is read (see
+    `numba_vulkan.vkimpl.expression_element`).
+
+    Parameters
+    ----------
+    name : str
+        The reduction: a key of ``_REDUCTIONS``.
+    acc : numba.types.Type
+        Type in which sums, products and means are formed.
+    """
+
+    name: str
+    acc: types.Type
+    # Tells numba_vulkan.vkimpl how to compute the shape and the elements.
+    reduces_axis = True
+
+    @property
+    def __name__(self):
+        """Name used in the expression's type name."""
+        return f"{self.name}(axis)"
+
+    def element_function(self, ndim):
+        """The function that computes one element.
+
+        Parameters
+        ----------
+        ndim : int
+            Dimensions of the operand.
+
+        Returns
+        -------
+        function
+            Takes the operand, the axis (made non-negative, and checked
+            when the shape was computed) and one index per dimension of
+            the result.
+        """
+        return _element_function(self, ndim)
+
+
+# Loops that compute one element of a reduction along an axis; ``{item}``
+# reads the element at position ``k`` along it, ``ACC`` is the accumulator
+# type.
+_ELEMENT_BODIES = {
+    "sum": """
+    total = ACC(0)
+    for k in range(n):
+        total += {item}
+    return total""",
+    "prod": """
+    total = ACC(1)
+    for k in range(n):
+        total *= {item}
+    return total""",
+    "mean": """
+    total = ACC(0)
+    for k in range(n):
+        total += {item}
+    return total / ACC(n)""",
+    "min": """
+    if n == 0:
+        raise ValueError("zero-size array has no minimum")
+    k = 0
+    best = {item}
+    for k in range(1, n):
+        value = {item}
+        if value < best or value != value:  # noqa: PLR0124 - NaN
+            best = value
+    return best""",
+    "max": """
+    if n == 0:
+        raise ValueError("zero-size array has no maximum")
+    k = 0
+    best = {item}
+    for k in range(1, n):
+        value = {item}
+        if value > best or value != value:  # noqa: PLR0124 - NaN
+            best = value
+    return best""",
+    "argmin": """
+    if n == 0:
+        raise ValueError("attempt to get argmin of an empty sequence")
+    k = 0
+    best, where = {item}, 0
+    for k in range(1, n):
+        value = {item}
+        if value < best:
+            best, where = value, k
+    return where""",
+    "argmax": """
+    if n == 0:
+        raise ValueError("attempt to get argmax of an empty sequence")
+    k = 0
+    best, where = {item}, 0
+    for k in range(1, n):
+        value = {item}
+        if value > best:
+            best, where = value, k
+    return where""",
+    "any": """
+    for k in range(n):
+        if {item}:
+            return True
+    return False""",
+    "all": """
+    for k in range(n):
+        if not {item}:
+            return False
+    return True""",
+}
+
+
+@functools.cache
+def _element_function(op, ndim):
+    """Write and define the function of `AxisReduction.element_function`."""
+    outer = [f"i{d}" for d in range(ndim - 1)]
+    position = []
+    for d in range(ndim):
+        # Index along operand axis d: an index of the result before the
+        # reduced axis, k on it, and the next index of the result after it.
+        if d == 0:
+            position.append("(k if axis == 0 else i0)")
+        elif d == ndim - 1:
+            position.append(f"(k if axis == {d} else i{d - 1})")
+        else:
+            position.append(
+                f"(i{d} if axis > {d} else (k if axis == {d} else i{d - 1}))"
+            )
+    body = _ELEMENT_BODIES[op.name].format(item=f"a[{', '.join(position)}]")
+    source = f"def element(a, axis, {', '.join(outer)}):\n    n = a.shape[axis]{body}\n"
+    scope = {"ACC": op.acc}
+    exec(source, scope)  # noqa: S102 - generated from the table above
+    return scope["element"]
+
+
+def _result_type(name, dtype):
+    """Element type of a reduction of `dtype` elements."""
+    if name in ("sum", "prod"):
+        return _accumulator(dtype)
+    if name == "mean":
+        return dtype if isinstance(dtype, types.Float) else types.float64
+    if name in ("argmin", "argmax"):
+        return types.intp
+    if name in ("any", "all"):
+        return types.boolean
+    return dtype
+
+
+@functools.cache
+def _axis_reducer(name):
+    """An intrinsic that makes the expression of a reduction along an axis."""
+
+    @intrinsic(target=TARGET)
+    def reduce_along(typingctx, a, axis):
+        acc = _result_type("mean" if name == "mean" else "sum", a.dtype)
+        op = AxisReduction(name, acc)
+        dtype = _result_type(name, a.dtype)
+        exprty = VulkanExpr(op, (a, types.intp), dtype, a.ndim - 1)
+
+        def codegen(context, builder, sig, args):
+            proxy = cgutils.create_struct_proxy(sig.return_type)(context, builder)
+            proxy.operand0 = args[0]
+            proxy.operand1 = context.cast(builder, args[1], sig.args[1], types.intp)
+            return proxy._getvalue()
+
+        return exprty(a, axis), codegen
+
+    return reduce_along
+
+
+_REDUCTIONS = {
+    "sum": (_sum, [np.sum]),
+    "prod": (_prod, [np.prod]),
+    "mean": (_mean, [np.mean]),
+    "min": (_min, [np.min, min]),
+    "max": (_max, [np.max, max]),
+    "argmin": (_argmin, [np.argmin]),
+    "argmax": (_argmax, [np.argmax]),
+    "any": (_any, [np.any]),
+    "all": (_all, [np.all]),
+}
+
+
+def _register(name, factory, functions):
+    """Register a reduction as an array method and for its functions.
+
+    Parameters
+    ----------
+    name : str
+        Name of the array method.
+    factory : callable
+        Builds the implementation for an array type.
+    functions : list of callable
+        Functions that mean the same when called with one array.
+    """
+
+    def typer(a, axis=None):
+        """Select the implementation for arrays and expressions of this target."""
+        if not isinstance(a, (VulkanArray, VulkanExpr)):
+            return None
+        if axis is None or isinstance(axis, (types.NoneType, types.Omitted)):
+            return factory(a)
+        if not isinstance(axis, types.Integer):
+            raise errors.TypingError(f"{name}(axis=...) takes an integer axis")
+        ndim = a.ndim
+        if ndim == 1:
+            return factory(a)  # a scalar, as in NumPy; it checks the axis
+        reduce_along = _axis_reducer(name)
+
+        def impl(a, axis=None):
+            # Checked where the shape is computed: raising here would make
+            # the array of the expression depend on the branch taken.
+            return reduce_along(a, axis)
+
+        return impl
+
+    def builtin_typer(a):
+        """``min(a)`` and ``max(a)`` on an array: over all elements."""
+        if isinstance(a, (VulkanArray, VulkanExpr)):
+            return factory(a)  # its axis keeps its default
+        return None
+
+    overload_method(VulkanArray, name, target=TARGET)(typer)
+    overload_method(VulkanExpr, name, target=TARGET)(typer)
+    for function in functions:
+        if function in (min, max):
+            overload(function, target=TARGET, strict=False)(builtin_typer)
+        else:
+            overload(function, target=TARGET)(typer)
+
+
+for _name, (_factory, _functions) in _REDUCTIONS.items():
+    _register(_name, _factory, _functions)
+
+
+@overload(np.dot, target=TARGET)
+def _dot(a, b):
+    """Build ``np.dot`` for two 1-d arrays."""
+    if not (
+        isinstance(a, (VulkanArray, VulkanExpr))
+        and isinstance(b, (VulkanArray, VulkanExpr))
+    ):
+        return None
+    if a.ndim != 1 or b.ndim != 1:
+        return None
+    acc = np.result_type(str(a.dtype), str(b.dtype)).type
+
+    def impl(a, b):
+        if a.shape[0] != b.shape[0]:
+            raise ValueError("incompatible array sizes for np.dot(a, b)")
+        total = acc(0)
+        for k in range(a.shape[0]):
+            total += flat_item(a, k) * flat_item(b, k)
+        return total
+
+    return impl