pixpoint 8.0.0__py3-none-any.whl

pixpoint/__init__.py

@@ -0,0 +1,86 @@
+"""
+PixPoint: GPU/CPU measurement of integrated, two- and three-point correlations.
+
+Three estimators on HEALPix maps, named after the statistic they deliver:
+:class:`Integrated3PCF` (per-patch aperture masses and 2PCFs, the i3PCF
+ingredients), :class:`Global2PCF` (footprint-wide 2PCFs),
+:class:`Global3PCF` (footprint-wide 3PCF multipoles) and :class:`OnePoint`
+(PDFs and moments of aperture fields: ``g`` for the aperture mass, ``k`` for a
+filtered scalar field).  Methods are TreeCorr
+field codes (``g`` = shear, ``k`` = scalar field): ``gg``, ``kk``, ``kg``,
+``three_by_two_pt``, ``ggg``, ``kkk``, ``kgg``, ``kkg``.  Every call returns a
+labelled result (:class:`IntegratedResult`, :class:`TwoPointResult`,
+:class:`ThreePointResult`, :class:`OnePointResult`).
+"""
+
+from .correlation_helpers import (
+    Q_crittenden,
+    Q_schneider,
+    U_crittenden,
+    U_schneider,
+    U_tophat,
+    calculate_all_zetas,
+    zeta_a_g,
+    zeta_a_minus,
+    zeta_a_plus,
+    zeta_a_t,
+    zeta_g_g,
+    zeta_g_minus,
+    zeta_g_plus,
+    zeta_g_t,
+)
+from .correlations import Integrated3PCF
+from .global_2pcf import Global2PCF
+from .global_3pcf import Global3PCF
+from .onepoint import OnePoint
+from .pipeline import MapFileLoader, MapLoader
+from .result_writer import OnePointWriter, ResultWriter, result_arrays
+from .results import (
+    Component,
+    IntegratedResult,
+    OnePointResult,
+    ThreePointResult,
+    TwoPointResult,
+)
+from .treecode import DEFAULT_RESOLUTION_FACTOR
+from .utils import pixel2RaDec, select_patch_centers
+from .zeta_writer import ZetaWriter
+
+__version__ = "8.0.0"
+__author__ = "David Gebauer"
+__email__ = "git@gebauer.ai"
+
+__all__ = [
+    "Component",
+    "Global2PCF",
+    "Global3PCF",
+    "Integrated3PCF",
+    "IntegratedResult",
+    "OnePoint",
+    "OnePointResult",
+    "OnePointWriter",
+    "ResultWriter",
+    "result_arrays",
+    "ThreePointResult",
+    "TwoPointResult",
+    "DEFAULT_RESOLUTION_FACTOR",
+    "MapLoader",
+    "MapFileLoader",
+    "Q_crittenden",
+    "Q_schneider",
+    "U_crittenden",
+    "U_schneider",
+    "U_tophat",
+    "calculate_all_zetas",
+    "zeta_g_plus",
+    "zeta_g_minus",
+    "zeta_a_plus",
+    "zeta_a_minus",
+    "zeta_g_g",
+    "zeta_a_g",
+    "zeta_g_t",
+    "zeta_a_t",
+    "pixel2RaDec",
+    "select_patch_centers",
+    "ZetaWriter",
+]

pixpoint/_numba.py

@@ -0,0 +1,23 @@
+"""Numba entry points with types mypy can check.
+
+``numba.prange`` is typed as a class without ``__iter__``; to a type checker
+it is ``range`` (which is also what it is with ``NUMBA_DISABLE_JIT=1``).  At
+run time the jitted functions see the real ``numba.prange``.
+"""
+
+from typing import TYPE_CHECKING, Any, Callable
+
+import numba
+from numba import njit
+
+if TYPE_CHECKING:
+    prange = range
+else:
+    from numba import prange
+
+#: ``numba.get_num_threads``, typed (numba < 0.68 leaves it untyped; going
+#: through ``Any`` is correct for both).
+_numba: Any = numba
+get_num_threads: Callable[[], int] = _numba.get_num_threads
+
+__all__ = ["get_num_threads", "njit", "prange"]

