Repository navigation
Avoid atomics in deterministic actuator and sparse contact kernels - #1932
Closed
johnnynunez wants to merge 7 commits into
Closed
johnnynunez wants to merge 7 commits into
johnnynunez wants to merge 7 commits into
Conversation
With `DeterminismType.ATOMICS`, `_qfrc_actuator` scattered `actuator_moment * actuator_force` into `qfrc_actuator` with a float atomic compiled in Warp's `RUN_TO_RUN` mode, which records and sorts every contribution. On the humanoid benchmark with 8192 worlds that launch took about 406 us per step against 19 us for the default kernel. The deterministic variant now lets the thread of the first actuator of each world accumulate all the actuators of that world in ascending order with plain loads and stores. Each element has one writer and a fixed summation order, and the compiled kernel has no atomic left, so Warp neither allocates scatter buffers nor sorts. The update is written as `x = x + y` because Warp compiles `+=` on an array element to an atomic. The launch now takes about 33 us per step in the same benchmark, and 36 us instead of 176 us with three humanoids and 64 worlds. The factory signature, its `RUN_TO_RUN` module options and the default kernel are unchanged: the generated CUDA source of the default variant is the same as before. Add a test that compares the deterministic forces with MuJoCo and with the default kernel for one and two worlds, on a model where every degree of freedom is driven by several actuators.
`_transmission` took the address of each sparse `actuator_moment` row from `wp.atomic_add(moment_nnz, worldid, rownnz)`. The returned value is used, so with `DeterminismType.ATOMICS` Warp's `RUN_TO_RUN` mode ran the kernel twice around a prefix scan to hand out the addresses in thread order. On the humanoid benchmark with 8192 worlds that launch took about 381 us per step against 46 us for the default kernel. The width of a row only depends on the transmission: the joint type, the number of non-zeros of the tendon, `nv` for body transmissions, and the degrees of freedom that move the sites for site and slider-crank transmissions. `put_model` now derives these widths and stores their exclusive prefix sum in the new Model field `actuator_moment_rowadr`. In deterministic steps `transmission` passes this array as `moment_nnz` instead of a zeroed counter per world, and the deterministic variant reads the number of non-zeros that precede the row of its actuator from it. This gives the same layout as before (rows in actuator order) with no atomic in the compiled kernel and without allocating the counter. The launch now takes about 50 us per step in the same benchmark, and 42 us instead of 271 us with three humanoids and 64 worlds. The factory signature, its `RUN_TO_RUN` module options, the kernel arguments and the default kernel are unchanged: the default variant still takes its addresses from the counter and its generated CUDA source is the same as before. Add a test with every transmission type that checks the row addresses against the row widths written by the kernel and against the Model field, and compares lengths and moments with MuJoCo and with the default kernel for one and two worlds.
`_efc_contact_jac_sparse` splits the non-zeros of a contact across lanes, and each lane adds its part of `J * qvel` to `efc.Jqvel`. With `DeterminismType.ATOMICS` these float atomics were compiled in Warp's `RUN_TO_RUN` mode, which records and sorts every contribution. With three humanoids and 64 worlds the launch took about 216 us per step against 43 us for the default kernel. The deterministic variant now only writes the Jacobian, so its compiled kernel has no atomic left. A new kernel, `_efc_contact_jqvel_sparse`, then computes `J * qvel` for each contact row with one thread per row that sums the row in storage order. It has no atomics, so it runs in the default mode. The two launches now take about 47 us and 20 us per step. The factory signature, its `RUN_TO_RUN` module options and the default kernel are unchanged: the generated CUDA source of the default variant is the same as before. Sparse models with flexes use `_efc_contact_jac_sparse_flex`, which already sums each row in one thread, and the dense kernels are left alone. Add a test that checks `efc.Jqvel` of the contact rows against the product of the assembled Jacobian and `qvel`, and compares `efc.vel` with MuJoCo and with the default kernel, for both cones and one and two worlds.
This was referenced Nov 10, 2026
The deterministic variant of forward._qfrc_actuator had the thread of the
first actuator of each world add all the actuators in order, which is
latency bound: 8x slower than the default kernel on three_humanoids.
put_model now builds the transpose of the static sparsity of
actuator_moment (the rows of smooth._transmission only depend on the
transmission): actuator_momentT_rowadr, actuator_momentT_elem and
max_actuator_momentT_rownnz. The deterministic variant is one of:
- If no dof is driven by more than one actuator, the default kernel
with a plain load/add/store instead of the atomic: each element has
one writer, so there is nothing to order.
- Otherwise a second kernel with one thread per dof, which adds the
actuators of its dof in ascending order. The moment of an actuator is
at a fixed offset of its row; rows copied from MuJoCo by put_data,
which leave out zero moments, have fewer dofs than in the model and
are searched instead. The address of each row is read from Data, so
rows placed by the kernel with atomics also work.
The default kernel is not changed: its generated CUDA source is
identical.
GPU time per step of the kernel in mjwarp-testspeed (Nsight Systems,
8192 worlds, RTX PRO 6000; default -> deterministic before -> after):
three_humanoids 5.53 -> 45.66 -> 5.68 us
humanoid 1.61 -> 9.02 -> 1.71 us
unitree_g1_flat 1.81 -> 12.69 -> 1.93 us
franka_emika_panda 1.86 -> 5.54 -> 1.97 us (32768 worlds)
myoarm 142.89 -> 392.01 -> 91.67 us (63 tendon actuators)
unitree_g1_flat with a second actuator per joint
2.70 -> 35.74 -> 3.24 us
johnnynunez
added a commit
to johnnynunez/mujoco_warp
that referenced
this pull request
Nov 10, 2026
Three steps of the deterministic solver were serial per world: - The list of rows whose quadratic state changed was built by one thread per world that scanned a flag per constraint row. The rows now set one bit each, 31 rows per word, and `_quad_changed_ids` lists the set bits in row order. - With sparse Jacobians, `qfrc_constraint` was added row by row by one block per world. `init_context` now lists, for each dof, the JTDAJ groups whose rows contain it (a count, a prefix sum in two levels, a fill and a sort of each list), and one thread per dof adds its rows in that order. Solvers without those groups keep the block per world. - The incremental sparse Hessian uses blocks twice as wide, and `_JTDACJ_sparse` does too while all its blocks are resident, which shortens the chain of rows of a world. SolverContext gains `quad_changed_flag`, `qfrc_adr` and `qfrc_entry`. The default kernels generate the same CUDA source as before. GPU time per step with ATOMICS and 8192 worlds (Nsight Systems, RTX PRO 6000, on top of google-deepmind#1932 and google-deepmind#1915), before -> after: three humanoids qfrc_constraint, with its index 673 -> 386 us _quad_changed_ids 86 -> 27 us incremental sparse Hessian 128 -> 114 us ALOHA clutter qfrc_constraint, with its index 594 -> 211 us _JTDACJ_sparse 532 -> 473 us humanoid _quad_changed_ids 25 -> 8 us Add tests of the list of changed rows and of the sort of the dof lists.
This was referenced Nov 10, 2026
Open
Contributor
Author
|
Benchmarks on main 60dbbdf:
"fails": Warp's scatter buffer for a kernel that another PR converts exceeds int32 (here the solver's, #1933). With NONE this PR is within 0.2% of main on the three 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,_qfrc_actuator,_transmissionand_efc_contact_jac_sparserun their atomics in Warp'sRUN_TO_RUNmode. Their deterministic variants now compile no atomic:actuator_momentrows take their address from a prefix sum built input_model, a second kernel sumsJ * qvelof each contact row in a fixed order, and_qfrc_actuatorgathers per dof through the transposed sparsity ofactuator_moment(three new Model fields), or keeps the default body when no dof has two actuators. The default kernels generate the same CUDA source as before.With 8192 humanoids,
_qfrc_actuatortakes 1.7 us of GPU time per step against 1.6 us for the default kernel.mjwarp-testspeedwith ATOMICS on main 5cfac7b, median of 3, RTX PRO 6000: humanoid at 8192 worlds 1.79M to 2.55M steps/s (NONE 9.58M); with #1933 and #1937, 9.53M, 0.6% below NONE.Gravity compensation was left to #1915, now merged. Conflicts textually with #1929.
cc @thowell @adenzler-nvidia