[MISC] Migrate rigid_solver.py to data_oriented.#2799
Open
hughperkins wants to merge 15 commits into
Open
Conversation
…stable_members=True) Checkpoint of the migration attempt: moves kernel_step_1, kernel_step_2, and update_geom_aabbs into RigidSolver as @qd.kernel methods, decorates the class with @qd.data_oriented(stable_members=True), exposes _contact_island_state and _collider_state as direct attributes, and updates the substep() / detection() call sites. This was originally failing against quadrants hp/data-oriented-ndarray-fix with: Missing argument __qd_links_state__qd_cinr_inertial. Unexpected argument links_state. The fix lives on quadrants branch hp/data-oriented-qd-func-dataclass (Option A: extend caller-side dataclass-arg expansion to handle data_oriented self.X attribute access).
Collaborator
Author
|
Note: this depends on a new Quadrants release of a new Quadrants branch |
The forward path was migrated to @qd.kernel methods on the data_oriented RigidSolver (self.step_1 / self.step_2). The only remaining caller of the top-level kernel_step_2 was the backward-grad call site; replace it with self.step_2.grad(is_backward=True) (the BoundQuadrantsCallable.grad descriptor uses the kernels adjoint correctly) and delete both old top-level kernel definitions.
…ep_1 - test_grad.py::test_diff_solver now calls rigid_solver.step_1(...) (method form) instead of the deleted top-level kernel_step_1. - test_quadrants.py: access RigidSolver.step_1._primal (via the class-level QuadrantsCallable descriptor) for the cache observation checks. The attribute chain works because QuadrantsCallable directly exposes _primal. Smoke-tested both forward and backward paths (scene.step() + rigid_solver.step_2.grad).
``qd.simt.subgroup`` recently moved to the suffix convention where the sized form is ``<op>_tiled(v, log2_size)`` and the no-suffix form ``<op>(v)`` operates over the full subgroup (quadrants commit d07644e4, "rename to _tiled suffix convention"). ``solver._kernel_mass_factor_solve_chol_subgroup_16`` was still calling the old ``reduce_all_add(dot, 4)`` two-arg form, which no longer exists on the new API (the 1-arg form ignores tile width and reduces over the full subgroup, which would give wrong results for the 16-thread parallel dot the kernel implements). Switch the two call sites to ``reduce_all_add_tiled(dot, 4)`` to preserve the intended 2**4 = 16 lane tile.
…nstraint_cost coop
…module-level globals Two qd.static expressions captured module-level mutable globals (`gs.backend`, `gs.qd_float`) whose values can vary across process runs but aren't part of the kernel's fastcache key — risking cross-process cache collisions where a kernel compiled for one backend / precision gets loaded for another. Fix by routing both through declared `static_rigid_sim_config` fields: * `constraint/solver.py::add_frictionloss_constraints`: `gs.backend != gs.metal` -> `static_rigid_sim_config.backend != gs.metal` * `constraint/solver.py::func_cholesky_solve_tiled`: `gs.qd_float == qd.f32` -> `static_rigid_sim_config.is_qd_float_f32` * `array_class.py::RigidSimStaticConfig`: new `is_qd_float_f32: bool` field, populated by `RigidSolver._build_static_config` and `KinematicSolver._build_static_config`. All `qd.static(...)` expressions in fastcache-enabled kernels must derive from declared kernel parameters (now documented in the fastcache user guide). Co-authored-by: Cursor <cursoragent@cursor.com>
Collaborator
Author
|
Benchmark results: 
|
Collaborator
Author
|
Genesis unit tests passing: 
|
hughperkins
changed the title
Collaborator
Author
|
apparently 0.8.1b was inadvertently built from main... |
…s/Genesis into hp/data-oriented-rigid-solver
Collaborator
Author
|
making a new quadrants release... |
Collaborator
Author
|
Note: for the delta in the benchmarks, got AI to compare main versus this branch using several runs of each, on the same machine. Results: |
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.
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.