Repository navigation
optimize flex performance - #1620
Conversation
…xis tests and add dimension flags to Model.
…rallel reduction kernels.
…n candidate count exceeds contact limits.
…tive degrees of freedom.
…recomputed topological adjacency bitmask.
d53011c to
5c28ebd
Compare
5c28ebd to
e161a57
Compare
smallquail
left a comment
There was a problem hiding this comment.
Summary in plain terms. The solver changes (equality-row pre-reduction, fused beta) are correct: I checked the algebra, the row ordering including islands, and the existing constraints.xml parity test covers them. The SAT prefilter and the parallel FPS are the right ideas and match the serial semantics. Three things to fix before merge, then nits:
- The new spatial sort key loses precision in float32 and can make the dedup scan stop early. Inline at the key computation.
- The adjacency mask is N^2 memory and the runtime check it replaces is still there, so nothing is saved. Pick one. Inline at the sweep and at the build.
- The benchmark table doesn't isolate the code change: the baseline ran with overflow and the configs changed in the same PR. Inline on the benchmark config.
Description accuracy: the "compact workspace" commit skips a wasted eager allocation rather than allocating "dynamically", and there are no weld exclusions in the self-collision path. Details inline.
| # Clamped monotonic 1D integer mapping with 1 um resolution. | ||
| p = cand_pos[i] | ||
| s = (p[0] + p[1] + p[2]) * wp.static(1.0 / math.sqrt(3.0)) | ||
| spatial_key = wp.clamp(wp.int64((s + 100.0) * 1000000.0), wp.int64(0), wp.int64(2147483647)) |
There was a problem hiding this comment.
This is float32 arithmetic on a value near 1e8 (s ~ +/-1 m gives (s+100)*1e6 ~ 1.01e8), where the ulp is 8. So the key is quantized to ~8 um, not 1 um, and the neighbour scan below breaks when keys differ by >= 1001. Two candidates 999 um apart along the projection can land 1016 apart in key space and never get compared, so a dominated duplicate survives. Mild consequence (an extra candidate, not a missing contact), but live with the shipped constants.
Suggest doing the offset in integer space: wp.int64(s * 1e6) + wp.int64(100000000). The float product then stays under 2e6 (ulp 0.25) and the +100 m clamp range is unchanged.
| my_key = sort_key[si] | ||
| my_group = my_key >> wp.int64(32) | ||
| my_spatial = my_key & wp.int64(0x7FFFFFFF) | ||
| eps_key = wp.int64(epsilon * 1000000.0) + wp.int64(1) |
There was a problem hiding this comment.
Exact today only because 1e-3 * 1e6 is an integer. With a fractional epsilon * 1e6 the truncation makes the window one unit too small (e.g. 5.5 gives keys 11 and 5 for points 5.4 apart). ceil(...) + 1 makes it safe for any epsilon.
| worldelemid += nsweep_in | ||
|
|
||
| if wp.static(is_self): | ||
| word = flex_selfcollide_mask[elem1, elem2 >> 5] |
There was a problem hiding this comment.
The mask encodes exactly the truth table of _exclude_self_collision (same vertex, or same body with b >= 0), which still runs a few lines below. So for self pairs the runtime check can never return True after this point (it is dead), and the mask has bought one 4-byte load in exchange for the table.
Either delete the runtime check on the self path (then the mask does something, and it should show up in a measurement) or delete the mask. Also, on "maintaining dynamic weld exclusions" in the description: there is no weld logic in this path; _exclude_self_collision does not read body_weldid.
| m.flex_vertflexid = flex_vertflexid | ||
| m.flex_shelladr = flex_shelladr | ||
|
|
||
| # Precompute flex_selfcollide_mask: topologically adjacent elements sharing vertices/bodies |
There was a problem hiding this comment.
This table is nflexelem^2 / 8 bytes of model memory. On the benchmark cloth (1682 elements) it is 0.36 MB and builds in 0.06 s, so it is invisible there, but it is 0.8 GB at 200x200 and 4 GB at 300x300, for a per-element neighbour set of ~12. A per-element neighbour list would be a few hundred bytes per element. Worth an ablation of SAT-only vs SAT+mask on the benchmarks before committing to the quadratic form.
| "map_iefc2efc": None, | ||
| "dof_islandid": None, | ||
| "efc_islandid": None, | ||
| # compact arrays |
There was a problem hiding this comment.
What this does: keep the generic field loop below from eagerly allocating full (nworld, nvmax_pad, nvmax_pad) arrays for cM, cqLD, cJ, ... that _allocate_compact_arrays immediately replaces. That is a real fix for transient peak memory (gigabytes at 8192 worlds), but the commit message and PR bullet describe something else ("dynamically based on active DOFs"). Suggest saying what it does, and deriving the skip list from _allocate_compact_arrays instead of two hand-maintained lists of fifteen; a new c* field silently regresses otherwise.
| flex_group_start_indices = wp.full(nmax_groups, -1, dtype=int) | ||
| flex_fps_min_dist = wp.empty(d.naconmax, dtype=float) | ||
| flex_num_groups = wp.zeros(1, dtype=int) | ||
| fps_scratch_dist = wp.empty((nmax_groups, 64), dtype=float) |
There was a problem hiding this comment.
64 is hard-coded here, in the three dim=(nmax_groups, 64) launches, and as the stride and reduce width inside five kernels. One named constant.
| min_d = d_val | ||
| sel_cidx = c_idx | ||
| elif d_val == min_d: | ||
| if _tie_break_fps(c_idx, sel_cidx, cand_elem): |
There was a problem hiding this comment.
_tie_break_fps orders on (elem1, elem2) only, so candidates from the same element pair compare equal and the winner depends on scan order, which now differs from the serial kernel. Adding curr_idx < sel_idx as the final key makes it a strict total order and the parallel result bit-identical to the serial one regardless of partitioning.
| mjw.kinematics(m, d) | ||
| mjw.collision(m, d) | ||
| nacon = int(d.nacon.numpy()[0]) | ||
| self.assertEqual(nacon, types.MJ_MAXCONPAIR * nworld) |
There was a problem hiding this comment.
This checks the cap, not the selection. Two cheap, strong tests would cover the new code properly: (a) on a self-colliding cloth, the contact set with the SAT filter equals the set without it (it is a pure prefilter, so any difference is a bug); (b) parallel FPS selects the same candidates as a serial/numpy reference.
| "nworld": 32, | ||
| "nconmax": 2200, | ||
| "njmax": 8000, | ||
| "nworld": 256, |
There was a problem hiding this comment.
The "Before" column ran with buffer overflow (dropped contacts mean less work and different physics), and this PR changes nworld, nconmax, njmax, and adds nccdmax in the same change. The 6-12x therefore mixes a code change, a batch-size change, and an overflow/no-overflow comparison. For a number that isolates the code: baseline with buffers raised until overflow is zero, same nworld as the optimized run, Jacobian setting stated.
|
I posted some comments, nice speedups! |
Removes the N^2 / 8 bitmask matrix and relies on the 2D triangle separating axis prefilter and existing register-based element vertex loop for self-collision exclusions. This avoids multi-gigabyte model allocations on fine meshes and speeds up both model construction and broadphase sweep execution.
|
the pr is updated Below is a summary of the changes responding to each review comment. 1. Spatial Sort Key Precision
Applied. In 2. Fractional Epsilon Safety in Neighbor Search
Applied. In 3 & 4. Adjacency Bitmask Redundancy & Quadratic Memory Footprint
Removed the quadratic bitmask entirely. We ran an ablation measuring SAT-only vs. SAT+mask:
We cleaned up 5. Compact Solver Workspace Eager Allocation
Applied. Defined a single canonical 6. FPS Block Size Named Constant
Applied. Defined 7. Strict Total Order for FPS Tie-Breaking
Applied. Appended 8. Unit Tests for SAT Prefilter and Parallel FPS Parity
Applied. In
|
smallquail
left a comment
There was a problem hiding this comment.
Thanks: all seven code items land cleanly, and settling the mask question with an ablation was exactly the right call.
test_parallel_fps_numpy_parityis failing on CI, and the test is what is wrong, not the kernel. Details inline on the assertion.- PR description is stale. Bullet 4 still describes the removed adjacency bitmask and "dynamic weld exclusions".
- Benchmark table is still not apples-to-apples (unchanged from the earlier comment on
benchmarks/cloth/__init__.py): the baseline ran through overflow, andnworld,nconmax,njmaxandnccdmaxchanged in the same PR. One added row, with the baseline buffers raised until overflow is zero, the samenworldas the optimized run, and the Jacobian setting stated, would make the headline number defensible. - Nit inline on
collision(): the newenable_satkeyword exists only for the test.
|
|
||
| warp_selected = sorted([c_idx for c_idx in group_cands if cand_active[c_idx] == 1]) | ||
| np_selected = sorted(selected) | ||
| self.assertEqual(warp_selected, np_selected) |
There was a problem hiding this comment.
The reference does not compute what the kernel computes, so exact equality cannot hold on this scene:
- The reference ranks by squared distance (
np.sum((a-b)**2), float32, numpy summation order); the kernels rank bywp.length(...)(float32sqrtofx*x+y*y+z*z). Same ordering in exact arithmetic, different rounding on ties. - Ties are decided by exact float equality in both, and this scene, an 8x8 lattice stacked exactly 1 cm over an identical lattice, has many candidates at geometrically identical distances. One tie breaks differently (312 vs 484 in the CI log) and the greedy chain diverges from there.
- The reference also has no
max_d <= 0.0stop, which the kernel has in_parallel_fps_resolve_max; on coincident positions the two would select different counts.
Suggest: perturb the fixture so exact ties cannot occur (qpos_noise= on test_data.fixture), mirror the kernel's sqrt and stop rule in the reference, and keep the exact-match assertion. With a strict total order on the tie-break and no exact ties, bit-identical selection is a sound invariant.
| m: Model, | ||
| d: Data, | ||
| awake_prev: Optional[wp.array] = None, | ||
| enable_sat: bool = True, |
There was a problem hiding this comment.
This keyword is on the public collision() signature only so test_flex_sat_prefilter_conservative can toggle the filter. A module-level flag in collision_flex.py that the test patches would keep it off the API surface.
|
Left some minor comments. |
|
Request: mirror the contact-selection changes in MuJoCo C, so the two implementations keep producing the same contacts. The motivation is behavior preservation. MuJoCo C and mjWarp should return the same flex contact set from the same state, and today the place where they can differ is the per-pair contact cap ( Three items:
Items 1 and 2 are the ones that affect physics; item 3 is bookkeeping. |
|
I left one request regarding C mujoco, I think we should make sure we don't deviate in behavior between the two. |
…o C parity. Narrowphase candidate contacts are passed directly to farthest-point sampling without 1 mm spatial deduplication, matching MuJoCo C contact filtering semantics. The broadphase radix sort of candidate filter keys is now executed only when FPS contact sampling is enabled. Updated sphere cloth collision tests assert exact contact count parity against MuJoCo C.
|
updated the pr description including updated benchmark results comparing the before and after commits with modified settings to avoid overflow in order to directly compare performance. results with the updated settings for the latest pr commit are also provided. |
Brings the deterministic-mode branch (google-deepmind#1422 plus the inertia-matrix race fix) up to current main, which it had not been synced with since google-deepmind#1530. Conflict resolutions: - constraint.py: keep google-deepmind#1422's `deterministic` factory argument alongside main's `grid_stride=False`. Main's passive adhesion (google-deepmind#1550) makes a contact active when it is penetrating or adhesive, so google-deepmind#1422's per-contact row counters (`_efc_contact_count`, `_efc_contact_count_flex`) now take `flg_adhesion` and `adhesion` and apply the same rule as `_efc_contact_init`; otherwise the precomputed row offsets would disagree with the rows actually written. - solver.py: main replaced `_JTDAJ_sparse` with `_JTDACJ_sparse`, which adds elliptic-cone curvature terms. The deterministic `_JTDAJ_sparse_det` gather covers only the quadratic J^T D J term, so it is used for non-elliptic cones and elliptic cones fall back to `_JTDACJ_sparse`. Dropped the flex diagonal preconditioner kernels that main reverted in google-deepmind#1576. - collision_flex.py: take main's rewritten flex pipeline (google-deepmind#1496, google-deepmind#1583, google-deepmind#1620); google-deepmind#1422's tie-break in the old `_filter_flex_candidates_sorted` has no counterpart there, and the flex cases in determinism_test.py pass without it. - io.py, types.py, pyproject.toml: union of both sides (`opt.deterministic` plus main's new options and the `warn_overflow` bitmask property; main's version and dependency floors).
This PR introduces targeted performance optimizations for flexible body simulation in MuJoCo Warp, focusing on broadphase collision filtering, farthest point contact sampling, constraint solver linesearch, and conjugate gradient search direction updates.
During the sweep-and-prune broadphase, 2D flex elements are prefiltered using a separating axis test between planar triangle bounding disks, eliminating non-colliding element pairs before narrowphase dispatch. Farthest point contact sampling is parallelized using block-cooperative reductions with a strict total-order tie-breaker and early termination when candidate counts are within capacity, matching MuJoCo C contact filtering semantics. In the constraint solver, quadratic costs for invariant equality constraint rows are pre-reduced prior to the linesearch loop, and the Fletcher-Reeves beta calculation is fused into a single vectorized reduction kernel. In data initialization, eager full-size tensor allocation is skipped for compacted solver arrays that are subsequently populated by compact workspace setup.
Benchmark results below isolate code speedups on identical configurations with zero buffer overflows using dense Jacobians:
Cloth (32 worlds, nconmax=14000, nccdmax=14000, njmax=3600, dense Jacobian, zero overflow):
Aloha Cloth (32 worlds, nconmax=8000, njmax=20000, opt.iterations=1000, dense Jacobian, zero overflow):
Standalone performance of the PR commit with updated benchmark suite settings (256 worlds, zero overflow):