pixpoint/backend/__init__.py

@@ -0,0 +1,22 @@
+"""
+Backend abstraction layer for CPU (NumPy/Numba) and GPU (CuPy/CUDA) execution.
+
+All correlation measurement kernels exist in two forms:
+  - **CPU**: Numba @njit functions with parallel=True for multi-core execution
+  - **GPU**: CuPy ElementwiseKernels for simple per-pair operations, and
+    compiled CUDA kernels (.cu files in cuda/) for fused/vectorised
+    tomographic operations, launched through :class:`raw.TypedKernel`
+    (every launch checked against the signature parsed from the source)
+
+One module per kernel family, CPU kernel and GPU wrapper side by side:
+:mod:`.aperture` (M_ap, M_g), :mod:`.density` (ξ_g, ξ_t), :mod:`.shear`
+(ξ±), :mod:`.three_by_two` (the fused 3x2pt path), :mod:`.degrade` (treecode
+rows), :mod:`.threepoint` (step 1 of Global3PCF); :mod:`.combinations` holds
+the tomographic layouts of the tiled kernels, :mod:`.raw` compiling and
+launching, :mod:`.core` the :class:`Backend` that holds the active kernel set
+and provides to_device()/to_numpy() for data movement.
+"""
+
+from .core import Backend, get_backend, safe_divide
+
+__all__ = ["Backend", "get_backend", "safe_divide"]

pixpoint/backend/aperture.py

