fastpose 1.0.0__py3-none-any.whl

fastpose/__init__.py

@@ -0,0 +1,6 @@
+import fastpose.estimators
+import fastpose.refiners
+import fastpose.scorers
+import fastpose.solvers
+
+__all__ = ["estimators", "refiners", "scorers", "solvers"]

fastpose/clean_cache.py

@@ -0,0 +1,170 @@
+"""Command-line removal of fastpose's cached Numba kernels.
+
+`fastpose-warmup` writes a lot: a full warmup produces hundreds of MB of
+`.nbc`/`.nbi` files, and nothing ever reclaims them. Superseded entries are not
+overwritten in place either - the index is keyed on each source file's
+timestamp and size and on the host CPU, so editing a kernel, upgrading numba or
+upgrading fastpose orphans the old entries rather than replacing them, and pip
+does not remove them on upgrade because they are generated at runtime and never
+appear in the wheel's RECORD.
+
+Numba may put them in any of three places depending on how fastpose is
+installed, which is why finding them by hand is awkward:
+
+    NUMBA_CACHE_DIR   if that variable is set (UserProvidedCacheLocator)
+    <package>/**/__pycache__   the usual case  (InTreeCacheLocator)
+    a user-wide cache directory, when the install directory is not writable,
+                      e.g. a system or container install (UserWideCacheLocator)
+
+This checks all three. It only ever deletes `.nbc` and `.nbi` files that sit in
+a cache directory belonging to a fastpose source directory: never `.py`, never
+`.pyc` (those are CPython's own and are cheap to rebuild), and never anything
+under another package. Deleting them is always safe - the kernels recompile on
+next use, or run `fastpose-warmup` to rebuild them up front.
+"""
+
+import argparse
+import os
+
+import fastpose
+
+# what numba writes: the per-signature index and the compiled data files
+CACHE_SUFFIXES = ('.nbi', '.nbc')
+
+
+def _package_dirs():
+    """Every directory of the installed fastpose package that holds sources."""
+    root = os.path.dirname(os.path.abspath(fastpose.__file__))
+    dirs = []
+    for dirpath, dirnames, filenames in os.walk(root):
+        dirnames[:] = [d for d in dirnames if d != '__pycache__']
+        if any(f.endswith('.py') for f in filenames):
+            dirs.append(dirpath)
+    return dirs
+
+
+def cache_dirs():
+    """Cache directories numba could have used for fastpose's kernels.
+
+    Mirrors the three locators in `numba.core.caching` that apply to a package
+    backed by real source files, using numba's own subpath helper so this keeps
+    agreeing with it. Directories that do not exist are still returned; callers
+    filter them.
+    """
+    from numba import config
+    from numba.core.caching import _CacheLocator
+
+    candidates = []
+    bases = []
+    if getattr(config, 'CACHE_DIR', None):
+        bases.append(config.CACHE_DIR)
+    try:
+        from numba.misc.appdirs import AppDirs
+        bases.append(AppDirs(appname='numba', appauthor=False).user_cache_dir)
+    except Exception:  # pragma: no cover - numba internals moved
+        pass
+
+    for source_dir in _package_dirs():
+        # in-tree: the __pycache__ beside the sources
+        candidates.append(os.path.join(source_dir, '__pycache__'))
+        # the redirected and user-wide variants, which hash the source path
+        marker = os.path.join(source_dir, 'x.py')
+        for base in bases:
+            subpath = _CacheLocator.get_suitable_cache_subpath(marker)
+            candidates.append(os.path.join(base, subpath))
+
+    seen = set()
+    unique = []
+    for path in candidates:
+        key = os.path.normcase(os.path.abspath(path))
+        if key not in seen:
+            seen.add(key)
+            unique.append(path)
+    return unique
+
+
+def find_cache_files():
+    """(path, size) of every fastpose kernel cache file, largest first."""
+    found = []
+    for directory in cache_dirs():
+        if not os.path.isdir(directory):
+            continue
+        for name in os.listdir(directory):
+            if not name.endswith(CACHE_SUFFIXES):
+                continue
+            path = os.path.join(directory, name)
+            if os.path.isfile(path):
+                found.append((path, os.path.getsize(path)))
+    found.sort(key=lambda item: -item[1])
+    return found
+
+
+def _human(num_bytes):
+    size = float(num_bytes)
+    for unit in ('B', 'KB', 'MB', 'GB'):
+        if size < 1024.0 or unit == 'GB':
+            return f'{size:.0f} {unit}' if unit == 'B' else f'{size:.1f} {unit}'
+        size /= 1024.0
+
+
+def clean(dry_run=False, verbose=False):
+    """Delete the cache files; returns (num_removed, bytes_removed)."""
+    files = find_cache_files()
+    removed = 0
+    freed = 0
+    for path, size in files:
+        if dry_run:
+            removed += 1
+            freed += size
+            if verbose:
+                print(f'  would remove {path} ({_human(size)})')
+            continue
+        try:
+            os.remove(path)
+        except OSError as exc:
+            print(f'  could not remove {path}: {exc}')
+            continue
+        removed += 1
+        freed += size
+        if verbose:
+            print(f'  removed {path} ({_human(size)})')
+    return removed, freed
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(
+        description="Remove fastpose's cached Numba kernels (.nbi/.nbc). They "
+                    "recompile on next use; `fastpose-warmup` rebuilds them "
+                    "up front.",
+    )
+    parser.add_argument(
+        '-n', '--dry-run', action='store_true',
+        help='Report what would be removed without deleting anything.',
+    )
+    parser.add_argument(
+        '-v', '--verbose', action='store_true',
+        help='List every file rather than just the total.',
+    )
+    parser.add_argument(
+        '--list-dirs', action='store_true',
+        help='Print the cache directories that were searched and exit.',
+    )
+    args = parser.parse_args(argv)
+
+    if args.list_dirs:
+        for directory in cache_dirs():
+            mark = '' if os.path.isdir(directory) else '  (does not exist)'
+            print(f'{directory}{mark}')
+        return 0
+
+    removed, freed = clean(dry_run=args.dry_run, verbose=args.verbose)
+    if not removed:
+        print('No cached kernels found.')
+        return 0
+    verb = 'Would remove' if args.dry_run else 'Removed'
+    print(f'{verb} {removed} file(s), {_human(freed)}.')
+    return 0
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())

