refactor solver#1395
Merged
Merged
Conversation
adenzler-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
Nice refactor — the _eval_constraint extraction and the monolithic/island unification look correct to me (traced the per-branch force/state/cost equivalence, the if cost != 0.0 gating, and the padding-island done=True change; solver_test.py passes locally). Two minor non-blocking nits in _eval_elliptic_cost_island, otherwise LGTM.
15f5394 to
c8631a5
Compare
adenzler-nvidia
approved these changes
Jun 3, 2026
adenzler-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
Both nits from the previous review are addressed:
_eval_ellipticnow takesmuas an argument, removing the double-compute._eval_elliptic_cost_islandreuses theu0/v0locals when buildingquad1.
The _eval_constraint extraction and monolithic/island unification look correct. LGTM — please resolve the conflicts with main before merging.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
small refactor to update solver.py and start unifying the monolithic and island solvers a bit
1. Shared
_eval_constraintfunctionThe inline constraint-evaluation logic that was duplicated in both the monolithic (
_update_constraint_efc) and island (_update_constraint_efc_island) kernels is extracted into a single@wp.funccalled_eval_constraint. It handles all constraint types (equality, friction, elliptic cone, inequality) and returns avec3of(force, state, cost). Both kernels now call this shared function instead of reimplementing the logic.2. Unify
_update_constraint_init_costacross monolithic and island pathsThe island-specific
update_constraint_init_cost_islandkernel is removed. Both paths now use the same_update_constraint_init_costkernel — the island path calls it with flattened (.reshape(-1)) arrays anddim=d.nworld * m.ntree.3. Unify
_solve_zero_search_dotacross monolithic and island pathsSame pattern as above: the island path previously used
ctx.search_dot.zero_(), which unconditionally zeroed all islands including converged ones. Now it reuses the monolithic_solve_zero_search_dotkernel (which skips done islands) via.reshape(-1).4. Reuse
_eval_ellipticin island linesearch_eval_elliptic_cost_islandwas inlining the elliptic cone cost evaluation. It now delegates to the existing_eval_elliptichelper, and similarly uses_eval_frictionloss_pt_3alphas/_eval_pt_direct_3alphasinstead of three separate calls.5. Prefix all non-public kernels/functions with
_Every kernel and host-level function that isn't part of the public API is renamed with a leading underscore (e.g.,
solve_init_search→_solve_init_search,linesearch_island→_linesearch_kernel_island,_solve_islands→_solve_island, etc.). ~30 kernels and functions are affected.this pr
193404a