Skip to content

perf(state): derive n_systems from cell shape to avoid GPU-CPU sync - #621

Merged
CompRhys merged 1 commit into
TorchSim:mainfrom
lil-lon:perf/n-systems-from-cell
Oct 26, 2026
Merged

CompRhys merged 1 commit into
TorchSim:mainfrom
lil-lon:perf/n-systems-from-cell

Conversation

@lil-lon

@lil-lon lil-lon commented Oct 24, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Derive SimState.n_systems from the leading dimension of cell instead of torch.unique on system_idx, so reading it no longer forces a GPU-CPU sync.
  • Use state.n_systems in the FIRE steps, as discussed in PERF: Remove GPU-CPU Synchronization in FIRE #609.
  • Validate constraints after __post_init__ batches the cell, so system constraint bounds are checked against the correct n_systems when a (3, 3) cell is passed. Add a test for this case.

Checklist

Before a pull request can be merged, the following items must be checked:

  • Doc strings have been added in the Google docstring format.
  • Run ruff on your code.
  • Tests have been added for any new functionality or bug fixes.

@CompRhys
CompRhys merged commit c89ad7d into TorchSim:main Oct 26, 2026
62 of 64 checks passed
@lil-lon
lil-lon deleted the perf/n-systems-from-cell branch October 26, 2026 10:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants