perf(rbd): fuse the solver color loops — one workgroup per env (stacked on #21)#22
Open
haixuanTao wants to merge 4 commits into
Open
perf(rbd): fuse the solver color loops — one workgroup per env (stacked on #21)#22haixuanTao wants to merge 4 commits into
haixuanTao wants to merge 4 commits into
Conversation
The narrow-phase kernels dispatched [max_len/64, num_batches, 1]: one 64-lane workgroup per batch rounded up from that batch's live pair count. A batched robot env has ~7 pairs, so at 2048+ envs the GPU ran thousands of workgroups at ~11% lane occupancy and narrow-phase scaled ~linearly with env count (0.02 ms at 1 env -> 18.6 ms at 4096). gpu_flatten_batches_dispatch (one thread, same style as the existing max-scan init kernels) now builds exclusive prefix offsets over the per-batch work-lists plus a flat [total/64, 1, 1] grid; the classify, deferred and PFM kernels walk 0..total and recover (batch, item) with a binary search over the offsets. Warps fill with real pairs from consecutive batches. Buffer layout is unchanged — only the dispatch shape and index math moved. RTX 5090, 12-DOF biped batch, dt=5ms, steady-state contacts: narrow-phase 2048 envs: 10.31 -> 5.09 ms 4096 envs: 18.57 -> 6.94 ms whole step 2048 envs: 22.45 -> 16.50 ms/step (91k -> 124k env-steps/s) Physics bit-identical (robot trajectory + cube settle unchanged to printed precision). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…s via workgroup barriers The contact solver ran reset_color + num_colors x (solve + inc_color) as sequential device-wide dispatches, three times per substep — ~230 dependent launches per step whose cost was launch/barrier latency, not lane math (the color kernels touch ~1-2 contacts per env; flattening their lanes was measured neutral). But the ordering the colors enforce is only ever *within* an env — bodies are never shared across batches — so the barrier only needs workgroup scope. gpu_warmstart_fused / gpu_step_gauss_seidel_fused run one 64-lane workgroup per batch: stage the batch's solver_vels in shared memory (which the workgroup barrier fences), walk colors 1..=num_colors with a barrier between, write back once. Each replaces a whole reset+N x (solve+inc) chain with a single dispatch; the per-color path remains as fallback for batches with more than FUSED_SOLVE_MAX_BODIES (64) bodies. num_colors reaches the kernels via a tiny uniform refreshed only when the Grow policy raises it. RTX 5090, 12-DOF biped batch, dt=5ms, steady-state, wgpu — on top of the flat narrow-phase dispatch: 2048 envs: 16.50 -> 8.68 ms/step (124k -> 236k env-steps/s) 4096 envs: 27.27 -> 13.72 ms/step (150k -> 298k env-steps/s) Cumulative vs pre-flatten baseline: 91k -> 298k env-steps/s at 4096 (3.3x). Physics bit-identical (robot trajectory + cube settle unchanged to printed precision) — same-color contacts touch disjoint bodies, so intra-color order cannot matter, and inter-color order is preserved. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ls + flat refit_leaves Round 2 of the one-workgroup-per-env fusion: - gpu_warmstart_fused absorbs gpu_apply_solver_vels_inc (increment added while staging velocities into shared memory, same bodies_len bound) and gpu_solver_update_constraints (each lane refreshes its own contacts' constraints before the color walk — a constraint is only read by the lane that updated it, so no extra barrier). Exactly fills the 8-storage-buffer budget. - gpu_step_gauss_seidel_fused_no_bias replaces the final color loop, with gpu_remove_cfm_and_bias_kernel absorbed as a prologue (same lane-locality argument). - gpu_lbvh_refit_leaves goes flat: the per-batch collider count is uniform, so batch/index recover by division. A 14-collider env used 14 of 64 lanes per workgroup, and per-leaf work is mesh-AABB computation — this pass was 17% of the profiled frame. Removes another ~12 dispatches per step in the fused path; the standalone kernels remain for the >FUSED_SOLVE_MAX_BODIES fallback. RTX 5090, 12-DOF biped batch, dt=5ms, wgpu, on top of the round-1 fusion: 2048 envs: 8.68 -> 6.52 ms/step (236k -> 314k env-steps/s) Physics bit-identical (robot trajectory endpoint + cube settle to printed precision). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…narrow phase Collision pairs carry env-local collider ids, and collider_parent is batch-strided like every other per-collider buffer (its construction comment says so: 'Env-local body slot; the kernels apply the per-batch stride') — but the classify and pfm_pfm kernels read it unsliced, resolving every batch's parents through batch 0's table. With identical environments the tables coincide and nothing observable goes wrong — which is why every bit-exactness check passed. With heterogeneous environments (the point of per-env MJCF insertion) contacts are silently mis-parented: a pair that is same-body in batch 0 gets skipped in batches where it isn't, and solved impulses can target the wrong bodies. Repro (now a stacking test): env0 = body with two glued boxes + a single box; env1 = single box dropped onto a two-glued-box body (equal body and collider counts, different parent tables). Pre-fix the falling box never rests on the stack; post-fix it settles on top (z=0.75). Identical-env scenes are bit-identical before/after, as the aliasing argument predicts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
haixuanTao
added a commit
to haixuanTao/nexus
that referenced
this pull request
Jul 17, 2026
…capture CUDA-graph capture of the physics step was failing with STREAM_CAPTURE_INVALIDATED on the unified stack (production captured fine; these are unified-added regressions from the dimforge#20 small-sort and dimforge#22 fused-color merges). Any allocation (cuMemAlloc) or pageable host→device copy inside the captured region invalidates the capture. Two offenders, localized with a cuStreamGetCaptureInfo probe after each op: 1. Small-sort fast path (radix_sort): allocated a fresh SortUniforms tensor every dispatch (`pass_uniforms.clear(); push(Tensor::scalar(..))`), bypassing the `uniforms_key` cache the radix path already uses for exactly this reason. Now cached/keyed by per_batch_max — allocate once, reuse. 2. num_colors_uniform sync (rbd_step): updated the fused-solver color-count uniform inside step() — first as a realloc, then (mistakenly) as an in-place write_buffer, both illegal during capture. Moved the sync to auto_resize_buffers, which runs OUTSIDE the captured region; max_colors only changes there (the ratchet) and in set_max_colors (setup), so the value is stable across a capture. With these + the 1-D fixed-grid fix (daf5a43), BIPED_GRAPH=1 now captures and replays on the unified stack (verified: 12 iters @2048 g1_29dof_agile terrain full-parity, falls within baseline noise). Unblocks the env-per-lane LᵀDL wall-clock win, whose per-kernel gain was hidden by eager launch overhead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W
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
2 participants
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.
Stacked on #21 (its commit is included as the base of this branch; only the top two commits are new).
Problem
After #21 packs the narrow-phase warps, the solver becomes the dominant pass — and its cost is not lane math. Each substep runs
reset_color + num_colors × (solve + inc_color)three times as sequential device-wide dispatches: ~230 dependent launches per step, each touching ~1–2 contacts per env for ~a microsecond of work. Flattening these kernels' lanes (same treatment as #21) measures neutral — they are latency-bound, not occupancy-bound.Change
The ordering the colors enforce only ever applies within an env — bodies are never shared across batches — so the barrier between colors needs workgroup scope, not a dispatch boundary.
Commit 1:
gpu_warmstart_fused/gpu_step_gauss_seidel_fusedrun one 64-lane workgroup per env: stage that env'ssolver_velsin shared memory (which the workgroup barrier fences), walk colors1..=num_colorswith a barrier between, write back once. Each replaces an entire reset+N×(solve+inc) chain with a single dispatch.num_colorsarrives via a one-word uniform refreshed only when the Grow policy raises it. Batches with more thanFUSED_SOLVE_MAX_BODIES(64) bodies fall back to the per-color chain automatically.Commit 2: the fused kernels absorb their per-substep neighbours —
apply_solver_vels_incfolds into the staging load,update_constraintsinto a prologue (a constraint is only read by the lane that updated it, so lane-locality replaces the barrier),remove_cfm_and_biasinto a_no_biasvariant — andgpu_lbvh_refit_leavesgoes flat (uniform per-batch collider count → batch by division; its per-leaf mesh-AABB work ran at 14-of-64 lanes).Results
RTX 5090, batched 12-DOF biped (mesh feet), dt = 5 ms, steady-state, wgpu — cumulative over #21:
Physics is bit-identical at every stage (robot trajectory and cube-settle checks to printed precision): same-color contacts touch disjoint bodies so intra-color order cannot matter, and inter-color order is preserved by the barriers.
🤖 Generated with Claude Code