# 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.
"""PyTorch bindings for single-system cluster-pair tile neighbor list.
Mirrors the ``cell_list`` / ``naive`` wrapper pattern: exposes low-level
component ops (``build_cluster_tile_list``, ``query_cluster_tile``,
``query_cluster_tile_coo``) that fill pre-allocated tensors, plus a high-level
convenience entry point ``cluster_tile_neighbor_list``. All torch-side work
(Morton sort, allocation, ``wp.from_torch`` conversion) lives in this
module; the Warp layer at ``nvalchemiops.neighbors.cluster_tile`` only
sees ``wp.array`` inputs.
Scope: single system, orthorhombic or triclinic PBC, float32, any
``N >= 0`` (the wrapper pads internally to a multiple of TILE_GROUP_SIZE).
"""
import torch
import warp as wp
from nvalchemiops.neighbors.cluster_tile import (
TILE_GROUP_SIZE,
estimate_max_tiles_per_group,
)
from nvalchemiops.neighbors.cluster_tile import (
build_cluster_tile_list as wp_build_cluster_tile_list,
)
from nvalchemiops.neighbors.cluster_tile import (
query_cluster_tile as wp_query_cluster_tile,
)
from nvalchemiops.neighbors.cluster_tile import (
query_cluster_tile_coo as wp_query_cluster_tile_coo,
)
from nvalchemiops.neighbors.cluster_tile.kernels import (
_get_query_cluster_tile_direct_csr_count_kernel,
_get_query_cluster_tile_direct_csr_fill_kernel,
)
from nvalchemiops.neighbors.neighbor_utils import (
_selective_fill_neighbor_matrix_tail as wp_selective_fill_neighbor_matrix_tail,
)
from nvalchemiops.neighbors.neighbor_utils import estimate_max_neighbors
from nvalchemiops.neighbors.neighbor_utils import (
fill_neighbor_matrix_tail as wp_fill_neighbor_matrix_tail,
)
from nvalchemiops.neighbors.neighbor_utils import (
selective_zero_num_neighbors_single as wp_selective_zero_num_neighbors_single,
)
from nvalchemiops.neighbors.output_args import (
_has_partial_or_pair_outputs,
_prepare_coo_pair_output_args,
)
from nvalchemiops.torch._warp_op_helpers import (
_capture_safe_inverse,
scoped_torch_warp_stream,
)
from nvalchemiops.torch.neighbors._autograd import (
_reconstruct_coo_geometry,
_reconstruct_matrix_geometry,
)
from nvalchemiops.torch.neighbors._fixed_cell import _FixedCellGeometry
from nvalchemiops.torch.neighbors.neighbor_utils import (
_check_neighbor_capacity,
_check_tile_buffer_capacity,
_compact_coo_prefix,
_normalize_compiled_single_segment_coo_count,
_prepare_compact_coo_geometry_buffers,
_validate_cluster_tile_matrix_outputs,
_validate_compact_coo_outputs,
_validate_segmented_coo_state,
)
from nvalchemiops.torch.types import get_wp_dtype
__all__ = [
"TILE_GROUP_SIZE",
"estimate_cluster_tile_list_sizes",
"allocate_cluster_tile_list",
"build_cluster_tile_list",
"query_cluster_tile",
"query_cluster_tile_coo",
"cluster_tile_neighbor_list",
]
def _wp_fixed_cluster_geometry(
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> dict[str, object]:
"""Convert optional prepared cluster geometry arrays for a Warp launcher."""
return {
"fixed_cell": fixed_cell,
"qr": wp.from_torch(
qr, dtype=wp.float32, requires_grad=False, return_ctype=True
)
if qr is not None
else None,
"axis_aligned": wp.from_torch(
axis_aligned, dtype=wp.bool, requires_grad=False, return_ctype=True
)
if axis_aligned is not None
else None,
"fractional_rounding_certified": wp.from_torch(
fractional_rounding_certified,
dtype=wp.bool,
requires_grad=False,
return_ctype=True,
)
if fractional_rounding_certified is not None
else None,
"qr_height_certified": wp.from_torch(
qr_height_certified, dtype=wp.bool, requires_grad=False, return_ctype=True
)
if qr_height_certified is not None
else None,
"bbox_cutoff_bounds": wp.from_torch(
bbox_cutoff_bounds, dtype=wp.float32, requires_grad=False, return_ctype=True
)
if bbox_cutoff_bounds is not None
else None,
}
def _fixed_cluster_op_args(
geometry: _FixedCellGeometry | None,
) -> tuple[
bool,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
]:
"""Return static fixed mode and readonly geometry tensors for custom ops."""
if geometry is None:
return False, None, None, None, None, None
return (
True,
geometry.qr,
geometry.axis_aligned,
geometry.fractional_rounding_certified,
geometry.qr_height_certified,
geometry.bbox_cutoff_bounds,
)
def _fixed_cluster_geometry_record(
inv_cell: torch.Tensor,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> _FixedCellGeometry | None:
"""Wrap query-op tensors in the private geometry record expected downstream."""
if not fixed_cell:
return None
return _FixedCellGeometry(
inv_cell=inv_cell,
cell=None,
cells_per_dimension=None,
neighbor_search_radius=None,
cells_per_system=None,
cell_offsets=None,
qr=qr,
axis_aligned=axis_aligned,
fractional_rounding_certified=fractional_rounding_certified,
qr_height_certified=qr_height_certified,
bbox_cutoff_bounds=bbox_cutoff_bounds,
)
def _wp_direct_csr_geometry_inputs(
device: torch.device,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> tuple[wp.array, wp.array, wp.array, wp.array, wp.array]:
"""Convert direct-CSR geometry inputs, using empty dynamic sentinels."""
device_str = str(device)
return (
wp.from_torch(qr, dtype=wp.float32, requires_grad=False)
if qr is not None
else wp.empty((0, 15), dtype=wp.float32, device=device_str),
wp.from_torch(axis_aligned, dtype=wp.bool, requires_grad=False)
if axis_aligned is not None
else wp.empty(0, dtype=wp.bool, device=device_str),
wp.from_torch(fractional_rounding_certified, dtype=wp.bool, requires_grad=False)
if fractional_rounding_certified is not None
else wp.empty(0, dtype=wp.bool, device=device_str),
wp.from_torch(qr_height_certified, dtype=wp.bool, requires_grad=False)
if qr_height_certified is not None
else wp.empty(0, dtype=wp.bool, device=device_str),
wp.from_torch(bbox_cutoff_bounds, dtype=wp.float32, requires_grad=False)
if bbox_cutoff_bounds is not None
else wp.empty((0, 3), dtype=wp.float32, device=device_str),
)
@torch.library.custom_op(
"nvalchemiops::_cluster_tile_direct_csr_count",
mutates_args=("row_counts",),
)
@scoped_torch_warp_stream
def _cluster_tile_direct_csr_count_op(
cutoff: float,
natom: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
row_counts: torch.Tensor,
launch_tiles: int,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
geometry = _wp_direct_csr_geometry_inputs(
sorted_pos_x.device,
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
)
wp.launch_tiled(
_get_query_cluster_tile_direct_csr_count_kernel(
batched=False, fixed_cell=fixed_cell
),
dim=[launch_tiles],
inputs=[
wp.from_torch(x, dtype=d, return_ctype=True)
for x, d in (
(sorted_pos_x, wp.float32),
(sorted_pos_y, wp.float32),
(sorted_pos_z, wp.float32),
(sorted_atom_index, wp.int32),
(cell.unsqueeze(0), wp.mat33f),
(inv_cell.unsqueeze(0), wp.mat33f),
)
]
+ list(geometry)
+ [
float(cutoff * cutoff),
int(natom),
wp.from_torch(num_tiles, dtype=wp.int32, return_ctype=True),
wp.from_torch(tile_row_group, dtype=wp.int32, return_ctype=True),
wp.from_torch(tile_col_group, dtype=wp.int32, return_ctype=True),
wp.empty(1, dtype=wp.int32, device=str(sorted_pos_x.device)),
wp.from_torch(row_counts, dtype=wp.int32, return_ctype=True),
],
device=str(sorted_pos_x.device),
block_dim=TILE_GROUP_SIZE,
)
@_cluster_tile_direct_csr_count_op.register_fake
def _(*args, **kwargs) -> None:
return None
@torch.library.custom_op(
"nvalchemiops::_cluster_tile_direct_csr_fill",
mutates_args=(
"cursors",
"coo_list",
"coo_shifts",
"neighbor_vectors",
"neighbor_distances",
),
)
@scoped_torch_warp_stream
def _cluster_tile_direct_csr_fill_op(
cutoff: float,
natom: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
cursors: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
neighbor_vectors: torch.Tensor,
neighbor_distances: torch.Tensor,
return_vectors: bool,
return_distances: bool,
max_pairs: int,
launch_tiles: int,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
geometry = _wp_direct_csr_geometry_inputs(
sorted_pos_x.device,
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
)
wp.launch_tiled(
_get_query_cluster_tile_direct_csr_fill_kernel(
batched=False,
fixed_cell=fixed_cell,
return_vectors=return_vectors,
return_distances=return_distances,
),
dim=[launch_tiles],
inputs=[
wp.from_torch(x, dtype=d, return_ctype=True)
for x, d in (
(sorted_pos_x, wp.float32),
(sorted_pos_y, wp.float32),
(sorted_pos_z, wp.float32),
(sorted_atom_index, wp.int32),
(cell.unsqueeze(0), wp.mat33f),
(inv_cell.unsqueeze(0), wp.mat33f),
)
]
+ list(geometry)
+ [
float(cutoff * cutoff),
int(natom),
int(max_pairs),
wp.from_torch(num_tiles, dtype=wp.int32, return_ctype=True),
wp.from_torch(tile_row_group, dtype=wp.int32, return_ctype=True),
wp.from_torch(tile_col_group, dtype=wp.int32, return_ctype=True),
wp.empty(1, dtype=wp.int32, device=str(sorted_pos_x.device)),
wp.from_torch(cursors, dtype=wp.int32, return_ctype=True),
wp.from_torch(coo_list, dtype=wp.int32, return_ctype=True),
wp.from_torch(coo_shifts, dtype=wp.int32, return_ctype=True),
wp.from_torch(neighbor_vectors, dtype=wp.vec3f, return_ctype=True),
wp.from_torch(neighbor_distances, dtype=wp.float32, return_ctype=True),
wp.empty((1, 1), dtype=wp.float32, device=str(sorted_pos_x.device)),
wp.empty(1, dtype=wp.float32, device=str(sorted_pos_x.device)),
wp.empty(1, dtype=wp.vec3f, device=str(sorted_pos_x.device)),
],
device=str(sorted_pos_x.device),
block_dim=TILE_GROUP_SIZE,
)
@_cluster_tile_direct_csr_fill_op.register_fake
def _(*args, **kwargs) -> None:
return None
@scoped_torch_warp_stream
def _cluster_tile_direct_csr_fill_pair_outputs(
cutoff: float,
natom: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
cursors: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
pair_fn: wp.Function,
pair_params: torch.Tensor,
pair_energies: torch.Tensor,
pair_forces: torch.Tensor,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
return_vectors: bool,
return_distances: bool,
max_pairs: int,
launch_tiles: int,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> None:
"""Fill compact CSR rows and eager raw pair-function outputs together."""
device = str(sorted_pos_x.device)
wp_pair_params = wp.from_torch(pair_params, dtype=wp.float32)
wp_neighbor_vectors = (
wp.from_torch(neighbor_vectors, dtype=wp.vec3f)
if neighbor_vectors is not None
else None
)
wp_neighbor_distances = (
wp.from_torch(neighbor_distances, dtype=wp.float32)
if neighbor_distances is not None
else None
)
(
vectors_arg,
distances_arg,
pair_params_arg,
energies_arg,
forces_arg,
) = _prepare_coo_pair_output_args(
wp.float32,
device,
max_pairs,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=wp_pair_params,
neighbor_vectors=wp_neighbor_vectors,
neighbor_distances=wp_neighbor_distances,
pair_energies=wp.from_torch(pair_energies, dtype=wp.float32),
pair_forces=wp.from_torch(pair_forces, dtype=wp.vec3f),
)
wp.launch_tiled(
_get_query_cluster_tile_direct_csr_fill_kernel(
batched=False,
fixed_cell=_fixed_cell_geometry is not None,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
),
dim=[launch_tiles],
inputs=[
wp.from_torch(sorted_pos_x, dtype=wp.float32),
wp.from_torch(sorted_pos_y, dtype=wp.float32),
wp.from_torch(sorted_pos_z, dtype=wp.float32),
wp.from_torch(sorted_atom_index, dtype=wp.int32),
wp.from_torch(cell.unsqueeze(0), dtype=wp.mat33f),
wp.from_torch(inv_cell.unsqueeze(0), dtype=wp.mat33f),
*_wp_direct_csr_geometry_inputs(
sorted_pos_x.device,
*_fixed_cluster_op_args(_fixed_cell_geometry),
),
float(cutoff * cutoff),
int(natom),
int(max_pairs),
wp.from_torch(num_tiles, dtype=wp.int32),
wp.from_torch(tile_row_group, dtype=wp.int32),
wp.from_torch(tile_col_group, dtype=wp.int32),
wp.empty(1, dtype=wp.int32, device=device),
wp.from_torch(cursors, dtype=wp.int32),
wp.from_torch(coo_list, dtype=wp.int32),
wp.from_torch(coo_shifts, dtype=wp.int32),
vectors_arg,
distances_arg,
pair_params_arg,
energies_arg,
forces_arg,
],
device=device,
block_dim=TILE_GROUP_SIZE,
)
@torch.library.custom_op(
"nvalchemiops::_cluster_tile_fill_neighbor_matrix_tail",
mutates_args=("neighbor_matrix",),
)
@scoped_torch_warp_stream
def _cluster_tile_fill_neighbor_matrix_tail_op(
num_neighbors: torch.Tensor,
neighbor_matrix: torch.Tensor,
n_rows: int,
max_neighbors: int,
fill_value: int,
) -> None:
if max_neighbors <= 0:
return
wp_fill_neighbor_matrix_tail(
wp.from_torch(
num_neighbors, dtype=wp.int32, requires_grad=False, return_ctype=True
),
int(n_rows),
int(max_neighbors),
int(fill_value),
wp.from_torch(
neighbor_matrix, dtype=wp.int32, requires_grad=False, return_ctype=True
),
str(neighbor_matrix.device),
)
@_cluster_tile_fill_neighbor_matrix_tail_op.register_fake
def _(
num_neighbors: torch.Tensor,
neighbor_matrix: torch.Tensor,
n_rows: int,
max_neighbors: int,
fill_value: int,
) -> None:
return None
@torch.library.custom_op(
"nvalchemiops::_cluster_tile_selective_zero_num_neighbors_single",
mutates_args=("num_neighbors",),
)
@scoped_torch_warp_stream
def _cluster_tile_selective_zero_num_neighbors_single_op(
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor,
) -> None:
wp_selective_zero_num_neighbors_single(
wp.from_torch(
num_neighbors, dtype=wp.int32, requires_grad=False, return_ctype=True
),
wp.from_torch(
rebuild_flags, dtype=wp.bool, requires_grad=False, return_ctype=True
),
str(num_neighbors.device),
)
@_cluster_tile_selective_zero_num_neighbors_single_op.register_fake
def _(
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor,
) -> None:
return None
@torch.library.custom_op(
"nvalchemiops::_cluster_tile_selective_fill_neighbor_matrix_tail",
mutates_args=("neighbor_matrix",),
)
@scoped_torch_warp_stream
def _cluster_tile_selective_fill_neighbor_matrix_tail_op(
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor,
neighbor_matrix: torch.Tensor,
max_neighbors: int,
fill_value: int,
) -> None:
if max_neighbors <= 0:
return
wp_selective_fill_neighbor_matrix_tail(
wp.from_torch(
num_neighbors, dtype=wp.int32, requires_grad=False, return_ctype=True
),
None,
wp.from_torch(
rebuild_flags, dtype=wp.bool, requires_grad=False, return_ctype=True
),
int(num_neighbors.shape[0]),
int(max_neighbors),
int(fill_value),
wp.from_torch(
neighbor_matrix, dtype=wp.int32, requires_grad=False, return_ctype=True
),
str(neighbor_matrix.device),
batched=False,
)
@_cluster_tile_selective_fill_neighbor_matrix_tail_op.register_fake
def _(
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor,
neighbor_matrix: torch.Tensor,
max_neighbors: int,
fill_value: int,
) -> None:
return None
# =============================================================================
# Sizing + allocation helpers (torch-side, not ``custom_op``-wrapped)
# =============================================================================
[docs]
def estimate_cluster_tile_list_sizes(
total_atoms: int,
max_tiles_per_group: int = 256,
) -> tuple[int, int, int, int]:
"""Estimate allocation sizes for the tile neighbor list state.
Any ``total_atoms >= 0`` is accepted. Nonempty inputs use
``n_padded = ceil(total_atoms / TILE_GROUP_SIZE) * TILE_GROUP_SIZE`` so the
kernels see a 32-aligned layout; empty input reserves one tile group.
Padding slots receive a sentinel maximum Morton code (see
``_compute_morton_kernel``), sort to the end, and are dropped by the
convert/COO kernels' ``i_sorted < natom`` filter.
Parameters
----------
total_atoms : int
Real atom count.
max_tiles_per_group : int, default 256
Sets the capacity of the tile-pair buffer shared by all row groups. The
allocated group count is ``ngroup = max(1, ceil(total_atoms / 32))``.
The capacity is ``ngroup * min(ngroup, max_tiles_per_group)`` entries.
Returns
-------
n_padded : int
Padded atom count. This is at least ``TILE_GROUP_SIZE``; nonempty inputs
are rounded up to a multiple of ``TILE_GROUP_SIZE``.
ngroup : int
Number of 32-atom groups: ``n_padded // TILE_GROUP_SIZE``.
ngroup_padded : int
Group-array pad length for in-bounds ``wp.tile_load`` at any
TILE-aligned offset. Multiple of ``TILE_GROUP_SIZE``; at least one
TILE slack over ``ngroup``.
max_tiles : int
Allocated tile-pair list capacity.
"""
if total_atoms < 0:
raise ValueError(f"total_atoms must be >= 0; got {total_atoms}")
n_padded = (
(total_atoms + TILE_GROUP_SIZE - 1) // TILE_GROUP_SIZE
) * TILE_GROUP_SIZE
if n_padded == 0:
n_padded = TILE_GROUP_SIZE # always reserve at least one tile
ngroup = n_padded // TILE_GROUP_SIZE
ngroup_padded = (
(ngroup + TILE_GROUP_SIZE - 1) // TILE_GROUP_SIZE
) * TILE_GROUP_SIZE
if ngroup_padded == ngroup:
ngroup_padded = ngroup + TILE_GROUP_SIZE
max_tiles = ngroup * min(ngroup, max_tiles_per_group)
return n_padded, ngroup, ngroup_padded, max_tiles
[docs]
def allocate_cluster_tile_list(
total_atoms: int,
device: torch.device,
dtype: torch.dtype = torch.float32,
max_tiles_per_group: int = 256,
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""Allocate all state tensors consumed by ``build_cluster_tile_list``.
Sizes each buffer according to the padded atom count and group count
returned by
:func:`nvalchemiops.torch.neighbors.cluster_tile.estimate_cluster_tile_list_sizes`.
Parameters
----------
total_atoms : int
Real atom count. Any value ``>= 0`` is accepted.
device : torch.device
Target device for all allocated tensors.
dtype : torch.dtype, optional
Floating-point dtype for position and bounding-box arrays.
Default is ``torch.float32``.
max_tiles_per_group : int, optional
Capacity factor for the intermediate tile-pair buffer. For ``g`` row
groups, the buffer holds ``g * min(g, max_tiles_per_group)`` tile pairs.
Increasing the value up to ``g`` uses more memory and accommodates more
candidate tile pairs. Defaults to 256. See
:ref:`cluster-tile-buffer-capacity` for sizing details.
Returns
-------
sorted_atom_index : torch.Tensor, shape (n_padded,), dtype=int32
Permutation that maps sorted rank to original atom index.
morton_codes : torch.Tensor, shape (n_padded,), dtype=int32
30-bit Morton codes for each atom in sorted order; padding slots
carry a sentinel value ``0x40000000``.
sorted_pos_x : torch.Tensor, shape (n_padded,), dtype=dtype
x-coordinates in Morton-sorted order.
sorted_pos_y : torch.Tensor, shape (n_padded,), dtype=dtype
y-coordinates in Morton-sorted order.
sorted_pos_z : torch.Tensor, shape (n_padded,), dtype=dtype
z-coordinates in Morton-sorted order.
group_ctr_x : torch.Tensor, shape (ngroup_padded,), dtype=dtype
x-component of each group bounding-box centre.
group_ctr_y : torch.Tensor, shape (ngroup_padded,), dtype=dtype
y-component of each group bounding-box centre.
group_ctr_z : torch.Tensor, shape (ngroup_padded,), dtype=dtype
z-component of each group bounding-box centre.
group_ext_x : torch.Tensor, shape (ngroup_padded,), dtype=dtype
x half-extent of each group bounding box.
group_ext_y : torch.Tensor, shape (ngroup_padded,), dtype=dtype
y half-extent of each group bounding box.
group_ext_z : torch.Tensor, shape (ngroup_padded,), dtype=dtype
z half-extent of each group bounding box.
num_tiles : torch.Tensor, shape (1,), dtype=int32
Atomic counter holding the number of emitted tile pairs.
tile_row_group : torch.Tensor, shape (max_tiles,), dtype=int32
Row group index for each emitted tile pair.
tile_col_group : torch.Tensor, shape (max_tiles,), dtype=int32
Column group index for each emitted tile pair.
See Also
--------
:func:`nvalchemiops.torch.neighbors.cluster_tile.estimate_cluster_tile_list_sizes` :
Returns the sizing integers used here.
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list` :
Fills the allocated buffers.
"""
n_padded, ngroup, ngroup_padded, max_tiles = estimate_cluster_tile_list_sizes(
total_atoms,
max_tiles_per_group=max_tiles_per_group,
)
# Scratch arrays sized at the padded layout so non-32-aligned
# ``total_atoms`` is handled inside ``_build_cluster_tile_list_op``.
# Padding slots are populated with sentinel Morton codes (see
# ``_compute_morton_kernel``) and dropped by the convert/coo
# kernels' ``i_sorted < natom`` filter.
sorted_atom_index = torch.empty(n_padded, dtype=torch.int32, device=device)
morton_codes = torch.empty(n_padded, dtype=torch.int32, device=device)
sorted_pos_x = torch.empty(n_padded, dtype=dtype, device=device)
sorted_pos_y = torch.empty(n_padded, dtype=dtype, device=device)
sorted_pos_z = torch.empty(n_padded, dtype=dtype, device=device)
group_ctr_x = torch.zeros(ngroup_padded, dtype=dtype, device=device)
group_ctr_y = torch.zeros(ngroup_padded, dtype=dtype, device=device)
group_ctr_z = torch.zeros(ngroup_padded, dtype=dtype, device=device)
group_ext_x = torch.zeros(ngroup_padded, dtype=dtype, device=device)
group_ext_y = torch.zeros(ngroup_padded, dtype=dtype, device=device)
group_ext_z = torch.zeros(ngroup_padded, dtype=dtype, device=device)
num_tiles = torch.zeros(1, dtype=torch.int32, device=device)
tile_row_group = torch.zeros(max_tiles, dtype=torch.int32, device=device)
tile_col_group = torch.zeros(max_tiles, dtype=torch.int32, device=device)
return (
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
)
# =============================================================================
# Internal helpers
# =============================================================================
def _cell_from_cell(cell: torch.Tensor) -> torch.Tensor:
"""Normalize a single-system cell to one contiguous ``(3, 3)`` tensor."""
if cell.ndim == 3:
if cell.shape[0] != 1:
raise ValueError(
f"single-system cluster_tile expects (1, 3, 3) cell; got {tuple(cell.shape)}"
)
cell_mat = cell[0]
elif cell.ndim == 2:
if cell.shape != (3, 3):
raise ValueError(
f"cell must be (3, 3) or (1, 3, 3); got {tuple(cell.shape)}"
)
cell_mat = cell
else:
raise ValueError(f"cell must be (3, 3) or (1, 3, 3); got {tuple(cell.shape)}")
return cell_mat.contiguous()
def _cell_invcell_from_cell(
cell: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Normalize a cell and compute its inverse with capture-safe checking."""
cell_mat = _cell_from_cell(cell)
inv_cell_mat = _capture_safe_inverse(cell_mat).contiguous()
return cell_mat, inv_cell_mat
def _cell_volume(cell: torch.Tensor) -> float:
"""Return ``abs(det(cell))`` for a ``(3, 3)`` or ``(1, 3, 3)`` cell."""
cell_mat = cell[0] if cell.ndim == 3 else cell
return float(torch.linalg.det(cell_mat.to(torch.float64)).abs().item())
@scoped_torch_warp_stream
def _mat33f_from_torch(mat: torch.Tensor):
"""Zero-copy view a ``(1, 3, 3)`` or ``(3, 3)`` torch tensor as a
``wp.array(dtype=wp.mat33f, shape=(1,))``.
These Warp kernels read the cell / inv_cell as length-1
``wp.array(dtype=wp.mat33f)`` and dereference ``cell[0]`` inside the
kernel body, matching the cell_list pattern. This avoids a per-call
host sync.
"""
if mat.ndim == 2:
mat = mat.unsqueeze(0)
return wp.from_torch(
mat.detach().contiguous().to(torch.float32),
dtype=wp.mat33f,
requires_grad=False,
return_ctype=True,
)
# =============================================================================
# Component ops (torch.library.custom_op wrappers)
# =============================================================================
@scoped_torch_warp_stream
def _build_cluster_tile_list_impl(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
use_rebuild_flags: bool,
compute_inv_cell: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
"""Compute Morton codes + argsort + SoA gather in torch, then run
bbox reduction + tile enumeration on the warp side.
Triclinic-safe via the (cell, inv_cell) pair; orthorhombic cells
just have a diagonal inv_cell and produce the same result as the
pre-triclinic implementation.
See Also
--------
nvalchemiops.neighbors.cluster_tile.build_cluster_tile_list : warp launcher
"""
device = positions.device
if compute_inv_cell:
computed_inv_cell, info = torch.linalg.inv_ex(cell, check_errors=False)
torch._assert_async(info == 0, "cell matrix must be non-singular")
inv_cell.copy_(computed_inv_cell)
N = positions.shape[0]
if N == 0:
return
wp_device = str(device)
wp_dtype = get_wp_dtype(positions.dtype)
if not use_rebuild_flags:
num_tiles.zero_()
wp_cell = _mat33f_from_torch(cell)
wp_inv_cell = _mat33f_from_torch(inv_cell)
# ---- Steps 1-3: Morton codes + argsort + gather (torch-side) ----
# Fractional coords via inv_cell; wrap into [0, 1) so the bucket
# produces a deterministic 30-bit Morton code. Orthorhombic cells
# collapse this back to the cheaper diagonal multiply.
#
# Pad to ``n_padded = ceil(N / TILE_GROUP_SIZE) * TILE_GROUP_SIZE`` so the
# kernels see a 32-aligned layout. Padding slots get a sentinel
# max 30-bit Morton code (0x3FFFFFFF) so radix sort places them at
# the end of the sorted layout, where the convert/coo kernels'
# ``i_sorted < natom`` filter naturally drops any pair involving
# them. Padding positions copy the last real atom (any in-cell
# position is safe) so the SoA gather has well-defined reads.
n_padded = int(sorted_atom_index.shape[0])
if N < n_padded:
padded_positions = torch.empty(
(n_padded, 3),
dtype=positions.dtype,
device=positions.device,
)
padded_positions[:N].copy_(positions)
padded_positions[N:].copy_(positions[-1:].expand(n_padded - N, 3))
else:
padded_positions = positions
frac = padded_positions @ inv_cell.T
frac = frac - torch.floor(frac)
bucket = (frac * 1024.0).clamp(0, 1023).to(torch.int32)
ix, iy, iz = bucket.unbind(dim=-1)
def _spread(x: torch.Tensor) -> torch.Tensor:
x = x & 0x3FF
x = (x | (x << 16)) & 0x030000FF
x = (x | (x << 8)) & 0x0300F00F
x = (x | (x << 4)) & 0x030C30C3
x = (x | (x << 2)) & 0x09249249
return x
codes32 = (_spread(iz) << 2) | (_spread(iy) << 1) | _spread(ix)
if N < n_padded:
# Sentinel: one bit above any real 30-bit Morton code. Real
# codes saturate at 0x3FFFFFFF; 0x40000000 sorts after all of
# them. See ``_compute_morton_kernel`` for the matching
# warp-side value.
codes32[N:] = 0x40000000
morton_codes.copy_(codes32)
perm = torch.argsort(codes32)
sorted_atom_index.copy_(perm.to(torch.int32))
sorted_pos = padded_positions[perm].contiguous()
sorted_pos_x.copy_(sorted_pos[:, 0].contiguous())
sorted_pos_y.copy_(sorted_pos[:, 1].contiguous())
sorted_pos_z.copy_(sorted_pos[:, 2].contiguous())
# ---- Step 4: rank2group + group2tile (warp launcher) ----
wp_sorted_pos_x = wp.from_torch(
sorted_pos_x,
dtype=wp_dtype,
return_ctype=True,
)
wp_sorted_pos_y = wp.from_torch(
sorted_pos_y,
dtype=wp_dtype,
return_ctype=True,
)
wp_sorted_pos_z = wp.from_torch(
sorted_pos_z,
dtype=wp_dtype,
return_ctype=True,
)
wp_group_ctr_x = wp.from_torch(group_ctr_x, dtype=wp_dtype, return_ctype=True)
wp_group_ctr_y = wp.from_torch(group_ctr_y, dtype=wp_dtype, return_ctype=True)
wp_group_ctr_z = wp.from_torch(group_ctr_z, dtype=wp_dtype, return_ctype=True)
wp_group_ext_x = wp.from_torch(group_ext_x, dtype=wp_dtype, return_ctype=True)
wp_group_ext_y = wp.from_torch(group_ext_y, dtype=wp_dtype, return_ctype=True)
wp_group_ext_z = wp.from_torch(group_ext_z, dtype=wp_dtype, return_ctype=True)
wp_num_tiles = wp.from_torch(num_tiles, dtype=wp.int32, return_ctype=True)
wp_tile_row_group = wp.from_torch(tile_row_group, dtype=wp.int32, return_ctype=True)
wp_tile_col_group = wp.from_torch(
tile_col_group,
dtype=wp.int32,
return_ctype=True,
)
wp_rebuild_flags = wp.from_torch(
rebuild_flags,
dtype=wp.bool,
return_ctype=True,
)
wp_build_cluster_tile_list(
sorted_pos_x=wp_sorted_pos_x,
sorted_pos_y=wp_sorted_pos_y,
sorted_pos_z=wp_sorted_pos_z,
cell=wp_cell,
inv_cell=wp_inv_cell,
cutoff=float(cutoff),
num_tiles=wp_num_tiles,
tile_row_group=wp_tile_row_group,
tile_col_group=wp_tile_col_group,
wp_dtype=wp_dtype,
device=wp_device,
group_ctr_x_buffer=wp_group_ctr_x,
group_ctr_y_buffer=wp_group_ctr_y,
group_ctr_z_buffer=wp_group_ctr_z,
group_ext_x_buffer=wp_group_ext_x,
group_ext_y_buffer=wp_group_ext_y,
group_ext_z_buffer=wp_group_ext_z,
rebuild_flags=wp_rebuild_flags if use_rebuild_flags else None,
fixed_cell=fixed_cell,
qr=(
wp.from_torch(qr, dtype=wp.float32, return_ctype=True)
if qr is not None
else None
),
axis_aligned=(
wp.from_torch(axis_aligned, dtype=wp.bool, return_ctype=True)
if axis_aligned is not None
else None
),
fractional_rounding_certified=(
wp.from_torch(
fractional_rounding_certified, dtype=wp.bool, return_ctype=True
)
if fractional_rounding_certified is not None
else None
),
qr_height_certified=(
wp.from_torch(qr_height_certified, dtype=wp.bool, return_ctype=True)
if qr_height_certified is not None
else None
),
bbox_cutoff_bounds=(
wp.from_torch(bbox_cutoff_bounds, dtype=wp.float32, return_ctype=True)
if bbox_cutoff_bounds is not None
else None
),
)
@torch.library.custom_op(
"nvalchemiops::_build_cluster_tile_list",
mutates_args=(
"sorted_atom_index",
"morton_codes",
"sorted_pos_x",
"sorted_pos_y",
"sorted_pos_z",
"group_ctr_x",
"group_ctr_y",
"group_ctr_z",
"group_ext_x",
"group_ext_y",
"group_ext_z",
"inv_cell",
"num_tiles",
"tile_row_group",
"tile_col_group",
),
)
def _build_cluster_tile_list_op(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
use_rebuild_flags: bool,
compute_inv_cell: bool,
) -> None:
_build_cluster_tile_list_impl(
positions,
cutoff,
cell,
inv_cell,
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
rebuild_flags,
use_rebuild_flags,
compute_inv_cell,
False,
None,
None,
None,
None,
None,
)
@torch.library.custom_op(
"nvalchemiops::_build_cluster_tile_list_fixed",
mutates_args=(
"sorted_atom_index",
"morton_codes",
"sorted_pos_x",
"sorted_pos_y",
"sorted_pos_z",
"group_ctr_x",
"group_ctr_y",
"group_ctr_z",
"group_ext_x",
"group_ext_y",
"group_ext_z",
"num_tiles",
"tile_row_group",
"tile_col_group",
),
)
def _build_cluster_tile_list_fixed_op(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
use_rebuild_flags: bool,
qr: torch.Tensor,
axis_aligned: torch.Tensor,
fractional_rounding_certified: torch.Tensor,
qr_height_certified: torch.Tensor,
bbox_cutoff_bounds: torch.Tensor,
) -> None:
_build_cluster_tile_list_impl(
positions,
cutoff,
cell,
inv_cell,
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
rebuild_flags,
use_rebuild_flags,
False,
True,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
)
@_build_cluster_tile_list_op.register_fake
def _(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
use_rebuild_flags: bool,
compute_inv_cell: bool,
) -> None:
return None
@_build_cluster_tile_list_fixed_op.register_fake
def _(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
use_rebuild_flags: bool,
qr: torch.Tensor,
axis_aligned: torch.Tensor,
fractional_rounding_certified: torch.Tensor,
qr_height_certified: torch.Tensor,
bbox_cutoff_bounds: torch.Tensor,
) -> None:
return None
def _build_cluster_tile_list_normalized(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
*,
inv_cell: torch.Tensor | None = None,
rebuild_flags: torch.Tensor | None = None,
force_compute_inv_cell: bool = False,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build tiles and return the normalized cell and inverse used by queries."""
if positions.dtype != torch.float32:
raise TypeError("positions must be float32")
cell_mat = _cell_from_cell(cell).to(positions.dtype)
compute_inv_cell = _fixed_cell_geometry is None and (
force_compute_inv_cell or torch.compiler.is_compiling()
)
inv_cell_mat = (
torch.empty_like(cell_mat)
if inv_cell is None and compute_inv_cell
else (torch.linalg.inv(cell_mat).contiguous() if inv_cell is None else inv_cell)
)
dummy_rebuild_flags = torch.empty(1, dtype=torch.bool, device=positions.device)
build_op = (
_build_cluster_tile_list_fixed_op
if _fixed_cell_geometry is not None
else _build_cluster_tile_list_op
)
build_args = (
positions,
cutoff,
cell_mat,
inv_cell_mat,
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
rebuild_flags if rebuild_flags is not None else dummy_rebuild_flags,
rebuild_flags is not None,
)
if _fixed_cell_geometry is None:
build_op(*build_args, compute_inv_cell)
else:
geometry = _fixed_cell_geometry
if any(
value is None
for value in (
geometry.qr,
geometry.axis_aligned,
geometry.fractional_rounding_certified,
geometry.qr_height_certified,
geometry.bbox_cutoff_bounds,
)
):
raise RuntimeError("prepared cluster geometry cache is incomplete")
build_op(
*build_args,
geometry.qr,
geometry.axis_aligned,
geometry.fractional_rounding_certified,
geometry.qr_height_certified,
geometry.bbox_cutoff_bounds,
)
return cell_mat, inv_cell_mat
[docs]
def build_cluster_tile_list(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
morton_codes: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
group_ctr_x: torch.Tensor,
group_ctr_y: torch.Tensor,
group_ctr_z: torch.Tensor,
group_ext_x: torch.Tensor,
group_ext_y: torch.Tensor,
group_ext_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
*,
rebuild_flags: torch.Tensor | None = None,
) -> None:
"""Build cluster-tile neighbor list state into pre-allocated tensors.
Normalizes ``cell`` to a ``(3, 3)`` matrix and computes ``inv_cell``,
then runs Morton sort (torch) + Warp bounding-box reduction + Warp
tile-pair enumeration. Triclinic cells are supported. All output
tensors are filled in place.
Parameters
----------
positions : torch.Tensor, shape (N, 3), dtype=float32
Atomic coordinates wrapped to the primary cell. Non-32-aligned ``N``
is padded internally to ``ceil(N / TILE_GROUP_SIZE) * TILE_GROUP_SIZE``.
cutoff : float
Cutoff distance in Cartesian units used for tile-pair pruning.
cell : torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32
Any non-degenerate cell (orthorhombic or triclinic).
sorted_atom_index : torch.Tensor, shape (n_padded,), dtype=int32
Output permutation mapping sorted rank to original atom index.
Modified in-place.
morton_codes : torch.Tensor, shape (n_padded,), dtype=int32
Output scratch buffer for 30-bit Morton codes. Modified in-place.
sorted_pos_x : torch.Tensor, shape (n_padded,), dtype=float32
Output x-coordinates in Morton-sorted order. Modified in-place.
sorted_pos_y : torch.Tensor, shape (n_padded,), dtype=float32
Output y-coordinates in Morton-sorted order. Modified in-place.
sorted_pos_z : torch.Tensor, shape (n_padded,), dtype=float32
Output z-coordinates in Morton-sorted order. Modified in-place.
group_ctr_x : torch.Tensor, shape (ngroup_padded,), dtype=float32
Output x-component of group bounding-box centres. Modified in-place.
group_ctr_y : torch.Tensor, shape (ngroup_padded,), dtype=float32
Output y-component of group bounding-box centres. Modified in-place.
group_ctr_z : torch.Tensor, shape (ngroup_padded,), dtype=float32
Output z-component of group bounding-box centres. Modified in-place.
group_ext_x : torch.Tensor, shape (ngroup_padded,), dtype=float32
Output x half-extents of group bounding boxes. Modified in-place.
group_ext_y : torch.Tensor, shape (ngroup_padded,), dtype=float32
Output y half-extents of group bounding boxes. Modified in-place.
group_ext_z : torch.Tensor, shape (ngroup_padded,), dtype=float32
Output z half-extents of group bounding boxes. Modified in-place.
num_tiles : torch.Tensor, shape (1,), dtype=int32
Output atomic counter holding the number of emitted tile pairs.
Reset to zero internally before use. Modified in-place.
tile_row_group : torch.Tensor, shape (max_tiles,), dtype=int32
Output row group index for each emitted tile pair. Modified in-place.
tile_col_group : torch.Tensor, shape (max_tiles,), dtype=int32
Output column group index for each emitted tile pair. Modified in-place.
See Also
--------
:func:`nvalchemiops.torch.neighbors.cluster_tile.allocate_cluster_tile_list` :
Allocates all buffers consumed by this function.
:func:`nvalchemiops.neighbors.cluster_tile.build_cluster_tile_list` :
Warp-level launcher called internally.
"""
_build_cluster_tile_list_normalized(
positions,
cutoff,
cell,
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
rebuild_flags=rebuild_flags,
)
@torch.library.custom_op(
"nvalchemiops::_query_cluster_tile",
mutates_args=(
"neighbor_matrix",
"neighbor_matrix_shifts",
"num_neighbors",
"neighbor_vectors",
"neighbor_distances",
),
)
@scoped_torch_warp_stream
def _query_cluster_tile_op(
cutoff: float,
natom: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_vectors: torch.Tensor,
neighbor_distances: torch.Tensor,
n_tiles: int,
return_vectors: bool,
return_distances: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
# ``n_tiles`` (host-synced emitted-tile count from the caller) sets the
# launch dimension so we don't launch over the full allocated tile
# buffer. The kernel still guards ``tile >= num_tiles[0]`` defensively.
# Warp writes active pairs only. Clear padding inside the same opaque
# mutation boundary before launching the query.
if return_vectors:
neighbor_vectors.zero_()
if return_distances:
neighbor_distances.zero_()
device = sorted_pos_x.device
wp_device = str(device)
wp_dtype = get_wp_dtype(sorted_pos_x.dtype)
wp_cell = _mat33f_from_torch(cell)
wp_inv_cell = _mat33f_from_torch(inv_cell)
wp_query_cluster_tile(
sorted_atom_index=wp.from_torch(
sorted_atom_index, dtype=wp.int32, return_ctype=True
),
sorted_pos_x=wp.from_torch(sorted_pos_x, dtype=wp_dtype, return_ctype=True),
sorted_pos_y=wp.from_torch(sorted_pos_y, dtype=wp_dtype, return_ctype=True),
sorted_pos_z=wp.from_torch(sorted_pos_z, dtype=wp_dtype, return_ctype=True),
num_tiles=wp.from_torch(num_tiles, dtype=wp.int32, return_ctype=True),
tile_row_group=wp.from_torch(tile_row_group, dtype=wp.int32, return_ctype=True),
tile_col_group=wp.from_torch(
tile_col_group,
dtype=wp.int32,
return_ctype=True,
),
cell=wp_cell,
inv_cell=wp_inv_cell,
cutoff=float(cutoff),
natom=int(natom),
neighbor_matrix=wp.from_torch(
neighbor_matrix,
dtype=wp.int32,
return_ctype=True,
),
num_neighbors=wp.from_torch(
num_neighbors,
dtype=wp.int32,
return_ctype=True,
),
neighbor_matrix_shifts=wp.from_torch(
neighbor_matrix_shifts,
dtype=wp.int32,
return_ctype=True,
),
wp_dtype=wp_dtype,
device=wp_device,
n_tiles=int(n_tiles),
return_vectors=bool(return_vectors),
return_distances=bool(return_distances),
neighbor_vectors=(
wp.from_torch(neighbor_vectors, dtype=wp.vec3f, return_ctype=True)
if return_vectors
else None
),
neighbor_distances=(
wp.from_torch(neighbor_distances, dtype=wp_dtype, return_ctype=True)
if return_distances
else None
),
**_wp_fixed_cluster_geometry(
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
),
)
@_query_cluster_tile_op.register_fake
def _(
cutoff: float,
natom: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_vectors: torch.Tensor,
neighbor_distances: torch.Tensor,
n_tiles: int,
return_vectors: bool,
return_distances: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
return None
def _query_cluster_tile_normalized(
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
cell: torch.Tensor,
cutoff: float,
natom: int,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
*,
inv_cell_mat: torch.Tensor | None = None,
cutoff2: float | None = None,
neighbor_matrix2: torch.Tensor | None = None,
num_neighbors2: torch.Tensor | None = None,
neighbor_matrix_shifts2: torch.Tensor | None = None,
rebuild_flags: torch.Tensor | None = None,
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> None:
"""Convert the tile pair list to neighbor_matrix form in place.
Cluster-tile does not support partial neighbor lists; there is no
``target_indices`` kwarg. Use
:func:`nvalchemiops.torch.neighbors.cell_list.cell_list` or
:func:`nvalchemiops.torch.neighbors.naive.naive_neighbor_list` for
partial neighbor lists.
Parameters
----------
sorted_atom_index : torch.Tensor, shape (n_padded,), dtype=int32
Permutation mapping sorted rank to original atom index; output of
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list`.
sorted_pos_x : torch.Tensor, shape (n_padded,), dtype=float32
x-coordinates in Morton-sorted order.
sorted_pos_y : torch.Tensor, shape (n_padded,), dtype=float32
y-coordinates in Morton-sorted order.
sorted_pos_z : torch.Tensor, shape (n_padded,), dtype=float32
z-coordinates in Morton-sorted order.
num_tiles : torch.Tensor, shape (1,), dtype=int32
Device-side tile counter written by
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list`.
tile_row_group : torch.Tensor, shape (max_tiles,), dtype=int32
Row group indices of emitted tile pairs.
tile_col_group : torch.Tensor, shape (max_tiles,), dtype=int32
Column group indices of emitted tile pairs.
cell : torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32
Simulation cell matrix (orthorhombic or triclinic).
cutoff : float
Neighbor search cutoff radius in Cartesian units.
natom : int
True atom count (before padding).
neighbor_matrix : torch.Tensor, shape (natom, max_neighbors), dtype=int32
Output neighbor indices. Modified in-place.
num_neighbors : torch.Tensor, shape (natom,), dtype=int32
Output per-atom neighbor counts. Modified in-place.
neighbor_matrix_shifts : torch.Tensor, shape (natom, max_neighbors, 3), dtype=int32
Output per-pair periodic image shift vectors. Modified in-place.
cutoff2 : float, optional
Second cutoff for a dual-cutoff query; fills ``neighbor_matrix2`` /
``num_neighbors2`` / ``neighbor_matrix_shifts2`` when provided.
neighbor_matrix2 : torch.Tensor, shape (natom, max_neighbors), dtype=int32, optional
Second-cutoff output neighbor indices. Modified in-place.
num_neighbors2 : torch.Tensor, shape (natom,), dtype=int32, optional
Second-cutoff per-atom neighbor counts. Modified in-place.
neighbor_matrix_shifts2 : torch.Tensor, shape (natom, max_neighbors, 3), dtype=int32, optional
Second-cutoff per-pair shift vectors. Modified in-place.
rebuild_flags : torch.Tensor, shape (1,), dtype=bool, optional
When ``False``, all output buffers are left unchanged and the call
returns early.
return_vectors : bool, optional
Write per-pair Cartesian displacement vectors to ``neighbor_vectors``.
Default is ``False``.
return_distances : bool, optional
Write per-pair scalar distances to ``neighbor_distances``.
Default is ``False``.
pair_fn : wp.Function, optional
Module-scope Warp ``@wp.func`` of signature
``(r_ij, distance, pair_params, i, j) -> (energy, force)``.
pair_params : torch.Tensor, shape (natom, num_parameters), optional
Per-atom pair-function parameters; required with ``pair_fn``.
neighbor_vectors : torch.Tensor, shape (natom, max_neighbors, 3), optional
Output buffer for per-pair displacement vectors. Modified in-place.
neighbor_distances : torch.Tensor, shape (natom, max_neighbors), optional
Output buffer for per-pair scalar distances. Modified in-place.
pair_energies : torch.Tensor, shape (natom, max_neighbors), optional
Output buffer for per-pair energies; required with ``pair_fn``.
Modified in-place.
pair_forces : torch.Tensor, shape (natom, max_neighbors, 3), optional
Output buffer for per-pair forces; required with ``pair_fn``.
Modified in-place.
See Also
--------
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list` :
Produces the tile state consumed by this function.
:func:`nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile_coo` :
COO-format alternative that emits a flat pair list instead.
"""
_validate_cluster_tile_matrix_outputs(
device=sorted_pos_x.device,
dtype=sorted_pos_x.dtype,
natom=int(natom),
max_neighbors=int(neighbor_matrix.shape[1]),
cutoff2=cutoff2,
neighbor_matrix2=neighbor_matrix2,
num_neighbors2=num_neighbors2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
return_vectors=return_vectors,
return_distances=return_distances,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
allocate_missing=False,
)
cell_mat = _cell_from_cell(cell).to(sorted_pos_x.dtype)
if inv_cell_mat is None:
inv_cell_mat = torch.linalg.inv(cell_mat).contiguous()
else:
inv_cell_mat = inv_cell_mat.to(sorted_pos_x.dtype)
fixed_geometry_args = _fixed_cluster_op_args(_fixed_cell_geometry)
# Eager execution tightens the launch to the emitted tile count. Compiled
# execution validates on device and launches the full static capacity.
tile_capacity = int(tile_row_group.shape[0])
n_tiles = _check_tile_buffer_capacity(num_tiles, tile_capacity)
geometry_only = (
(return_vectors or return_distances)
and cutoff2 is None
and rebuild_flags is None
and pair_fn is None
and pair_params is None
and pair_energies is None
and pair_forces is None
)
feature_path = (
cutoff2 is not None
or rebuild_flags is not None
or _has_partial_or_pair_outputs(
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
)
topology_only = not _has_partial_or_pair_outputs(
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
if topology_only and (cutoff2 is not None or rebuild_flags is not None):
device = sorted_pos_x.device
dummy_matrix = torch.empty((1, 1), dtype=torch.int32, device=device)
dummy_counts = torch.empty(1, dtype=torch.int32, device=device)
dummy_shifts = torch.empty((1, 1, 3), dtype=torch.int32, device=device)
dummy_bool = torch.empty(1, dtype=torch.bool, device=device)
_query_cluster_tile_topology_op(
cell_mat,
inv_cell_mat,
int(natom),
float(cutoff),
float(cutoff2) if cutoff2 is not None else 0.0,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
neighbor_matrix2 if neighbor_matrix2 is not None else dummy_matrix,
num_neighbors2 if num_neighbors2 is not None else dummy_counts,
neighbor_matrix_shifts2
if neighbor_matrix_shifts2 is not None
else dummy_shifts,
rebuild_flags if rebuild_flags is not None else dummy_bool,
int(n_tiles),
bool(cutoff2 is not None),
bool(rebuild_flags is not None),
*fixed_geometry_args,
)
return
if geometry_only:
device = sorted_pos_x.device
dummy_vectors = torch.empty((1, 3), dtype=sorted_pos_x.dtype, device=device)
dummy_distances = torch.empty(1, dtype=sorted_pos_x.dtype, device=device)
_query_cluster_tile_op(
cutoff,
natom,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
neighbor_vectors if neighbor_vectors is not None else dummy_vectors,
neighbor_distances if neighbor_distances is not None else dummy_distances,
n_tiles,
bool(return_vectors),
bool(return_distances),
*fixed_geometry_args,
)
return
if feature_path:
if (
pair_fn is None
and pair_params is None
and pair_energies is None
and pair_forces is None
):
return _query_cluster_tile_optional_no_pair_fn_op(
cell_mat,
inv_cell_mat,
natom,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
neighbor_matrix2,
num_neighbors2,
neighbor_matrix_shifts2,
rebuild_flags,
neighbor_vectors,
neighbor_distances,
n_tiles,
cutoff2,
return_vectors,
return_distances,
*fixed_geometry_args,
)
if torch.compiler.is_compiling():
raise NotImplementedError(
"cluster_tile pair_fn outputs are eager-only because callable Warp "
"functions cannot cross a torch.library.custom_op schema boundary.",
)
# Pair outputs are exercised - bypass the torch custom op because it
# cannot carry a callable ``pair_fn``.
_query_cluster_tile_optional(
cell_mat,
inv_cell_mat,
natom,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
cutoff2=cutoff2,
neighbor_matrix2=neighbor_matrix2,
num_neighbors2=num_neighbors2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
rebuild_flags=rebuild_flags,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
n_tiles=n_tiles,
_fixed_cell_geometry=_fixed_cell_geometry,
)
return
_query_cluster_tile_op(
cutoff,
natom,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
torch.empty((1, 3), dtype=sorted_pos_x.dtype, device=sorted_pos_x.device),
torch.empty(1, dtype=sorted_pos_x.dtype, device=sorted_pos_x.device),
n_tiles,
False,
False,
*fixed_geometry_args,
)
[docs]
def query_cluster_tile(
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
cell: torch.Tensor,
cutoff: float,
natom: int,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
*,
cutoff2: float | None = None,
neighbor_matrix2: torch.Tensor | None = None,
num_neighbors2: torch.Tensor | None = None,
neighbor_matrix_shifts2: torch.Tensor | None = None,
rebuild_flags: torch.Tensor | None = None,
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> None:
"""Convert the tile pair list to neighbor-matrix form in place."""
_query_cluster_tile_normalized(
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
cell,
cutoff,
natom,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
cutoff2=cutoff2,
neighbor_matrix2=neighbor_matrix2,
num_neighbors2=num_neighbors2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
rebuild_flags=rebuild_flags,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
_fixed_cell_geometry=_fixed_cell_geometry,
)
query_cluster_tile.__doc__ = _query_cluster_tile_normalized.__doc__
@torch.library.custom_op(
"nvalchemiops::_query_cluster_tile_optional_no_pair_fn",
mutates_args=(
"neighbor_matrix",
"num_neighbors",
"neighbor_matrix_shifts",
"neighbor_matrix2",
"num_neighbors2",
"neighbor_matrix_shifts2",
"neighbor_vectors",
"neighbor_distances",
),
)
@scoped_torch_warp_stream
def _query_cluster_tile_optional_no_pair_fn_op(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
cutoff: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
neighbor_matrix2: torch.Tensor | None,
num_neighbors2: torch.Tensor | None,
neighbor_matrix_shifts2: torch.Tensor | None,
rebuild_flags: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
n_tiles: int,
cutoff2: float | None,
return_vectors: bool,
return_distances: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
# Selective queries must retain every output entry when the system is not
# rebuilt; a rebuilt system starts with zeroed geometry padding.
if rebuild_flags is None:
if return_vectors and neighbor_vectors is not None:
neighbor_vectors.zero_()
if return_distances and neighbor_distances is not None:
neighbor_distances.zero_()
else:
rebuild = rebuild_flags.reshape(-1)[0]
if return_vectors and neighbor_vectors is not None:
neighbor_vectors.copy_(
torch.where(
rebuild, torch.zeros_like(neighbor_vectors), neighbor_vectors
)
)
if return_distances and neighbor_distances is not None:
neighbor_distances.copy_(
torch.where(
rebuild,
torch.zeros_like(neighbor_distances),
neighbor_distances,
)
)
_query_cluster_tile_optional(
cell_mat,
inv_cell_mat,
natom,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
n_tiles=n_tiles,
cutoff2=cutoff2,
neighbor_matrix2=neighbor_matrix2,
num_neighbors2=num_neighbors2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
rebuild_flags=rebuild_flags,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=None,
pair_params=None,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=None,
pair_forces=None,
_fixed_cell_geometry=_fixed_cluster_geometry_record(
inv_cell_mat,
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
),
)
@_query_cluster_tile_optional_no_pair_fn_op.register_fake
def _(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
cutoff: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
neighbor_matrix2: torch.Tensor | None,
num_neighbors2: torch.Tensor | None,
neighbor_matrix_shifts2: torch.Tensor | None,
rebuild_flags: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
n_tiles: int,
cutoff2: float | None,
return_vectors: bool,
return_distances: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
return None
@torch.library.custom_op(
"nvalchemiops::_query_cluster_tile_topology",
mutates_args=(
"neighbor_matrix",
"num_neighbors",
"neighbor_matrix_shifts",
"neighbor_matrix2",
"num_neighbors2",
"neighbor_matrix_shifts2",
),
)
@scoped_torch_warp_stream
def _query_cluster_tile_topology_op(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
cutoff: float,
cutoff2: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
neighbor_matrix2: torch.Tensor,
num_neighbors2: torch.Tensor,
neighbor_matrix_shifts2: torch.Tensor,
rebuild_flags: torch.Tensor,
n_tiles: int,
use_cutoff2: bool,
use_rebuild_flags: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
_query_cluster_tile_optional(
cell_mat,
inv_cell_mat,
natom,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
n_tiles=n_tiles,
cutoff2=cutoff2 if use_cutoff2 else None,
neighbor_matrix2=neighbor_matrix2 if use_cutoff2 else None,
num_neighbors2=num_neighbors2 if use_cutoff2 else None,
neighbor_matrix_shifts2=neighbor_matrix_shifts2 if use_cutoff2 else None,
rebuild_flags=rebuild_flags if use_rebuild_flags else None,
return_vectors=False,
return_distances=False,
pair_fn=None,
pair_params=None,
neighbor_vectors=None,
neighbor_distances=None,
pair_energies=None,
pair_forces=None,
_fixed_cell_geometry=_fixed_cluster_geometry_record(
inv_cell_mat,
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
),
)
@_query_cluster_tile_topology_op.register_fake
def _(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
cutoff: float,
cutoff2: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
neighbor_matrix2: torch.Tensor,
num_neighbors2: torch.Tensor,
neighbor_matrix_shifts2: torch.Tensor,
rebuild_flags: torch.Tensor,
n_tiles: int,
use_cutoff2: bool,
use_rebuild_flags: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
return None
@scoped_torch_warp_stream
def _query_cluster_tile_optional(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
cutoff: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
*,
n_tiles: int,
cutoff2: float | None,
neighbor_matrix2: torch.Tensor | None,
num_neighbors2: torch.Tensor | None,
neighbor_matrix_shifts2: torch.Tensor | None,
rebuild_flags: torch.Tensor | None,
return_vectors: bool,
return_distances: bool,
pair_fn: wp.Function | None,
pair_params: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
pair_energies: torch.Tensor | None,
pair_forces: torch.Tensor | None,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> None:
"""Pair-output path: bypass the torch custom op + call warp directly.
Mirrors :func:`nvalchemiops.torch.neighbors.cell_list._query_cell_list_optional`.
Torch custom ops cannot carry a callable ``pair_fn`` across their
schema boundary, so the pair-output path drops down to the warp
launcher with ``wp.from_torch``-wrapped tensors directly. No
host-side ``num_tiles.item()`` sync: the warp kernel guards
per-tile via the device-side ``num_tiles`` array.
"""
device = sorted_pos_x.device
wp_device = str(device)
wp_dtype = get_wp_dtype(sorted_pos_x.dtype)
wp_vec_dtype = wp.vec3f
wp_cell = _mat33f_from_torch(cell_mat)
wp_inv_cell = _mat33f_from_torch(inv_cell_mat)
wp_pair_params = (
wp.from_torch(pair_params, dtype=wp_dtype, requires_grad=False)
if pair_params is not None
else None
)
wp_neighbor_vectors = (
wp.from_torch(neighbor_vectors, dtype=wp_vec_dtype, requires_grad=False)
if neighbor_vectors is not None
else None
)
wp_neighbor_distances = (
wp.from_torch(neighbor_distances, dtype=wp_dtype, requires_grad=False)
if neighbor_distances is not None
else None
)
wp_pair_energies = (
wp.from_torch(pair_energies, dtype=wp_dtype, requires_grad=False)
if pair_energies is not None
else None
)
wp_pair_forces = (
wp.from_torch(pair_forces, dtype=wp_vec_dtype, requires_grad=False)
if pair_forces is not None
else None
)
wp_query_cluster_tile(
sorted_atom_index=wp.from_torch(
sorted_atom_index, dtype=wp.int32, requires_grad=False
),
sorted_pos_x=wp.from_torch(sorted_pos_x, dtype=wp_dtype, requires_grad=False),
sorted_pos_y=wp.from_torch(sorted_pos_y, dtype=wp_dtype, requires_grad=False),
sorted_pos_z=wp.from_torch(sorted_pos_z, dtype=wp_dtype, requires_grad=False),
num_tiles=wp.from_torch(num_tiles, dtype=wp.int32, requires_grad=False),
tile_row_group=wp.from_torch(
tile_row_group, dtype=wp.int32, requires_grad=False
),
tile_col_group=wp.from_torch(
tile_col_group, dtype=wp.int32, requires_grad=False
),
cell=wp_cell,
inv_cell=wp_inv_cell,
cutoff=float(cutoff),
natom=int(natom),
neighbor_matrix=wp.from_torch(
neighbor_matrix, dtype=wp.int32, requires_grad=False
),
num_neighbors=wp.from_torch(num_neighbors, dtype=wp.int32, requires_grad=False),
neighbor_matrix_shifts=wp.from_torch(
neighbor_matrix_shifts, dtype=wp.int32, requires_grad=False
),
wp_dtype=wp_dtype,
device=wp_device,
n_tiles=int(n_tiles),
cutoff2=cutoff2,
neighbor_matrix2=(
wp.from_torch(neighbor_matrix2, dtype=wp.int32, requires_grad=False)
if neighbor_matrix2 is not None
else None
),
num_neighbors2=(
wp.from_torch(num_neighbors2, dtype=wp.int32, requires_grad=False)
if num_neighbors2 is not None
else None
),
neighbor_matrix_shifts2=(
wp.from_torch(neighbor_matrix_shifts2, dtype=wp.int32, requires_grad=False)
if neighbor_matrix_shifts2 is not None
else None
),
rebuild_flags=(
wp.from_torch(rebuild_flags, dtype=wp.bool, requires_grad=False)
if rebuild_flags is not None
else None
),
return_vectors=bool(return_vectors),
return_distances=bool(return_distances),
pair_fn=pair_fn,
pair_params=wp_pair_params,
neighbor_vectors=wp_neighbor_vectors,
neighbor_distances=wp_neighbor_distances,
pair_energies=wp_pair_energies,
pair_forces=wp_pair_forces,
**_wp_fixed_cluster_geometry(*_fixed_cluster_op_args(_fixed_cell_geometry)),
)
@torch.library.custom_op(
"nvalchemiops::_query_cluster_tile_coo",
mutates_args=("pair_counter", "coo_list", "coo_shifts"),
)
@scoped_torch_warp_stream
def _query_cluster_tile_coo_op(
cutoff: float,
natom: int,
max_pairs: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
pair_counter: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
n_tiles: int,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
# ``n_tiles`` (host-synced emitted-tile count from the caller) tightens
# the launch dimension; the kernel still guards per-tile defensively.
device = sorted_pos_x.device
wp_device = str(device)
wp_dtype = get_wp_dtype(sorted_pos_x.dtype)
wp_cell = _mat33f_from_torch(cell)
wp_inv_cell = _mat33f_from_torch(inv_cell)
wp_query_cluster_tile_coo(
sorted_atom_index=wp.from_torch(
sorted_atom_index, dtype=wp.int32, return_ctype=True
),
sorted_pos_x=wp.from_torch(sorted_pos_x, dtype=wp_dtype, return_ctype=True),
sorted_pos_y=wp.from_torch(sorted_pos_y, dtype=wp_dtype, return_ctype=True),
sorted_pos_z=wp.from_torch(sorted_pos_z, dtype=wp_dtype, return_ctype=True),
num_tiles=wp.from_torch(num_tiles, dtype=wp.int32, return_ctype=True),
tile_row_group=wp.from_torch(tile_row_group, dtype=wp.int32, return_ctype=True),
tile_col_group=wp.from_torch(
tile_col_group,
dtype=wp.int32,
return_ctype=True,
),
cell=wp_cell,
inv_cell=wp_inv_cell,
cutoff=float(cutoff),
natom=int(natom),
max_pairs=int(max_pairs),
pair_counter=wp.from_torch(
pair_counter,
dtype=wp.int32,
return_ctype=True,
),
coo_list=wp.from_torch(coo_list, dtype=wp.int32, return_ctype=True),
coo_shifts=wp.from_torch(coo_shifts, dtype=wp.int32, return_ctype=True),
wp_dtype=wp_dtype,
device=wp_device,
n_tiles=int(n_tiles),
**_wp_fixed_cluster_geometry(
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
),
)
@_query_cluster_tile_coo_op.register_fake
def _(
cutoff: float,
natom: int,
max_pairs: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
pair_counter: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
n_tiles: int,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
return None
@torch.library.custom_op(
"nvalchemiops::_query_cluster_tile_coo_segmented",
mutates_args=("pair_counter", "pair_counts", "coo_list", "coo_shifts"),
)
@scoped_torch_warp_stream
def _query_cluster_tile_coo_segmented_op(
cutoff: float,
natom: int,
max_pairs: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
pair_counter: torch.Tensor,
pair_offsets: torch.Tensor,
pair_counts: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
n_tiles: int,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
"""Run fixed-capacity selective COO conversion in place."""
device = sorted_pos_x.device
wp_device = str(device)
wp_dtype = get_wp_dtype(sorted_pos_x.dtype)
wp_query_cluster_tile_coo(
sorted_atom_index=wp.from_torch(
sorted_atom_index, dtype=wp.int32, return_ctype=True
),
sorted_pos_x=wp.from_torch(sorted_pos_x, dtype=wp_dtype, return_ctype=True),
sorted_pos_y=wp.from_torch(sorted_pos_y, dtype=wp_dtype, return_ctype=True),
sorted_pos_z=wp.from_torch(sorted_pos_z, dtype=wp_dtype, return_ctype=True),
num_tiles=wp.from_torch(num_tiles, dtype=wp.int32, return_ctype=True),
tile_row_group=wp.from_torch(tile_row_group, dtype=wp.int32, return_ctype=True),
tile_col_group=wp.from_torch(tile_col_group, dtype=wp.int32, return_ctype=True),
cell=_mat33f_from_torch(cell),
inv_cell=_mat33f_from_torch(inv_cell),
cutoff=float(cutoff),
natom=int(natom),
max_pairs=int(max_pairs),
pair_counter=wp.from_torch(pair_counter, dtype=wp.int32, return_ctype=True),
coo_list=wp.from_torch(coo_list, dtype=wp.int32, return_ctype=True),
coo_shifts=wp.from_torch(coo_shifts, dtype=wp.int32, return_ctype=True),
wp_dtype=wp_dtype,
device=wp_device,
n_tiles=int(n_tiles),
rebuild_flags=wp.from_torch(rebuild_flags, dtype=wp.bool, return_ctype=True),
pair_offsets=wp.from_torch(pair_offsets, dtype=wp.int32, return_ctype=True),
pair_counts=wp.from_torch(pair_counts, dtype=wp.int32, return_ctype=True),
**_wp_fixed_cluster_geometry(
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
),
)
@_query_cluster_tile_coo_segmented_op.register_fake
def _(
cutoff: float,
natom: int,
max_pairs: int,
cell: torch.Tensor,
inv_cell: torch.Tensor,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
rebuild_flags: torch.Tensor,
pair_counter: torch.Tensor,
pair_offsets: torch.Tensor,
pair_counts: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
n_tiles: int,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
return None
@torch.library.custom_op(
"nvalchemiops::_query_cluster_tile_coo_optional_no_pair_fn",
mutates_args=(
"pair_counter",
"coo_list",
"coo_shifts",
"neighbor_vectors",
"neighbor_distances",
),
)
@scoped_torch_warp_stream
def _query_cluster_tile_coo_optional_no_pair_fn_op(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
max_pairs: int,
cutoff: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
pair_counter: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
n_tiles: int,
return_vectors: bool,
return_distances: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
_query_cluster_tile_coo_optional(
cell_mat,
inv_cell_mat,
natom,
max_pairs,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
pair_counter,
coo_list,
coo_shifts,
n_tiles=n_tiles,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=None,
pair_params=None,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=None,
pair_forces=None,
_fixed_cell_geometry=_fixed_cluster_geometry_record(
inv_cell_mat,
fixed_cell,
qr,
axis_aligned,
fractional_rounding_certified,
qr_height_certified,
bbox_cutoff_bounds,
),
)
@_query_cluster_tile_coo_optional_no_pair_fn_op.register_fake
def _(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
max_pairs: int,
cutoff: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
pair_counter: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
n_tiles: int,
return_vectors: bool,
return_distances: bool,
fixed_cell: bool,
qr: torch.Tensor | None,
axis_aligned: torch.Tensor | None,
fractional_rounding_certified: torch.Tensor | None,
qr_height_certified: torch.Tensor | None,
bbox_cutoff_bounds: torch.Tensor | None,
) -> None:
return None
@scoped_torch_warp_stream
def _query_cluster_tile_coo_optional(
cell_mat: torch.Tensor,
inv_cell_mat: torch.Tensor,
natom: int,
max_pairs: int,
cutoff: float,
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
pair_counter: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
*,
n_tiles: int,
return_vectors: bool,
return_distances: bool,
pair_fn: wp.Function | None,
pair_params: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
pair_energies: torch.Tensor | None,
pair_forces: torch.Tensor | None,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> None:
"""Pair-output COO path: bypass the torch custom op.
``n_tiles`` (host-synced emitted-tile count) sets the launch dimension;
the kernel still guards per-tile defensively.
"""
device = sorted_pos_x.device
wp_device = str(device)
wp_dtype = get_wp_dtype(sorted_pos_x.dtype)
wp_vec_dtype = wp.vec3f
wp_pair_params = (
wp.from_torch(pair_params, dtype=wp_dtype, requires_grad=False)
if pair_params is not None
else None
)
wp_neighbor_vectors = (
wp.from_torch(neighbor_vectors, dtype=wp_vec_dtype, requires_grad=False)
if neighbor_vectors is not None
else None
)
wp_neighbor_distances = (
wp.from_torch(neighbor_distances, dtype=wp_dtype, requires_grad=False)
if neighbor_distances is not None
else None
)
wp_pair_energies = (
wp.from_torch(pair_energies, dtype=wp_dtype, requires_grad=False)
if pair_energies is not None
else None
)
wp_pair_forces = (
wp.from_torch(pair_forces, dtype=wp_vec_dtype, requires_grad=False)
if pair_forces is not None
else None
)
wp_query_cluster_tile_coo(
sorted_atom_index=wp.from_torch(
sorted_atom_index, dtype=wp.int32, requires_grad=False
),
sorted_pos_x=wp.from_torch(sorted_pos_x, dtype=wp_dtype, requires_grad=False),
sorted_pos_y=wp.from_torch(sorted_pos_y, dtype=wp_dtype, requires_grad=False),
sorted_pos_z=wp.from_torch(sorted_pos_z, dtype=wp_dtype, requires_grad=False),
num_tiles=wp.from_torch(num_tiles, dtype=wp.int32, requires_grad=False),
tile_row_group=wp.from_torch(
tile_row_group, dtype=wp.int32, requires_grad=False
),
tile_col_group=wp.from_torch(
tile_col_group, dtype=wp.int32, requires_grad=False
),
cell=_mat33f_from_torch(cell_mat),
inv_cell=_mat33f_from_torch(inv_cell_mat),
cutoff=float(cutoff),
natom=int(natom),
max_pairs=int(max_pairs),
pair_counter=wp.from_torch(pair_counter, dtype=wp.int32, requires_grad=False),
coo_list=wp.from_torch(coo_list, dtype=wp.int32, requires_grad=False),
coo_shifts=wp.from_torch(coo_shifts, dtype=wp.int32, requires_grad=False),
wp_dtype=wp_dtype,
device=wp_device,
n_tiles=int(n_tiles),
return_vectors=bool(return_vectors),
return_distances=bool(return_distances),
pair_fn=pair_fn,
pair_params=wp_pair_params,
neighbor_vectors=wp_neighbor_vectors,
neighbor_distances=wp_neighbor_distances,
pair_energies=wp_pair_energies,
pair_forces=wp_pair_forces,
**_wp_fixed_cluster_geometry(*_fixed_cluster_op_args(_fixed_cell_geometry)),
)
[docs]
def query_cluster_tile_coo(
sorted_atom_index: torch.Tensor,
sorted_pos_x: torch.Tensor,
sorted_pos_y: torch.Tensor,
sorted_pos_z: torch.Tensor,
num_tiles: torch.Tensor,
tile_row_group: torch.Tensor,
tile_col_group: torch.Tensor,
cell: torch.Tensor,
cutoff: float,
natom: int,
max_pairs: int,
pair_counter: torch.Tensor,
coo_list: torch.Tensor,
coo_shifts: torch.Tensor,
*,
rebuild_flags: torch.Tensor | None = None,
pair_offsets: torch.Tensor | None = None,
pair_counts: torch.Tensor | None = None,
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
) -> None:
"""Convert the tile pair list to flat COO format in place.
Cluster-tile does not support partial neighbor lists; there is no
``target_indices`` kwarg. Optional pair outputs use flat COO buffers of
length ``max_pairs`` and are written in the same order as ``coo_list``.
Selective segmented COO is topology-only and cannot combine with pair
outputs.
Parameters
----------
sorted_atom_index : torch.Tensor, shape (n_padded,), dtype=int32
Permutation mapping sorted rank to original atom index; output of
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list`.
sorted_pos_x : torch.Tensor, shape (n_padded,), dtype=float32
x-coordinates in Morton-sorted order.
sorted_pos_y : torch.Tensor, shape (n_padded,), dtype=float32
y-coordinates in Morton-sorted order.
sorted_pos_z : torch.Tensor, shape (n_padded,), dtype=float32
z-coordinates in Morton-sorted order.
num_tiles : torch.Tensor, shape (1,), dtype=int32
Device-side tile counter written by
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list`.
tile_row_group : torch.Tensor, shape (max_tiles,), dtype=int32
Row group indices of emitted tile pairs.
tile_col_group : torch.Tensor, shape (max_tiles,), dtype=int32
Column group indices of emitted tile pairs.
cell : torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32
Simulation cell matrix (orthorhombic or triclinic).
cutoff : float
Neighbor search cutoff radius in Cartesian units.
natom : int
True atom count (before padding).
max_pairs : int
Allocated capacity of the COO output buffers.
pair_counter : torch.Tensor, shape (1,), dtype=int32
Atomic counter that accumulates the number of written pairs.
Reset to zero internally before use. Modified in-place.
coo_list : torch.Tensor, shape (max_pairs, 2), dtype=int32
Output flat pair list; each row is ``(i, j)``. Modified in-place.
coo_shifts : torch.Tensor, shape (max_pairs, 3), dtype=int32
Output periodic image shift vectors for each pair. Modified in-place.
rebuild_flags : torch.Tensor, shape (1,), dtype=bool, optional
Enables selective fixed-capacity COO. Requires ``pair_offsets`` and
``pair_counts``. A false flag preserves the supplied topology buffers;
a true flag rebuilds their active prefix.
pair_offsets : torch.Tensor, shape (2,), dtype=int32, optional
Fixed COO segment boundaries ``[0, max_pairs]``. Must be supplied with
``rebuild_flags`` and ``pair_counts``.
pair_counts : torch.Tensor, shape (1,), dtype=int32, optional
Number of active pairs in the segment. Modified in-place.
return_vectors : bool, optional
Write per-pair Cartesian displacement vectors to ``neighbor_vectors``.
Default is ``False``.
return_distances : bool, optional
Write per-pair scalar distances to ``neighbor_distances``.
Default is ``False``.
pair_fn : wp.Function, optional
Module-scope Warp ``@wp.func`` of signature
``(r_ij, distance, pair_params, i, j) -> (energy, force)``.
pair_params : torch.Tensor, shape (natom, num_parameters), optional
Per-atom pair-function parameters; required with ``pair_fn``.
neighbor_vectors : torch.Tensor, shape (max_pairs, 3), optional
Output buffer for per-pair displacement vectors. Modified in-place.
neighbor_distances : torch.Tensor, shape (max_pairs,), optional
Output buffer for per-pair scalar distances. Modified in-place.
pair_energies : torch.Tensor, shape (max_pairs,), optional
Output buffer for per-pair energies; required with ``pair_fn``.
Modified in-place.
pair_forces : torch.Tensor, shape (max_pairs, 3), optional
Output buffer for per-pair forces; required with ``pair_fn``.
Modified in-place.
See Also
--------
:func:`nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list` :
Produces the tile state consumed by this function.
:func:`nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile` :
Row-padded matrix-format alternative.
"""
segmented_inputs = (rebuild_flags, pair_offsets, pair_counts)
segmented = rebuild_flags is not None
if any(value is not None for value in segmented_inputs) and not all(
value is not None for value in segmented_inputs
):
raise ValueError(
"rebuild_flags, pair_offsets, and pair_counts must be supplied together"
)
if segmented:
if _has_partial_or_pair_outputs(
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
):
raise ValueError(
"cluster_tile selective rebuild cannot be combined with pair outputs"
)
if coo_list.shape != (max_pairs, 2):
raise ValueError("coo_list must have shape (max_pairs, 2)")
physical_capacity = _validate_segmented_coo_state(
device=coo_list.device,
num_systems=1,
neighbor_list=coo_list.transpose(0, 1),
neighbor_list_shifts=coo_shifts,
pair_offsets=pair_offsets,
pair_counts=pair_counts,
rebuild_flags=rebuild_flags,
)
if physical_capacity != max_pairs:
raise ValueError("max_pairs must equal coo_list.shape[0]")
cell_mat, inv_cell_mat = _cell_invcell_from_cell(cell)
cell_mat = cell_mat.to(sorted_pos_x.dtype)
inv_cell_mat = inv_cell_mat.to(sorted_pos_x.dtype)
fixed_geometry_args = _fixed_cluster_op_args(_fixed_cell_geometry)
# Eager execution tightens the launch to the emitted tile count. Compiled
# execution validates on device and launches the full static capacity.
tile_capacity = int(tile_row_group.shape[0])
n_tiles = _check_tile_buffer_capacity(num_tiles, tile_capacity)
pair_counter.zero_()
if segmented:
_query_cluster_tile_coo_segmented_op(
cutoff,
natom,
max_pairs,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
rebuild_flags,
pair_counter,
pair_offsets,
pair_counts,
coo_list,
coo_shifts,
n_tiles,
*fixed_geometry_args,
)
if torch.compiler.is_compiling():
_normalize_compiled_single_segment_coo_count(
pair_offsets=pair_offsets,
pair_counts=pair_counts,
rebuild_flags=rebuild_flags,
physical_capacity=physical_capacity,
)
return
if _has_partial_or_pair_outputs(
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
):
if (
pair_fn is None
and pair_params is None
and pair_energies is None
and pair_forces is None
):
return _query_cluster_tile_coo_optional_no_pair_fn_op(
cell_mat,
inv_cell_mat,
natom,
max_pairs,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
pair_counter,
coo_list,
coo_shifts,
neighbor_vectors,
neighbor_distances,
n_tiles,
return_vectors,
return_distances,
*fixed_geometry_args,
)
if torch.compiler.is_compiling():
raise NotImplementedError(
"cluster_tile COO pair_fn outputs are eager-only because callable "
"Warp functions cannot cross a torch.library.custom_op schema boundary.",
)
_query_cluster_tile_coo_optional(
cell_mat,
inv_cell_mat,
natom,
max_pairs,
cutoff,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
pair_counter,
coo_list,
coo_shifts,
n_tiles=n_tiles,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
_fixed_cell_geometry=_fixed_cell_geometry,
)
return
_query_cluster_tile_coo_op(
cutoff,
natom,
max_pairs,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
pair_counter,
coo_list,
coo_shifts,
n_tiles,
*fixed_geometry_args,
)
@scoped_torch_warp_stream
def _cluster_tile_neighbor_list_impl(
positions: torch.Tensor,
cutoff: float | None = None,
cell: torch.Tensor | None = None,
max_neighbors: int | None = None,
fill_value: int | None = None,
format: str = "matrix",
max_pairs: int | None = None,
cutoff2: float | None = None,
rebuild_flags: torch.Tensor | None = None,
return_state: bool = False,
# matrix-format outputs
neighbor_matrix: torch.Tensor | None = None,
neighbor_matrix_shifts: torch.Tensor | None = None,
num_neighbors: torch.Tensor | None = None,
neighbor_matrix2: torch.Tensor | None = None,
neighbor_matrix_shifts2: torch.Tensor | None = None,
num_neighbors2: torch.Tensor | None = None,
# coo-format outputs
neighbor_list: torch.Tensor | None = None,
neighbor_list_shifts: torch.Tensor | None = None,
pair_offsets: torch.Tensor | None = None,
pair_counts: torch.Tensor | None = None,
pair_counter: torch.Tensor | None = None,
# scratch buffers
sorted_atom_index: torch.Tensor | None = None,
morton_codes: torch.Tensor | None = None,
sorted_pos_x: torch.Tensor | None = None,
sorted_pos_y: torch.Tensor | None = None,
sorted_pos_z: torch.Tensor | None = None,
group_ctr_x: torch.Tensor | None = None,
group_ctr_y: torch.Tensor | None = None,
group_ctr_z: torch.Tensor | None = None,
group_ext_x: torch.Tensor | None = None,
group_ext_y: torch.Tensor | None = None,
group_ext_z: torch.Tensor | None = None,
num_tiles: torch.Tensor | None = None,
tile_row_group: torch.Tensor | None = None,
tile_col_group: torch.Tensor | None = None,
# Optional matrix outputs / pair-fn surface
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
*,
max_tiles_per_group: int | None = None,
eager_rebuild_count: int | None = None,
inv_cell: torch.Tensor | None = None,
force_compute_inv_cell: bool = False,
_fixed_cell_geometry: _FixedCellGeometry | None = None,
_live_cell: torch.Tensor | None = None,
) -> tuple[torch.Tensor, ...]:
"""Build and query a cluster-pair tile neighbor list in one call.
Single-system PyTorch binding for the cluster-pair tile algorithm.
Runs Morton sort, Warp bounding-box reduction, and tile enumeration,
then emits the result in one of three formats selected by ``format=``.
Supports orthorhombic and triclinic cells alike via
``_wrap_triclinic``. Cluster-tile is CUDA float32 only.
Parameters
----------
positions : torch.Tensor, shape (N, 3), dtype=float32
Atomic coordinates. Any ``N >= 0``; non-32-aligned ``N`` is
supported via internal padding to
``ceil(N / TILE_GROUP_SIZE) * TILE_GROUP_SIZE``. Padding slots
use sentinel Morton codes and are filtered out by the
convert/coo kernels.
cutoff : float, optional
Cutoff distance in Cartesian units. Must be positive.
cutoff2 : float, optional
Cutoff for the second matrix. It is normally the outer cutoff and may
equal ``cutoff``. Either order is accepted because the tile buffer is
sized for the larger value. Cannot be combined with pair outputs or
COO/tile formats.
cell : torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32
Any non-degenerate cell (orthorhombic or triclinic). Required on every
call.
max_neighbors : int, optional
Falls back to ``estimate_max_neighbors`` using the larger active cutoff.
Matrix format only.
fill_value : int, optional
Matrix sentinel; defaults to ``N``.
format : {"matrix", "coo", "tile"}, default "matrix"
Output representation:
- ``"matrix"``: returns
``(neighbor_matrix, num_neighbors, neighbor_matrix_shifts)`` —
the dense ``(N, max_neighbors)`` row-padded form used by
``cell_list`` and ``naive``.
- ``"coo"``: compact calls return
``(neighbor_list, neighbor_ptr, neighbor_list_shifts)`` and append
exact distances, vectors, or both when requested. The geometry order
is distances followed by vectors. The trimmed flat pair list is
placed directly into source-owned CSR rows; pair order within a row
is unspecified. With ``rebuild_flags``,
fixed-capacity segmented COO returns
``(neighbor_list, pair_offsets, pair_counts,
neighbor_list_shifts)`` instead.
- ``"tile"``: returns the native cluster-pair tile state as a
7-tuple
``(num_tiles, tile_row_group, tile_col_group,
sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z)``.
No convert kernel is run. Intended for downstream kernels
that consume the tile-pair list directly with shared-memory
tile loads. Tile pairs are group-level half-fill: every
emitted pair has ``tile_col_group[t] >= tile_row_group[t]``.
The consumer chooses atom-level fill.
max_pairs : int, optional
Upper bound for COO output; defaults to ``N * max_neighbors``.
max_tiles_per_group : int, optional
Capacity factor for an internally allocated intermediate tile-pair
buffer. For ``g`` row groups, the buffer holds
``g * min(g, max_tiles_per_group)`` tile pairs. Increasing the value up
to ``g`` uses more memory and accommodates more candidate tile pairs.
Eager calls estimate the value when it is ``None``. Caller-owned tile
arrays determine the actual capacity. See
:ref:`cluster-tile-buffer-capacity` for sizing details.
rebuild_flags : torch.Tensor, shape (1,), dtype=bool, optional
Selective rebuild flag. An eager all-true call may bootstrap omitted
state. When false, caller-owned fixed topology and tile state are
required before launch. COO calls additionally require
``pair_offsets`` and ``pair_counts``. The explicit single-system COO
route supports ``torch.compile(fullgraph=True)`` only after eager
bootstrap with fixed-shape validated buffers.
return_state : bool, default=False
With ``rebuild_flags``, append ``(num_tiles, tile_row_group,
tile_col_group)`` after topology outputs. These are the same tensor
objects as caller-supplied state buffers.
neighbor_matrix, num_neighbors, neighbor_matrix_shifts : optional
Pre-allocated matrix-format outputs. All-or-nothing only across
the trio; supply all three or none.
neighbor_list, neighbor_list_shifts, pair_counter : optional
Pre-allocated compact COO outputs with shapes ``(2, max_pairs)``,
``(max_pairs, 3)``, and ``(1,)``, dtype int32 and on the same device
as ``positions``. Each buffer may be supplied independently; omitted
buffers are allocated.
pair_offsets, pair_counts : optional
Fixed selective-COO segment metadata with shapes ``(2,)`` and
``(1,)`` int32. Requires ``rebuild_flags``, caller-owned
``neighbor_list``, and ``neighbor_list_shifts``.
sorted_atom_index, morton_codes, sorted_pos_x, sorted_pos_y, sorted_pos_z, group_ctr_x, group_ctr_y, group_ctr_z, group_ext_x, group_ext_y, group_ext_z, num_tiles, tile_row_group, tile_col_group : torch.Tensor, optional
Pre-allocated scratch buffers (shapes as returned by
``allocate_cluster_tile_list``). All-or-nothing: either
provide every scratch buffer or none. The trigger is
``sorted_atom_index``. Reuse is safe — ``num_tiles`` is reset
each call and every other scratch tensor is either fully
overwritten or only read in regions the kernels just wrote.
return_vectors, return_distances : bool, default ``False``
Return per-pair Cartesian displacements / scalar distances and write
their active prefixes to ``neighbor_vectors`` /
``neighbor_distances`` when those buffers are supplied.
Matrix format uses ``(N, max_neighbors, ...)`` buffers; COO format
uses flat ``(max_pairs, ...)`` buffers.
pair_fn : callable, optional
Module-scope Warp ``@wp.func`` of signature
``(r_ij, distance, pair_params, i, j) -> (energy, force)``. Writes
per-pair energies / forces to ``pair_energies`` /
``pair_forces``. Matrix format uses row-padded buffers; COO
format uses flat buffers written in pair-list order.
pair_params : torch.Tensor, shape ``(num_atoms, num_parameters)``, optional
Per-atom pair-function parameters; required with ``pair_fn``.
neighbor_vectors, neighbor_distances : torch.Tensor, optional
OUTPUT buffers for per-pair displacements / distances. Without
autograd reconstruction, matrix and compact COO formats allocate them
when omitted. Compact COO accepts flat capacity buffers, leaves their
inactive tails unspecified, and returns newly allocated exact geometry
that does not alias them. Geometry buffers must not require gradients.
When geometry is reconstructed for autograd, supplied buffers receive
detached value snapshots while the returned geometry uses separate
differentiable tensors. Omitted geometry buffers are not allocated.
pair_energies, pair_forces : torch.Tensor, optional
OUTPUT buffers for per-pair energies / forces. Matrix format
allocates them when omitted; COO format requires caller-owned flat
buffers.
Returns
-------
tuple of torch.Tensor
Shape depends on ``format``:
- ``"matrix"`` (default): ``(neighbor_matrix, num_neighbors,
neighbor_matrix_shifts)``, with optional ``(*, distances)`` and/or
``(*, vectors)`` appended when ``return_distances`` /
``return_vectors`` is True, and optional ``(*, pair_energies,
pair_forces)`` when ``pair_fn`` is set. With ``cutoff2``, returns
the primary group followed by the secondary cutoff group. Selective
calls with ``return_state=True`` append ``(num_tiles,
tile_row_group, tile_col_group)``, yielding six tensors for one
cutoff or nine for two cutoffs. Differentiable reconstructed
geometry does not alias supplied output buffers; otherwise returned
geometry is the supplied or internally allocated buffer.
- ``"coo"``: compact calls return ``(neighbor_list, neighbor_ptr,
neighbor_list_shifts)`` and append ``distances`` and/or ``vectors``
in that order when requested without a pair callback. Exact outputs
do not alias reusable capacity buffers. Pair callbacks keep the
topology-only return and write their buffers in place. Selective
calls return fixed-capacity
``(neighbor_list, pair_offsets, pair_counts,
neighbor_list_shifts)`` and append tile state when
``return_state=True``.
- ``"tile"``: ``(num_tiles, tile_row_group, tile_col_group,
sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z)``.
Notes
-----
- Cluster-tile is CUDA float32 only; float64 ``positions`` is rejected.
- Cluster-tile does not support partial neighbor lists (no
``target_indices`` kwarg).
- ``torch.compile(fullgraph=True)`` supports tile and matrix output,
including dual-cutoff matrices and differentiable geometry. On PyTorch
2.10 or newer, nonselective exact COO output also supports fullgraph and
may return a different pair count on each call. A compiled call that
allocates scratch internally requires a positive static
``max_tiles_per_group``; complete caller-owned scratch may be supplied
instead. Compiled capacity failures use asynchronous device assertions.
Pair callbacks remain eager-only.
- Build differentiable losses from returned geometry, not from supplied
output buffers, which are non-differentiable value snapshots when
reconstruction is required.
- The unified
:func:`nvalchemiops.torch.neighbors.neighbor_list` entry point may
select this binding automatically when the selector guards and cost
model prefer it; pass ``method="cluster_tile"`` to force it.
See Also
--------
nvalchemiops.torch.neighbors.batch_cluster_tile_neighbor_list :
Batched companion entry point.
nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list :
Lower-level build step exposed for caching across queries.
nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile :
Lower-level query step.
"""
if cutoff is None or cell is None:
missing = [
name
for name, value in (("cutoff", cutoff), ("cell", cell))
if value is None
]
raise ValueError("missing required arguments: " + ", ".join(missing))
if positions.dtype != torch.float32:
raise TypeError("positions must be float32")
if format not in ("matrix", "coo", "tile"):
raise ValueError(
f"format must be 'matrix' | 'coo' | 'tile'; got {format!r}",
)
if max_tiles_per_group is not None and (
not isinstance(max_tiles_per_group, int)
or isinstance(max_tiles_per_group, bool)
or max_tiles_per_group <= 0
):
raise ValueError("max_tiles_per_group must be a positive integer")
has_pair_outputs = (
bool(return_vectors)
or bool(return_distances)
or pair_fn is not None
or pair_params is not None
or neighbor_vectors is not None
or neighbor_distances is not None
or pair_energies is not None
or pair_forces is not None
)
dual_cutoff = cutoff2 is not None
selective = rebuild_flags is not None
is_compiling = torch.compiler.is_compiling()
if is_compiling and format == "coo" and not selective and pair_fn is not None:
raise NotImplementedError(
"torch.compile(fullgraph=True) cluster-tile pair_fn is not supported"
)
if selective and not is_compiling:
if eager_rebuild_count is None:
eager_rebuild_count = int(torch.count_nonzero(rebuild_flags).item())
eager_rebuild_requested = eager_rebuild_count > 0
else:
eager_rebuild_requested = False
if return_state and not selective:
raise ValueError("return_state=True requires rebuild_flags")
if not selective and (pair_offsets is not None or pair_counts is not None):
raise ValueError("pair_offsets and pair_counts require rebuild_flags")
if has_pair_outputs and format == "tile":
raise NotImplementedError(
"Pair outputs (return_vectors / return_distances / pair_fn) "
"are not supported with format='tile'. Use format='matrix' "
"or format='coo'.",
)
if dual_cutoff:
if format != "matrix":
raise ValueError(
"cluster_tile cutoff2 is supported only with format='matrix'"
)
if has_pair_outputs:
raise ValueError(
"cluster_tile cutoff2 cannot be combined with pair outputs"
)
if selective:
if format not in ("matrix", "coo"):
raise ValueError(
"cluster_tile selective rebuild is supported only with "
"format='matrix' or format='coo'"
)
if has_pair_outputs:
raise ValueError(
"cluster_tile selective rebuild cannot be combined with pair outputs"
)
required = {
"num_tiles": num_tiles,
"tile_row_group": tile_row_group,
"tile_col_group": tile_col_group,
}
if format == "matrix":
required.update(
{
"neighbor_matrix": neighbor_matrix,
"num_neighbors": num_neighbors,
"neighbor_matrix_shifts": neighbor_matrix_shifts,
}
)
else:
required.update(
{
"neighbor_list": neighbor_list,
"pair_offsets": pair_offsets,
"pair_counts": pair_counts,
"neighbor_list_shifts": neighbor_list_shifts,
}
)
if dual_cutoff:
required.update(
{
"neighbor_matrix2": neighbor_matrix2,
"num_neighbors2": num_neighbors2,
"neighbor_matrix_shifts2": neighbor_matrix_shifts2,
}
)
missing = [name for name, value in required.items() if value is None]
bootstrap_coo = (
format == "coo" and eager_rebuild_requested and bool(return_state)
)
if missing and not bootstrap_coo:
raise ValueError(
"rebuild_flags requires previous cluster_tile state: "
+ ", ".join(missing)
)
N = positions.shape[0]
device = positions.device
if max_neighbors is None:
max_neighbors = (
int(neighbor_matrix.shape[1])
if format == "matrix" and neighbor_matrix is not None
else max(
estimate_max_neighbors(
cutoff if cutoff2 is None else max(float(cutoff), float(cutoff2))
),
TILE_GROUP_SIZE,
)
)
elif (
format == "matrix"
and neighbor_matrix is not None
and (neighbor_matrix.shape != (N, int(max_neighbors)))
):
raise ValueError("neighbor_matrix must have shape (N, max_neighbors)")
if format == "matrix":
_validate_cluster_tile_matrix_outputs(
device=device,
dtype=positions.dtype,
natom=N,
max_neighbors=int(max_neighbors),
cutoff2=cutoff2,
neighbor_matrix2=neighbor_matrix2,
num_neighbors2=num_neighbors2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
return_vectors=return_vectors,
return_distances=return_distances,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
allocate_missing=True,
)
if fill_value is None:
fill_value = N
if selective and format == "matrix":
matrix_state = {
"neighbor_matrix": neighbor_matrix,
"num_neighbors": num_neighbors,
"neighbor_matrix_shifts": neighbor_matrix_shifts,
"num_tiles": num_tiles,
"tile_row_group": tile_row_group,
"tile_col_group": tile_col_group,
}
for name, value in matrix_state.items():
if value is None or value.device != device or value.dtype != torch.int32:
raise ValueError(f"{name} must be an int32 tensor on positions.device")
if (
num_neighbors.shape != (N,)
or neighbor_matrix_shifts.shape != (N, int(max_neighbors), 3)
or num_tiles.shape != (1,)
or tile_row_group.ndim != 1
or tile_col_group.shape != tile_row_group.shape
):
raise ValueError("selective matrix state has an invalid shape")
if dual_cutoff and (
neighbor_matrix2 is None
or num_neighbors2 is None
or neighbor_matrix_shifts2 is None
or neighbor_matrix2.dtype != torch.int32
or num_neighbors2.dtype != torch.int32
or neighbor_matrix_shifts2.dtype != torch.int32
or neighbor_matrix2.device != device
or num_neighbors2.device != device
or neighbor_matrix_shifts2.device != device
or neighbor_matrix2.shape != (N, int(max_neighbors))
or num_neighbors2.shape != (N,)
or neighbor_matrix_shifts2.shape != (N, int(max_neighbors), 3)
):
raise ValueError("selective secondary matrix state has an invalid shape")
if selective and format == "coo" and eager_rebuild_requested:
coo_state = (
neighbor_list,
pair_offsets,
pair_counts,
neighbor_list_shifts,
num_tiles,
tile_row_group,
tile_col_group,
)
if any(value is not None for value in coo_state) and any(
value is None for value in coo_state
):
raise ValueError(
"selective COO bootstrap requires complete caller-owned state "
"or no persistent state"
)
if neighbor_list is None:
capacity = max_pairs if max_pairs is not None else N * max_neighbors
neighbor_list = torch.empty(
(2, capacity), dtype=torch.int32, device=positions.device
)
neighbor_list_shifts = torch.empty(
(capacity, 3), dtype=torch.int32, device=positions.device
)
pair_offsets = torch.tensor(
[0, capacity], dtype=torch.int32, device=positions.device
)
pair_counts = torch.zeros(1, dtype=torch.int32, device=positions.device)
if selective and format == "coo":
max_pairs_from_buffers = _validate_segmented_coo_state(
device=positions.device,
num_systems=1,
neighbor_list=neighbor_list,
neighbor_list_shifts=neighbor_list_shifts,
pair_offsets=pair_offsets,
pair_counts=pair_counts,
rebuild_flags=rebuild_flags,
)
if max_pairs is not None and int(max_pairs) != max_pairs_from_buffers:
raise ValueError("max_pairs must equal neighbor_list.shape[1]")
if selective and not is_compiling and not eager_rebuild_requested:
if format == "coo":
outputs = (
neighbor_list,
pair_offsets,
pair_counts,
neighbor_list_shifts,
)
elif dual_cutoff:
outputs = (
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
neighbor_matrix2,
num_neighbors2,
neighbor_matrix_shifts2,
)
else:
outputs = (neighbor_matrix, num_neighbors, neighbor_matrix_shifts)
if return_state:
return (*outputs, num_tiles, tile_row_group, tile_col_group)
return outputs
geometry_requested = (
(return_vectors or return_distances)
and pair_fn is None
and pair_params is None
and pair_energies is None
and pair_forces is None
)
requires_reconstruction = (
geometry_requested
and torch.is_grad_enabled()
and (
positions.requires_grad
or (_live_cell if _live_cell is not None else cell).requires_grad
)
)
snapshot_vectors = neighbor_vectors is not None
snapshot_distances = neighbor_distances is not None
if format == "coo" and not selective:
if max_pairs is None:
max_pairs = N * max_neighbors
_validate_compact_coo_outputs(
device=device,
capacity=int(max_pairs),
neighbor_list=neighbor_list,
neighbor_list_shifts=neighbor_list_shifts,
pair_counter=pair_counter,
)
if geometry_requested:
neighbor_vectors, neighbor_distances = (
_prepare_compact_coo_geometry_buffers(
device=device,
dtype=positions.dtype,
capacity=int(max_pairs),
return_vectors=bool(return_vectors),
return_distances=bool(return_distances),
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
allocate_missing=not requires_reconstruction,
)
)
# Candidate tiles must cover both radii. The query then filters each matrix
# with its own cutoff.
build_cutoff = cutoff if cutoff2 is None else max(float(cutoff), float(cutoff2))
# Allocate scratch if caller didn't supply. ``sorted_atom_index`` is the
# all-or-nothing sentinel.
if sorted_atom_index is None:
previous_tile_state = (num_tiles, tile_row_group, tile_col_group)
if max_tiles_per_group is not None:
max_tiles_per_group = int(max_tiles_per_group)
elif is_compiling and not selective:
raise ValueError(
"compiled cluster_tile_neighbor_list requires max_tiles_per_group "
"when scratch buffers are not provided"
)
elif selective and torch.compiler.is_compiling():
if tile_row_group is None:
raise ValueError(
"compiled selective cluster_tile_neighbor_list requires "
"tile_row_group"
)
groups = max(1, (N + TILE_GROUP_SIZE - 1) // TILE_GROUP_SIZE)
max_tiles_per_group = max(1, int(tile_row_group.shape[0]) // groups)
else:
cell_volume = float(_cell_volume(cell))
max_tiles_per_group = estimate_max_tiles_per_group(
N,
build_cutoff,
cell_volume,
)
(
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
) = allocate_cluster_tile_list(
N,
device,
dtype=positions.dtype,
max_tiles_per_group=max_tiles_per_group,
)
if selective:
previous_num_tiles, previous_row, previous_col = previous_tile_state
if previous_num_tiles is not None:
num_tiles = previous_num_tiles
if previous_row is not None:
tile_row_group = previous_row
if previous_col is not None:
tile_col_group = previous_col
cell_mat, inv_cell_mat = _build_cluster_tile_list_normalized(
positions.detach() if requires_reconstruction else positions,
build_cutoff,
cell.detach() if requires_reconstruction else cell,
sorted_atom_index,
morton_codes,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
group_ctr_x,
group_ctr_y,
group_ctr_z,
group_ext_x,
group_ext_y,
group_ext_z,
num_tiles,
tile_row_group,
tile_col_group,
inv_cell=inv_cell,
rebuild_flags=rebuild_flags if selective else None,
force_compute_inv_cell=force_compute_inv_cell,
_fixed_cell_geometry=_fixed_cell_geometry,
)
if format == "tile":
# Raw-tile callers must not receive a silently truncated tile list.
tile_capacity = int(tile_row_group.shape[0])
_check_tile_buffer_capacity(num_tiles, tile_capacity)
return (
num_tiles,
tile_row_group,
tile_col_group,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
)
if format == "coo":
if selective:
max_pairs = int(neighbor_list.shape[1])
elif max_pairs is None:
max_pairs = N * max_neighbors
# Kernels write row-major (max_pairs, 2); transpose to the public
# (2, num_pairs) layout on return. Pre-allocation kwargs accept the
# public (2, max_pairs) layout.
if neighbor_list is None:
coo_buf = torch.empty(
(max_pairs, 2),
dtype=torch.int32,
device=device,
)
else:
# Caller passed (2, max_pairs). Use a transposed view; the
# kernel writes row-major into this buffer.
coo_buf = neighbor_list.transpose(0, 1)
if neighbor_list_shifts is None:
neighbor_list_shifts = torch.empty(
(max_pairs, 3),
dtype=torch.int32,
device=device,
)
if pair_counter is None:
pair_counter = torch.zeros(1, dtype=torch.int32, device=device)
if not selective:
if is_compiling:
version = torch.__version__.split("+", maxsplit=1)[0]
major, minor = (int(part) for part in version.split(".")[:2])
if (major, minor) < (2, 10):
raise RuntimeError(
"torch.compile(fullgraph=True) exact cluster-tile COO requires PyTorch 2.10 or newer"
)
copy_vectors = bool(
geometry_requested and return_vectors and not requires_reconstruction
)
copy_distances = bool(
geometry_requested and return_distances and not requires_reconstruction
)
if not geometry_requested:
prefix_vectors = torch.empty(
(1, 3), dtype=positions.dtype, device=device
)
prefix_distances = torch.empty(1, dtype=positions.dtype, device=device)
if geometry_requested:
prefix_vectors = neighbor_vectors
prefix_distances = neighbor_distances
tile_capacity = int(tile_row_group.shape[0])
n_tiles = _check_tile_buffer_capacity(num_tiles, tile_capacity)
launch_tiles = tile_capacity if is_compiling else n_tiles
row_counts = torch.zeros(N, dtype=torch.int32, device=device)
_cluster_tile_direct_csr_count_op(
cutoff,
N,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
row_counts,
launch_tiles,
*_fixed_cluster_op_args(_fixed_cell_geometry),
)
neighbor_ptr = torch.cat(
(
torch.zeros(1, dtype=torch.int32, device=device),
torch.cumsum(row_counts, dim=0, dtype=torch.int32),
)
)
total = neighbor_ptr[-1:]
resolved_count = _check_neighbor_capacity(
total,
int(max_pairs),
kind="coo",
)
pair_counter.copy_(total)
cursors = neighbor_ptr[:-1].clone()
if pair_fn is None:
_cluster_tile_direct_csr_fill_op(
cutoff,
N,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
cursors,
coo_buf,
neighbor_list_shifts,
prefix_vectors,
prefix_distances,
copy_vectors,
copy_distances,
int(max_pairs),
launch_tiles,
*_fixed_cluster_op_args(_fixed_cell_geometry),
)
else:
if pair_params is None or pair_energies is None or pair_forces is None:
raise ValueError(
"pair_fn requires pair_params, pair_energies, and pair_forces"
)
_cluster_tile_direct_csr_fill_pair_outputs(
cutoff,
N,
cell_mat,
inv_cell_mat,
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
cursors,
coo_buf,
neighbor_list_shifts,
pair_fn,
pair_params,
pair_energies,
pair_forces,
neighbor_vectors,
neighbor_distances,
return_vectors,
return_distances,
int(max_pairs),
launch_tiles,
_fixed_cell_geometry=_fixed_cell_geometry,
)
nl, nls, exact_vectors, exact_distances = _compact_coo_prefix(
total,
coo_buf,
neighbor_list_shifts,
prefix_vectors,
prefix_distances,
int(max_pairs),
-1 if is_compiling else resolved_count,
copy_vectors,
copy_distances,
)
if geometry_requested:
if requires_reconstruction:
exact_distances, exact_vectors = _reconstruct_coo_geometry(
positions,
_live_cell if _live_cell is not None else cell,
nl,
nls,
)
if return_vectors and snapshot_vectors:
neighbor_vectors[: exact_vectors.shape[0]].copy_(
exact_vectors.detach()
)
if return_distances and snapshot_distances:
neighbor_distances[: exact_distances.shape[0]].copy_(
exact_distances.detach()
)
outputs = (nl, neighbor_ptr, nls)
if return_distances:
outputs = (*outputs, exact_distances)
if return_vectors:
outputs = (*outputs, exact_vectors)
return outputs
return nl, neighbor_ptr, nls
query_cluster_tile_coo(
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
cell,
cutoff,
N,
int(max_pairs),
pair_counter,
coo_buf,
neighbor_list_shifts,
rebuild_flags=rebuild_flags,
pair_offsets=pair_offsets,
pair_counts=pair_counts,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
_fixed_cell_geometry=_fixed_cell_geometry,
)
if selective:
if not torch.compiler.is_compiling():
pair_capacity = int(pair_offsets[1].item()) - int(
pair_offsets[0].item()
)
_check_neighbor_capacity(
pair_counts,
pair_capacity,
kind="coo",
)
outputs = (
neighbor_list,
pair_offsets,
pair_counts,
neighbor_list_shifts,
)
if return_state:
return (*outputs, num_tiles, tile_row_group, tile_col_group)
return outputs
# ``format == "matrix"``: skip-prefill matrix path with tail fill.
if neighbor_matrix is None:
neighbor_matrix = torch.empty(
(N, max_neighbors),
dtype=torch.int32,
device=device,
)
if num_neighbors is None:
num_neighbors = torch.zeros(N, dtype=torch.int32, device=device)
elif selective:
_cluster_tile_selective_zero_num_neighbors_single_op(
num_neighbors,
rebuild_flags,
)
else:
num_neighbors.zero_()
if neighbor_matrix_shifts is None:
neighbor_matrix_shifts = torch.empty(
(N, max_neighbors, 3),
dtype=torch.int32,
device=device,
)
if dual_cutoff:
if neighbor_matrix2 is None:
neighbor_matrix2 = torch.empty(
(N, max_neighbors), dtype=torch.int32, device=device
)
if num_neighbors2 is None:
num_neighbors2 = torch.zeros(N, dtype=torch.int32, device=device)
elif selective:
_cluster_tile_selective_zero_num_neighbors_single_op(
num_neighbors2,
rebuild_flags,
)
else:
num_neighbors2.zero_()
if neighbor_matrix_shifts2 is None:
neighbor_matrix_shifts2 = torch.empty(
(N, max_neighbors, 3), dtype=torch.int32, device=device
)
# Pair-output buffer allocation: caller may omit any of the four OUTPUT
# buffer kwargs. Reconstructed geometry needs no internal snapshot buffer;
# it is returned directly and copied only to caller-owned storage.
# Required-presence rules
# (``return_vectors`` ⇒ ``neighbor_vectors``,
# ``pair_fn`` ⇒ ``pair_{energies,forces}_buffer``) are enforced by
# the warp launcher.
if return_vectors and neighbor_vectors is None and not requires_reconstruction:
neighbor_vectors = torch.empty(
(N, max_neighbors, 3),
dtype=positions.dtype,
device=device,
)
if return_distances and neighbor_distances is None and not requires_reconstruction:
neighbor_distances = torch.empty(
(N, max_neighbors),
dtype=positions.dtype,
device=device,
)
if pair_fn is not None:
if pair_energies is None:
pair_energies = torch.empty(
(N, max_neighbors),
dtype=positions.dtype,
device=device,
)
if pair_forces is None:
pair_forces = torch.empty(
(N, max_neighbors, 3),
dtype=positions.dtype,
device=device,
)
_query_cluster_tile_normalized(
sorted_atom_index,
sorted_pos_x,
sorted_pos_y,
sorted_pos_z,
num_tiles,
tile_row_group,
tile_col_group,
cell,
cutoff,
N,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
inv_cell_mat=inv_cell_mat,
cutoff2=cutoff2,
neighbor_matrix2=neighbor_matrix2,
num_neighbors2=num_neighbors2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
rebuild_flags=rebuild_flags,
return_vectors=return_vectors and not requires_reconstruction,
return_distances=return_distances and not requires_reconstruction,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=None if requires_reconstruction else neighbor_vectors,
neighbor_distances=None if requires_reconstruction else neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
_fixed_cell_geometry=_fixed_cell_geometry,
)
# Skip-prefill tail fill: write ``fill_value`` into the unused columns
# of ``neighbor_matrix``. Pairs with the always-write-shifts kernel
# above to eliminate the per-step ``neighbor_matrix.fill_`` and
# ``neighbor_matrix_shifts.zero_`` ops.
if max_neighbors > 0:
if selective:
_cluster_tile_selective_fill_neighbor_matrix_tail_op(
num_neighbors,
rebuild_flags,
neighbor_matrix,
int(max_neighbors),
int(fill_value),
)
if dual_cutoff:
_cluster_tile_selective_fill_neighbor_matrix_tail_op(
num_neighbors2,
rebuild_flags,
neighbor_matrix2,
int(max_neighbors),
int(fill_value),
)
else:
_cluster_tile_fill_neighbor_matrix_tail_op(
num_neighbors,
neighbor_matrix,
int(N),
int(max_neighbors),
int(fill_value),
)
if dual_cutoff:
_cluster_tile_fill_neighbor_matrix_tail_op(
num_neighbors2,
neighbor_matrix2,
int(N),
int(max_neighbors),
int(fill_value),
)
_check_neighbor_capacity(num_neighbors, int(max_neighbors))
if dual_cutoff and num_neighbors2 is not None:
_check_neighbor_capacity(num_neighbors2, int(max_neighbors))
if requires_reconstruction:
distances, vectors = _reconstruct_matrix_geometry(
positions,
_live_cell if _live_cell is not None else cell,
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
)
if return_vectors and snapshot_vectors:
neighbor_vectors.copy_(vectors.detach())
if return_distances and snapshot_distances:
neighbor_distances.copy_(distances.detach())
if dual_cutoff:
outputs = (
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
neighbor_matrix2,
num_neighbors2,
neighbor_matrix_shifts2,
)
else:
outputs = (neighbor_matrix, num_neighbors, neighbor_matrix_shifts)
if geometry_requested:
if return_distances:
outputs = (
*outputs,
distances if requires_reconstruction else neighbor_distances,
)
if return_vectors:
outputs = (
*outputs,
vectors if requires_reconstruction else neighbor_vectors,
)
if return_state:
return (*outputs, num_tiles, tile_row_group, tile_col_group)
return outputs
[docs]
def cluster_tile_neighbor_list(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
max_neighbors: int | None = None,
fill_value: int | None = None,
format: str = "matrix",
max_pairs: int | None = None,
cutoff2: float | None = None,
rebuild_flags: torch.Tensor | None = None,
return_state: bool = False,
neighbor_matrix: torch.Tensor | None = None,
neighbor_matrix_shifts: torch.Tensor | None = None,
num_neighbors: torch.Tensor | None = None,
neighbor_matrix2: torch.Tensor | None = None,
neighbor_matrix_shifts2: torch.Tensor | None = None,
num_neighbors2: torch.Tensor | None = None,
neighbor_list: torch.Tensor | None = None,
neighbor_list_shifts: torch.Tensor | None = None,
pair_offsets: torch.Tensor | None = None,
pair_counts: torch.Tensor | None = None,
pair_counter: torch.Tensor | None = None,
sorted_atom_index: torch.Tensor | None = None,
morton_codes: torch.Tensor | None = None,
sorted_pos_x: torch.Tensor | None = None,
sorted_pos_y: torch.Tensor | None = None,
sorted_pos_z: torch.Tensor | None = None,
group_ctr_x: torch.Tensor | None = None,
group_ctr_y: torch.Tensor | None = None,
group_ctr_z: torch.Tensor | None = None,
group_ext_x: torch.Tensor | None = None,
group_ext_y: torch.Tensor | None = None,
group_ext_z: torch.Tensor | None = None,
num_tiles: torch.Tensor | None = None,
tile_row_group: torch.Tensor | None = None,
tile_col_group: torch.Tensor | None = None,
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
*,
max_tiles_per_group: int | None = None,
) -> tuple[torch.Tensor, ...]:
"""Build and query a cluster-pair tile neighbor list in one call."""
return _cluster_tile_neighbor_list_impl(
positions,
cutoff,
cell,
max_neighbors=max_neighbors,
fill_value=fill_value,
format=format,
max_pairs=max_pairs,
cutoff2=cutoff2,
rebuild_flags=rebuild_flags,
return_state=return_state,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_matrix_shifts,
num_neighbors=num_neighbors,
neighbor_matrix2=neighbor_matrix2,
neighbor_matrix_shifts2=neighbor_matrix_shifts2,
num_neighbors2=num_neighbors2,
neighbor_list=neighbor_list,
neighbor_list_shifts=neighbor_list_shifts,
pair_offsets=pair_offsets,
pair_counts=pair_counts,
pair_counter=pair_counter,
sorted_atom_index=sorted_atom_index,
morton_codes=morton_codes,
sorted_pos_x=sorted_pos_x,
sorted_pos_y=sorted_pos_y,
sorted_pos_z=sorted_pos_z,
group_ctr_x=group_ctr_x,
group_ctr_y=group_ctr_y,
group_ctr_z=group_ctr_z,
group_ext_x=group_ext_x,
group_ext_y=group_ext_y,
group_ext_z=group_ext_z,
num_tiles=num_tiles,
tile_row_group=tile_row_group,
tile_col_group=tile_col_group,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
max_tiles_per_group=max_tiles_per_group,
)
cluster_tile_neighbor_list.__doc__ = _cluster_tile_neighbor_list_impl.__doc__
