[johnnynunez]

Implements the CONTACTS and CONSTRAINT sorting hooks from #1687, keeping the flags independent and upstream atomics unchanged. The review follow-up uses explicit fields and lazy comparisons, preserves efc_address in permutations, resets reused contact addresses and initializes imported collision IDs.

At 831a476, 2,359 CPU and 2,366 GPU tests pass; pre-commit passes. Standalone assembly still sorts external contacts. This does not guarantee complete stepping determinism.

Adhesion and sensor work is separate in #1860 and #1861 (drafts); FLEXSTRAIN already has #1793. Original joint work with @mar-yan24.

[johnnynunez]

Test and benchmark results on this branch as pushed (RTX PRO 6000, CUDA 13.3).

Full test suite: uv run pytest -n 4 -> 1108 passed, 23 skipped.

Determinism: the included determinism_test.py verifies qacc/qpos/qvel are bitwise identical across repeated runs over 1000 graph-captured steps.

Forward-step throughput with opt.deterministic left off (the default), mjwarp-testspeed benchmarks/humanoid/humanoid.xml --nstep 500 (steps/second, this branch vs current main):

nworld	main	this PR (det off)	delta
1	4230	4205	-0.6%
256	621692	618133	-0.6%
1024	2231204	2223241	-0.4%
4096	6178023	6177181	-0.0%

This is the goal of this PR: bitwise-deterministic stepping as an opt-in flag, with the default path unchanged within run-to-run noise (the wp.static(deterministic) dispatch compiles the off path to the original atomic kernels). The cost of deterministic=on itself, measured on three_humanoids, is characterized in #1422 (comment); optimizing that overhead (tile-parallel scans, JTDAJ pair index) is deferred to a follow-up PR.

(Edited: an earlier version of this comment claimed the deterministic gathers were faster than the atomic reductions; that comparison was against the naive deterministic kernels, not the atomics — see the linked comment for the corrected numbers.)

[thowell]

awesome work!

[johnnynunez]

awesome work!

I already have differentiability in the other PR. I’d also like to add these two features, since I think they’re both important for the research community.

Let me know if you have any comments or if you spot anything wrong. I tried to make the implementation as clean as possible.

[adenzler-nvidia]

This is awesome stuff indeed, I really like it.

Question about the performance numbers though - I though we might as well do the precomputation for JTDAJ always based on your claims, but I cannot reproduce that it is faster. I used three_humanoids to test that, what model did you measure this on?

[adenzler-nvidia]

same question about your testspeed numbers - what are you actually comparing there? Is this just to show that deterministic = off did not get any slower? Would love to have a good set of numbers to characterize the overhead of deterministic = on.

[johnnynunez]

Good catch, and apologies, it was my own test environment with the previous PRs.... Let me explain:

The perf claim in the original description was wrong, and you were right not to be able to reproduce it. The "-35% / -71%" figures came from the optimization history in mar-yan24#8: we compared the segmented-gather version of the deterministic reductions against the first naive O(nefc)-scan version of the same deterministic kernels, on a local 20-body contact pile scene. So we describe how much the det-mode kernels improved during development, not a win over the atomic baseline. The deterministic gathers are slower than the atomics everywhere, as you'd expect... there is no case where you'd want the precomputation with deterministic off. I've rewritten the PR description accordingly.

Here are fresh numbers on three_humanoids so we're on the same model (mjwarp-testspeed, graph captured, nstep=1000, njmax=160, RTX PRO 6000 Blackwell sm_120, CUDA 13.3):

--function step (worlds·steps/s):

nworld	upstream main	det=off	det=on	det on-vs-off
1	1588	1510	1042	−31%
256	201k	205k	100k	−51%
1024	664k	662k	199k	−70%
8192	1420k	1388k	259k	−81%

--function solve (solver only):

nworld	upstream main	det=off	det=on	det on-vs-off
1	8051	8049	6089	−24%
256	1693k	1704k	1200k	−30%
1024	5088k	5089k	3346k	−34%
8192	7247k	7449k	4486k	−40%

On your second question: yes, the testspeed numbers in the earlier comments were exactly that, checking that deterministic=off doesn't regress. The off path compiles to the original atomic kernels via wp.static dispatch, and the table bears that out: det=off matches upstream main within run-to-run noise at every world count (I re-ran the nworld=1 pair 3x to confirm the small deltas are noise).

For det=on, the step-level overhead is larger than the solve-level overhead because the count-scan-emit row allocation and the contact sort also run outside the solver, and the per-world exclusive scans serialize at high world counts. If the high-nworld overhead matters for the RL use case, the clearest headroom is making those scans tile-parallel and giving the JTDAJ gather a precomputed (row, col) pair index.... happy to pursue that in this PR or a follow-up, whichever you prefer.

[adenzler-nvidia]

Thanks for the update, makes a lot of sense.

I'd say the JTDAJ follow-up is definitely material for a follow-up PR, and this PR should focus on correctness while not regressing the "off" path, which you already showed.

[johnnynunez]

Thanks for the update, makes a lot of sense.

I'd say the JTDAJ follow-up is definitely material for a follow-up PR, and this PR should focus on correctness while not regressing the "off" path, which you already showed.

