Repository navigation
Conversation
Adds the batch contract (torchref/base/electron_density/batching.py): a (B, nx, ny, nz) map plus CSR-style batch_offsets, or dense (B, N, ...) atoms flattened by normalize_atoms. The portable splat shifts each atom's flat voxel index into its member's slab; unbatched calls are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The Rust entry points take an i64 offsets buffer (ABI 2): forward runs each member's x-plane blocks over its own slab, backward resolves each atom's member and reads that member's slab of the gradient. The per-member kernels are unchanged, so a batch is bitwise the stack of single-member calls. The double-backward re-derivation passes the batch on, and returns zero gradients for a call with no atoms. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Each kernel takes a per-atom member id behind a HAS_BATCH constexpr and offsets its map pointer by member * nvox in int64. Per-atom arithmetic follows the input dtype, so the cuda_triton row now accepts float64 (uniform dtypes only); the voxel-index decode stays float32. Gradient sums accumulate per lane and reduce once after the loop. Not yet run on a CUDA host. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The shaders take a member-id buffer and a has_batch flag and work in their member's slab of the grid. Indices stay 32-bit, so the wrapper refuses a batch of 2**31 voxels or more instead of wrapping. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
build_electron_density takes dense (B, N, ...) atoms or ragged ones with batch_offsets_iso / batch_offsets_aniso and returns (B, nx, ny, nz). The reciprocal extractor returns (B, N) from a (B, Nx, Ny, Nz) grid, and both map-symmetry operators accept one leading batch axis. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
build_density_map and compute_structure_factors take dense or ragged batched atoms (and ragged imaginary atoms via imaginary_batch_offsets) and return (B, n_reflections) structure factors through either symmetry route. Implied batch offsets are built on the host, so batching costs no device sync. Changelog entries for the batched path, float64 Triton and torchref-kernels 0.2.0. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cb23c0d690
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| [package] | ||
| name = "torchref_kernels" | ||
| version = "0.1.0" | ||
| version = "0.2.0" |
There was a problem hiding this comment.
For installations using the declared dependencies, pyproject.toml still requires torchref-kernels>=0.1.0,<0.2, which excludes this new 0.2.0 package. Meanwhile, torchref/utils/native.py now requires ABI version 2 and rejects the permitted 0.1.0 build. Consequently, a normal installation cannot load the native CPU kernels and falls back to the much slower portable implementation; explicitly installing both projects also produces conflicting dependency requirements. Update the dependency range alongside the package and ABI versions.
AGENTS.md reference: AGENTS.md:L300-L305
Useful? React with 👍 / 👎.
| grad_xyz = torch.zeros_like(xyz) | ||
| grad_third = torch.zeros_like(third) | ||
| grad_occ = torch.zeros_like(occ) |
There was a problem hiding this comment.
When CUDA callers supply a non-contiguous dense tensor, such as coordinates constructed with xyz.T.contiguous().T, the refactored forward saves that original layout rather than the contiguous copy passed to the kernel. These zeros_like allocations preserve its strides, but the backward kernels write gradients using contiguous offsets such as atom * 3 + component. The returned coordinate gradients—and anisotropic gradients for similarly laid-out u—are therefore silently permuted even though the forward result is correct. Allocate the gradient buffers with contiguous memory format or save contiguous inputs as the previous implementation did.
Useful? React with 👍 / 👎.
The splat entry points changed (ABI 2), and the <0.2 pin made pip install -e . replace the freshly built 0.2.0 wheel with 0.1.0 from PyPI, so CI tested against kernels this TorchRef refuses to load. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The portable splat's scatter_add sums in a different order when atoms are regrouped into a batch on some hosts (Linux CI differed by one ulp), so it gets the same rounding tolerance as the atomic GPU kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Summary
Adds an optional leading batch over models that share one crystal (same cell, space group, grid and HKL), such as ensemble members, timepoints, conformers or multistarts. The batch is supported in the density splat kernels, the base functions and
SfFFT. The input shape picks the path: unbatched input behaves exactly as before, and every kernel runs it as a single member.torchref/base/electron_density/batching.py): either dense(B, N, ...)atoms (A/Bmay be shared as(N, 5)), or flat atoms with CSR-stylebatch_offsets (B+1,). In the flat form members can have different atom counts, including zero. Maps come back as(B, nx, ny, nz)and structure factors as(B, n_refl).HAS_BATCHconstexpr, with an int64 slab offset.has_batchflag; refuses a batch of 2³¹ voxels or more because its indices are 32-bit.cuda_tritonrow accepts float64 when all inputs share one dtype. The voxel-index decode stays float32, since its exactness bound is a float32 one. Gradient sums now accumulate per lane and reduce once after the loop; this slightly changes the float32 summation order.build_electron_density(batch_offsets_iso/batch_offsets_aniso),ReciprocalSymmetryExtractor.extract_from_grid, both map-symmetry operators, andSfFFT.build_density_map/compute_structure_factors(includingimaginary_batch_offsets).ModelFT, refinement and I/O are unchanged.torchref-kernels0.2.0. This matches thekernel-expansionbranch, and the two should be released together. Merging them conflicts only inkernels/src/lib.rsandCargo.toml, where both sets of modules need registering.Performance (Apple Silicon, float32, splat + FFT + backward, loop vs batched)
At B=1 and for large structures (2000 atoms), batched is at parity with the loop, within timer noise. CUDA has not been measured.
Test plan
tests/unit/structure_factor/test_batched.py: compares batched against per-member results for forward, gradients and second order on every backend available on the host; ragged vs dense input;build_electron_density, the extractor, map symmetry (direct and interpolated) andSfFFT(late and early symmetry, with and without the imaginary term); error paths.tests/unit/structure_factor,tests/unit/base,tests/unit/symmetry,tests/unit/model, dtype conformance, kernel fixes and import smoke all pass.pip install ./kernels, thenpytest tests/unit/structure_factor/test_batched.pyandpytest tests/unit/structure_factor tests/unit/base tests/integration/test_triton_vs_eager_targets.py -m "gpu or cuda". Pass criteria: allcuda-*cases pass, includingcuda-f64-cuda_triton;test_triton_kernel_is_actually_dispatched[*-float64]passes; noTorchRefDegradationWarning.🤖 Generated with Claude Code