fastpose/cuda/__init__.py

@@ -0,0 +1,104 @@
+"""CUDA backend: batch-parallel LO-RANSAC on the GPU.
+
+The CPU engine draws one hypothesis at a time (or a small batch across numba
+threads) and is bounded by how fast a single core can solve and score. This
+backend instead keeps thousands of hypotheses in flight: a round of `batch`
+minimal samples is solved one-thread-per-hypothesis, scored one-block-per-
+hypothesis with the block reducing over all correspondences, and locally
+optimized one-block-per-candidate. Only a handful of scalars cross the PCIe
+bus per round.
+
+That makes it a win in exactly the regime the CPU driver is worst at - many
+iterations and many matches - and a loss on small problems, where a round of
+kernel launches (~30-60us) costs more than the few hundred microseconds the
+CPU driver would have needed in total.
+
+Layout
+------
+Everything that is not specific to one estimation problem is shared:
+
+    ransac.py     the batched driver - round loop, local-optimization gate,
+                  adaptive termination, packed readback
+    scoring.py    the block-reduction truncated (MSAC) scorer
+    lm.py         the whole LM loop inside one kernel launch
+    reductions.py the tree reductions and the damped Cholesky solve
+    problem.py    `CudaProblem`, the seam between the two
+    problems/     one module per estimation problem: the solve kernel's
+                  scratch, and the per-point device functions the scorer and
+                  the LM call
+
+What is shared with the CPU path and what is not
+------------------------------------------------
+The minimal-solver math is *the same source*: `solvers/essential.py` builds
+its kernels through `build_five_point_kernels(jit)` and `problems/essential.py`
+instantiates that factory with `cuda.jit(device=True)` instead of `njit`.
+There is no second copy of the Sturm/Danilevsky chain to keep in sync, and the
+same holds for every other problem.
+
+Two things could not be shared, both because numba's runtime is host-only, so
+neither `np.empty` nor `.reshape` compiles in device code:
+
+- scratch is passed in pre-shaped (see the note in `solvers/essential.py`);
+  the CUDA kernels allocate it with `cuda.local.array`, which the hardware
+  interleaves across threads, so the accesses coalesce for free.
+- the scorer and the LM accumulate are *reductions* here, not serial loops,
+  so their drivers are written fresh in `scoring.py` and `lm.py` rather than
+  re-decorated. They are checked against the CPU kernels point-for-point in
+  `tests/test_cuda.py`.
+
+Availability
+------------
+`is_available()` reports whether a usable CUDA device is present; `require()`
+raises with the reason if not. Import of this package never fails and never
+initializes a context, so `import fastpose` stays cheap on machines with no
+GPU.
+"""
+
+_UNAVAILABLE_REASON = None
+_AVAILABLE = None
+
+
+def _probe():
+    global _AVAILABLE, _UNAVAILABLE_REASON
+    if _AVAILABLE is not None:
+        return _AVAILABLE
+    try:
+        from numba import cuda
+    except ImportError as exc:  # pragma: no cover - numba is a hard dependency
+        _AVAILABLE, _UNAVAILABLE_REASON = False, f"numba.cuda unimportable: {exc}"
+        return False
+    try:
+        if not cuda.is_available():
+            _AVAILABLE = False
+            _UNAVAILABLE_REASON = (
+                "numba.cuda.is_available() is False - no CUDA driver or no "
+                "visible device")
+            return False
+        # is_available() only checks the driver; touching a device is what
+        # catches a driver/toolkit mismatch or an already-exhausted GPU
+        cuda.current_context()
+    except Exception as exc:
+        _AVAILABLE, _UNAVAILABLE_REASON = False, f"{type(exc).__name__}: {exc}"
+        return False
+    _AVAILABLE, _UNAVAILABLE_REASON = True, None
+    return True
+
+
+def is_available():
+    # True when a CUDA device can actually be used, not merely imported
+    return _probe()
+
+
+def unavailable_reason():
+    # human-readable reason is_available() returned False, or None
+    _probe()
+    return _UNAVAILABLE_REASON
+
+
+def require():
+    # raises RuntimeError with the probe's reason; used by the device='cuda'
+    # entry points so a missing GPU fails with something actionable
+    if not _probe():
+        raise RuntimeError(
+            "device='cuda' requires a working CUDA device: "
+            f"{unavailable_reason()}")