thanks! Let me do that... Really appreciate it the feedback..
Last push the CPU failures are fixed (the host-side overflow check called wp.get_stream() on CPU devices, now guarded behind is_cuda) and the kernel analyzer is clean, pre-commit and the full determinism suite pass locally on both CPU and GPU. The workflows just need approval to run on this PR.

[johnnynunez]

Follow-up PR with the first round of deterministic-mode overhead reductions is up as a draft: #1425 (JTDAJ merge-intersection gather + parallel contact world boundaries; det=on overhead at 8192 worlds goes from -81% to -66%, bitwise determinism unchanged). It's stacked on this branch and I'll rebase it once this lands.

[thowell]

@johnnynunez @mar-yan24

couple of questions:

• do we have comparisons with warp's built-in determinism Warp determinism NVIDIA/warp#1355?
• what features are not currently support with determinism in this pr (eg, rendering, island solver, flex)?
• how do these changes affect memory and compile time?

thanks!

[johnnynunez]

@johnnynunez @mar-yan24

couple of questions:

• do we have comparisons with warp's built-in determinism Warp determinism NVIDIA/warp#1355?
• what features are not currently support with determinism in this pr (eg, rendering, island solver, flex)?
• how do these changes affect memory and compile time?

thanks!

@eric-heiden @adenzler-nvidia are the expert here… i think that in that thread @eric-heiden answered me that probably we get better performance if we do it in in subset of kernels.

related to the other questions, let me check in detail, anyways i haver ready the other PR getting better performance than here.

[johnnynunez]

Following up on the first question properly — I went and actually ran NVIDIA/warp#1355 on mujoco_warp rather than leaving it at "not merged." I built mmacklin/warp@warp-deterministic locally (reports as 1.15.0.dev0), ran mjwarp main (the unmodified atomic solver, none of this PR's changes) against it, and just flipped Warp's global wp.config.deterministic = RUN_TO_RUN before importing mjwarp so every module compiles with the deterministic atomic rewrite. Same GPU as before (RTX PRO 6000 Blackwell sm_120), humanoid.xml, sparse Jacobian, Newton.

The headline: NVIDIA/warp#1355 does make unmodified mjwarp bitwise-reproducible — so the approach genuinely works end to end. But applied wholesale to the solver it isn't practical on this model/hardware yet, for two concrete reasons that both surfaced as hard errors.

Matched run at 256 worlds (bit_stable = sha1(qacc,qpos,qvel) identical across two in-process runs):

config	bit-stable	graph capture	µs/step	worlds·steps/s	used_MB
main, atomics (det off)	no	yes	368	6.96e5	1038
main + warp#1355 RUN_TO_RUN	yes	no (eager fallback)	13765	1.86e4	45132
this PR, det=off	no	yes	377	6.80e5	1038
this PR, det=on	yes	yes	655	3.91e5	1038

So at 256 worlds this PR is ~1.8x over baseline with graph capture intact and no extra memory, whereas NVIDIA/warp#1355 is ~37x with graph capture lost and +44 GB of scratch.

Why NVIDIA/warp#1355 struggles specifically on the solver:

1. Scatter-buffer sizing on many-atomics-per-thread kernels. The generic scatter-sort-reduce sizes a temp buffer of roughly num_threads × deterministic_max_records per atomic target. mjwarp's _M_sparse (CRB) and especially _JTDAJ_sparse issue many atomic_adds per thread into the Hessian, so the default bound overflows right away: Deterministic scatter buffer overflow in kernel '_JTDAJ_sparse'. Bumping deterministic_max_records to 64 restores correctness but the buffer becomes ~44 GB at just 256 worlds, and at 8192 worlds the record count overflows int32 (2,733,721,343 > 2^31) so the scatter array can't be allocated at all.

2. Graph-capture incompatibility with the solver loop. mjwarp runs Newton inside a CUDA graph conditional-while (capture_while). Warp determinism NVIDIA/warp#1355 allocates its deterministic scratch inside the launch, and conditional graph bodies disallow allocation: Conditional body graph contains an unsupported operation (memory allocation). So even where correctness holds, the step falls back to eager — that's the 37x above.

This PR sidesteps both by exploiting that J is fixed for the duration of a solve: it builds one compact per-dof column index (~87 MB at 8192 worlds) once per step and gathers in a fixed order, which stays capture-safe and costs ~1.8x at 256 worlds / ~3x at 8192 rather than 37x.

My read is unchanged but now backed by numbers: the two are complementary — NVIDIA/warp#1355 is the zero-rewrite general tool, this PR is the domain-specific fast path for the solver. The genuinely useful next step is feeding these _JTDAJ_sparse / _M_sparse cases back to the Warp team as a stress test for scatter-buffer sizing and for allocation inside conditional graphs; if those two are addressed, swapping individual hand-written gathers for the generic path becomes worth re-measuring. Happy to share the harness if useful.

[johnnynunez]

