[MISC] Use new GlobalState in narrowphase save 1309 lines#2791
Draft
hughperkins wants to merge 11 commits into
Draft
[MISC] Use new GlobalState in narrowphase save 1309 lines#2791hughperkins wants to merge 11 commits into
hughperkins wants to merge 11 commits into
Conversation
Introduces ``array_class.GlobalState`` bundling every rigid-solver dataclass that narrowphase kernels need (rigid_global_info, dofs/links/joints/geoms/verts/edges/entities/etc. plus collider/mpr/gjk/sdf state), and migrates every ``@qd.func`` / ``@qd.kernel`` in ``narrowphase.py`` from a long list of dataclass parameters down to a single ``global_state: array_class.GlobalState`` argument. Why now: Quadrants Genesis-Embodied-AI#698 enables passing dataclass sub-structs into ``qd.func`` (see ``docs/source/user_guide/compound_types.md``), so callees can keep their existing narrow typed signatures while callers just pass ``global_state.X``. Adding new state no longer requires editing every kernel signature in the chain. Highlights: - Replaced ``geoms_init_AABB`` and ``diff_contact_input`` standalone parameters with ``global_state.rigid_global_info.geoms_init_AABB`` and ``global_state.gjk_state.diff_contact_input`` (or ``global_state.collider_state.diff_contact_input`` in the diff/grad kernel where the caller previously passed the collider one). - Internal narrowphase calls collapse to ``func(..., global_state, ...)``; external helper calls (``mpr``, ``gjk``, ``diff_gjk``, ``capsule_contact``, ``box_contact``, ``contact``, ``sdf``, ``utils``) keep their existing signatures and now receive ``global_state.X`` attribute accesses. - ``Collider`` gains ``_build_global_state(mpr_state, gjk_state)`` so the per-call-site variability (e.g. ``_contact0_*`` vs ``_multicontact_*`` MPR/GJK state) stays explicit. - Removed obsolete ``FIXME: Passing nested data structure as input argument is not supported for now.`` comments.
Many call sites and signatures shrank after the GlobalState refactor (and in the surrounding factory functions in array_class.py) so a few that previously required one-arg-per-line now fit in 120 chars. Ran ``ruff format`` with ``skip-magic-trailing-comma=true`` once and accepted the resulting collapses. Pure formatting; no semantic changes.
hughperkins
changed the title
Extend the previous narrowphase migration to the helpers it directly invokes inside the collider package: mpr.py, gjk.py, diff_gjk.py, capsule_contact.py, box_contact.py, and contact.py. Each function that took individual ``array_class.*`` sub-dataclasses now takes a single ``global_state: array_class.GlobalState``; Quadrants prunes the unused fields. Call sites in those files and in collider.py are updated to pass ``global_state`` (or ``self._build_global_state(...)`` from the Collider class) where they previously assembled the long parameter list. Helpers in sdf.py and utils.py are intentionally left untouched in this PR -- they have genuine callers outside the collider package (legacy_coupler, mpm_solver, kinematic_tactile sensor) and would force a wider refactor. Net effect: ~700 fewer lines across the migrated files with no functional change.
Drive-by formatting: ruff format with ``skip-magic-trailing-comma=true`` collapses imports and signatures that comfortably fit within 120 columns. Touches the non-migrated collider helpers only (broadphase, epa, gjk_utils, multi_contact, support_field, utils); no behavior changes.
``global_state`` is referenced thousands of times across the migrated collider files; the 13-character prefix on every ``global_state.<sub_struct>`` push a number of calls past the 120-col limit even when they used to fit on one line with bare arg names. Shrink the parameter name to ``gst`` (zero existing bindings anywhere in genesis), which collapses 15 of the 20 calls that the migration originally split, with no functional change.
hughperkins
changed the title
``errno`` is a solver-wide per-env error-flag tensor. Inside the collider it was a pure passthrough -- every call site forwarded ``self._solver._errno`` and every kernel that took it as a parameter only ever indexed into it. Fold it into ``GlobalState`` and let ``_build_global_state`` bind ``errno=self._solver._errno`` once, then drop the explicit parameter from the 13 collider kernels (and 25 call sites). No behavior change. Pyright: -7 errors net (600 -> 593).
Commit 096bc55 dropped the explicit ``errno: qd.Tensor`` parameter from six narrowphase entry points (``_func_narrowphase_multicontact_mixed``, ``_func_narrowphase_contact0``, ``func_narrow_phase_convex_vs_convex``, ``func_narrow_phase_convex_specializations``, ``func_narrow_phase_any_vs_terrain``, ``func_narrow_phase_nonconvex_vs_nonterrain``) but left the matching ``self._solver._errno`` positional argument in place at every collider call site. Strip them.
… params. ``mpr_state`` and ``gjk_state`` are algorithm-local scratch buffers, not solver-wide state: the main, ``contact0`` and ``multicontact`` narrowphase entry points each need their own instance (sized for that entry point's parallelism). They were folded into ``GlobalState`` only because that was the convenient way to plumb them down the call chain; in return, ``GlobalState`` had to be rebuilt at every narrowphase call site via ``_build_global_state(mpr_state=..., gjk_state=...)``. Promote them back out as explicit ``mpr_state: MPRState`` / ``gjk_state: GJKState`` parameters on the 43 ``@qd.func`` / ``@qd.kernel`` functions in the collider call graph that actually need them (computed as the transitive closure of "directly references ``gst.mpr_state`` / ``gst.gjk_state``"), and pass the right scratch buffer positionally at each entry point. Drop the two fields from ``GlobalState`` and simplify ``_build_global_state`` to take no arguments. This is a prerequisite for caching ``GlobalState`` as a singleton on the Collider; with the scratch buffers gone, the remaining ``GlobalState`` is constant for the Collider's lifetime.
Now that ``GlobalState`` no longer carries call-site-specific scratch buffers (those moved to explicit ``mpr_state`` / ``gjk_state`` parameters in the previous commit), every field it bundles is a stable reference for the lifetime of the Collider. Build it once via a ``functools.cached_property`` and reuse the same instance at all 14 collider call sites. The ``cached_property`` (rather than building in ``__init__``) is deliberate: ``self._solver.constraint_solver`` is wired up later in the solver bring-up, so eager construction in ``Collider.__init__`` would dereference ``None``. Deferring to first use sidesteps that ordering constraint without introducing a second init phase.
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
1 participant
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.