fix(electrostatics): compile explicit single-system batches

[zubatyuk]

ALCHEMI Toolkit-Ops Pull Request

Description

Fix torch.compile for a logically single Ewald/PME/slab system when callers provide an explicit batch_idx=zeros(N).

Without this change, the backward path treats explicit B=1 input as a segmented batch and reduces atom values with index_add or scatter operations into a tensor whose leading dimension is one. PyTorch Inductor's CPU vector-atomic code generation can reject that size-one destination, even though the equivalent unbatched call (batch_idx=None) compiles because it uses a direct mean.

The same atom-to-system pattern appears in energy cotangents, double backward, cell and virial gradients, charge corrections, direct-k Ewald, and slab correction. The fix centralizes those operations and specializes only the structural num_systems == 1 case:

if num_systems == 1:
    return values.sum(dim=0, keepdim=True)

result = values.new_zeros((num_systems, *values.shape[1:]))
return result.index_add(0, batch_idx, values)

Explicit batch_idx remains intact, so batched custom-op and Warp-kernel dispatch does not change. Multi-system inputs continue to use segmented reductions.

The reciprocal PME autograd chain also materializes the Hermitian-weighted RFFT cotangent before passing it to the opaque convolution backward:

backward_args=lambda g, f: (
    f[0],
    g[0] * 1.0,
    # ...
)

This prevents Inductor from scheduling the interior-frequency weighting after the custom backward has already consumed the cotangent.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

Merge this PR before #132: fix(electrostatics): avoid per-system backward recomputation.

Changes Made

• Add shared atom-to-system sum/mean, system-to-atom broadcast, and mean-adjoint helpers with explicit single-system and empty-input behavior.
• Route real-space, reciprocal-space, PME, slab, direct-k, correction, cell-gradient, virial, first-backward, and double-backward reductions through the shared single-system policy.
• Materialize the reciprocal PME convolution cotangent at the custom-op boundary to preserve RFFT adjoint ordering under Inductor.
• Add CPU/CUDA eager and compiled coverage for explicit B=1 PME, Ewald, and slab energy, position/charge/cell gradients, direct outputs, and supported higher-order derivatives.
• Add exact helper coverage for empty inputs, one-atom input, and heterogeneous B=2 segmentation with unequal atom counts and non-uniform cotangents.
• Replace the quadratic test-only n-fold boxcar convolution oracle with an equivalent linear-time prefix-sum implementation.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Focused CPU/CUDA validation passed 302 tests covering helper contracts, compile paths, first- and second-order derivatives, batch consistency, q(R), and empty inputs. The supplied explicit-B=1 reproducer reports successful compilation for both batch_idx=zeros(N) and batch_idx=None.

make interrogate and git diff --check also pass.

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Faster B-spline test

Optimized test_weight_n_fold_convolution_parity for spline orders 5 and 6 by replacing repeated np.convolve calls with an equivalent prefix-sum calculation. These cases now take about 1–2 ms instead of 4–10 seconds, with less than 1e-14 numerical difference. This improves test runtime only; production PME code is unchanged.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR fixes torch.compile failures for explicit single-system batches (batch_idx=zeros(N)) by centralizing all atom-to-system and system-to-atom reductions through four new shared helpers in _util.py, each with a fast num_systems == 1 direct-reduction path that avoids the index_add/scatter_add operations that Inductor's CPU vector-atomic code generation rejects for size-one destinations. A secondary fix passes is_compiled through to the PME convolve custom-op boundary and conditionally multiplies the cotangent by 1.0 to prevent Inductor from hoisting Hermitian RFFT interior-frequency weighting past the opaque backward.

• _util.py: Adds _sum_atom_values_by_system, _mean_atom_values_by_system, _broadcast_system_values_to_atoms, and _distribute_system_mean_cotangent_to_atoms; all callers across real-space, reciprocal-space, PME, slab, direct-k, correction, cell-gradient, virial, and double-backward paths are rerouted through these helpers.
• pme.py: Extracts _pme_convolve_backward_args from the inline lambda, adds is_compiled as a registered-autograd passthrough parameter (n_forward_inputs bumped from 8 → 9), and applies the cotangent materialization only when compiling.
• Tests: 302 new checks cover helper contracts, compile graph inspection, first/second-order gradients, and batch consistency; the spline test oracle is rewritten with an O(N) prefix-sum to cut runtime from ~10 s to ~1 ms.

Important Files Changed

Filename	Overview
nvalchemiops/torch/interactions/electrostatics/_util.py	Adds four well-guarded helper functions; empty-input, single-system, and multi-system branches all look correct; .clone() on the expand in _distribute_system_mean_cotangent_to_atoms correctly prevents aliasing.
nvalchemiops/torch/interactions/electrostatics/pme.py	Extracts _pme_convolve_backward_args, correctly bumps n_forward_inputs to 9, preserves diff_input_positions, and gates cotangent materialization on is_compiled.
nvalchemiops/torch/interactions/electrostatics/ewald.py	_atom_ranges and _apply_reciprocal_corrections both add a num_systems == 1 fast path; the early-return in _atom_ranges correctly computes ends from batch_idx.shape[0].
nvalchemiops/torch/interactions/electrostatics/_ewald_real_chain.py	Removes duplicate _atom_counts helper and rewires all scatter/index_add reductions through shared helpers; _scatter_dedcell_cache correctly accepts ctx.batch_idx=None for the single-system path.
nvalchemiops/torch/interactions/electrostatics/_ewald_recip_chain.py	Removes duplicate _atom_counts, rewires _per_system_cotangent, _distribute_to_atoms, and _atom_cotangent through shared helpers.
nvalchemiops/torch/interactions/electrostatics/_slab_chain.py	Replaces bincount-based helpers with shared utilities; _slab_double_backward_layout now handles the empty-atom case implicitly.
nvalchemiops/torch/interactions/electrostatics/_ewald_direct.py	Adds num_systems == 1 guards for q_over_v and virial q_tot computations.
test/interactions/electrostatics/test_util.py	Adds contract tests for all four helpers including compile-graph inspection asserting absence of atomic scatter ops for B=1.
test/math/test_spline.py	Replaces O(N^2) FFT-convolve oracle with an O(N) prefix-sum sliding-window implementation; mathematically equivalent, substantially faster.

_{Reviews (5): Last reviewed commit: "Merge remote-tracking branch 'refs/remot..." | Re-trigger Greptile}

[laserkelvin]

/ok to test 6d283cd

[laserkelvin]

One performance suggestion; left inline.

[laserkelvin]

/ok to test 2fabd9a