Rebased onto current main to clear the merge conflict, which pulled in mujoco v3.9.0 (#1428), the flex collision changes (#1421), the aloha cloth benchmark update (#1416) and per-world skybox (#1429).

The only textual conflict was in constraint.py and resolved cleanly: main added the has_flex branch around contact initialization, so the deterministic count-scan-emit row allocation is threaded through both the flex and non-flex contact init kernels.

Merging those changes surfaced a real crash in deterministic mode that I also fixed here.
The flex collision work added a new contact.elem field and made contact.flex, contact.elem and contact.vert lazily allocated, so they are zero-length null arrays whenever a model has no flex contacts.
The deterministic contact permutation kernel _permute_contacts dereferenced src_flex_in and src_vert_in unconditionally, which read from a null pointer and raised CUDA 700 illegal memory access for any model without flex (for example humanoid) whenever opt.deterministic was set.
compute-sanitizer pinned it precisely: Invalid global read of size 4 bytes at _permute_contacts, Access to 0x0 is out of bounds.

The fix guards the flex, elem and vert gathers behind a runtime shape check so the kernel never dereferences an empty array, and it also carries the new elem field through the sort so flex contact element ids stay consistent after reordering.
The determinism test helper had the same latent assumption on the empty arrays and is updated to skip empty contact fields and to track elem.

Verified with compute-sanitizer memcheck (0 errors) and the full determinism, constraint and collision_driver suites on both GPU and CPU: determinism 48 passed, constraint plus collision_driver 119 passed.

[johnnynunez]

Rebased onto current main via merge (the conflicts were against the new flexstrain constraints, the compact sleep solve, and the stable-state fast path; the deterministic gathers cover the standard solve and the compact sleep path keeps upstream's kernels) and pushed e4d4a0a, which closes coverage gaps found while stress-testing the deterministic machinery:

• equality_flexstrain rows allocate through count-scan-emit instead of raw atomics; row placement raced once more than a handful of strain constraints were active per world
• the deterministic flex-contact count kernel now shares its rownnz computation with _efc_contact_init_flex through a common wp.func; the two could disagree and overlap adjacent J storage
• tendon-limit rows no longer double-reserve efc_nnz, which could raise spurious deterministic-mode overflow errors on tightly sized buffers
• contact sort keys are int64 with worldid as the leading radix; SDF contacts get unique geomcollisionids from the initpoint index and flex contacts get content-stable keys built from element and vertex ids, so key ties no longer fall back to racy emission order
• deterministic mode raises on collision-side buffer overflows (nacon, broadphase pairs, CCD, flex candidates) the same way it already does for nefc
• Newton termination uses an order-fixed reduction, matching the CG path

Determinism tests grew from 48 to 55; the full suite passes (1250 passed, 18 skipped).

[thowell]

@mar-yan24 @johnnynunez if you are interested in completing the determinism feature please rebase this pr on main containing #1687 and focus the updated pr on the contact sorting and constraint sorting. these sorting routines should follow the implemented determinism api outline. please let us know if you have any questions. thanks!

[thowell]

lets move this code to collision_driver_test.py

[thowell]

we can drop solver=Newton since this is the default solver

[thowell]

nit: formatting

[thowell]

_import_mj_contacts?

[thowell]

a few more recommendations focused on sort.py and sort_test.py

mujoco_warp/_src/sort.py

1. mujoco_warp/_src/sort.py (lines 148–168): Single compare evaluation and simplified destination index in _merge

   • Issue: In each binary-search iteration, compare(keys_in, worldid, other, source) is unconditionally evaluated first, and then right-half threads (not left) evaluate compare(keys_in, worldid, source, other) a second time, doubling key comparisons and introducing intra-warp branch divergence when width < 32. Additionally, dest = start + (i - wp.where(left, start, middle)) + lo - other_start simplifies algebraically to i - middle + lo in both branches (start + (i - start) + lo - middle when left and start + (i - middle) + lo - start when not left).
   • Recommendation: Select lhs and rhs via wp.where so compare is called once per binary-search step without branch divergence, remove other_start, and write directly to order_out[worldid, i - middle + lo].
   • Example:

     lo = wp.where(left, middle, start)
     hi = wp.where(left, end, middle)
     source = order_in[worldid, i]
     while lo < hi:
       mid = (lo + hi) // 2
       other = order_in[worldid, mid]
       lhs = wp.where(left, other, source)
       rhs = wp.where(left, source, other)
       # Stable merge of *identical* keys, not a key based on emission order.
       if compare(keys_in, worldid, lhs, rhs) == left:
         lo = mid + 1
       else:
         hi = mid
     order_out[worldid, i - middle + lo] = source

2. mujoco_warp/_src/sort.py (lines 51–58, 188–215, 434–474): Replace eager 43-column key materialization (_contact_key_field) with a lazy @wp.struct (_ContactKeys)

   • Issue: contact_permutation allocates a (1, d.naconmax, 43) uint32 key buffer and launches _contact_key_field up to 17 separate times to pre-convert every scalar, vector, and matrix contact field before sorting, even though most contact comparisons resolve on worldid, geom, or pos.
   • Recommendation: Mirror _ConstraintKeys by defining a @wp.struct _ContactKeys holding views of d.contact fields and comparing fields lazily in _contact_less. Because geom is compared before flex/elem/vert, ga == gb is already established when reaching the flex fields, so rigid-rigid contacts (ga[0] >= 0 and ga[1] >= 0) skip reading flex, elem, and vert entirely, and _contact_key_field (along with all 17 kernel launches and the 43 * naconmax temporary buffer) can be deleted.
   • Example:

     @wp.struct
     class _ContactKeys:
       worldid: wp.array[int]
       geom: wp.array[wp.vec2i]
       flex: wp.array[wp.vec2i]
       elem: wp.array[wp.vec2i]
       vert: wp.array[wp.vec2i]
       pos: wp.array2d[wp.uint32]
       frame: wp.array3d[wp.uint32]
       dist: wp.array[wp.uint32]
       dim: wp.array[int]
       type: wp.array[int]
       includemargin: wp.array[wp.uint32]
       friction: wp.array2d[wp.uint32]
       solref: wp.array2d[wp.uint32]
       solreffriction: wp.array2d[wp.uint32]
       solimp: wp.array2d[wp.uint32]
       adhesion: wp.array[wp.uint32]
       geomcollisionid: wp.array[int]
     
     
     @wp.func
     def _contact_less(keys: _ContactKeys, worldid: int, a: int, b: int) -> bool:
       if keys.worldid[a] != keys.worldid[b]:
         return keys.worldid[a] < keys.worldid[b]
       ga = keys.geom[a]
       gb = keys.geom[b]
       for i in range(2):
         if ga[i] != gb[i]:
           return ga[i] < gb[i]
       # Rigid writers do not initialize flex feature fields; ignore those slots when geom >= 0.
       if ga[0] < 0 or ga[1] < 0:
         for i in range(2):
           if ga[i] < 0 and keys.flex[a][i] != keys.flex[b][i]:
             return keys.flex[a][i] < keys.flex[b][i]
         for i in range(2):
           if ga[i] < 0 and keys.elem[a][i] != keys.elem[b][i]:
             return keys.elem[a][i] < keys.elem[b][i]
         for i in range(2):
           if ga[i] < 0 and keys.vert[a][i] != keys.vert[b][i]:
             return keys.vert[a][i] < keys.vert[b][i]
       for i in range(3):
         if keys.pos[a, i] != keys.pos[b, i]:
           return _float_key(keys.pos[a, i]) < _float_key(keys.pos[b, i])
       for i in range(3):
         for j in range(3):
           if keys.frame[a, i, j] != keys.frame[b, i, j]:
             return _float_key(keys.frame[a, i, j]) < _float_key(keys.frame[b, i, j])
       if keys.dist[a] != keys.dist[b]:
         return _float_key(keys.dist[a]) < _float_key(keys.dist[b])
       if keys.dim[a] != keys.dim[b]:
         return keys.dim[a] < keys.dim[b]
       if keys.type[a] != keys.type[b]:
         return keys.type[a] < keys.type[b]
       if keys.includemargin[a] != keys.includemargin[b]:
         return _float_key(keys.includemargin[a]) < _float_key(keys.includemargin[b])
       for i in range(5):
         if keys.friction[a, i] != keys.friction[b, i]:
           return _float_key(keys.friction[a, i]) < _float_key(keys.friction[b, i])
       for i in range(2):
         if keys.solref[a, i] != keys.solref[b, i]:
           return _float_key(keys.solref[a, i]) < _float_key(keys.solref[b, i])
       for i in range(2):
         if keys.solreffriction[a, i] != keys.solreffriction[b, i]:
           return _float_key(keys.solreffriction[a, i]) < _float_key(keys.solreffriction[b, i])
       for i in range(5):
         if keys.solimp[a, i] != keys.solimp[b, i]:
           return _float_key(keys.solimp[a, i]) < _float_key(keys.solimp[b, i])
       if keys.adhesion[a] != keys.adhesion[b]:
         return _float_key(keys.adhesion[a]) < _float_key(keys.adhesion[b])
       return keys.geomcollisionid[a] < keys.geomcollisionid[b]

3. mujoco_warp/_src/sort.py (lines 174–185, 413–432): View rows as wp.uint32 in row_view and type _gather concretely

   • Issue: row_view accesses private Warp internals (value.dtype._wp_scalar_type_) and returns arrays of varying scalar dtypes (int32, float32), requiring wp.array3d[Any] annotations on _gather and compiling separate _gather kernel specializations per scalar type even though _gather only copies raw 32-bit words.
   • Recommendation: Always view arrays as wp.uint32 in row_view and annotate value_in / value_out in _gather with concrete wp.array3d[wp.uint32].
   • Example:

     @wp.kernel
     def _gather(
       # In:
       count_in: wp.array[int],
       order_in: wp.array2d[int],
       value_in: wp.array3d[wp.uint32],
       # Out:
       value_out: wp.array3d[wp.uint32],
     ):
       worldid, row, col = wp.tid()
       if row < count_in[worldid] and row < order_in.shape[1]:
         value_out[worldid, row, col] = value_in[worldid, order_in[worldid, row], col]
     
     
     def row_view(value: wp.array, nworld: int, capacity: int) -> wp.array:
       """View a contiguous array as scalar rows without moving its storage."""
       view = value.view(wp.uint32)
       width = view.size // max(nworld * capacity, 1)
       # Keep the backing allocation alive when a captured gather retains the view.
       return view.reshape((nworld, capacity, width))

4. mujoco_warp/_src/sort.py (lines 80–117): Check != on raw uint32 bits before computing _float_key in _constraint_less

   • Issue: _constraint_less computes _float_key on both operands before testing inequality across J and the 6 scalar tie-breaker fields (pos, margin, D, vel, aref, frictionloss), repeating the same 4-line block 8 times. Because _float_key is a bijection on uint32, _float_key(a) != _float_key(b) holds if and only if a != b.
   • Recommendation: Compare raw uint32 equality first and only evaluate _float_key when the two values differ.
   • Example:

     if keys.pos[worldid, a] != keys.pos[worldid, b]:
       return _float_key(keys.pos[worldid, a]) < _float_key(keys.pos[worldid, b])
     if keys.margin[worldid, a] != keys.margin[worldid, b]:
       return _float_key(keys.margin[worldid, a]) < _float_key(keys.margin[worldid, b])
     if keys.D[worldid, a] != keys.D[worldid, b]:
       return _float_key(keys.D[worldid, a]) < _float_key(keys.D[worldid, b])
     if keys.vel[worldid, a] != keys.vel[worldid, b]:
       return _float_key(keys.vel[worldid, a]) < _float_key(keys.vel[worldid, b])
     if keys.aref[worldid, a] != keys.aref[worldid, b]:
       return _float_key(keys.aref[worldid, a]) < _float_key(keys.aref[worldid, b])
     if keys.frictionloss[worldid, a] != keys.frictionloss[worldid, b]:
       return _float_key(keys.frictionloss[worldid, a]) < _float_key(keys.frictionloss[worldid, b])
     return False

5. mujoco_warp/_src/sort.py (lines 376–398, 504–508): Remove redundant loop in _remap_blocks and redundant if d.naconmax: branch in constraints

   • Issue:
     • In _remap_blocks, the loop for i in range(1, nrow_in[worldid, block]): first = wp.min(first, inverse_in[worldid, start + i]) is unnecessary because multi-row groups (CONNECT, WELD, and contacts) are assigned identical (type, subtype, id/rank) primary keys with strictly increasing component = 0 .. nrow - 1, so start (component == 0) always maps to the minimum post-sort index inverse_in[worldid, start].
     • In constraints, the if d.naconmax: ... else: ... branch is unnecessary because contact_permutation(d) already handles d.naconmax == 0.
   • Recommendation: Index inverse_in[worldid, efc_jtdaj_adr_in[worldid, block]] directly in _remap_blocks (dropping nrow_in) and call contact_order = contact_permutation(d) unconditionally in constraints.
   • Example:

     @wp.kernel
     def _remap_blocks(
       # Data in:
       efc_jtdaj_adr_in: wp.array2d[int],
       efc_jtdaj_nblock_in: wp.array[int],
       # In:
       count_in: wp.array[int],
       inverse_in: wp.array2d[int],
       # Data out:
       efc_jtdaj_adr_out: wp.array2d[int],
       # Out:
       nblock_out: wp.array[int],
     ):
       worldid, block = wp.tid()
       if block == 0:
         nblock_out[worldid] = wp.where(count_in[worldid] > 0, efc_jtdaj_nblock_in[worldid], 0)
       if count_in[worldid] == 0 or block >= efc_jtdaj_nblock_in[worldid]:
         return
       efc_jtdaj_adr_out[worldid, block] = inverse_in[worldid, efc_jtdaj_adr_in[worldid, block]]

6. mujoco_warp/_src/sort.py (lines 240–398): Use canonical Model / Data parameter names and # Data in: / # Data out: section headers

   • Issue: Kernels in sort.py pass Model and Data fields under generic # In: / # Out: names (type_in, id_in, worldid_in, efc_address_in, dim_in, is_pyramidal, address_out, rownnz_out, adr_in, nblock_in), bypassing kernel_analyzer validation of Model and Data field types and ordering.
   • Recommendation: Use the canonical Model and Data parameter names (opt_cone, contact_dim_in, contact_efc_address_in, contact_worldid_in, efc_type_in, efc_id_in, efc_J_rownnz_in, efc_jtdaj_adr_in, efc_jtdaj_nblock_in, contact_efc_address_out, efc_J_rownnz_out, efc_J_rowadr_out, efc_J_colind_out, efc_J_out, efc_jtdaj_adr_out) under # Model:, # Data in:, and # Data out:.
   • Example:

     @wp.kernel
     def _contact_constraint_primary(
       # Model:
       opt_cone: int,
       # Data in:
       contact_dim_in: wp.array[int],
       contact_efc_address_in: wp.array2d[int],
       contact_worldid_in: wp.array[int],
       efc_type_in: wp.array2d[int],
       efc_id_in: wp.array2d[int],
       nacon_in: wp.array[int],
       # In:
       count_in: wp.array[int],
       order_in: wp.array2d[int],
       # Out:
       keys_out: wp.array2d[wp.vec4i],
     ):

---

mujoco_warp/_src/sort_test.py (and contact_address_test.py, imported_contact_test.py)

7. mujoco_warp/_src/contact_address_test.py, mujoco_warp/_src/imported_contact_test.py, and mujoco_warp/_src/sort_test.py: Consolidate contact address and imported contact tests into sort_test.py

   • Issue: contact_address_test.py and imported_contact_test.py do not correspond to modules of the same name and compile flex collision/constraint kernels in separate test processes (~60s each), while testing contact efc_address clearing/remapping and put_data contact initialization that are part of the canonical sorting pipeline.
   • Recommendation: Move the tests from contact_address_test.py and imported_contact_test.py into sort_test.py (reusing the flex fixture already compiled in sort_test.py) and remove the two standalone test files.

8. mujoco_warp/_src/sort_test.py (lines 99–101, 112–114, 262, 278–280, 314, 331–332, 399–405, 460–494, 564–624, 734–755): Avoid getattr / setattr in tests

   • Issue: sort_test.py uses getattr and setattr 22 times across test methods and helpers to access known attributes on d, d.contact, d.efc, mjd, and mujoco.mjtEq.
   • Recommendation: Use explicit attribute references or dicts mapping field names to array attributes (e.g. {"force": d.efc.force, "state": d.efc.state, "island": d.efc.island} and {"connect": mujoco.mjtEq.mjEQ_CONNECT, "weld": mujoco.mjtEq.mjEQ_WELD}).
   • Example:

     for array in (d.efc.force, d.efc.state, d.efc.island):
       array.assign(np.tile(np.arange(array.shape[1]) + 10, (nworld, 1)))

9. mujoco_warp/_src/sort_test.py (lines 39–40, 306–308, 458–488): Replace np.random.default_rng with deterministic permutations

   • Issue: test_constraint_scalar_keys, test_interleaved_constraints, and test_contact_permutation use np.random.default_rng to generate test inputs and permutations.
   • Recommendation: Use fixed deterministic index permutations (such as reversals and np.roll shifts) instead of np.random.default_rng.
   • Example:

     base = np.array([5, 2, 7, 0, 6, 1, 4, 3])
     values = np.stack([bits[np.roll(base, w)] for w in range(nworld)])

10. mujoco_warp/_src/sort_test.py (lines 424–435): Remove if wp.get_device().is_cuda: guard in test_forward_capture

    • Issue: test_forward_capture conditionally skips wp.ScopedCapture() when not on a CUDA device instead of using @absltest.skipIf(not wp.get_device().is_cuda, ...) or capturing unconditionally when CUDA is required.
    • Recommendation: Capture and replay wp.ScopedCapture() unconditionally in test_forward_capture (or decorate with @absltest.skipIf(not wp.get_device().is_cuda, "ScopedCapture requires CUDA")).
    • Example:

      with wp.ScopedCapture() as capture:
        mjw.forward(m, d)
      for i in range(6):
        restore()
        if i % 2:
          wp.capture_launch(capture.graph)
        else:
          mjw.forward(m, d)

11. mujoco_warp/_src/sort_test.py (lines 272–303, 396–413, 497–644): Deduplicate JTDAJ block contiguity assertions, hoist .numpy() transfers out of loops, and simplify test_equality_forces

    • Issue:
      • test_interleaved_constraints, test_forward_capture, and test_equality_forces repeat the same sparse JTDAJ block contiguity check (starts = d.efc.jtdaj_adr.numpy()[world, :nb], sizes = d.efc.jtdaj_nrow.numpy()[world, :nb], assert_array_equal(starts, np.cumsum(sizes) - sizes)), and call .numpy() inside per-world loops.
      • test_equality_forces is parameterized over stage=("solved", "forward") (24 subtests) where stage="solved" manually shuffles d.efc in NumPy and calls constraint._sort_constraints, even though the XML already interleaves <weld> and <connect> equalities so mjw.forward with DeterminismType.CONSTRAINT exercises the exact same group-reordering path end-to-end.
    • Recommendation: Extract a shared _assert_jtdaj_blocks(test, d, world, n) helper, hoist .numpy() calls outside per-world loops, and drop the stage parameterization in test_equality_forces (halving its subtest count from 24 to 12).
    • Example:

      def _assert_jtdaj_blocks(test, d, world: int, n: int):
        nblock = int(d.efc.jtdaj_nblock.numpy()[world])
        starts = d.efc.jtdaj_adr.numpy()[world, :nblock]
        widths = d.efc.jtdaj_nrow.numpy()[world, :nblock]
        np.testing.assert_array_equal(starts, np.cumsum(widths) - widths)
        test.assertEqual(int(np.sum(widths)), n)
        return starts, widths

12. mujoco_warp/_src/sort_test.py (lines 183, 208, 347, 510): Format inline MJCF XMLs with one element per line and omit default attributes

    • Issue: Inline XMLs in sort_test.py include unused <compiler angle="radian"/> (in test_forward_capture, which has no angles), default mass="1" on geoms/flexes, and long single-line <mesh> / <flexcomp> tags.
    • Recommendation: Remove unused <compiler angle="radian"/> and default mass="1" attributes, and format inline XMLs concisely.

[thowell]

lets also test graph capture here

mujoco_warp/_src/sort_test.py: Add direct CUDA graph capture tests for contact and constraint sorting across varying permutations and counts

• Issue: test_forward_capture captures mjw.forward(m, d) as a whole (guarded by if wp.get_device().is_cuda:) and replays it on the exact same restored initial state, while test_contact_permutation, test_unused_contact_addresses, and test_interleaved_constraints invoke collision_driver._sort_contacts(m, d) and constraint._sort_constraints(m, d) eagerly without wp.ScopedCapture(). Because sort.gather relies on captured wp.clone(value) copies and value.view(...).reshape(...) views, replaying only on an unchanged input state does not verify that a captured sort graph reads updated contact/constraint permutations or updated device counts (d.nacon, d.nefc) on subsequent wp.capture_launch replays.
• Recommendation: Capture collision_driver._sort_contacts(m, d) and constraint._sort_constraints(m, d) directly with wp.ScopedCapture(), replay wp.capture_launch(capture.graph) across distinct input permutations and active counts (d.nacon, d.nefc), and remove the if wp.get_device().is_cuda: guard in test_forward_capture.
• Example:

  # Contact sort graph capture: replay across distinct permutations and active counts
  with wp.ScopedCapture() as capture:
    collision_driver._sort_contacts(m, d)
  for shift in (1, 3, 5):
    permutation = np.roll(np.arange(count)[::-1], shift)
    for array, value in reference:
      shuffled = value.copy()
      shuffled[:count] = value[permutation]
      array.assign(shuffled)
    wp.capture_launch(capture.graph)
    for array, value in reference:
      np.testing.assert_array_equal(array.numpy(), value)
  
  # Constraint sort graph capture: replay across distinct contact/constraint permutations
  with wp.ScopedCapture() as capture:
    constraint.make_constraint(m, d)
  for shift in (1, 2, 3):
    order = np.roll(np.arange(count)[::-1], shift)
    for array, value in contact_fields:
      shuffled = value.copy()
      shuffled[:count] = value[order]
      array.assign(shuffled)
    wp.capture_launch(capture.graph)
    for name, value in snapshot().items():
      np.testing.assert_array_equal(value, reference[name], err_msg=str(name))

[thowell]

@johnnynunez what is the status of this pr? thanks!

[johnnynunez]

@thowell Updated in 3e9aab1: merged main and addressed the second review round and the nine inline threads. Tests moved to sort_test.py as in your later comment, with capture replays over permutations and counts. Main needed three additions: surfacevel, dof_frictionadr and flex strain blocks follow the sort. Humanoid, 8192 worlds, 10 alternating pairs, main vs PR: 9.132M vs 9.126M steps/s (-0.06%), make_constraint +0.43%; three_humanoids +0.02%. I kept mass="1" in test_equality_forces: without it cfrc_int misses 1e-5. Locally 2725 GPU and 2709 CPU tests pass; CI needs approval.

[thowell]

Thanks @johnnynunez! A few items to resolve before merging:

1. mujoco_warp/_src/sort_test.py (lines 1303--1416): Fix flaky/failing test_contact_rank_paths and test_sparse_offsets_paths

• Issue: test_contact_rank_paths and test_sparse_offsets_paths run a fresh mjw.forward(m, d) (with DeterminismType.NONE) inside each iteration of the path-comparison loop and then assert exact equality on all vars(d.efc) arrays. Because sort.constraints does not touch solver workspace d.efc.Ma (and mjw.forward populates d.efc.Ma via nondeterministic float32 atomics), d.efc.Ma differs by ~1 ULP across the two independent forward calls, failing test_contact_rank_paths2 and test_sparse_offsets_paths2.
• Recommendation: Build the unsorted state once before the loop, snapshot initial = [(a, wp.clone(a)) for a in _arrays(d)], and restore d before running each sort.constraints(m, d) path so both paths operate on identical inputs.

    constraint.make_constraint(m, d)
    before = d.efc.J.numpy().copy()
    initial = [(a, wp.clone(a)) for a in _arrays(d)]
    results = []
    for row_rank_njmax in (sort._ROW_RANK_NJMAX, 0):
      for target, value in initial:
        wp.copy(target, value)
      with mock.patch.object(sort, "_ROW_RANK_NJMAX", row_rank_njmax):
        sort.constraints(m, d)

2. mujoco_warp/_src/sort.py (lines 18--19, 50--86, 284--334, 996--1002, 1032--1080, 1200--1204) & mujoco_warp/_src/sort_test.py (lines 19, 1208): Remove getattr / setattr, SimpleNamespace, and Any

• Issue: Commit c7b7d71 reintroduced getattr, setattr, SimpleNamespace, and string-name field tuples (_CONTACT_FIELDS, _CONSTRAINT_FIELDS) to populate _ContactRows and _ConstraintRows (and copy in test_contact_bucket_ties), and uses Any in _rank, _merge, _contact_keys, and _contact_rows.
• Recommendation: Assign _ContactRows and _ConstraintRows fields directly (d.efc.J_rownnz already has shape (nworld, 0) in dense mode), allocate copy_rows directly as a _ContactRows instance (removing SimpleNamespace, _CONTACT_FIELDS, _CONSTRAINT_FIELDS, getattr, and setattr), and resolve key_type = compare.input_types["keys"] in _rank and _merge (plus | _BlockKeys.cls on permutation) to drop Any.

@cache_kernel
def _rank(compare: Callable):
  key_type = compare.input_types["keys"]

  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
  def kernel(
    # In:
    count_in: wp.array[int],
    keys_in: key_type,
    width: int,
    # Out:
    order_out: wp.array2d[int],
    inverse_out: wp.array2d[int],
  ):
    ...

def _contact_keys(contact: types.Contact | _ContactRows.cls) -> _ContactKeys.cls:
  ...


def _empty_contact_rows(contact: types.Contact) -> _ContactRows.cls:
  rows = _ContactRows()
  rows.dist = wp.empty_like(contact.dist)
  rows.pos = wp.empty_like(contact.pos)
  rows.frame = wp.empty_like(contact.frame)
  rows.includemargin = wp.empty_like(contact.includemargin)
  rows.friction = wp.empty_like(contact.friction)
  rows.solref = wp.empty_like(contact.solref)
  rows.solreffriction = wp.empty_like(contact.solreffriction)
  rows.solimp = wp.empty_like(contact.solimp)
  rows.dim = wp.empty_like(contact.dim)
  rows.geom = wp.empty_like(contact.geom)
  rows.flex = wp.empty_like(contact.flex)
  rows.elem = wp.empty_like(contact.elem)
  rows.vert = wp.empty_like(contact.vert)
  rows.efc_address = wp.empty_like(contact.efc_address)
  rows.worldid = wp.empty_like(contact.worldid)
  rows.type = wp.empty_like(contact.type)
  rows.geomcollisionid = wp.empty_like(contact.geomcollisionid)
  rows.adhesion = wp.empty_like(contact.adhesion)
  rows.surfacevel = wp.empty_like(contact.surfacevel)
  return rows

3. mujoco_warp/_src/sort_test.py (lines 1092--1094, 1135--1146, 1238--1240): Replace np.random.default_rng with deterministic expressions

• Issue: test_permutation_blocks, test_contact_paths, and test_sparse_move reintroduce np.random.default_rng.
• Recommendation: Use deterministic modular index arithmetic and np.roll permutations instead of np.random.default_rng.

    # In test_permutation_blocks:
    primary = ((np.arange(nworld * capacity * 4, dtype=np.int32) * 7 + 3) % 3).reshape(nworld, capacity, 4)

    # In test_sparse_move:
    colind = ((np.arange(nworld * capacity, dtype=np.int32) * 17 + 5) % 100).reshape(nworld, capacity)
    J = (np.arange(1, nworld * capacity + 1, dtype=np.float32) * 0.125).reshape(nworld, capacity)

4. mujoco_warp/_src/collision_driver.py (lines 1155--1157, 1238--1239): Inline sort.contacts(d) and remove _sort_contacts

• Issue: _sort_contacts(m: Model, d: Data) is now a 1-line wrapper (sort.contacts(d)) with an unused m argument.
• Recommendation: Remove _sort_contacts and call sort.contacts(d) directly in collision (matching sort.constraints in constraint.py) and in sort_test.py.

  if m.opt.deterministic & DeterminismType.CONTACTS:
    sort.contacts(d)

5. mujoco_warp/_src/collision_flex.py (lines 2875--2876): Remove duplicate contact_efc_address_out reset

• Issue: main (Reserve sparse Jacobian storage before allocating constraint rows #1793) already resets contact_efc_address_out[id_] to -1 at lines 2892--2893 of _write_filtered_contacts, so keeping lines 2875--2876 writes -1 to global memory twice per flex contact.
• Recommendation: Remove lines 2875--2876 so collision_flex.py matches main.

    contact_dist_out[id_] = cand_dist[i]
    contact_pos_out[id_] = cand_pos[i]

6. mujoco_warp/_src/sort.py (lines 659--692, 1181): Zero stale d.efc.Jqvel on non-contact rows

• Issue: Non-contact constraint kernels compute Jqvel in registers without writing d.efc.Jqvel, so packing/unpacking d.efc.Jqvel moves stale prior-step values from nondeterministic pre-sort slots into sorted non-contact rows.
• Recommendation: Zero efc_Jqvel_out[worldid, row] when row_type < int(types.ConstraintType.CONTACT_FRICTIONLESS) in _constraint_primary.

  row_type = efc_type_in[worldid, row]
  id = efc_id_in[worldid, row]
  if row_type < int(types.ConstraintType.CONTACT_FRICTIONLESS):
    efc_Jqvel_out[worldid, row] = 0.0

[johnnynunez]

Thanks @thowell, all six are resolved in 6e3c073:

1. The path tests build the unsorted state once and restore a snapshot before each path.
2. sort.py has no getattr/setattr, SimpleNamespace, Any or field tuples; structs are assigned field by field.
3. No default_rng in the three tests.
4. _sort_contacts is removed; collision calls sort.contacts(d).
5. collision_flex.py equals main again.
6. _constraint_primary zeroes efc.Jqvel on non-contact rows; test_noncontact_jqvel covers it.

The sorted output is bitwise identical to the previous head.

[thowell]

lgtm, thanks @johnnynunez!

[thowell]

@johnnynunez we just need to resolve the merge conflict

Cause

PR #1422 adds from mujoco_warp._src import sort directly above from mujoco_warp._src import util_pkg in mujoco_warp/_src/determinism_test.py, while main (#1927) removed the unused util_pkg import on that adjacent line.

Conflict Hunk

from mujoco_warp._src import render_util
from mujoco_warp._src import set_const
<<<<<<< feature/determinism
from mujoco_warp._src import sort
from mujoco_warp._src import util_pkg
=======
>>>>>>> main
from mujoco_warp._src import warp_util

Resolution

Keep from mujoco_warp._src import sort and omit from mujoco_warp._src import util_pkg:

from mujoco_warp._src import render_util
from mujoco_warp._src import set_const
from mujoco_warp._src import sort
from mujoco_warp._src import warp_util

(Whether resolved via GitHub's web conflict editor, git merge upstream/main, or git rebase upstream/main, this 1-line import change is the only conflict across the entire PR.)

[johnnynunez]

Done and thank your for your patient @thowell

[thowell]

@johnnynunez thanks for getting this feature over the finish line 🚀

[thowell]

@mar-yan24 thanks for your contributions to this feature! 🚀