tack-core 0.7.0__py3-none-any.whl

tack/__init__.py

@@ -0,0 +1,294 @@
+# SPDX-FileCopyrightText: Copyright (c) Kitware Inc.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tack: a Python-first compute framework.
+
+Kernels are Python functions decorated with ``@tack.kernel``, compiled when
+first called for the device they run on: the CPU (through LLVM), Metal, CUDA,
+HIP or Level Zero. Arrays (``tack.empty``, ``tack.array``, ...) hold typed
+values on a device, and a kernel runs where its arrays are. ``tack.dataset``,
+in the tack-vis package, adds meshes, fields on them, and the algorithms of
+scientific visualization.
+"""
+
+# Enable namespace package merging across tack-core and tack-vis
+from importlib.metadata import PackageNotFoundError as _NotInstalled
+from importlib.metadata import version as _installed_version
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+del extend_path
+
+try:
+    __version__ = _installed_version("tack-core")
+except _NotInstalled:          # a source tree that was never installed
+    __version__ = "0+unknown"
+del _NotInstalled, _installed_version
+
+from tack.lang.array import (
+    Array,
+    ExportedMemory,
+    SoaArray,
+    Texture3D,
+    arange,
+    array,
+    array_from_ptr,
+    asarray,
+    concat,
+    copyto,
+    empty,
+    empty_like,
+    from_dlpack,
+    full,
+    full_like,
+    memory_space,
+    ones,
+    ones_like,
+    soa,
+    zeros,
+    zeros_like,
+)
+from tack.lang.constant import constant
+from tack.lang.device_class import device_class
+from tack.lang.func import func
+from tack.lang.inspect_kernel import inspect
+from tack.lang.kernel import kernel
+from tack.lang.types import (
+    ElementType,
+    ScalarType,
+    f32,
+    f64,
+    get_default_float,
+    i8,
+    i16,
+    i32,
+    i64,
+    mat,
+    mat2,
+    mat3,
+    mat4,
+    set_default_float,
+    u8,
+    u16,
+    u32,
+    u64,
+    vec,
+    vec2,
+    vec3,
+    vec4,
+)
+from tack.runtime.dispatch import (
+    Device,
+    device,
+    device_count,
+    get_default_device,
+    set_debug,
+    set_default_device,
+    sync,
+    sync_async,
+)
+
+
+# Shared memory, barrier, thread_id — only usable inside @tack.kernel
+def shared(dtype, size):
+    """Allocate threadgroup shared memory. Only usable inside a @tack.kernel."""
+    raise RuntimeError("shared() can only be used inside a @tack.kernel")
+
+def shared_like(array, size):
+    """Allocate shared memory with the same dtype as an array. Only usable inside a @tack.kernel."""
+    raise RuntimeError("shared_like() can only be used inside a @tack.kernel")
+
+def select(mask, a, b):
+    """``a`` where ``mask`` is true, else ``b``, component by component.
+
+    Only usable inside a kernel or device function. ``mask`` is a scalar
+    or a comparison vector such as ``v > 0.0``; ``a`` and ``b`` are scalars
+    or vectors of one shape, a scalar standing for every component.
+    """
+    raise RuntimeError("tack.select() can only be used inside a @tack.kernel")
+
+
+def barrier():
+    """Threadgroup synchronization barrier. Only usable inside a @tack.kernel."""
+    raise RuntimeError("barrier() can only be used inside a @tack.kernel")
+
+def thread_id():
+    """Thread index within workgroup. Only usable inside a @tack.kernel."""
+    raise RuntimeError("thread_id() can only be used inside a @tack.kernel")
+
+# Atomic operations — only usable inside @tack.kernel
+def atomic_add(array, index, value):
+    """Atomic add: array[index] += value. Only usable inside a @tack.kernel."""
+    raise RuntimeError("atomic_add() can only be used inside a @tack.kernel")
+
+def atomic_min(array, index, value):
+    """Atomic min: array[index] = min(array[index], value). Only usable inside a @tack.kernel."""
+    raise RuntimeError("atomic_min() can only be used inside a @tack.kernel")
+
+def atomic_max(array, index, value):
+    """Atomic max: array[index] = max(array[index], value). Only usable inside a @tack.kernel."""
+    raise RuntimeError("atomic_max() can only be used inside a @tack.kernel")
+
+# ndindex for multi-dimensional parallel iteration
+def ndindex(*args):
+    """Multi-dimensional parallel iteration, as ``np.ndindex``.
+
+    Used in kernels: ``for i, j in tack.ndindex(w, h)``; each argument is a
+    size or a ``(start, end)`` pair. Cannot be called from Python directly.
+    """
+    raise RuntimeError("ndindex() can only be used inside a @tack.kernel")
+
+
+def identity(n):
+    """The n-by-n identity matrix, as ``np.identity``. Only usable inside a @tack.kernel."""
+    raise RuntimeError("identity() can only be used inside a @tack.kernel")
+
+def block_sum(value):
+    """Sum a value across all threads in the workgroup. Only usable inside a @tack.kernel."""
+    raise RuntimeError("block_sum() can only be used inside a @tack.kernel")
+
+def block_max(value):
+    """Max of a value across all threads in the workgroup. Only usable inside a @tack.kernel."""
+    raise RuntimeError("block_max() can only be used inside a @tack.kernel")
+
+def block_min(value):
+    """Min of a value across all threads in the workgroup. Only usable inside a @tack.kernel."""
+    raise RuntimeError("block_min() can only be used inside a @tack.kernel")
+
+def local_array(dtype, size):
+    """Allocate a per-thread local array. Only usable inside a @tack.kernel.
+
+    Usage: arr = tack.local_array(tack.f32, 8)
+    Then: arr[i] = value, value = arr[i]
+    """
+    raise RuntimeError("local_array() can only be used inside a @tack.kernel")
+
+def local_array_like(array, size):
+    """Allocate a per-thread local array with the same dtype as an array. Only usable inside a @tack.kernel."""
+    raise RuntimeError("local_array_like() can only be used inside a @tack.kernel")
+
+def texture3d(source_array, shape=None, interp='linear'):
+    """Create a 3D texture from an array for hardware-accelerated sampling.
+
+    The texture copies the array's data now. Later writes to the array --
+    kernel stores, ``from_numpy``, external or DLPack writes -- are not
+    visible to it until ``tex.update()`` copies the array again. That holds
+    on every backend, including CPU.
+
+    Args:
+        source_array: A Tack array with f32 dtype holding W*H*D elements,
+               x varying fastest.
+        shape: (W, H, D) tuple for the 3D dimensions.  If the array's shape
+               is already 3D, this can be omitted.
+        interp: Interpolation mode. Only 'linear' (the default) is
+               implemented.
+    """
+    if shape is None:
+        shape = source_array.shape
+    if len(shape) != 3:
+        raise ValueError("texture3d requires a 3D shape")
+    return Texture3D(source_array, shape_3d=shape, interp=interp)
+
+
+#: Public submodules, loaded on first use, so ``import tack`` alone reaches
+#: ``tack.arrays.CountingArray``. ``dataset`` and ``interop`` are tack-vis's.
+#: How to install tack-vis, which is not on PyPI: from the repository, at this
+#: release's tag (a development install has no tag to name).
+_TACK_VIS_INSTALL = (
+    "pip install 'tack-vis @ git+https://github.com/Kitware/tack"
+    + ("" if "+" in __version__ or "dev" in __version__ else f"@v{__version__}")
+    + "#subdirectory=packages/tack-vis'")
+
+_SUBMODULES = frozenset({"algorithms", "arrays", "cache", "communication", "math", "random",
+                         "dataset", "interop"})
+
+
+def __getattr__(name):
+    if name in _SUBMODULES:
+        import importlib
+        try:
+            return importlib.import_module(f"{__name__}.{name}")
+        except ModuleNotFoundError as error:
+            if error.name != f"{__name__}.{name}":
+                raise
+            raise AttributeError(f"module 'tack' has no attribute {name!r}: tack.{name} is "
+                                 f"in the tack-vis package, which is not installed. It is "
+                                 f"not on PyPI; install it from the repository:\n"
+                                 f"    {_TACK_VIS_INSTALL}") from error
+    raise AttributeError(f"module 'tack' has no attribute {name!r}")
+
+
+__all__ = [
+    "Array",
+    "Device",
+    "ElementType",
+    "ExportedMemory",
+    "ScalarType",
+    "SoaArray",
+    "Texture3D",
+    "arange",
+    "array",
+    "array_from_ptr",
+    "asarray",
+    "atomic_add",
+    "atomic_max",
+    "atomic_min",
+    "barrier",
+    "block_max",
+    "block_min",
+    "block_sum",
+    "concat",
+    "constant",
+    "copyto",
+    "device",
+    "device_class",
+    "device_count",
+    "empty",
+    "empty_like",
+    "f32",
+    "f64",
+    "from_dlpack",
+    "full",
+    "full_like",
+    "func",
+    "get_default_device",
+    "get_default_float",
+    "i8",
+    "i16",
+    "i32",
+    "i64",
+    "identity",
+    "inspect",
+    "kernel",
+    "local_array",
+    "local_array_like",
+    "mat",
+    "mat2",
+    "mat3",
+    "mat4",
+    "memory_space",
+    "ndindex",
+    "ones",
+    "ones_like",
+    "select",
+    "set_debug",
+    "set_default_device",
+    "set_default_float",
+    "shared",
+    "shared_like",
+    "soa",
+    "sync",
+    "sync_async",
+    "texture3d",
+    "thread_id",
+    "u8",
+    "u16",
+    "u32",
+    "u64",
+    "vec",
+    "vec2",
+    "vec3",
+    "vec4",
+    "zeros",
+    "zeros_like",
+]

tack/algorithms/__init__.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) Kitware Inc.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""tack.algorithms — GPU algorithm primitives.
+
+Portable, backend-agnostic building blocks for common parallel patterns:
+scan, sort, segmented reductions, copy, fill, stats. Their inputs are any
+arrays a kernel indexes -- stored Tack arrays or the implicit arrays of
+``tack.arrays`` -- and their outputs stored arrays; they run as kernels on
+their inputs' device (CPU, Metal, CUDA, HIP, Level Zero), reading back at
+most a scalar result.
+
+Visualization algorithms are in ``tack.dataset`` (tack-vis).
+"""
+
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+
+from tack.algorithms.copy import copy, fill_value
+from tack.algorithms.scan import exclusive_scan, inclusive_scan
+from tack.algorithms.segmented import (
+    segment_broadcast,
+    segment_ids,
+    segment_reduce,
+    segment_scan,
+)
+from tack.algorithms.sort import argsort, gather, reduce_by_key, sort_by_key, unique
+from tack.algorithms.stats import (
+    absmax,
+    count_nonzero,
+    dot,
+    histogram,
+    norm,
+    std,
+    var,
+)
+
+__all__ = [
+    "absmax",
+    "argsort",
+    "copy",
+    "count_nonzero",
+    "dot",
+    "exclusive_scan",
+    "fill_value",
+    "gather",
+    "histogram",
+    "inclusive_scan",
+    "norm",
+    "reduce_by_key",
+    "segment_broadcast",
+    "segment_ids",
+    "segment_reduce",
+    "segment_scan",
+    "sort_by_key",
+    "std",
+    "unique",
+    "var",
+]

