Repository navigation
Gather applied forces per dof in deterministic steps - #1937
Closed
johnnynunez wants to merge 4 commits into
Closed
johnnynunez wants to merge 4 commits into
johnnynunez wants to merge 4 commits into
Conversation
With `opt.deterministic=ATOMICS`, `_apply_ft_scatter` ran its float atomic in Warp's RUN_TO_RUN mode, which records and sorts every contribution on each step. `apply_ft` already has a kernel without atomics for dense wrenches, `_apply_ft_gather`: one thread per dof adds the bodies of its subtree in order. Deterministic steps now use it for sparse wrenches too, so `_apply_ft_scatter` needs no deterministic variant. Its default kernel generates the same CUDA source as before. mjwarp-testspeed with ATOMICS and 8192 worlds, median of 3, RTX PRO 6000, on top of google-deepmind#1932, google-deepmind#1933 and google-deepmind#1915 (steps per second): humanoid 4.33M -> 9.11M (NONE 9.24M) three humanoids 0.54M -> 2.01M (NONE 2.14M) The gather takes 14.0 us per step on three humanoids and 7.3 us on humanoid, against 7.4 and 3.5 us for the default scatter. The determinism tests now check that deterministic steps do not launch `_apply_ft_scatter`, and `test_apply_ft_subtrees` also runs with ATOMICS.
This was referenced Nov 10, 2026
Deterministic steps map the applied wrenches to dof space with _apply_ft_gather: a thread per dof read the wrench of every body of its subtree, twice the cost of the default scatter when no body has one. A thread per body now sets the bit of each body with a wrench with an integer or (_apply_ft_mask, the same loads as a scatter thread that returns), plus one word per world that tells whether it has any. The gather (_apply_ft_gather_mask) reads that word and returns for a world without wrenches; otherwise it adds the bodies of its subtree whose bit is set, in body order: the same sum as before. Dense wrenches (fluid forces) keep _apply_ft_gather, and the default scatter is not edited. 8192 worlds, us of GPU time per step (default scatter / deterministic before / after = mask + gather): humanoid 3.5 / 7.3 / 3.5 + 1.1 unitree_g1_flat 5.5 / 12.4 / 5.4 + 1.2 three_humanoids 7.4 / 13.9 / 7.6 + 2.0 The mask is a zeroed array of 2 + (nbody - 1) / 31 ints per world. Measured and rejected: a thread per world or per word that tests its bodies one after the other (15 to 29 us), and several dofs per gather thread (1.5 -> 2.3 us with 4).
Contributor
Author
|
Benchmarks on main 60dbbdf:
"fails": Warp's scatter buffer for the solver kernels exceeds int32 (#1933). With NONE this PR is within 0.1% of main on both benchmarks. |
Collaborator
|
@johnnynunez plan for now is to utilize warp's built-in deterministic atomics functionality |
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.
With
opt.deterministic=ATOMICS,_apply_ft_scatter(#1891) runs its float atomic in Warp'sRUN_TO_RUNmode, which records and sorts every contribution. Deterministic steps now gather instead: a thread per body sets the bit of each body with a wrench (_apply_ft_mask) and a thread per dof adds the flagged bodies of its subtree in body order (_apply_ft_gather_mask). Dense wrenches keep_apply_ft_gather._apply_ft_scatterkeeps no deterministic variant, like_solve_LD_sparse_fused, and its default kernel generates the same CUDA source as before.mjwarp-testspeedwith ATOMICS at 8192 worlds, median of 3, on top of #1932 and #1933: humanoid goes from 4.42M to 9.53M steps/s (NONE 9.59M) and three humanoids from 0.55M to 2.08M (NONE 2.14M). Mask and gather take 7.3 + 2.0 us of GPU time per step on three humanoids, against 7.4 us for the default scatter.The determinism tests check that deterministic steps do not launch
_apply_ft_scatter.cc @thowell @adenzler-nvidia