@@ -0,0 +1,387 @@
+"""Aperture statistics M_ap (shear) and M_g (density) per patch.
+
+CPU: Numba kernels.  GPU: per-pixel ``ElementwiseKernel`` contributions and
+the block-reduced tomographic kernels of ``cuda/aperture_tomo.cu``.
+"""
+
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+from .._numba import njit, prange
+from .raw import _CUPY_FASTMATH_OPTIONS, KernelFile, _has_raw_cuda_compiler, c_real
+
+# Beyond this many tomographic bins the fused aperture kernels (one block per
+# patch, the bins looped inside it) hold too many per-thread accumulators --
+# 2*NZ each -- and the per-(patch, bin) kernel is launched instead.
+_MAX_FUSED_APERTURE_BINS = 16
+# Map arguments the aperture kernels address through explicit strides
+# (planar, strided and AoS views are all passed without a copy).
+_STRIDED = ("g1", "g2", "weights", "values")
+
+
+# ── Aperture statistics kernels ──────────────────────────────────────────
+# Compute the aperture-filtered field for each sky patch:
+#   M = A_patch · Σ(w · field · Q(θ)) / Σ(w)
+# where Q(θ) is the compensated filter function.
+
+
+@njit(fastmath=True, parallel=True, cache=True)
+def _cpu_aperture_density_kernel(
+    Q_inds: np.ndarray,
+    Q_val: np.ndarray,
+    Q_offsets: np.ndarray,
+    map_values: np.ndarray,
+    weights: np.ndarray,
+    Q_patch_area: np.ndarray,
+    out_aperture: np.ndarray,
+) -> None:
+    """Galaxy mean density M_g within an aperture: weighted δ_g convolved
+    with the compensated filter Q(θ)."""
+    n_patches = Q_offsets.shape[0] - 1
+    # float64 explicitly, not ``x[0] * 0.0``.  That expression types as
+    # float64 under numba (python-float promotion) but as float32 under
+    # NUMBA_DISABLE_JIT=1, which pytest-env forces -- so the suite was
+    # measuring a float32 accumulation of a kernel that ships accumulating in
+    # float64.  Pinning it keeps the shipped numbers and makes the tests
+    # exercise them.  (It is deliberately not ``acc_dtype``: these four
+    # kernels reduce over a whole aperture disc or angular bin, where a
+    # float32 accumulator loses precision for nothing.)
+    zero = np.float64(0)
+
+    for patch_idx in prange(n_patches):
+        start = Q_offsets[patch_idx]
+        stop = Q_offsets[patch_idx + 1]
+        sum_w = zero
+        sum_wdelta_q = zero
+        for i in range(start, stop):
+            pix_idx = Q_inds[i]
+            weight = weights[pix_idx]
+            sum_w += weight
+            sum_wdelta_q += weight * map_values[pix_idx] * Q_val[i]
+        # 0, not NaN, for an empty aperture -- like every other M_g path
+        if sum_w == 0:
+            out_aperture[patch_idx] = 0.0
+        else:
+            out_aperture[patch_idx] = Q_patch_area[patch_idx] * sum_wdelta_q / sum_w
+
+
+@njit(fastmath=True, parallel=True, cache=True)
+def _cpu_aperture_shear_kernel(
+    Q_inds: np.ndarray,
+    Q_cos: np.ndarray,
+    Q_sin: np.ndarray,
+    Q_val: np.ndarray,
+    Q_offsets: np.ndarray,
+    g1: np.ndarray,
+    g2: np.ndarray,
+    weights: np.ndarray,
+    Q_patch_area: np.ndarray,
+    out_aperture: np.ndarray,
+) -> None:
+    """Aperture mass M_ap: tangential shear γ_t convolved with Q(θ).
+
+    γ_t = -γ₁·cos(2φ) - γ₂·sin(2φ) is the tangential shear component
+    relative to the patch centre.
+    """
+    n_patches = Q_offsets.shape[0] - 1
+    # float64 explicitly, not ``x[0] * 0.0``.  That expression types as
+    # float64 under numba (python-float promotion) but as float32 under
+    # NUMBA_DISABLE_JIT=1, which pytest-env forces -- so the suite was
+    # measuring a float32 accumulation of a kernel that ships accumulating in
+    # float64.  Pinning it keeps the shipped numbers and makes the tests
+    # exercise them.  (It is deliberately not ``acc_dtype``: these four
+    # kernels reduce over a whole aperture disc or angular bin, where a
+    # float32 accumulator loses precision for nothing.)
+    zero = np.float64(0)
+
+    for patch_idx in prange(n_patches):
+        start = Q_offsets[patch_idx]
+        stop = Q_offsets[patch_idx + 1]
+        sum_w = zero
+        sum_wgt_q = zero
+        for i in range(start, stop):
+            pix_idx = Q_inds[i]
+            weight = weights[pix_idx]
+            # Tangential shear projection
+            gt = -g1[pix_idx] * Q_cos[i] - g2[pix_idx] * Q_sin[i]
+            sum_w += weight
+            sum_wgt_q += weight * gt * Q_val[i]
+        out_aperture[patch_idx] = Q_patch_area[patch_idx] * sum_wgt_q / sum_w
+
+
+# ── CuPy ElementwiseKernel builders (GPU per-element operations) ──────────
+
+
+def _build_cupy_aperture_density_kernel(module: Any) -> Any:
+    """GPU kernel: per-pixel contribution to galaxy mean density M_g.
+
+    The filter geometry arrays (type letter ``Q``) may be narrower than the
+    map type ``T``; the ``(T)`` promotion at use is exact for float32 →
+    float64.
+    """
+    return module.ElementwiseKernel(
+        "raw I Q_inds, raw Q Q_val, raw T map_values, raw T weights",
+        "T out_num, T out_den",
+        """
+        const I idx = Q_inds[i];
+        const T w = weights[idx];
+        out_num = w * map_values[idx] * (T)Q_val[i];
+        out_den = w;
+        """,
+        "gpu_aperture_density_kernel",
+        options=_CUPY_FASTMATH_OPTIONS,
+    )
+
+
+def _build_cupy_aperture_shear_kernel(module: Any) -> Any:
+    """GPU kernel: per-pixel contribution to aperture mass M_ap via γ_t·Q(θ).
+
+    The filter geometry arrays (type letter ``Q``) may be narrower than the
+    map type ``T``; the ``(T)`` promotion at use is exact for float32 →
+    float64.
+    """
+    return module.ElementwiseKernel(
+        "raw I Q_inds, raw Q Q_cos, raw Q Q_sin, raw Q Q_val,"
+        " raw T g1, raw T g2, raw T weights",
+        "T out_num, T out_den",
+        """
+        const I idx = Q_inds[i];
+        const T w = weights[idx];
+        const T gt = -g1[idx] * (T)Q_cos[i] - g2[idx] * (T)Q_sin[i];
+        out_num = w * gt * (T)Q_val[i];
+        out_den = w;
+        """,
+        "gpu_aperture_shear_kernel",
+        options=_CUPY_FASTMATH_OPTIONS,
+    )
+
+
+# ── Block-reduced aperture RawKernels (all tomo bins in one launch) ───────
+# Replace the per-pixel ElementwiseKernel + add.reduceat path with a single
+# launch over (patch, tomo bin) blocks; the wrappers fall back (return
+# False) when no raw compiler is available so callers can use the legacy
+# ElementwiseKernel path.
+
+
+def _aperture_tomo_prepare_planar(module: Any, arr: Any) -> Tuple[Any, int, int]:
+    """Return (array, bin stride, row stride) in elements for a 2D view.
+
+    The kernel takes both strides, so nothing is ever copied: planar
+    ``(nz, npix)`` arrays, strided views such as ``shear[:, 0]`` of an
+    ``(nz, 2, npix)`` array, and the transposed AoS buffers
+    ``shear_aos[:, :, 0].T`` of an ``(npix, nz, 2)`` array all work.
+    """
+    del module  # no copy is needed for any layout
+    return (
+        arr,
+        int(arr.strides[0] // arr.itemsize),
+        int(arr.strides[1] // arr.itemsize),
+    )
+
+
+# (runtime id, device) -> (runtime, SM count).  The runtime object is kept
+# alive by the value so its id cannot be reused by a different one.
+_SM_COUNT_CACHE: dict[Tuple[int, int], Tuple[Any, Optional[int]]] = {}
+
+
+def _multiprocessor_count(module: Any) -> Optional[int]:
+    """SM count of the active device, or None when it cannot be queried.
+
+    Only used to size a launch grid, so a stand-in module without
+    ``cuda.runtime`` (the emulated cupy of the tests) simply imposes no
+    constraint.
+    """
+    runtime = getattr(getattr(module, "cuda", None), "runtime", None)
+    if runtime is None:
+        return None
+    try:
+        device = int(runtime.getDevice())
+    except Exception:  # pragma: no cover - no device behind the module
+        return None
+    key = (id(runtime), device)
+    if key not in _SM_COUNT_CACHE:
+        try:
+            props = runtime.getDeviceProperties(device)
+            count: Optional[int] = int(props["multiProcessorCount"])
+        except Exception:  # pragma: no cover - older runtime without the field
+            count = None
+        _SM_COUNT_CACHE[key] = (runtime, count)
+    return _SM_COUNT_CACHE[key][1]
+
+
+def _use_fused_aperture(module: Any, ntomo: int, npatches: int) -> bool:
+    """Whether to launch the one-block-per-patch aperture kernel.
+
+    It reads the aperture disc geometry once per pixel instead of once per
+    tomographic bin: 1.72x on the kernel at nside 512 and 1.83x measured in
+    situ inside ``gg`` (the result is bitwise identical
+    either way).  Two things take it back to the
+    per-(patch, bin) kernel: a bin set wide enough for the per-thread
+    accumulators to spill, and a patch set small enough that one block per
+    patch underfills the device -- the fused grid is ``ntomo`` times
+    smaller, which is the point of it.
+    """
+    if ntomo <= 0 or ntomo > _MAX_FUSED_APERTURE_BINS:
+        return False
+    n_sm = _multiprocessor_count(module)
+    if n_sm is not None and npatches < 2 * n_sm:
+        return False
+    return True
+
+
+def _build_cupy_aperture_tomo_shear_kernel(module: Any) -> Any:
+    """Builder for the GPU block-reduced tomographic aperture-mass kernel.
+
+    Two kernels behind one launch contract (see ``aperture_tomo.cu``): the
+    fused one-block-per-patch kernel is the default and reads the disc
+    geometry once per pixel instead of once per tomographic bin; the
+    per-(patch, bin) kernel is the fallback for wide bin sets and small
+    patch counts.  The two are bitwise identical, so which one ran is never
+    observable in the result.  ``ntomo`` is a template argument of the
+    fused kernel and therefore part of the cache key; it is small and
+    stable, so compiling on demand per value is fine.
+    """
+    kernels = KernelFile(
+        module, "aperture_tomo.cu", "Aperture-shear tomo", strided=_STRIDED
+    )
+
+    def _cupy_aperture_tomo_shear_kernel(
+        g1: Any,
+        g2: Any,
+        weights: Any,
+        q_inds: Any,
+        q_cos: Any,
+        q_sin: Any,
+        q_val: Any,
+        q_offsets: Any,
+        q_patch_area: Any,
+        out_num: Any,
+        out_den: Any,
+    ) -> bool:
+        if not _has_raw_cuda_compiler(module):
+            return False
+        if g1.ndim != 2 or g2.ndim != 2 or weights.ndim != 2:
+            return False
+
+        targs: Tuple[Any, ...] = (c_real(weights.dtype), c_real(q_cos.dtype))
+        npatches = int(q_offsets.shape[0] - 1)
+        ntomo = int(g1.shape[0])
+        fused = _use_fused_aperture(module, ntomo, npatches)
+        raw_kernel = (
+            kernels("gpu_aperture_shear_tomo_fused", targs + (ntomo,))
+            if fused
+            else kernels("gpu_aperture_shear_tomo", targs)
+        )
+        if raw_kernel is None:
+            return False
+
+        g1, g1_stride, g1_elem = _aperture_tomo_prepare_planar(module, g1)
+        g2, g2_stride, g2_elem = _aperture_tomo_prepare_planar(module, g2)
+        if (g1_stride, g1_elem) != (g2_stride, g2_elem):
+            g1 = module.ascontiguousarray(g1)
+            g2 = module.ascontiguousarray(g2)
+            g1_stride = g2_stride = int(g1.shape[1])
+            g1_elem = g2_elem = 1
+        weights, w_stride, w_elem = _aperture_tomo_prepare_planar(module, weights)
+
+        blocks = (
+            (max(1, npatches), 1, 1) if fused else (max(1, npatches), max(1, ntomo), 1)
+        )
+        # The fused kernel knows NZ at compile time and takes no `ntomo`.
+        tail: Tuple[int, ...] = (npatches,) if fused else (npatches, ntomo)
+        raw_kernel(
+            blocks,
+            (256,),
+            (
+                g1,
+                g2,
+                g1_stride,
+                g1_elem,
+                weights,
+                w_stride,
+                w_elem,
+                q_inds,
+                q_cos,
+                q_sin,
+                q_val,
+                q_offsets,
+                q_patch_area,
+                out_num,
+                out_den,
+                *tail,
+            ),
+        )
+        return True
+
+    return _cupy_aperture_tomo_shear_kernel
+
+
+def _build_cupy_aperture_tomo_density_kernel(module: Any) -> Any:
+    """Builder for the GPU block-reduced tomographic aperture-density kernel.
+
+    Same two-kernel contract as the aperture-mass builder above: the fused
+    one-block-per-patch kernel by default, the per-(patch, bin) kernel as
+    the fallback, bitwise identical results either way.
+    """
+    kernels = KernelFile(
+        module, "aperture_tomo.cu", "Aperture-density tomo", strided=_STRIDED
+    )
+
+    def _cupy_aperture_tomo_density_kernel(
+        values: Any,
+        weights: Any,
+        q_inds: Any,
+        q_val: Any,
+        q_offsets: Any,
+        q_patch_area: Any,
+        out_num: Any,
+        out_den: Any,
+    ) -> bool:
+        if not _has_raw_cuda_compiler(module):
+            return False
+        if values.ndim != 2 or weights.ndim != 2:
+            return False
+
+        targs: Tuple[Any, ...] = (c_real(weights.dtype), c_real(q_val.dtype))
+        npatches = int(q_offsets.shape[0] - 1)
+        ntomo = int(values.shape[0])
+        fused = _use_fused_aperture(module, ntomo, npatches)
+        raw_kernel = (
+            kernels("gpu_aperture_density_tomo_fused", targs + (ntomo,))
+            if fused
+            else kernels("gpu_aperture_density_tomo", targs)
+        )
+        if raw_kernel is None:
+            return False
+
+        values, v_stride, v_elem = _aperture_tomo_prepare_planar(module, values)
+        weights, w_stride, w_elem = _aperture_tomo_prepare_planar(module, weights)
+
+        blocks = (
+            (max(1, npatches), 1, 1) if fused else (max(1, npatches), max(1, ntomo), 1)
+        )
+        # The fused kernel knows NZ at compile time and takes no `ntomo`.
+        tail: Tuple[int, ...] = (npatches,) if fused else (npatches, ntomo)
+        raw_kernel(
+            blocks,
+            (256,),
+            (
+                values,
+                v_stride,
+                v_elem,
+                weights,
+                w_stride,
+                w_elem,
+                q_inds,
+                q_val,
+                q_offsets,
+                q_patch_area,
+                out_num,
+                out_den,
+                *tail,
+            ),
+        )
+        return True
+
+    return _cupy_aperture_tomo_density_kernel

pixpoint/backend/combinations.py

@@ -0,0 +1,177 @@
+"""Tomographic combination layouts of the combination-tiled pair kernels.
+
+A tiled kernel produces a *canonical* set of combinations (the row-major
+upper triangle of one bin set, or the full product of two); a requested list
+is classified against it once and, when it is a subset, gathered from a
+canonical temporary.
+"""
+
+from typing import Any, Callable, Optional, Tuple, cast
+
+import numpy as np
+
+from ..utils import live_object_serial
+from .raw import _KERNEL_CACHE_MISS
+
+_MAX_VECTOR_TOMO_BINS = 64
+# Combination-tiled pair kernels keep every accumulator in a register; the
+# tiles are used only while the per-thread accumulator count (numerators +
+# weight sums over all combinations) stays within this budget.  Beyond it
+# the per-(bin, row) kernels are launched instead (unpacked geometry only).
+_MAX_TILED_ACCUMULATORS = 64
+# Budget of the single-pass 3x2pt tile (xi+- + xi_g + xi_t accumulators).
+# Measured on an A100 up to 120 (4 source x 6 lens bins, all combinations),
+# where it is still ahead of three separate tiles; untested beyond.
+_MAX_TILED_3X2PT_ACCUMULATORS = 120
+
+
+class _CombinationLayout:
+    """How a requested list of tomographic combinations maps onto the
+    canonical layout generated by the combination-tiled kernels."""
+
+    __slots__ = ("mode", "canonical_ncomb", "rows")
+
+    def __init__(
+        self, mode: str, canonical_ncomb: int, rows: Optional[np.ndarray]
+    ) -> None:
+        self.mode = mode  # "triangle" | "auto" | "cartesian"
+        self.canonical_ncomb = canonical_ncomb  # rows the tile produces
+        self.rows = rows  # canonical row per request, None = identity
+
+
+class TiledWrapper:
+    """A pair-kernel launch wrapper whose combination-tiled kernel can be
+    switched off: ``tiled = False`` forces the per-(bin, row) fallback
+    (tests and benchmarks compare the two)."""
+
+    def __init__(self, launch: Callable[..., bool]) -> None:
+        self.tiled = True
+        self._launch = launch
+
+    def __call__(self, *args: Any) -> bool:
+        return self._launch(*args)
+
+
+class SinglePassWrapper:
+    """The launch wrapper of the single-pass 3x2pt tile.  ``mode``:
+    ``"auto"`` (within ``max_accumulators``), ``"single"`` (always, for
+    benchmarks) or ``"off"``; it returns which of (ξ±, ξ_g, ξ_t) it computed."""
+
+    def __init__(
+        self, launch: Callable[..., Tuple[bool, bool, bool]], max_accumulators: int
+    ) -> None:
+        self.mode = "auto"
+        self.max_accumulators = max_accumulators
+        self._launch = launch
+
+    def __call__(self, *args: Any) -> Tuple[bool, bool, bool]:
+        return self._launch(*args)
+
+
+_COMBINATION_LAYOUT_CACHE: dict[
+    tuple[int, int, str, int, int], Optional[_CombinationLayout]
+] = {}
+
+
+def _combination_layout(
+    comb_i: Any, comb_j: Any, kind: str, n1: int, n2: int = 0
+) -> Optional[_CombinationLayout]:
+    """Classify ``(comb_i, comb_j)`` against the tiled kernels' layouts.
+
+    ``kind="triangle"``: n1 bins, canonical = row-major upper triangle
+    (a request (i, j) with i > j is the same symmetric combination as
+    (j, i)); a request consisting of exactly the auto combinations gets
+    ``mode="auto"`` (its own, smaller tile).  ``kind="cartesian"``:
+    canonical = row-major n1 x n2 product.  ``rows`` is None when the
+    request *is* the canonical layout, else the canonical row of every
+    requested combination.  Returns None for out-of-range bins.
+
+    The (tiny) device arrays are compared once per array object and the
+    result cached by identity: the orchestrator caches and reuses them.
+
+    By ``live_object_serial`` and not ``id()``.  The orchestrator caches the
+    *canonical* combination arrays, but an explicit ``ggl_bin_combinations``
+    selection builds fresh ones on every call and drops them, so a second
+    selection can land on the recycled addresses of the first -- and the
+    layout carries ``rows``, the canonical row of every requested
+    combination, so a stale hit scatters the results into the previous
+    selection's rows.  The collision needs *both* ids to recycle at once and
+    CPython's LIFO free lists tend to swap rather than match them (it did not
+    fire in 4000 alternating calls here), but it is allocator luck, not a
+    guarantee.
+    """
+    key = (
+        live_object_serial(comb_i),
+        live_object_serial(comb_j),
+        kind,
+        int(n1),
+        int(n2),
+    )
+    cached = _COMBINATION_LAYOUT_CACHE.get(key, _KERNEL_CACHE_MISS)
+    if cached is not _KERNEL_CACHE_MISS:
+        return cast(Optional[_CombinationLayout], cached)
+
+    def to_host(a: Any) -> np.ndarray:
+        return np.asarray(a.get() if hasattr(a, "get") else a).astype(np.int64)
+
+    ci, cj = to_host(comb_i), to_host(comb_j)
+    layout: Optional[_CombinationLayout]
+    if kind == "triangle":
+        if (
+            ci.size == 0
+            or ci.min() < 0
+            or cj.min() < 0
+            or ci.max() >= n1
+            or cj.max() >= n1
+        ):
+            layout = None
+        else:
+            lo, hi = np.minimum(ci, cj), np.maximum(ci, cj)
+            rows = lo * n1 - (lo * (lo - 1)) // 2 + (hi - lo)
+            canonical = (n1 * (n1 + 1)) // 2
+            identity = rows.size == canonical and np.array_equal(
+                rows, np.arange(canonical)
+            )
+            if (
+                not identity
+                and ci.size == n1
+                and np.array_equal(ci, np.arange(n1))
+                and np.array_equal(cj, ci)
+            ):
+                layout = _CombinationLayout("auto", int(n1), None)
+            else:
+                layout = _CombinationLayout(
+                    "triangle", int(canonical), None if identity else rows
+                )
+    elif kind == "cartesian":
+        if (
+            ci.size == 0
+            or ci.min() < 0
+            or cj.min() < 0
+            or ci.max() >= n1
+            or cj.max() >= n2
+        ):
+            layout = None
+        else:
+            rows = ci * n2 + cj
+            canonical = int(n1 * n2)
+            identity = rows.size == canonical and np.array_equal(
+                rows, np.arange(canonical)
+            )
+            layout = _CombinationLayout(
+                "cartesian", canonical, None if identity else rows
+            )
+    else:
+        raise ValueError(f"unknown combination layout kind {kind!r}")
+    if len(_COMBINATION_LAYOUT_CACHE) > 64:
+        _COMBINATION_LAYOUT_CACHE.clear()
+    _COMBINATION_LAYOUT_CACHE[key] = layout
+    return layout
+
+
+def _scatter_canonical_rows(
+    module: Any, layout: _CombinationLayout, tmp: Any, out: Any
+) -> None:
+    """out[r] = tmp[layout.rows[r]] (tmp holds the canonical rows)."""
+    rows = module.asarray(layout.rows)
+    out[...] = tmp[rows]