Proposal: replace quadratic island discovery with a GPU-appropriate linear-memory path

[teerthsharma]

The feature and motivation

MuJoCo Warp's current island discovery materializes both tree_tree[nworld, ntree, ntree] and stack_scratch[nworld, ntree * ntree] as int32 tensors. Those two scratch tensors alone cost

8 * nworld * ntree^2 bytes.

At the current upstream/main commit (c50dd837), the allocations and traversal are in island.py.

nworld	ntree	two current scratch tensors
8192	64	256 MiB
1024	256	512 MiB
64	1024	512 MiB
1	4096	128 MiB

Native MuJoCo has the same quadratic-memory problem in google-deepmind/mujoco#3388 and an open minimum-root DSU implementation in google-deepmind/mujoco#3396. MuJoCo Warp depends on MuJoCo for model/data bindings but has its own Warp island kernels, so the native change does not remove these GPU allocations. The native DSU is a useful correctness baseline, but I do not think it should be ported blindly: a sequential worker per world can underutilize a GPU when nworld is small and ntree is large, while a parallel GPU connectivity algorithm adds atomics, convergence, determinism, and CUDA graph-capture tradeoffs.

The exact object is a finite hypergraph: dynamic trees are vertices and each active constraint row induces a hyperedge containing the trees in the nonzero support of that Jacobian row. Island discovery asks only for its connected-component partition, equivalently its H_0, with the minimum tree id as a canonical component label. This gives several exact linear-memory candidates without importing unrelated manifold machinery.

Exact candidates screened

I independently screened three propositions against a bipartite tree/constraint BFS oracle. These are correctness experiments, not performance evidence.

1. Incidence-star DSU. Replacing every nonempty hyperedge by a star through any representative preserves connected components. Minimum-root DSU is therefore exact regardless of representative, row order, endpoint order, or duplicates. A direct implementation needs parent[nworld, ntree] plus incidence scratch and is the simplest production baseline. It matched 20,010 special, random, and metamorphic cases in an isolated CPU test.

2. Warp-tiled minimum-label propagation. Initialize every active tree with its id. In each synchronous round, cooperatively reduce the minimum label of each constraint hyperedge, scatter it to incident trees with integer atomic_min, then pointer-jump. Labels decrease monotonically within their original component and converge to that component's minimum. A prototype using a 32-lane wp.launch_tiled block and wp.tile_min matched 332 CUDA cases on an RTX 4060 Laptop GPU, including a 129-tree chain, 64-way incidence, reversed rows, duplicate incidence, a two-lobe bridge/cycle corpus, and repeated determinism trials. It took at most eight rounds in that screening matrix. The prototype uses a host convergence check and fixed padded arity, so neither its timings nor its capture behavior should be treated as a production result.

3. Static quotient plus dynamic Warp connectivity. If a subset of incidence is certified static, contract its DSU components, map each dynamic hyperedge to the quotient, solve connectivity there, and pull labels back through the quotient map. Contracting a connected static subgraph preserves H_0. This matched 5,000 random static/dynamic splits in a CPU quotient test; the first 120 also matched when the quotient was solved by the tiled CUDA kernel. This looks like the most interesting DSU-plus-Warp extension, but it is only exact with careful invalidation: per-world d.eq_active, contacts, limits, flex constraints, and sleep/wake state can change active incidence. I suggest treating quotient caching as a follow-up rather than part of the first replacement.

Proposed first scope

I suggest a benchmark-first implementation that compares:

• direct-incidence minimum-root DSU as the linear-memory correctness baseline;
• 32-lane tiled minimum-label propagation for more parallelism; and
• a same-memory, non-tiled label-propagation implementation as a null comparison.

The first PR should select the best full-rebuild implementation from measured evidence. A static quotient/cache can follow behind an exact invalidation and full-rebuild fallback. A small-model threshold retaining the dense implementation is also reasonable if it is measurably faster and the threshold accounts for available device memory.

Validation should require exact parity for nisland, tree_island, nidof, and every downstream island mapping, not merely equal component counts. The matrix should include sparse and dense Jacobians, repeated scalar rows, equality and contact incidence involving flex bodies, static endpoints, inactive trees, per-world eq_active, sleep/wake transitions, duplicate and reversed incidence, and hyperedges wider than one hardware warp. It should exercise ordinary CUDA execution, graph capture, and replay.

Performance results should report peak device memory, construction time, convergence rounds, kernel time, fwd_position, and full-step time together with nworld, ntree, active incidence count and arity, GPU, and dtype. The decisive memory gate is that no O(nworld * ntree^2) temporary remains. Representative MuJoCo models should be accompanied by generated topology cases that independently vary component diameter, hyperedge width, redundant cycles, and bridge deletion.

Would the maintainers be open to this benchmark-first direction? In particular:

1. Should the initial PR implement both direct DSU and the tiled candidate for comparison, or start with direct DSU only?
2. For a tiled implementation, would fixed rounds, a conditionally replayed CUDA graph, or another convergence strategy fit the project's capture requirements best?
3. Does a separately reviewed static-quotient follow-up sound useful, provided every mutable incidence source has an explicit invalidation test and fallback?

Related GPU connectivity work includes Adaptive Work-Efficient Connected Components on the GPU and ConnectIt. Warp already exposes tiled launches and tile reductions, but the most suitable schedule here should be established by repository benchmarks rather than assumed from those general results.

[thowell]

@teerthsharma thank you for your interest in contributing to mujoco warp!

when adding or changing an existing feature we like to measure the performance of relevant benchmark scenes in order to inform development. in this case for island discovery the relevant benchmark is the aloha clutter scene. dont think any other benchmarks currently utilize island discovery -- however it would be great to improve the performance of this feature to the point where its computational cost is negligible compared to the cost of step / forward and that we turn on this feature by default like mujoco. this benchmark can be run with

uv run python3 benchmarks/run.py -f aloha_clutter

ideally an updated island discovery routine would have both increased steps per second and reduced memory utilization. an additional consideration is the simplicity/complexity of the implementation compared to its performance.

if you create an implementation that is an improvement given these considerations please considering opening a pr. please let us know if you have any questions. thanks!