fastpose/cuda/backend.py

@@ -0,0 +1,63 @@
+"""CUDA side of the `jit` shim in `fastpose/jit_backend.py`.
+
+Kept in its own module so importing it is the only thing that touches
+`numba.cuda`; `fastpose.cuda.__init__` stays importable (and cheap) on a
+machine with no GPU.
+"""
+
+import contextlib
+import warnings
+
+from numba import cuda
+from numba.core.errors import NumbaPerformanceWarning
+
+# Threads per block for the reduction kernels (scorer, LM accumulate). Must be
+# a power of two - the tree reductions halve the stride - and at least a warp.
+# 128 measured a reasonable default: large enough to hide the global-memory
+# latency of streaming the correspondence columns, small enough that a block
+# per hypothesis still fills the machine at modest batch sizes.
+THREADS_PER_BLOCK = 128
+
+# Threads per block for the one-thread-per-hypothesis solve kernel. The 5-point
+# solver carries ~6.9 KB of per-thread local memory (see solvers.py), so a
+# small block keeps the per-SM local-memory footprint sane.
+SOLVE_THREADS_PER_BLOCK = 64
+
+
+def cuda_jit(fastmath=False, inline=False):
+    # GPU instantiation of a kernel written against the shim. Device functions
+    # are inlined into the kernel that calls them, so there is no separate
+    # on-disk cache entry for them - `cache=True` goes on the kernels instead.
+    return cuda.jit(device=True, fastmath=fastmath, inline=inline)
+
+
+@contextlib.contextmanager
+def quiet_low_occupancy():
+    """Silence numba's low-occupancy warning for the launches inside.
+
+    numba warns on every launch whose grid is under 128 blocks, and several of
+    this backend's launches are deliberately smaller: sampling and solving are
+    one thread per hypothesis (32 and 64 blocks at the default batch), local
+    optimization refines a single candidate per round (one block, see the note
+    in cuda/ransac.py), and the last round of an adaptive run scores however
+    many iterations are left. Only the main scorer's grid is a function of the
+    batch, and that one is 4096 blocks by default. So the warning says nothing
+    actionable here.
+
+    Filtered by message rather than by category, so a NumbaPerformanceWarning
+    that *is* actionable - the host-array copy one, say - still comes through.
+    The leading `.*` is not decoration: numba's error classes render `str()`
+    with ANSI highlighting, and the filter is matched against that, so a
+    pattern anchored at 'Grid size' never fires.
+
+    `warnings` filters are process-global state, so this is scoped to the
+    driver's own launches and is one more reason `CudaRansacEstimator` is not
+    thread-safe.
+    """
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            'ignore',
+            message=r'.*Grid size \d+ will likely result in GPU '
+                    r'under-utilization',
+            category=NumbaPerformanceWarning)
+        yield