[MISC] Use qd.ndrange(axes=) to collapse layout-flip duplications.#2861
[MISC] Use qd.ndrange(axes=) to collapse layout-flip duplications.#2861hughperkins wants to merge 1 commit into
Conversation
Folds 9 ``if qd.static(constraint_layout_transposed): ... else: ...`` blocks in
the rigid constraint solver — where the two branches differ only in the
``qd.ndrange(...)`` axis order — into single canonical-indexed loops:
for i_d, i_b in qd.ndrange(
n_dofs, _B,
axes=qd.static((1, 0) if static_rigid_sim_config.constraint_layout_transposed else None),
):
... # body unchanged
``axes=`` on ``qd.ndrange`` is canonical-preserving: the loop variables stay
bound to canonical axes regardless of the permutation, so the bodies stop
being copy-pasted and the visit order is still controlled at compile time.
The identity / ``axes=None`` fast-path in the Quadrants AST builder emits
byte-identical IR to the pre-rename codegen for the canonical branch.
Sites migrated:
solver.py
* _func_update_efc_force
* func_update_gradient_tiled (gradient half)
* _initialize_Jaref_parallel
* initialize_Ma
* func_solve_init / from_warmstart
* func_solve_init / assign_search
solver_breakdown.py
* _func_update_constraint_forces
* _func_update_qfrc_constraint_per_dof
* _func_update_gradient_no_solve
Left alone (4 remaining ``constraint_layout_transposed`` branches): these
dispatch to genuinely different algorithms (cooperative-warp vs serial), not
axis-order flips:
* solver.py: contact-add per-friction vs per-contact path
* solver.py: fused coop update vs serial func_update_constraint_batch
* solver_breakdown.py: linesearch refine coop vs serial
* solver_breakdown.py: update_constraint_cost coop vs serial
Net -57 LOC; one source of truth per kernel body. No functional change.
Requires Quadrants with ``axes=`` on qd.ndrange (currently on the
``hp/ndrange-layout`` branch; targeted for the next release after 0.8.0).
![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: efdb5952fb
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for i_c, i_b in qd.ndrange( | ||
| len_constraints, _B, axes=qd.static((1, 0) if static_rigid_sim_config.constraint_layout_transposed else None) |
There was a problem hiding this comment.
In environments created from this repo's current dependency spec, quadrants is pinned to 0.8.0, while these new qd.ndrange(..., axes=...) calls require the unreleased Quadrants API noted in the commit message. Any rigid constraint kernel that reaches this helper with the pinned wheel will fail at Quadrants parsing/compilation with an unexpected axes keyword instead of running, so this needs either a dependency bump to a released version that supports axes or to keep the previous explicit branches until that release is available.
Useful? React with 👍 / 👎.
1 participant
Folds 9
if qd.static(constraint_layout_transposed): ... else: ...blocks in the rigid constraint solver — where the two branches differ only in theqd.ndrange(...)axis order — into single canonical-indexed loops:axes=onqd.ndrangeis canonical-preserving: the loop variables stay bound to canonical axes regardless of the permutation, so the bodies stop being copy-pasted and the visit order is still controlled at compile time. The identity /axes=Nonefast-path in the Quadrants AST builder emits byte-identical IR to the pre-rename codegen for the canonical branch.Sites migrated:
solver.py
solver_breakdown.py
Left alone (4 remaining
constraint_layout_transposedbranches): these dispatch to genuinely different algorithms (cooperative-warp vs serial), not axis-order flips:Net -57 LOC; one source of truth per kernel body. No functional change.
Requires Quadrants with
axes=on qd.ndrange (currently on thehp/ndrange-layoutbranch; targeted for the next release after 0.8.0).Description
Related Issue
Resolves Genesis-Embodied-AI/Genesis#
Motivation and Context
How Has This Been / Can This Be Tested?
Screenshots (if appropriate):
Checklist:
Submitting Code Changessection of CONTRIBUTING document.