PERF: Remove GPU-CPU Synchronization in FIRE

tests/test_math.py

@@ -4,6 +4,7 @@
 # ruff: noqa: SLF001
 
 import numpy as np
+import pytest
 import scipy
 import torch
 from numpy.testing import assert_allclose
@@ -15,6 +16,29 @@
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+@pytest.mark.parametrize(
+    ("n_systems", "expected"),
+    [(None, [208, 212, 420]), (3, [208, 212, 420]), (5, [208, 212, 420, 0, 0])],
+)
+def test_batched_vdot_known_size(n_systems, expected):
+    x = torch.arange(18, dtype=DTYPE, device=device).reshape(6, 3)
+    y = x.flip(0)
+    indices = torch.tensor([0, 0, 1, 2, 2, 2], device=device)
+    result = fm.batched_vdot(x, y, indices, n_systems=n_systems)
+    expected = torch.tensor(expected, dtype=DTYPE, device=device)
+    torch.testing.assert_close(result, expected)
+
+
+def test_batched_vdot_known_size_empty():
+    x = torch.empty((0, 3), dtype=DTYPE, device=device)
+    indices = torch.empty(0, dtype=torch.int64, device=device)
+    expected = torch.zeros(2, dtype=DTYPE, device=device)
+
+    result = fm.batched_vdot(x, x, indices, n_systems=2)
+
+    torch.testing.assert_close(result, expected)
+
+
 class TestExpmFrechet:
     """Test suite for expm_frechet for 3x3 matrices."""
 

tests/test_optimizers.py

@@ -747,7 +747,7 @@ def test_fire_nan_velocities_dont_affect_other_systems(
 
     # Evolve 10 steps so system 0 has non-trivial FIRE state (dt, alpha, n_pos)
     for _ in range(10):
-        state = ts.fire_step(state=state, model=lj_model)
+        state = ts.fire_step(state=state, model=lj_model, fire_flavor=fire_flavor)
 
     # Clone, then inject NaN into system 1 of one copy
     state_clean = copy.deepcopy(state)
@@ -759,8 +759,8 @@ def test_fire_nan_velocities_dont_affect_other_systems(
         state_mixed.cell_velocities[1] = float("nan")
 
     # One step each
-    state_clean = ts.fire_step(state=state_clean, model=lj_model)
-    state_mixed = ts.fire_step(state=state_mixed, model=lj_model)
+    state_clean = ts.fire_step(state=state_clean, model=lj_model, fire_flavor=fire_flavor)
+    state_mixed = ts.fire_step(state=state_mixed, model=lj_model, fire_flavor=fire_flavor)
 
     # System 0 must be identical regardless of system 1's NaN velocities
     sys0 = state_clean.system_idx == 0
@@ -785,6 +785,41 @@ def test_fire_nan_velocities_dont_affect_other_systems(
         )
 
 
+@pytest.mark.parametrize("cell_filter", [ts.CellFilter.unit, ts.CellFilter.frechet])
+def test_vv_fire_nan_cell_velocities_with_finite_atomic_velocities(
+    ar_double_sim_state: SimState,
+    lj_model: ModelInterface,
+    cell_filter: ts.CellFilter,
+) -> None:
+    """Cell NaNs reset that system's cell velocity independently of atomic NaNs."""
+    ar_double_sim_state.cell *= 0.85
+    ar_double_sim_state.positions *= 0.85
+    state = ts.fire_init(
+        ar_double_sim_state, lj_model, fire_flavor="vv_fire", cell_filter=cell_filter
+    )
+    state.velocities.fill_(0.01)
+    state.cell_velocities.fill_(0.02)
+
+    # A single NaN should reset the whole cell velocity for that system.
+    expected = state.clone()
+    expected.cell_velocities[1] = 0
+    state.cell_velocities[1, 0, 0] = float("nan")
+
+    expected = ts.fire_step(state=expected, model=lj_model, fire_flavor="vv_fire")
+    actual = ts.fire_step(state=state, model=lj_model, fire_flavor="vv_fire")
+
+    for attr in (
+        "positions",
+        "velocities",
+        "cell",
+        "cell_velocities",
+        "dt",
+        "alpha",
+        "n_pos",
+    ):
+        torch.testing.assert_close(getattr(actual, attr), getattr(expected, attr))
+
+
 @pytest.mark.parametrize("fire_flavor", get_args(FireFlavor))
 def test_unit_cell_fire_optimization(
     ar_supercell_sim_state: SimState, lj_model: ModelInterface, fire_flavor: FireFlavor

torch_sim/math.py

@@ -746,14 +746,21 @@ def matrix_log_33(
 
 
 def batched_vdot(
-    x: torch.Tensor, y: torch.Tensor, batch_indices: torch.Tensor
+    x: torch.Tensor,
+    y: torch.Tensor,
+    batch_indices: torch.Tensor,
+    *,
+    n_systems: int | None = None,
 ) -> torch.Tensor:
     """Computes batched vdot (sum of element-wise product) for groups of vectors.
 
     Args:
         x: Tensor of shape [N_total_entities, D] (e.g., forces, velocities).
         y: Tensor of shape [N_total_entities, D].
         batch_indices: Tensor of shape [N_total_entities] indicating batch membership.
+        n_systems: Known output size. When provided, batch_indices are assumed to
+            lie in [0, n_systems), avoiding scalar validation and size inference
+            that synchronize CUDA tensors with the CPU.
 
     Returns:
         Tensor: shape [n_systems] where each element is the sum(x_i * y_i)
@@ -769,10 +776,12 @@ def batched_vdot(
     ):
         raise ValueError(f"Invalid input shapes: {x.shape=}, {batch_indices.shape=}")
 
-    if batch_indices.min() < 0:
-        raise ValueError("batch_indices must be non-negative")
+    if n_systems is None:
+        if batch_indices.min() < 0:
+            raise ValueError("batch_indices must be non-negative")
+        n_systems = int(batch_indices.max()) + 1
 
-    output = torch.zeros(int(batch_indices.max()) + 1, dtype=x.dtype, device=x.device)
+    output = torch.zeros(n_systems, dtype=x.dtype, device=x.device)
     output.scatter_add_(dim=0, index=batch_indices, src=(x * y).sum(dim=1))
 
     return output