Blame for tensor_pruning/core.py

"""NumPy interface to the OpenMP selector."""

from contextlib import contextmanager
import ctypes as ct
from functools import lru_cache
import mmap
import operator
import os
from pathlib import Path
import shlex
import subprocess
import tempfile
from threading import Lock

import numpy as np


class Projection(ct.Structure):
    _fields_ = [
        ("weights", ct.c_void_p), ("moments", ct.c_void_p), ("mask", ct.c_void_p),
        ("heads", ct.c_size_t), ("channels", ct.c_size_t), ("inputs", ct.c_size_t),
        ("dtype", ct.c_int),
    ]


_BUILD_LOCK = Lock()
_DTYPES = {"float16": 2, "bfloat16": 3, "float32": 4, "float64": 8}


@lru_cache(maxsize=1)
def _compile():
    compiler = shlex.split(os.environ.get("CXX", "g++"))
    source = Path(__file__).with_name("kernel.cpp")
    with tempfile.TemporaryDirectory(prefix="tensor-pruning-build-") as directory:
        library = Path(directory) / "kernel.so"
        command = [*compiler, "-std=c++17", "-O3", "-fopenmp", "-shared", "-fPIC",
                   str(source), "-o", str(library)]
        try:
            subprocess.run(command, check=True, capture_output=True, text=True)
        except (OSError, subprocess.CalledProcessError) as exc:
            message = getattr(exc, "stderr", None) or str(exc)
            raise RuntimeError(f"An OpenMP C++17 compiler is required. {message}") from exc
        function = ct.CDLL(str(library)).tensor_prune_many
    function.argtypes = [ct.POINTER(Projection), ct.c_size_t, ct.c_void_p,
                         ct.c_size_t, ct.c_int, ct.c_uint64]
    function.restype = ct.c_int
    return function


def _kernel():
    with _BUILD_LOCK:
        return _compile()


def _inputs(projections, input_second_moments):
    arrays = tuple(np.asarray(w) for w in projections)
    if not arrays:
        raise ValueError("at least one query projection is required")
    moments = (None,) * len(arrays) if input_second_moments is None else tuple(input_second_moments)
    if len(moments) != len(arrays):
        raise ValueError("provide one input-second-moment vector per projection")
    checked = []
    for w, v in zip(arrays, moments):
        if w.ndim != 3 or not all(w.shape):
            raise ValueError("projections must have nonempty [head, channel, input] shape")
        if w.shape[-1] > np.iinfo(np.uint32).max:
            raise ValueError("input axis exceeds the 32-bit channel index limit")
        if w.dtype.name not in _DTYPES or not w.dtype.isnative:
            raise TypeError("use native-endian float16, bfloat16, float32, or float64 weights")
        if not w.flags.c_contiguous or not w.flags.aligned:
            raise ValueError("projections must be C-contiguous and aligned")
        if v is None:
            v = np.ones(w.shape[-1], dtype=np.float64)
        else:
            v = np.asarray(v)
            if v.dtype.kind not in "fiu" or v.shape != (w.shape[-1],):
                raise ValueError("each moment vector must be real with shape [input]")
            v = np.require(v, dtype=np.float64, requirements=["C", "A"])
            if not np.all(np.isfinite(v) & (v >= 0)):
                raise ValueError("input second moments must be finite and nonnegative")
        checked.append(v)
    return arrays, checked


def _options(total, keep, threads, seed):
    keep = operator.index(keep)
    if not 0 <= keep <= total:
        raise ValueError("keep must be between zero and the total number of weights")
    threads = operator.index(threads)
    if not 1 <= threads <= np.iinfo(np.int32).max:
        raise ValueError("threads must be a positive integer")
    return keep, threads, operator.index(seed) % (1 << 64)


@contextmanager
def _partition_workspace(count, scratch_dir):
    if count == 0:
        yield None
    elif scratch_dir is None:
        indices = np.empty(count, dtype=np.uint32)
        yield indices.ctypes.data
    else:
        with tempfile.TemporaryFile(dir=scratch_dir) as file:
            size = count * np.dtype(np.uint32).itemsize
            os.posix_fallocate(file.fileno(), 0, size)
            with mmap.mmap(file.fileno(), size) as region:
                yield ct.addressof(ct.c_char.from_buffer(region))


def _run(arrays, moments, masks, total, keep, threads, seed, scratch_dir):
    descriptors = (Projection * len(arrays))(*[
        Projection(w.ctypes.data, v.ctypes.data, mask.ctypes.data, *w.shape, _DTYPES[w.dtype.name])
        for w, v, mask in zip(arrays, moments, masks)
    ])
    native = _kernel()
    with _partition_workspace(total if 0 < keep < total else 0, scratch_dir) as scratch:
        status = native(descriptors, len(descriptors), scratch, keep, threads, seed)
    if status == 1:
        raise ValueError("weights and intermediate energies must be finite")
    if status == 2:
        raise MemoryError("native partition allocation failed")
    if status:
        raise RuntimeError(f"native selector failed with status {status}")


def prune_queries(projections, keep, *, input_second_moments=None, threads=1,
                  seed=0, scratch_dir=None):
    """Return one boolean mask per [head, channel, input] projection.

    All projections participate in a single retention budget, even if they
    have different shapes or dtypes. Masks contain exactly keep true entries.
    Existing zeros count as eligible slots. Inputs can be read-only memmaps.
    Set scratch_dir to use a disk-backed partition workspace.
    """
    arrays, moments = _inputs(projections, input_second_moments)
    total = sum(w.size for w in arrays)
    keep, threads, seed = _options(total, keep, threads, seed)
    masks = tuple(np.empty(w.shape, dtype=bool) for w in arrays)
    _run(arrays, moments, masks, total, keep, threads, seed, scratch_dir)
    return masks