tack/algorithms/copy.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: Copyright (c) Kitware Inc.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Array copy and fill utilities, each a kernel on its arrays' device.
+
+Usage:
+    from tack import algorithms
+    from tack.algorithms.copy import copy_with_offset
+
+    algorithms.copy(src, dst, n)
+    algorithms.fill_value(array, value, n)
+    copy_with_offset(src, dst, dst_offset, n)
+"""
+
+import tack
+
+
+@tack.kernel
+def _copy_kernel(src, dst, n):
+    for i in range(n):
+        dst[i] = src[i]
+
+
+@tack.kernel
+def _fill_kernel(dst, val, n):
+    for i in range(n):
+        dst[i] = val
+
+
+@tack.kernel
+def _copy_offset_kernel(src, dst, dst_offset, n):
+    """Copy n elements from src into dst starting at dst_offset."""
+    for i in range(n):
+        dst[dst_offset + i] = src[i]
+
+
+def copy(src, dst, n):
+    """Copy n elements from src array to dst array with a kernel."""
+    _copy_kernel(src, dst, n)
+
+
+def copy_with_offset(src, dst, dst_offset, n):
+    """Copy n elements from src into dst starting at dst_offset with a kernel."""
+    _copy_offset_kernel(src, dst, dst_offset, n)
+
+
+def fill_value(dst, value, n):
+    """Fill n elements of dst array with a scalar value with a kernel."""
+    _fill_kernel(dst, value, n)

tack/algorithms/scan.py

@@ -0,0 +1,292 @@
+# SPDX-FileCopyrightText: Copyright (c) Kitware Inc.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Parallel prefix sum (scan) on tack arrays.
+
+Reduce, then scan, by chunks: one thread per chunk of ``_CHUNK``
+elements totals it, the chunk totals are scanned the same way, and each
+chunk then writes its running sums from its offset. That is a few launches
+per factor of ``_CHUNK`` -- three levels for a million elements -- where
+the Blelloch scan this replaced made two per factor of two, 44 for a
+million, and most of a filter's time went to them. No shared memory or
+barriers -- works on all backends.
+
+Usage:
+    import tack
+    from tack import algorithms
+
+    total = algorithms.exclusive_scan(counts, offsets, n)
+    algorithms.inclusive_scan(input_array, output_array, n)
+"""
+
+import tack
+from tack.runtime.dispatch import on_inputs_device
+
+_CHUNK = 256  # elements one thread scans in turn
+# A level of at most this many elements is scanned by one thread in one
+# launch: below it the recursion's launches cost more than the work (the
+# chunk totals of a million elements, 3907 of them, took two more levels
+# and five launches of a few microseconds' work each).
+_SERIAL_MAX = 4096
+
+
+@tack.kernel
+def _chunk_sums(src, dst, sums, n, chunk, nchunks):
+    """Copy each chunk into ``dst`` -- converting to its dtype on the store --
+    and total it in that dtype."""
+    for c in range(nchunks):
+        start = c * chunk
+        end = min(start + chunk, n)
+        dst[start] = src[start]
+        total = dst[start]
+        for i in range(start + 1, end):
+            dst[i] = src[i]
+            total += dst[i]
+        sums[c] = total
+
+
+@tack.kernel
+def _chunk_scan(data, offsets, n, chunk, nchunks, inclusive):
+    """Each chunk's running sums, in place, starting from its offset."""
+    for c in range(nchunks):
+        running = offsets[c]
+        for i in range(c * chunk, min(c * chunk + chunk, n)):
+            value = data[i]
+            if inclusive == 1:
+                running += value
+                data[i] = running
+            else:
+                data[i] = running
+                running += value
+
+
+# On backends with workgroups the chunk kernels above run few threads with
+# long serial loops; there a tile of _TILE elements, _PER_LANE to each of a
+# workgroup's 256 lanes, is scanned by the workgroup: each lane totals its
+# elements in order, the lane totals are scanned in shared memory, and each
+# lane then writes running sums from its offset. Same levels and dtype rules
+# as the chunks; between lanes the additions run in tree order.
+
+_TILE = tack.constant(2048)
+_PER_LANE = tack.constant(8)
+
+
+@tack.kernel
+def _tile_sums(src, dst, sums, n, ntiles):
+    """Copy each tile into ``dst``, converting on the store, and total it there."""
+    lanes = tack.shared_like(dst, 256)
+    for i in range(ntiles * 256):
+        t = tack.thread_id()
+        start = (i // 256) * _TILE + t * _PER_LANE
+        lanes[t] = 0
+        if start < n:
+            dst[start] = src[start]
+            total = dst[start]
+            for g in range(start + 1, min(start + _PER_LANE, n)):
+                dst[g] = src[g]
+                total += dst[g]
+            lanes[t] = total
+        tack.barrier()
+        step = 1
+        while step < 256:
+            below = lanes[t - step] if t >= step else lanes[t] * 0
+            tack.barrier()
+            if t >= step:
+                lanes[t] = lanes[t] + below
+            tack.barrier()
+            step = step * 2
+        if t == 0:
+            sums[i // 256] = lanes[255]
+
+
+@tack.kernel
+def _tile_apply(data, offsets, n, ntiles, inclusive):
+    """Each tile's running sums, in place, from its offset."""
+    lanes = tack.shared_like(data, 256)
+    for i in range(ntiles * 256):
+        t = tack.thread_id()
+        start = (i // 256) * _TILE + t * _PER_LANE
+        lanes[t] = 0
+        if start < n:
+            total = data[start]
+            for g in range(start + 1, min(start + _PER_LANE, n)):
+                total += data[g]
+            lanes[t] = total
+        tack.barrier()
+        step = 1
+        while step < 256:
+            below = lanes[t - step] if t >= step else lanes[t] * 0
+            tack.barrier()
+            if t >= step:
+                lanes[t] = lanes[t] + below
+            tack.barrier()
+            step = step * 2
+        # The lanes before this one, read from its neighbour rather than
+        # subtracted, so floating sums are not rounded twice.
+        running = offsets[i // 256]
+        if t > 0:
+            running = running + lanes[t - 1]
+        for g in range(start, min(start + _PER_LANE, n)):
+            value = data[g]
+            if inclusive == 1:
+                running += value
+                data[g] = running
+            else:
+                data[g] = running
+                running += value
+
+
+@tack.kernel
+def _scan_serial(src, dst, total, n, inclusive, close):
+    """One thread scans ``[0, n)`` into ``dst`` and leaves the sum in
+    ``total[0]`` -- and with ``close`` in ``dst[n]`` too, the end of the
+    offsets -- every value converted to ``dst``'s dtype on the copy so the
+    sums form in it, as the chunk kernels' do."""
+    for i in range(1):
+        dst[0] = src[0]
+        running = dst[0] * 0
+        for k in range(n):
+            dst[k] = src[k]
+            value = dst[k]
+            if inclusive == 1:
+                running += value
+                dst[k] = running
+            else:
+                dst[k] = running
+                running += value
+        total[0] = running
+        if close == 1:
+            dst[n] = running
+
+
+def _count(n, *arrays):
+    """Check the element count: the scan reads and writes [0, n)."""
+    n = int(n)
+    for f in arrays:
+        if not 0 <= n <= f.size:
+            raise ValueError(
+                f"n={n} is outside [0, {f.size}] for an array of {f.size} elements")
+    return n
+
+
+def _total(array, n):
+    """array[n - 1] as a Python number: one element read on the host, which
+    on every backend is a copy of just that element and no launch."""
+    return array[n - 1]
+
+
+@tack.kernel
+def _close_from(dst, n, offsets, sums, nchunks):
+    # dst[n]: the last chunk's offset plus its sum, the whole's total.
+    for i in range(1):
+        dst[n] = offsets[nchunks - 1] + sums[nchunks - 1]
+
+
+@tack.kernel
+def _close_empty(dst):
+    for i in range(1):
+        dst[0] = dst[0] * 0
+
+
+def _scan(src, dst, n, inclusive, want_total=True, close=False):
+    """Scan ``src`` into ``dst`` (which may be ``src``) in ``dst``'s dtype and return
+    the total: chunk sums, then -- recursively -- their exclusive scan as the
+    chunks' offsets, then each chunk's own running sums. A few launches per
+    level, and a level per factor of ``_CHUNK``. Without ``want_total`` the
+    total is never read to the host, so nothing waits for the device; the
+    result is ``None``."""
+    from tack.runtime.dispatch import get_backend
+
+    tiles = get_backend().supports_workgroups
+    size = int(_TILE) if tiles else _CHUNK
+    if n == 0:
+        return 0 if want_total else None
+    if n <= (size if tiles else _SERIAL_MAX):
+        # One launch; with workgroups a level of one tile is one thread's
+        # worth anyway. The total is read as one element, not a launch.
+        total = tack.empty((1,), dst.dtype)
+        _scan_serial(src, dst, total, n, 1 if inclusive else 0, 1 if close else 0)
+        return total[0] if want_total else None
+    nchunks = (n + size - 1) // size
+    sums = tack.empty((nchunks,), dst.dtype)
+    if tiles:
+        _tile_sums(src, dst, sums, n, nchunks)
+    else:
+        _chunk_sums(src, dst, sums, n, _CHUNK, nchunks)
+    offsets = tack.empty((nchunks,), dst.dtype)
+    total = _scan(sums, offsets, nchunks, inclusive=False, want_total=want_total)
+    if tiles:
+        _tile_apply(dst, offsets, n, nchunks, 1 if inclusive else 0)
+    else:
+        _chunk_scan(dst, offsets, n, _CHUNK, nchunks, 1 if inclusive else 0)
+    if close:
+        _close_from(dst, n, offsets, sums, nchunks)
+    return total
+
+
+@on_inputs_device
+def exclusive_scan(input_array, output_array, n, total=True, close=False):
+    """Compute exclusive prefix sum on the inputs' device.
+
+    output[i] = sum(input[0..i-1]), output[0] = 0.
+
+    The scan runs in the output array's dtype; input values convert to it
+    on the copy, so any dtype the backend allocates works.
+
+    Args:
+        input_array: array with input values.
+        output_array: array for the output offsets.
+        n: number of elements, at most either array's size. Zero writes
+            nothing and returns 0.
+        total: whether to return the total. Reading it waits for the
+            device (on Metal, for every launch queued so far); a caller
+            that needs only the offsets passes ``total=False`` and gets
+            ``None``, and the scan stays queued.
+        close: also write the total at ``output[n]``, so ``output`` is the
+            ``n + 1`` offsets of CSR rows (``output[0]`` is 0 when ``n`` is
+            0). In the scan's own launch where it is one, as for any array
+            of at most a tile or 4096 elements; ``output`` needs ``n + 1``
+            elements.
+
+    Returns:
+        The total of all input elements in the output's dtype, as a Python
+        int or float; ``None`` with ``total=False``.
+    """
+    n = _count(n, input_array, output_array)
+    if close and output_array.size < n + 1:
+        raise ValueError(f"exclusive_scan with close writes output[{n}]: the output has "
+                         f"{output_array.size} elements")
+    if n == 0:
+        # Nothing to scan; the empty sum is 0, as for Array.sum().
+        if close:
+            _close_empty(output_array)
+        return 0 if total else None
+    # Scan in the output's dtype, as the inclusive scan does: an i32 work
+    # buffer truncated float inputs and wrapped wider integers.
+    return _scan(input_array, output_array, n, inclusive=False, want_total=total, close=close)
+
+
+@on_inputs_device
+def inclusive_scan(input_array, output_array, n, total=True):
+    """Compute inclusive prefix sum on the inputs' device.
+
+    output[i] = sum(input[0..i]).
+
+    The scan runs in the output array's dtype; input values convert to it
+    on the copy, so any dtype the backend allocates works.
+
+    Args:
+        input_array: array with input values.
+        output_array: array for the output sums.
+        n: number of elements, at most either array's size. Zero writes
+            nothing and returns 0.
+        total: whether to return the total; see ``exclusive_scan``.
+
+    Returns:
+        The total of all input elements in the output's dtype, as a Python
+        int or float; ``None`` with ``total=False``.
+    """
+    n = _count(n, input_array, output_array)
+    if n == 0:
+        return 0 if total else None
+    return _scan(input_array, output_array, n, inclusive=True, want_total=total)