Spaces:
Sleeping
Sleeping
| """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} | |
| 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) | |
| 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 | |