Repository navigation
Sort contacts and constraint rows with fewer moves - #1940
Open
johnnynunez wants to merge 17 commits into
Open
johnnynunez wants to merge 17 commits into
johnnynunez wants to merge 17 commits into
Conversation
Fill the CONTACTS and CONSTRAINT hooks introduced by google-deepmind#1687 without changing atomic reductions or the NONE path. Canonicalize contact features and geometry, permute constraint rows with their sparse Jacobian metadata, and repair contact and JTDAJ references independently of incoming contact order. Add permutation, CPU parity, interleaved-world, mixed-dimension, overflow, flex, SDF, and CUDA capture regressions. Credit Mark Yang (@mar-yan24) for the original joint determinism work that motivated this follow-up.
Order equality rows by subtype and ID while preserving each connect/weld group's semantic component offsets. Keep physical Jacobian tie-breaks for repeated flex IDs and leave force reconstruction unchanged. Add CPU-parity regressions for solved-state sorting and full forward with asymmetric loads, interleaved equality types, dense/sparse Jacobians, heterogeneous worlds, reversed groups, and JTDAJ block support.
Copy retained NumPy snapshots before mutating Warp storage so CPU assertions compare independent data rather than aliased views. Use moderate binary-exact equality loads and velocities without relaxing the standard parity tolerance or changing production sorting.
Compare contacts lazily through a `_ContactKeys` struct instead of materializing a 43-column key buffer, evaluate the comparison once per merge step, test raw bits before computing float keys, gather rows as `uint32` words, and name kernel parameters after their Model and Data fields. `_remap_blocks` reads the head row of each block directly and `constraints` always ranks contacts. Call `sort.constraints` from `make_constraint` instead of going through `_sort_constraints`. Clear only the contact address components that assembly does not write, so the addresses of active contacts are no longer written twice. Move the contact-address and imported-contact regressions into `sort_test.py`, replace attribute reflection and random permutations with explicit fields and fixed permutations, and replay captured contact and constraint sorts across input orders and active counts. JTDAJ blocks are asserted to tile every row except the DOF friction rows, which have no block on current main. `test_contact_permutation` still fails at this commit because main now stores `contact.surfacevel`; the next commit adapts sorting to the merged main.
Merging main brought three things that sorting has to carry. `contact.surfacevel` is computed before the contact sort, so it is permuted with the other contact fields and used as a tie breaker. `dof_frictionadr` holds the row of each DOF friction constraint for the sparse Newton Hessian diagonal, so it is remapped through the row permutation. Flex strain rows are now allocated as one JTDAJ block per cell, so they keep their mode order within the cell like connect and weld components, which keeps the head row of each block first. Add regressions for the friction addresses and the flex strain blocks, and a surface velocity tie-break case.
The deterministic sort launched one kernel per merge pass and one clone and one gather per contact or constraint field: 80 sorting launches per step on the humanoid benchmark with 8192 worlds. Contacts are now counted per world, grouped and placed by rank, moving every field in the same launch. The second contact sort of the constraint rows is gone: one launch checks whether the contact buffer is already in canonical order, and only ranks the contacts if it is not. The row sort starts from a counting launch over blocks of rows, and the row fields and the Jacobian move through one packed copy that skips the rows that are already in place. Buffers larger than the static limits fall back to the merge passes. Sorting launches per step go from 80 to 15 (humanoid) and from 99 to 23 (three humanoids). With 8192 worlds, CONTACTS|CONSTRAINT costs 18.4% of the steps per second of the humanoid benchmark instead of 33.6%, and 14.9% of three humanoids instead of 29.1%. The sorted output is bitwise identical to the previous implementation.
sort_test.py: test_contact_rank_paths and test_sparse_offsets_paths build the unsorted state once, snapshot Data and restore it before each path, so both paths sort the same input instead of the output of two forward calls, whose efc.Ma differs by one ULP. sort.py: the fields of _ContactRows and _ConstraintRows are assigned directly, the contact copy is allocated as a _ContactRows instance by _empty_contact_rows, and _rank and _merge take the type of their keys from compare.input_types["keys"]. This removes getattr, setattr, SimpleNamespace, Any, _CONTACT_FIELDS and _CONSTRAINT_FIELDS. test_contact_bucket_ties clones the contact into a types.Contact. sort_test.py: test_permutation_blocks, test_contact_paths and test_sparse_move use modular index arithmetic instead of np.random.default_rng. collision_driver.py: collision calls sort.contacts(d) and _sort_contacts is removed. sort_test.py and the spies of determinism_test.py follow. collision_flex.py: the duplicate reset of contact_efc_address_out is removed, so the file matches main. sort.py: _constraint_primary zeroes efc.Jqvel of the rows that are not contact rows. Assembly never writes it for them, so sorted rows carried the values of earlier slots. test_noncontact_jqvel covers it.
Under opt.deterministic = CONTACTS | CONSTRAINT the sort moved most of what it ordered twice. Contacts were grouped into a copy and placed back, so every contact field moved out of the contact buffer and back into it. Contact rows were allocated in the order of the unsorted buffer, so nearly every row was ordered and moved with its Jacobian, by a counting launch and merge passes over all the rows of a world, and the JTDAJ blocks by merge passes of their own. Contact rows are placed before they are filled. sort.contact_rows runs in make_constraint between the launch that allocates the contact rows and the launches that fill them: in a canonical contact buffer each contact walks the contacts of its world and takes the addresses that the sort would give its rows, their nonzeros and their JTDAJ block. sort.constraints then finds the contact rows in place. Only the head rows are ordered. _constraint_unsorted counts the rows that precede a sorted tail of contact rows and tells whether any of them is out of order. _constraint_rank orders them as one block with 16 threads per world, only in the worlds that need it. The launches that move the row fields also move the dense Jacobian chunks, the packed nonzeros move up to the last row that is not in place, and the JTDAJ blocks take their order from one counting launch (_block_place). The narrowphase writes to a contact buffer, so each contact field moves once. Under the CONTACTS bit sort.contact_buffer returns a second types.Contact, the primitive, convex and SDF launches write to it through collision_core.contact_outputs, and sort.contacts places each contact from there at its canonical position in Data.contact. Contacts are ranked within their world while the buffer holds at most 256 contacts per world (was 128). Limits. The launches are chosen from static sizes and options, so a captured step replays them: 256 contacts per world for the contact buffer and the placed rows (_BUCKET_NCONMAX), 256 rows per world for the head as one block (_HEAD_NJMAX) and for the count of the blocks (_BLOCK_RANK_NJMAX). Above them, without the CONTACTS bit, in an incremental sleeping pass, with a contact filter callback or with flexes, the previous launches run. The sorted output does not change: 12 rollouts under opt.deterministic = ALL (humanoid and three_humanoids, dense and sparse, both cones, 1 to 8192 worlds) are bitwise equal to those of the previous commit. Steps per second with CONTACTS | CONSTRAINT against NONE of the same tree (mjwarp-testspeed, median of 3 runs, spread at most 0.7%, RTX PRO 6000), previous commit -> this commit: humanoid, 8192 worlds -18.7% -> -7.4% three_humanoids, 8192 worlds -14.8% -> -4.1% unitree_g1_flat, 8192 worlds -11.3% -> -5.0% aloha_clutter, 2048 worlds -8.6% -> -4.3% franka_emika_panda, 32768 worlds -4.5% -> -4.7% Launches from sort.py per step at 8192 worlds: humanoid 15 of 113.6 -> 13 of 111.6, three_humanoids 23 of 136 -> 17 of 130. A step under NONE launches nothing from sort.py (98.6 and 113 launches) and runs within 0.2% of the previous commit in the five benchmarks. Tradeoffs. Under the CONTACTS bit collision detection allocates a second set of contact arrays (every field but the addresses and the flex features), in place of the copy that the sort allocated. The limits are static: worlds above them keep the previous cost. franka_emika_panda (njmax 5) does not gain. test_noncontact_jqvel runs on the launches for every row, on the head ordered as one block and by merge passes, and on placed contact rows. The tests that compare two paths build one state and restore it before each path.
Conflict in determinism_test.py: main removed the unused util_pkg import next to the sort import added here. The sort import is kept.
johnnynunez
marked this pull request as ready for review
November 12, 2026 00:37
This was referenced Nov 12, 2026
Contributor
Author
|
Benchmarks on main 60dbbdf, which now has #1422:
With NONE this PR is within 0.2% of main on the five benchmarks. |
Conflicts in collision_core.py and collision_driver.py. `CollisionContext` keeps the `contact` field next to the fields main added. `collision` writes the narrowphase contacts to the sort buffer only when it sorts them itself: with `sort_contacts=False` (google-deepmind#1960) they stay in `d.contact`, where the sleeping path reads them between its two passes and sorts them once afterwards.
This branch has not been deployed
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
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.
Follow-up to #1422, which is now merged.
With
CONTACTS|CONSTRAINT, #1422 moves most contacts and constraint rows twice. Here contact rows take their final address before they are filled, only the rows ahead of them are ordered, and the narrowphase writes to a second set of contact arrays. Output is bitwise identical.Loss vs NONE on main 639faba (
mjwarp-testspeed): humanoid 19.2% to 8.0%, three humanoids 14.9% to 4.1%, Unitree G1 12.2% to 5.2%, Franka 4.6% to 4.8%.Memory: +64 MB on aloha clutter.
Related: #1923, #1924.
cc @thowell @adenzler-nvidia