# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Shared helpers for frontend neighbor-list auto dispatch."""
import math
import os
from collections.abc import Iterable
from functools import lru_cache
from typing import Literal
import warp as wp
from nvalchemiops.neighbors.neighbor_utils import DTYPE_INFO_ALL, empty_sentinel
__all__ = [
"AUTO_BASE_DEFAULTS",
"DEFAULT_BATCH_MAX_NBINS",
"DEFAULT_SINGLE_MAX_NBINS",
"FEATURE_BATCHED",
"FEATURE_CUDA",
"FEATURE_POSITIONS_FLOAT32",
"NEIGHBOR_LIST_STRATEGIES",
"auto_base_constants",
"fp64_naive_penalty",
"finalize_neighbor_list_method",
"get_select_neighbor_list_method_cost_kernel",
"neighbor_list_strategy_run_args",
"optional_outputs_mask",
"estimate_neighbor_list_costs",
"suggest_neighbor_list_method",
]
NeighborListMethod = Literal["naive", "cell_list", "cluster_tile"]
CellListStrategy = Literal["atom_centric", "pair_centric"]
AtomCentricPath = Literal["direct", "sorted"]
NaiveStrategy = Literal["scalar", "tile"]
# Cost-model constants (env-overridable, NVALCHEMI_NEIGHLIST_* style).
#
# The selector compares three families using quantities computed from the same
# geometry kernel:
# naive ~ per-system direct candidate work, image count, output work
# cell_list ~ cell-grid candidate work plus expected output work
# cluster_tile ~ fully-periodic tile output work when semantically compatible
#
# These are intentionally coarse: the selector is a routing guard, not a timing
# predictor. It uses volume and expected pairs so density/cutoff regimes drive
# the decision rather than raw atom-count thresholds.
AUTO_BASE_DEFAULTS = {
"NVALCHEMI_NEIGHLIST_CELL_SHELL": 27.0,
"NVALCHEMI_NEIGHLIST_CELL_SETUP": 1192450.0,
}
DEFAULT_SINGLE_MAX_NBINS = 524288
"""Default single-system cell-list cell cap."""
DEFAULT_BATCH_MAX_NBINS = 8192
"""Default per-system batched cell-list cell cap."""
FEATURE_CUDA = 1 << 0
"""Selector feature flag: the active device supports CUDA kernels."""
FEATURE_POSITIONS_FLOAT32 = 1 << 1
"""Selector feature flag: positions are float32, required by cluster-tile auto."""
FEATURE_BATCHED = 1 << 2
"""Selector feature flag: the call routes through a batched frontend."""
_MAX_WARP_LINEAR_LAUNCH = 2**31 - 1
_EPS = 1.0e-30
_INF = 3.0e30
# Reference stencil the ``shell`` knob is normalised against (3x3x3).
_SHELL_REFERENCE = 27.0
# Cap on the per-axis stencil radius in the cost model. A cell whose face
# distance is far below the cutoff gives a huge radius; the cost is already
# prohibitive long before this cap, and it keeps every stencil product within
# int64 (and the float-to-int radius conversion in range).
_MAX_STENCIL_RADIUS = 524288.0
# Saturation for the stencil offset count used in launch-size arithmetic.
_MAX_STENCIL_OFFSETS = 16777216
# Naive per-candidate scan and per-pair output write. The scan weight is on a
# float32 basis; float64 is handled by _FP64_NAIVE_PENALTY below.
_NAIVE_SCAN = 0.002
_NAIVE_WRITE = 1.3
# Scalar and tile are indistinguishable per candidate, so break the tie toward
# tile by a margin below the measurement noise rather than leaving the choice
# to enumeration order.
_NAIVE_SCALAR_TIEBREAK = 1.01
_NAIVE_FLOOR = 340700.0
_CELL_ATOM_SCAN = 0.1
_CELL_ATOM_WRITE = 1.5
_CELL_PAIR_WRITE = 0.437
_CELL_PAIR_LAUNCH = 0.3
_CELL_PAIR_SCAN = 0.1
# Pair-centric sorts positions into cell order first, a per-atom pass that
# atom-centric has no equivalent of.
_CELL_PAIR_PER_ATOM = 50.0
# Pair-centric enumerates (cell, neighbour-offset) blocks across the whole
# grid, so the vacuum in a free-boundary bounding box is work it performs;
# atom-centric starts from atoms and never visits an empty cell.
_APERIODIC_PAIR_VACUUM = 2.0
# Naive evaluates N^2 distances in the position dtype, while cell_list is
# dominated by an index-bound sort, so the crossover between them moves with
# the device's FP64 rate for float64 inputs only.
_FP64_NAIVE_PENALTY = 5.0
# SM architectures with ~1:2 FP64 throughput. Unlisted architectures and CPUs
# take the penalised path, which is the conservative direction.
_FAST_FP64_ARCHS = frozenset({70, 72, 80, 90, 100})
# Cluster-tile build is a Morton encode, radix sort, cluster assignment and
# tile scan, so its per-atom term is well above pair-centric's single sort
# pass; its tile queries emit pairs from shared memory, so its per-pair write
# is below pair-centric's.
_CLUSTER_TILE_PIPELINE = 1200000.0
_CLUSTER_TILE_PER_ATOM = 750.0
_CLUSTER_TILE_WRITE = 0.2
# Batched overheads are dominated by per-atom work: the wrap/fill launches
# preceding the search run over atoms, so summed across systems they scale with
# the total atom count however it is partitioned. Only naive and pair-centric
# carry these; atom-centric launches once per batch. The naive per-atom term is
# position-dtype arithmetic, so it takes the same float64 penalty as the scan
# and is stated on the same float32 basis.
_BATCH_NAIVE_PER_SYSTEM = 0.0
_BATCH_NAIVE_PER_ATOM = 20.0
_BATCH_PAIR_PER_SYSTEM = 270.0
_BATCH_PAIR_PER_ATOM = 10.0
_SPHERE_VOLUME_FACTOR = 4.1887902047863905
_OPTION_CUTOFF2 = 1 << 0
_OPTION_HALF_FILL = 1 << 1
_OPTION_RETURN_NEIGHBOR_LIST = 1 << 2
_OPTION_TARGET_INDICES = 1 << 3
_OPTION_RETURN_VECTORS = 1 << 4
_OPTION_RETURN_DISTANCES = 1 << 5
_OPTION_USE_PAIR_FN = 1 << 6
_OPTION_REBUILD_FLAGS = 1 << 8
_OPTION_WRAP_POSITIONS_FALSE = 1 << 9
_OPTION_BITS = {
"cutoff2": _OPTION_CUTOFF2,
"half_fill": _OPTION_HALF_FILL,
"return_neighbor_list": _OPTION_RETURN_NEIGHBOR_LIST,
"target_indices": _OPTION_TARGET_INDICES,
"return_vectors": _OPTION_RETURN_VECTORS,
"return_distances": _OPTION_RETURN_DISTANCES,
"use_pair_fn": _OPTION_USE_PAIR_FN,
"rebuild_flags": _OPTION_REBUILD_FLAGS,
"wrap_positions_false": _OPTION_WRAP_POSITIONS_FALSE,
}
_OPTION_ALIASES = {
"compute_vectors": "return_vectors",
"vectors": "return_vectors",
"neighbor_vectors": "return_vectors",
"compute_distances": "return_distances",
"distances": "return_distances",
"neighbor_distances": "return_distances",
"pair_fn": "use_pair_fn",
"pair_params": "use_pair_fn",
"pair_energies": "use_pair_fn",
"pair_forces": "use_pair_fn",
"wrap_positions": "wrap_positions_false",
}
_FLAG_NAMES = (
"invalid_input",
"naive_unsafe",
"cell_list_unsafe",
"cluster_tile_unsupported_options",
"cluster_tile_unsupported_geometry",
"cluster_tile_noncontiguous_batch",
"cluster_tile_image_multiplicity",
"pair_centric_unsafe",
"naive_tile_unsafe",
)
# Fine-grained, directly-runnable neighbor-list strategies. These are the names
# returned by :func:`estimate_neighbor_list_costs` / :func:`suggest_neighbor_list_method`
# and accepted by the public ``neighbor_list(method=...)`` frontends; each maps
# 1:1 to a base method plus its sub-options.
NEIGHBOR_LIST_STRATEGIES = (
"naive_scalar",
"naive_tile",
"cell_list_atom_centric",
"cell_list_pair_centric",
"cluster_tile",
)
# name -> (base_method, native_strategy, cell_list_strategy, atom_centric_path).
_STRATEGY_RUN_ARGS: dict[str, tuple[str, str, str, str]] = {
"naive_scalar": ("naive", "scalar", "auto", "auto"),
"naive_tile": ("naive", "tile", "auto", "auto"),
"cell_list_atom_centric": ("cell_list", "auto", "atom_centric", "direct"),
"cell_list_pair_centric": ("cell_list", "auto", "pair_centric", "sorted"),
"cluster_tile": ("cluster_tile", "auto", "auto", "auto"),
}
def neighbor_list_strategy_run_args(strategy: str) -> tuple[str, str, str, str]:
"""Resolve a strategy name to its run arguments.
Parameters
----------
strategy : str
A name from :data:`NEIGHBOR_LIST_STRATEGIES`, optionally ``batch_``
prefixed.
Returns
-------
tuple of str
``(method, native_strategy, cell_list_strategy, atom_centric_path)``.
The ``batch_`` prefix, if present, is carried onto ``method``.
"""
batched = strategy.startswith("batch_")
base = strategy[len("batch_") :] if batched else strategy
if base not in _STRATEGY_RUN_ARGS:
raise ValueError(f"unknown neighbor-list strategy {strategy!r}")
method, native_strategy, cell_list_strategy, atom_centric_path = _STRATEGY_RUN_ARGS[
base
]
if batched:
method = "batch_" + method
return method, native_strategy, cell_list_strategy, atom_centric_path
def _auto_env_float(name: str) -> float:
raw = os.environ.get(name)
try:
return float(raw) if raw is not None else float(AUTO_BASE_DEFAULTS[name])
except (TypeError, ValueError):
return float(AUTO_BASE_DEFAULTS[name])
def auto_base_constants() -> tuple[float, float]:
"""Return the current auto-dispatch constants.
Returns
-------
shell : float
Multiplicative estimate of cell-list neighbor-shell work.
setup : float
Per-system setup floor for cell-list work.
Notes
-----
The values are read from ``NVALCHEMI_NEIGHLIST_CELL_SHELL`` and
``NVALCHEMI_NEIGHLIST_CELL_SETUP`` when present, otherwise from
:data:`AUTO_BASE_DEFAULTS`.
"""
return (
_auto_env_float("NVALCHEMI_NEIGHLIST_CELL_SHELL"),
_auto_env_float("NVALCHEMI_NEIGHLIST_CELL_SETUP"),
)
def fp64_naive_penalty(device: wp.Device) -> float:
"""Return the float64 scan penalty for ``device``.
Parameters
----------
device : warp.context.Device
Device the selector will run on.
Returns
-------
float
Multiplier applied to the naive per-candidate scan cost when positions
are float64. ``1.0`` on architectures with fast FP64.
Notes
-----
Overridden by ``NVALCHEMI_NEIGHLIST_FP64_NAIVE_PENALTY`` when set, which is
the escape hatch for an architecture missing from
:data:`_FAST_FP64_ARCHS`. A value that is not a finite positive number is
ignored, so a malformed override cannot remove float64 naive strategies
from the selector's ranking.
See Also
--------
auto_base_constants : The other env-tunable selector constants.
"""
raw = os.environ.get("NVALCHEMI_NEIGHLIST_FP64_NAIVE_PENALTY")
if raw is not None:
try:
override = float(raw)
except (TypeError, ValueError):
override = None
if override is not None and math.isfinite(override) and override > 0.0:
return override
arch = getattr(device, "arch", None)
if getattr(device, "is_cuda", False) and arch is not None:
if int(arch) in _FAST_FP64_ARCHS:
return 1.0
return float(_FP64_NAIVE_PENALTY)
def _wp_scalar_from_cell_dtype(cell_dtype: type) -> type:
"""Return the Warp scalar dtype for a mat33 cell dtype."""
for scalar_dtype, (_, mat_dtype) in DTYPE_INFO_ALL.items():
if cell_dtype == mat_dtype:
return scalar_dtype
raise ValueError(f"Unsupported cell dtype: {cell_dtype!r}")
def _safe_int(value) -> int:
"""Convert host scalar-like values to ``int``."""
return int(value.item() if hasattr(value, "item") else value)
def _safe_float(value) -> float:
"""Convert host scalar-like values to ``float``."""
return float(value.item() if hasattr(value, "item") else value)
def _option_is_active(value: object) -> bool:
"""Return whether a public-style selector option should set a bit."""
if value is None or value is False:
return False
return True
def optional_outputs_mask(
optional_outputs: Iterable[str] | None = None,
**public_kwargs: object,
) -> int:
"""Encode public neighbor-list option names into a selector bitmask.
Parameters
----------
optional_outputs : iterable of str, optional
Public-style option names. Supported names include ``"cutoff2"``,
``"half_fill"``, ``"return_neighbor_list"``, ``"target_indices"``,
``"return_vectors"``, ``"return_distances"``, ``"use_pair_fn"``, and
``"rebuild_flags"``. Aliases matching common public buffers such as
``"neighbor_vectors"`` and ``"pair_fn"`` are accepted.
**public_kwargs
Public neighbor-list keyword names mapped to their values. Truthy or
non-``None`` values set the corresponding feasibility bit. For
``wrap_positions``, a value of ``False`` sets
``"wrap_positions_false"``.
Returns
-------
int
Bitmask consumed by the Warp selector kernel.
"""
mask = 0
for name in optional_outputs or ():
canonical = _OPTION_ALIASES.get(str(name), str(name))
try:
mask |= _OPTION_BITS[canonical]
except KeyError as exc:
raise ValueError(f"Unknown optional output option: {name!r}") from exc
for name, value in public_kwargs.items():
canonical = _OPTION_ALIASES.get(name, name)
if canonical == "wrap_positions_false":
if value is False:
mask |= _OPTION_BITS[canonical]
continue
if canonical not in _OPTION_BITS:
raise ValueError(f"Unknown optional output option: {name!r}")
if _option_is_active(value):
mask |= _OPTION_BITS[canonical]
return int(mask)
def finalize_neighbor_list_method(costs, flags) -> list[tuple[str, float]]:
"""Reduce selector costs and flags to feasible strategies, cheapest-first.
Parameters
----------
costs : array-like, shape (5,)
``(naive_scalar, naive_tile, cell_atom, cell_pair, cluster_tile)``.
flags : array-like, shape (9,)
Invalid-input and per-strategy feasibility flags.
Returns
-------
list of (str, float)
Feasible strategy names (from :data:`NEIGHBOR_LIST_STRATEGIES`) paired
with their relative estimated cost (lower is faster), sorted ascending.
Raises
------
ValueError
If the selector reported invalid input.
RuntimeError
If no strategy satisfies the safety and feasibility guards.
"""
raw_flags = [_safe_int(flags[i]) for i in range(min(len(flags), len(_FLAG_NAMES)))]
# Pad to the full flag set so a short ``flags`` argument cannot IndexError on
# the fixed indices used below; a missing flag defaults to 0 (feasible).
while len(raw_flags) < len(_FLAG_NAMES):
raw_flags.append(0)
if raw_flags[0] != 0:
raise ValueError("invalid batch_ptr, cutoff, or max_nbins for auto dispatch")
raw_costs = [_safe_float(costs[i]) for i in range(min(len(costs), 5))]
while len(raw_costs) < 5:
raw_costs.append(_INF)
naive_ok = raw_flags[1] == 0
cell_ok = raw_flags[2] == 0
cluster_ok = (
raw_flags[3] == 0
and raw_flags[4] == 0
and raw_flags[5] == 0
and raw_flags[6] == 0
)
candidates: list[tuple[str, float]] = []
if naive_ok:
candidates.append(("naive_scalar", raw_costs[0]))
if raw_flags[8] == 0:
candidates.append(("naive_tile", raw_costs[1]))
if cell_ok:
candidates.append(("cell_list_atom_centric", raw_costs[2]))
if raw_flags[7] == 0:
candidates.append(("cell_list_pair_centric", raw_costs[3]))
if cluster_ok:
candidates.append(("cluster_tile", raw_costs[4]))
feasible = [(name, cost) for name, cost in candidates if cost < _INF * 0.5]
if not feasible:
raise RuntimeError(
"No neighbor-list method fits the current safety and feasibility guards"
)
feasible.sort(key=lambda item: item[1])
return feasible
@lru_cache(maxsize=None)
def get_select_neighbor_list_method_cost_kernel(wp_dtype: type) -> wp.Kernel:
"""Return the shared guarded auto-dispatch cost kernel.
Parameters
----------
wp_dtype : type
Warp scalar dtype for cell matrices, usually ``wp.float32`` or
``wp.float64``.
Returns
-------
wp.Kernel
Kernel computing method costs, sub-option costs, and feasibility flags.
"""
if wp_dtype not in DTYPE_INFO_ALL:
raise ValueError(f"Unsupported dtype: {wp_dtype!r}")
_, mat_dtype = DTYPE_INFO_ALL[wp_dtype]
@wp.kernel(enable_backward=False)
def _kernel(
batch_ptr: wp.array(dtype=wp.int32),
batch_idx: wp.array(dtype=wp.int32),
batch_idx_is_provided: bool,
cell: wp.array(dtype=mat_dtype),
pbc_single: wp.array(dtype=wp.bool),
pbc_batch: wp.array2d(dtype=wp.bool),
pbc_is_batched: bool,
cutoff: wp.float32,
shell: wp.float32,
setup: wp.float32,
fp64_scan_penalty: wp.float32,
max_nbins: wp.int32,
max_launch_size: wp.int64,
option_mask: wp.int32,
feature_mask: wp.int32,
target_count: wp.int32,
costs: wp.array(dtype=wp.float32),
flags: wp.array(dtype=wp.int32),
) -> None:
"""Compute dispatch costs, sub-option costs, and safety flags.
Parameters
----------
batch_ptr : wp.array, shape (num_systems + 1,), dtype=wp.int32
Cumulative per-system atom counts.
batch_idx : wp.array, shape (total_atoms,), dtype=wp.int32
Optional dense system ids. Zero-size sentinel when absent.
batch_idx_is_provided : bool
Whether ``batch_idx`` should be validated for contiguous layout.
cell : wp.array, shape (num_systems,), dtype=wp.mat33*
Per-system cell matrices.
pbc_single : wp.array, shape (3,), dtype=wp.bool
Shared PBC flags. Zero-size sentinel when ``pbc_is_batched=True``.
pbc_batch : wp.array, shape (num_systems, 3), dtype=wp.bool
Per-system PBC flags. Zero-size sentinel when
``pbc_is_batched=False``.
pbc_is_batched : bool
Selects ``pbc_batch`` vs ``pbc_single``.
cutoff : float
Neighbor-list cutoff.
shell : float
Cell-list shell work multiplier.
setup : float
Cell-list per-system setup cost.
fp64_scan_penalty : float
Multiplier applied to the naive scan cost for float64 positions.
max_nbins : int
Cell-list per-system cap used by the selected frontend.
max_launch_size : int
Conservative one-dimensional launch-size guard.
option_mask : int
Public-option feasibility bits from :func:`optional_outputs_mask`.
feature_mask : int
Device/dtype/frontend feature bits.
target_count : int
Number of compact central rows when ``target_indices`` is active.
Zero means "all atoms".
costs : wp.array, shape (5,), dtype=wp.float32
OUTPUT: ``(naive_scalar, naive_tile, cell_atom, cell_pair,
cluster_tile)``.
flags : wp.array, shape (9,), dtype=wp.int32
OUTPUT: invalid input and method/sub-option feasibility flags.
Returns
-------
None
The kernel has no return value; results are written in place to the
``costs`` and ``flags`` output arrays.
Notes
-----
- Thread launch: ``max(num_systems, total_atoms when batch_idx is
supplied)``.
- Modifies: ``costs`` and ``flags`` by atomic updates.
See Also
--------
finalize_neighbor_list_method : Reduce these costs/flags to a method.
optional_outputs_mask : Builds the ``option_mask`` consumed here.
"""
tid = wp.tid()
num_systems = wp.int32(cell.shape[0])
has_cuda = (feature_mask & wp.int32(FEATURE_CUDA)) != wp.int32(0)
is_float32 = (feature_mask & wp.int32(FEATURE_POSITIONS_FLOAT32)) != wp.int32(0)
is_batched = (feature_mask & wp.int32(FEATURE_BATCHED)) != wp.int32(0)
has_cutoff2 = (option_mask & wp.int32(_OPTION_CUTOFF2)) != wp.int32(0)
half_fill = (option_mask & wp.int32(_OPTION_HALF_FILL)) != wp.int32(0)
return_list = (
option_mask & wp.int32(_OPTION_RETURN_NEIGHBOR_LIST)
) != wp.int32(0)
has_target_indices = (
option_mask & wp.int32(_OPTION_TARGET_INDICES)
) != wp.int32(0)
has_vectors = (option_mask & wp.int32(_OPTION_RETURN_VECTORS)) != wp.int32(0)
has_distances = (option_mask & wp.int32(_OPTION_RETURN_DISTANCES)) != wp.int32(
0
)
use_pair_fn = (option_mask & wp.int32(_OPTION_USE_PAIR_FN)) != wp.int32(0)
has_rebuild_flags = (option_mask & wp.int32(_OPTION_REBUILD_FLAGS)) != wp.int32(
0
)
wrap_positions_false = (
option_mask & wp.int32(_OPTION_WRAP_POSITIONS_FALSE)
) != wp.int32(0)
has_pair_outputs = has_vectors or has_distances or use_pair_fn
total_atoms = batch_ptr[num_systems]
if tid == 0:
if cutoff <= wp.float32(0.0) or max_nbins <= wp.int32(0):
wp.atomic_max(flags, 0, wp.int32(1))
if batch_ptr[0] != wp.int32(0) or total_atoms < wp.int32(0):
wp.atomic_max(flags, 0, wp.int32(1))
if has_cutoff2:
wp.atomic_max(flags, 2, wp.int32(1))
if batch_idx_is_provided and tid < total_atoms:
atom_system = batch_idx[tid]
if atom_system < wp.int32(0) or atom_system >= num_systems:
wp.atomic_max(flags, 0, wp.int32(1))
wp.atomic_max(flags, 5, wp.int32(1))
else:
if tid < batch_ptr[atom_system] or tid >= batch_ptr[atom_system + 1]:
wp.atomic_max(flags, 5, wp.int32(1))
if tid >= num_systems:
return
start = batch_ptr[tid]
end = batch_ptr[tid + 1]
n_atoms = end - start
if n_atoms < wp.int32(0):
wp.atomic_max(flags, 0, wp.int32(1))
n_atoms = wp.int32(0)
cell_i = cell[tid]
cell_volume = (
cell_i[0, 0] * (cell_i[1, 1] * cell_i[2, 2] - cell_i[1, 2] * cell_i[2, 1])
- cell_i[0, 1] * (cell_i[1, 0] * cell_i[2, 2] - cell_i[1, 2] * cell_i[2, 0])
+ cell_i[0, 2] * (cell_i[1, 0] * cell_i[2, 1] - cell_i[1, 1] * cell_i[2, 0])
)
volume = wp.float32(wp.abs(cell_volume))
pbc_x = wp.bool(False)
pbc_y = wp.bool(False)
pbc_z = wp.bool(False)
if pbc_is_batched:
pbc_x = pbc_batch[tid, 0]
pbc_y = pbc_batch[tid, 1]
pbc_z = pbc_batch[tid, 2]
else:
pbc_x = pbc_single[0]
pbc_y = pbc_single[1]
pbc_z = pbc_single[2]
any_pbc = pbc_x or pbc_y or pbc_z
all_pbc = pbc_x and pbc_y and pbc_z
n_float = wp.float32(n_atoms)
active_atoms = n_atoms
if has_target_indices:
active_atoms = wp.int32(
wp.ceil(
wp.float32(n_atoms)
* wp.float32(max(target_count, wp.int32(0)))
/ wp.float32(max(total_atoms, wp.int32(1)))
)
)
active_atoms = min(active_atoms, n_atoms)
active_float = wp.float32(active_atoms)
n_sq = n_float * n_float
# Naive candidate count: each active central atom is tested against
# all ``n`` atoms, so active_n_work ~ N^2 distance checks. n_pairs_cap =
# n*(n-1) is the exact upper bound on real pairs.
active_n_work = active_float * n_float
n_pairs_cap = wp.max(
n_float * wp.float32(max(n_atoms - wp.int32(1), 0)), wp.float32(0.0)
)
if volume <= wp.float32(_EPS):
wp.atomic_max(flags, 2, wp.int32(1))
wp.atomic_max(flags, 4, wp.int32(1))
if any_pbc:
wp.atomic_max(flags, 1, wp.int32(1))
wp.atomic_add(costs, 0, n_sq)
wp.atomic_add(costs, 1, n_sq)
wp.atomic_add(costs, 2, wp.float32(_INF))
wp.atomic_add(costs, 3, wp.float32(_INF))
wp.atomic_add(costs, 4, wp.float32(_INF))
return
inverse_cell_transpose = wp.transpose(wp.inverse(cell_i))
face_distance_x = wp.float32(1.0) / wp.float32(
wp.length(inverse_cell_transpose[0])
)
face_distance_y = wp.float32(1.0) / wp.float32(
wp.length(inverse_cell_transpose[1])
)
face_distance_z = wp.float32(1.0) / wp.float32(
wp.length(inverse_cell_transpose[2])
)
min_face_distance = wp.min(
face_distance_x, wp.min(face_distance_y, face_distance_z)
)
if min_face_distance <= wp.float32(_EPS):
wp.atomic_max(flags, 2, wp.int32(1))
wp.atomic_max(flags, 4, wp.int32(1))
return
shift_x = wp.int32(0)
shift_y = wp.int32(0)
shift_z = wp.int32(0)
if any_pbc:
if pbc_x:
shift_x = wp.int32(
wp.ceil(wp.float32(wp.length(inverse_cell_transpose[0])) * cutoff)
)
if pbc_y:
shift_y = wp.int32(
wp.ceil(wp.float32(wp.length(inverse_cell_transpose[1])) * cutoff)
)
if pbc_z:
shift_z = wp.int32(
wp.ceil(wp.float32(wp.length(inverse_cell_transpose[2])) * cutoff)
)
shift_count = (
wp.int64(shift_x) * wp.int64(2 * shift_y + 1) * wp.int64(2 * shift_z + 1)
+ wp.int64(shift_y) * wp.int64(2 * shift_z + 1)
+ wp.int64(shift_z)
+ wp.int64(1)
)
if shift_count <= wp.int64(0):
wp.atomic_max(flags, 1, wp.int32(1))
shift_count = wp.int64(1)
if shift_count > max_launch_size:
wp.atomic_max(flags, 1, wp.int32(1))
if shift_count * wp.int64(max(total_atoms, wp.int32(1))) > max_launch_size:
wp.atomic_max(flags, 1, wp.int32(1))
if (
shift_count
* wp.int64(max(n_atoms, wp.int32(1)))
* wp.int64(max(num_systems, wp.int32(1)))
> max_launch_size
):
wp.atomic_max(flags, 1, wp.int32(1))
# Naive viability bound: above ~8e9 candidate pairs
# (active_n_work * shift_count) naive stops being competitive with
# cell_list, so mark it unsafe.
if active_n_work * wp.float32(shift_count) > wp.float32(8.0e9):
wp.atomic_max(flags, 1, wp.int32(1))
cutoff_volume = cutoff * cutoff * cutoff
# Floor against float32 underflow for tiny cutoffs so the cell-cost
# division below (volume / cutoff_volume) cannot produce inf/NaN.
cutoff_volume = wp.max(cutoff_volume, wp.float32(_EPS))
density = n_float / wp.max(volume, wp.float32(_EPS))
# Expected output work, shared by every method: neighbors per atom
# (density * cutoff-sphere volume) times the active central atoms,
# capped by the exact pair count for non-periodic and halved for half_fill.
expected_neighbors = density * wp.float32(_SPHERE_VOLUME_FACTOR) * cutoff_volume
expected_pairs = active_float * expected_neighbors
if not any_pbc:
expected_pairs = wp.min(expected_pairs, wp.min(n_pairs_cap, active_n_work))
if half_fill:
expected_pairs = expected_pairs * wp.float32(0.5)
# This kernel accumulates per system, so constants here are charged
# per system; the batched naive path launches once for the whole batch,
# so its floor is amortised instead.
system_share = wp.float32(1.0) / wp.float32(max(num_systems, wp.int32(1)))
# Supplied by the host for this device; applies to float64 only.
precision_scale = wp.float32(1.0)
if not is_float32:
precision_scale = fp64_scan_penalty
pair_vacuum = wp.float32(1.0)
if not any_pbc:
pair_vacuum = wp.float32(_APERIODIC_PAIR_VACUUM)
naive_setup = wp.float32(_NAIVE_FLOOR) * system_share
if is_batched:
naive_setup = (
naive_setup
+ wp.float32(_BATCH_NAIVE_PER_SYSTEM)
+ wp.float32(_BATCH_NAIVE_PER_ATOM) * precision_scale * n_float
)
# Naive cost = per-candidate scan + per-pair output write + setup.
# The scalar path pays for global loads where the tile path reuses
# shared memory, so it carries the heavier scan weight.
scalar_cost = (
wp.float32(_NAIVE_SCAN * _NAIVE_SCALAR_TIEBREAK)
* precision_scale
* active_n_work
* wp.float32(shift_count)
+ wp.float32(_NAIVE_WRITE) * expected_pairs
+ naive_setup
)
tile_cost = (
wp.float32(_NAIVE_SCAN)
* precision_scale
* active_n_work
* wp.float32(shift_count)
+ wp.float32(_NAIVE_WRITE) * expected_pairs
+ naive_setup
)
if (
(not has_cuda)
or has_pair_outputs
or (any_pbc and is_batched and wrap_positions_false)
):
wp.atomic_max(flags, 8, wp.int32(1))
wp.atomic_add(costs, 0, scalar_cost)
wp.atomic_add(costs, 1, tile_cost)
cells_per_dimension = wp.vec3i(0, 0, 0)
ADAPTIVE_MIN_CELLS = wp.int32(4)
# Clamp the float->int cast so an extreme face_distance/cutoff ratio
# cannot overflow int32; the halving loop below reduces the count to
# <= max_nbins regardless, so capping per axis here is lossless.
max_cells_f = wp.float32(max_nbins)
cells_per_dimension[0] = max(
wp.int32(wp.min(face_distance_x / cutoff, max_cells_f)), 1
)
cells_per_dimension[1] = max(
wp.int32(wp.min(face_distance_y / cutoff, max_cells_f)), 1
)
cells_per_dimension[2] = max(
wp.int32(wp.min(face_distance_z / cutoff, max_cells_f)), 1
)
for dim in range(3):
pbc_dim = pbc_z
if dim == 0:
pbc_dim = pbc_x
elif dim == 1:
pbc_dim = pbc_y
if pbc_dim or cells_per_dimension[dim] > wp.int32(1):
while cells_per_dimension[dim] < ADAPTIVE_MIN_CELLS:
cells_per_dimension[dim] = cells_per_dimension[dim] * wp.int32(2)
# Accumulate the cell-count product in int64 so a large grid cannot wrap
# around int32 and slip past the guards below; the halving loop reduces
# it to <= max_nbins, which always fits int32.
total_cells_i64 = (
wp.int64(cells_per_dimension[0])
* wp.int64(cells_per_dimension[1])
* wp.int64(cells_per_dimension[2])
)
while total_cells_i64 > wp.int64(max_nbins):
for dim in range(3):
cells_per_dimension[dim] = max(
cells_per_dimension[dim] // wp.int32(2), wp.int32(1)
)
total_cells_i64 = (
wp.int64(cells_per_dimension[0])
* wp.int64(cells_per_dimension[1])
* wp.int64(cells_per_dimension[2])
)
total_cells_i32 = wp.int32(total_cells_i64)
if total_cells_i32 <= wp.int32(0):
wp.atomic_max(flags, 2, wp.int32(1))
total_cells_i32 = wp.int32(1)
if (
wp.int64(total_cells_i32) * wp.int64(max(num_systems, wp.int32(1)))
> max_launch_size
):
wp.atomic_max(flags, 2, wp.int32(1))
radius_x = wp.int32(
wp.min(
wp.ceil(cutoff * wp.float32(cells_per_dimension[0]) / face_distance_x),
wp.float32(_MAX_STENCIL_RADIUS),
)
)
radius_y = wp.int32(
wp.min(
wp.ceil(cutoff * wp.float32(cells_per_dimension[1]) / face_distance_y),
wp.float32(_MAX_STENCIL_RADIUS),
)
)
radius_z = wp.int32(
wp.min(
wp.ceil(cutoff * wp.float32(cells_per_dimension[2]) / face_distance_z),
wp.float32(_MAX_STENCIL_RADIUS),
)
)
radius_x = max(radius_x, wp.int32(0))
radius_y = max(radius_y, wp.int32(0))
radius_z = max(radius_z, wp.int32(0))
# Candidates per source atom = stencil cells * atoms per cell, both
# taken from the grid that is actually built rather than from a fixed
# stencil and an idealised cell count.
span_x = wp.int64(2) * wp.int64(radius_x) + wp.int64(1)
span_y = wp.int64(2) * wp.int64(radius_y) + wp.int64(1)
span_z = wp.int64(2) * wp.int64(radius_z) + wp.int64(1)
stencil_cells = wp.float32(span_x * span_y * span_z)
cost_cells = wp.max(wp.float32(total_cells_i32), wp.float32(1.0))
grid_work = (
(shell / wp.float32(_SHELL_REFERENCE))
* stencil_cells
* active_n_work
/ cost_cells
)
# The grid build is a single batched launch, so ``setup`` is amortised
# rather than charged once per system.
cell_atom_cost = (
setup * system_share
+ wp.float32(_CELL_ATOM_SCAN) * grid_work
+ wp.float32(_CELL_ATOM_WRITE) * expected_pairs
)
n_outer = (
wp.int64(radius_x) * span_y * span_z
+ wp.int64(radius_y) * span_z
+ wp.int64(radius_z)
)
if not half_fill:
n_outer = span_x * span_y * span_z - wp.int64(1)
n_outer = wp.min(n_outer, wp.int64(_MAX_STENCIL_OFFSETS))
# Pair-centric launches one 64-thread block per (cell, neighbor-offset);
# n_outer = stencil offsets (forward half only when half_fill).
# pair_launch ~ total threads; pair_blocks ~ block count.
pair_launch = wp.int64(total_cells_i32) * (n_outer + wp.int64(1)) * wp.int64(64)
atom_blocks = wp.int64(max(n_atoms, wp.int32(1)) + wp.int32(63)) // wp.int64(64)
pair_blocks = wp.int64(total_cells_i32) * (n_outer + wp.int64(1))
if (not has_cuda) or pair_blocks < atom_blocks:
wp.atomic_max(flags, 7, wp.int32(1))
pair_launch_effective = wp.min(pair_launch, max_launch_size)
# Pair-centric: cheaper per pair (0.55 write, 0.5 scan) but pays a
# per-block launch term (0.025 * pair_launch_effective) for its many
# small blocks (capped when coarsening is required).
# Pair-centric launches per (cell, neighbour-offset), so every system
# adds its own blocks and its floor stays per-system.
cell_pair_cost = (
setup * system_share
+ wp.float32(_CELL_PAIR_PER_ATOM) * n_float
+ (
wp.float32(_BATCH_PAIR_PER_SYSTEM)
+ wp.float32(_BATCH_PAIR_PER_ATOM) * n_float
)
* wp.float32(is_batched)
+ wp.float32(_CELL_PAIR_WRITE) * expected_pairs
+ wp.float32(_CELL_PAIR_LAUNCH) * wp.float32(pair_launch_effective)
+ wp.float32(_CELL_PAIR_SCAN) * grid_work * pair_vacuum
)
wp.atomic_add(costs, 2, cell_atom_cost)
wp.atomic_add(costs, 3, cell_pair_cost)
if (
half_fill
or has_target_indices
or use_pair_fn
or has_rebuild_flags
or (has_cutoff2 and (return_list or has_vectors or has_distances))
):
wp.atomic_max(flags, 3, wp.int32(1))
if (not has_cuda) or (not is_float32) or (not all_pbc):
wp.atomic_max(flags, 4, wp.int32(1))
if cutoff > wp.float32(0.5) * min_face_distance:
wp.atomic_max(flags, 6, wp.int32(1))
if is_batched and batch_idx_is_provided:
# The per-atom validation above sets flag 5 when the dense labels
# do not match the contiguous ranges implied by batch_ptr.
pass
cluster_large_enough = wp.bool(False)
if is_batched:
cluster_large_enough = total_atoms >= wp.int32(4096)
else:
cluster_large_enough = expected_neighbors >= wp.float32(512.0)
if not cluster_large_enough:
wp.atomic_max(flags, 4, wp.int32(1))
# Cluster-tile one-time pipeline floor (Morton encode, sort, cluster
# build, tile scan), added once per batch rather than per-system.
if tid == 0:
wp.atomic_add(costs, 4, wp.float32(_CLUSTER_TILE_PIPELINE))
# Per-atom build plus the tile-query output write.
cluster_cost = (
wp.float32(_CLUSTER_TILE_WRITE) * expected_pairs
+ wp.float32(_CLUSTER_TILE_PER_ATOM) * n_float
+ setup
)
wp.atomic_add(costs, 4, cluster_cost)
return _kernel
def _validate_selector_inputs(
batch_ptr: wp.array,
cell: wp.array,
pbc: wp.array,
cutoff: float,
batch_idx: wp.array | None,
) -> int:
"""Validate shared selector inputs and return ``num_systems``."""
if batch_ptr.dtype != wp.int32:
raise ValueError("batch_ptr must be a wp.array with dtype=wp.int32")
if batch_ptr.ndim != 1:
raise ValueError("batch_ptr must be a 1-D wp.array")
if batch_ptr.shape[0] < 2:
raise ValueError("batch_ptr must have length at least 2")
num_systems = int(batch_ptr.shape[0] - 1)
if cell.shape[0] != num_systems:
raise ValueError("cell must contain one matrix per system")
if pbc.dtype != wp.bool:
raise ValueError("pbc must be a wp.array with dtype=wp.bool")
if pbc.ndim not in (1, 2):
raise ValueError("pbc must have shape (3,) or (num_systems, 3)")
if pbc.ndim == 1 and pbc.shape[0] != 3:
raise ValueError("pbc must have shape (3,) or (num_systems, 3)")
if pbc.ndim == 2 and (pbc.shape[0] != num_systems or pbc.shape[1] != 3):
raise ValueError("pbc must have shape (3,) or (num_systems, 3)")
device = str(batch_ptr.device)
if str(cell.device) != device or str(pbc.device) != device:
raise ValueError("batch_ptr, cell, and pbc must be on the same device")
if batch_idx is not None:
if batch_idx.dtype != wp.int32 or batch_idx.ndim != 1:
raise ValueError("batch_idx must be a 1-D wp.array with dtype=wp.int32")
if str(batch_idx.device) != device:
raise ValueError("batch_idx must be on the same device as batch_ptr")
if not float(cutoff) > 0.0:
raise ValueError("cutoff must be positive")
return num_systems
def _default_feature_mask(batch_ptr: wp.array, cell: wp.array, num_systems: int) -> int:
"""Return feature bits inferred by the shared Warp selector."""
mask = 0
if "cuda" in str(batch_ptr.device).lower():
mask |= FEATURE_CUDA
if _wp_scalar_from_cell_dtype(cell.dtype) == wp.float32:
mask |= FEATURE_POSITIONS_FLOAT32
if num_systems > 1:
mask |= FEATURE_BATCHED
return mask
[docs]
def estimate_neighbor_list_costs(
batch_ptr: wp.array,
cell: wp.array,
pbc: wp.array,
cutoff: float,
*,
batch_idx: wp.array | None = None,
max_nbins: int | None = None,
max_launch_size: int = _MAX_WARP_LINEAR_LAUNCH,
optional_outputs: Iterable[str] | None = None,
option_mask: int = 0,
feature_mask: int | None = None,
target_count: int | None = None,
) -> list[tuple[str, float]]:
"""Report feasible neighbor-list strategies and their estimated cost.
Parameters
----------
batch_ptr : wp.array, shape (num_systems + 1,), dtype=wp.int32
Cumulative atom counts. The final entry is the total atom count.
cell : wp.array, shape (num_systems,), dtype=wp.mat33*
Per-system cell matrices. Non-periodic callers should pass a
synthesized bounding-box cell.
pbc : wp.array, shape (3,) or (num_systems, 3), dtype=wp.bool
Shared or per-system periodic-boundary flags.
cutoff : float
Neighbor cutoff. For dual-cutoff routing, pass the larger cutoff.
batch_idx : wp.array, optional
Dense per-atom batch ids. When provided, the selector validates that
the labels match the contiguous ranges implied by ``batch_ptr`` before
allowing auto cluster-tile.
max_nbins : int, optional
Per-system cell-list cap. Defaults to the single-system cap when
``num_systems == 1`` and the batched cap otherwise.
max_launch_size : int, default=2**31 - 1
Conservative launch-size guard used by current neighbor-list kernels.
optional_outputs : iterable of str, optional
Public-style neighbor-list option names, encoded with
:func:`optional_outputs_mask`.
option_mask : int, default=0
Pre-encoded option bits to OR with ``optional_outputs``.
feature_mask : int, optional
Device/dtype/frontend feature bits. When omitted, CUDA and cell dtype
are inferred from the Warp arrays.
target_count : int, optional
Number of compact central rows requested by ``target_indices``. When
omitted, the selector scores all atoms. Frontends should pass
``len(target_indices)`` when the public ``target_indices`` kwarg is
active.
Returns
-------
list of (str, float)
Feasible strategies (from :data:`NEIGHBOR_LIST_STRATEGIES`) and their
relative estimated cost (lower is faster), sorted cheapest-first.
Batched inputs (``num_systems > 1``) return ``batch_`` prefixed names.
Notes
-----
The returned costs are *relative* (arbitrary units): only their ordering is
meaningful, so compare them to each other, not to a wall-clock time. The
model approximates algorithmic work (candidate pairs, neighbors written,
launch overhead) and is **hardware-independent** -- the true crossover
between strategies shifts with the device, so when the top costs are within a
small factor the predicted best may be marginally slower than a close
runner-up; benchmark the top few on your hardware in that case.
This launches one Warp kernel over systems (and over atoms when validating
``batch_idx`` contiguity) and reads back five costs plus nine flags, so it
is **host-only**: call it outside ``torch.compile`` / ``jax.jit`` and pass
the chosen name as an explicit ``method=`` to run compiled.
"""
num_systems = _validate_selector_inputs(batch_ptr, cell, pbc, cutoff, batch_idx)
if max_nbins is None:
max_nbins = (
DEFAULT_SINGLE_MAX_NBINS if num_systems == 1 else DEFAULT_BATCH_MAX_NBINS
)
if int(max_nbins) <= 0:
raise ValueError("max_nbins must be positive")
device = batch_ptr.device
costs = wp.zeros(5, dtype=wp.float32, device=device)
flags = wp.zeros(len(_FLAG_NAMES), dtype=wp.int32, device=device)
shell, setup = auto_base_constants()
fp64_scan_penalty = fp64_naive_penalty(device)
pbc_is_batched = pbc.ndim == 2
pbc_single = pbc if not pbc_is_batched else empty_sentinel(1, wp.bool, device)
pbc_batch = pbc if pbc_is_batched else empty_sentinel(2, wp.bool, device)
batch_idx_is_provided = batch_idx is not None
batch_idx_arg = (
batch_idx if batch_idx is not None else empty_sentinel(1, wp.int32, device)
)
wp_dtype = _wp_scalar_from_cell_dtype(cell.dtype)
feature_bits = _default_feature_mask(batch_ptr, cell, num_systems)
if feature_mask is not None:
feature_bits |= int(feature_mask)
options = int(option_mask) | optional_outputs_mask(optional_outputs)
if options & _OPTION_TARGET_INDICES and target_count is None:
raise ValueError(
"target_count is required when target_indices is included in "
"optional_outputs"
)
target_count_int = 0 if target_count is None else int(target_count)
launch_dim = max(
num_systems, int(batch_idx.shape[0]) if batch_idx is not None else 0
)
wp.launch(
get_select_neighbor_list_method_cost_kernel(wp_dtype),
dim=launch_dim,
inputs=[
batch_ptr,
batch_idx_arg,
batch_idx_is_provided,
cell,
pbc_single,
pbc_batch,
pbc_is_batched,
wp.float32(float(cutoff)),
wp.float32(float(shell)),
wp.float32(float(setup)),
wp.float32(float(fp64_scan_penalty)),
wp.int32(int(max_nbins)),
wp.int64(int(max_launch_size)),
wp.int32(options),
wp.int32(feature_bits),
wp.int32(target_count_int),
costs,
flags,
],
device=device,
)
strategies = finalize_neighbor_list_method(costs.numpy(), flags.numpy())
if num_systems > 1:
strategies = [("batch_" + name, cost) for name, cost in strategies]
return strategies
[docs]
def suggest_neighbor_list_method(*args, **kwargs) -> str:
"""Return the cheapest feasible neighbor-list strategy name.
Thin wrapper over :func:`estimate_neighbor_list_costs` returning only the
top-ranked strategy name. Accepts the same arguments and carries the same
host-only sync caveat: call outside ``torch.compile`` / ``jax.jit`` and
pass the result as an explicit ``method=`` argument.
Parameters
----------
*args
Positional arguments forwarded to
:func:`estimate_neighbor_list_costs`.
**kwargs
Keyword arguments forwarded to
:func:`estimate_neighbor_list_costs`.
Returns
-------
str
Name of the cheapest feasible strategy, e.g. ``"cell_list_atom_centric"``
or ``"batch_naive_tile"``.
See Also
--------
estimate_neighbor_list_costs : Full ranked list of feasible strategies with costs.
"""
return estimate_neighbor_list_costs(*args, **kwargs)[0][0]
