Fuse dense elliptic curvature into the Hessian tile by adenzler-nvidia · Pull Request #1879 · google-deepmind/mujoco_warp

adenzler-nvidia · 2026-11-02T14:25:06Z

Dense Newton builds the elliptic cone curvature in a separate kernel, _update_gradient_JTCJ_dense, that strides over every contact in the batch for every upper-triangle Hessian entry and accumulates into the Hessian with atomics. Every solver pass pays that scan, including the late passes in which almost every world has already converged, and each contact's Jacobian rows are reloaded once per Hessian entry.

The cone curvature of a contact is diagonal plus rank one, so it can be written as a sum of squared, weighted combinations of the contact's own Jacobian rows: the normal direction minus the projected tangent, and an orthonormal (Householder) basis of the tangent space orthogonal to the sliding direction. Each cone-state row maps to one such virtual row and weight, and the existing per-world J'WJ tile product of _update_gradient_JTDAJ_dense_tiled then assembles the whole Hessian. This mirrors the sparse path, which already folds the elliptic curvature into its per-group assembly (#1521).

The constraint update already loads the contact, friction and tangent residuals when it classifies a row into the cone state, so with dense elliptic Newton it also writes the row's virtual-row coefficients, weight and the contact's first row into three solver-context arrays. The arrays are allocated with wp.empty and exist only for dense elliptic Newton.
The dense JTDAJ tile kernel (full and compact) tile-loads the Jacobian chunk, takes cone weights from the context, and rewrites only the cone-state rows of chunks that contain any, reading a row's state from efc.state. Pyramidal and other dense models compile the unchanged code, apart from clamping nefc to njmax.
_update_gradient_JTCJ_dense and its SM-count based launch sizing are removed, along with dof_tri_row/dof_tri_col on Model and cdof_tri_row/cdof_tri_col on Data, which only it read. The dense Hessian is built by one per-world kernel without atomics, and converged worlds exit at the start of the block.
Cone curvature is selected by constraint state alone, as in the sparse path and MuJoCo. The old dense kernel also skipped contacts with dist >= includemargin, but contact rows only exist inside the include margin (adhesive contacts use the full margin), so that check never applied.

test_elliptic_dense_hessian now launches the tile kernel with and without the elliptic term and checks the difference against the independent NumPy reference for mixed condim 3/4/6 contacts, for one and two heterogeneous worlds. test_constrained_solve_equivalence_all_active also covers dense Jacobians with sliding contacts; for elliptic cones it takes a single Newton step from the warmstart, where qacc depends on the cone Hessian of the compact kernel. test_elliptic_truncated_contact covers a contact truncated by njmax, and test_elliptic_dense_njmax_below_chunk an njmax below the row chunk size.

The three context arrays add 32 bytes per constraint row for dense elliptic Newton only and are allocated with wp.empty: the weights are padded to njmax_pad like efc.D, and rows past nefc are masked before a cone weight is read. The constraint update now gives the rows of an elliptic contact truncated by an njmax or naconmax overflow zero force and the satisfied state instead of leaving both stale, so every cone-state row has its coefficients written in the same pass.

Benchmarks on an RTX PRO 6000 Blackwell (MIG 2g.48gb slice) against 24e6d81b, interleaved runs, zero overflow and all worlds converged. The elliptic benchmarks in benchmarks/run.py (median of 4 main and 3 PR runs); aloha_clutter uses the sparse Jacobian and is unaffected:

benchmark	main	this PR	change
aloha_pot	1,860,494	2,214,322	+19.0%
franka_threading	90,308	95,607	+5.9%
aloha_sdf	204,495	213,097	+4.2%
aloha_clutter	265,207	265,064	-0.1%

mjwarp-testspeed, nworld 8192, nstep 1000, median of 5 interleaved runs, time per step:

model	change
humanoid, `opt.cone=ELLIPTIC opt.jacobian=DENSE`	-49.5%
unitree_g1 flat (replay), `opt.cone=ELLIPTIC opt.jacobian=DENSE`	-38.9%

Solver iteration counts are unchanged.

Re-measured after the review changes on a full RTX PRO 6000 Blackwell against main at d3f27d64, mjwarp-testspeed with nworld 8192 and nstep 1000, median of 5 interleaved runs, steps per second:

model	main	this PR	change
humanoid, `opt.cone=ELLIPTIC opt.jacobian=DENSE`	4,062,204	7,136,690	+75.7%
unitree_g1 flat (replay), `opt.cone=ELLIPTIC opt.jacobian=DENSE`	1,920,536	3,070,868	+59.9%
aloha_pot	3,330,435	3,817,170	+14.6%
humanoid, `opt.jacobian=DENSE` (pyramidal)	8,458,887	8,460,065	+0.0%

Dense Newton built the elliptic cone curvature in a separate contact-major kernel that strided over every contact in the batch for every upper-triangle Hessian entry and accumulated into the Hessian with atomics. Every solver pass paid that full scan, including passes in which almost every world had already converged, and each contact's rows were reloaded once per Hessian entry. The cone curvature of a contact is diagonal plus rank one, so it can be written as a sum of squared, weighted combinations of the contact's own Jacobian rows: the normal direction minus the projected tangent, plus an orthonormal (Householder) basis of the tangent space orthogonal to the sliding direction. Each cone-state row maps to one such virtual row and weight, and the existing per-world J'WJ tile product of the dense JTDAJ kernel then assembles the full Hessian. This mirrors the sparse path, which already folds elliptic curvature into its per-group assembly. The constraint update already loads the contact, friction and tangent residuals when it classifies a row into the cone state, so it now also writes the row's virtual-row coefficients, weight and the contact's first row into solver-context arrays that exist only for dense elliptic Newton. The Hessian kernel tile-loads the Jacobian chunk, takes the cone weights from the context, and rewrites only the cone-state rows of chunks that contain any. The dense Hessian is now built by one per-world kernel without atomics, converged worlds exit at the start of the block, and the standalone dense JTCJ kernel and its launch sizing are removed. Cone curvature is selected by constraint state alone, as in the sparse path and MuJoCo. The dense Hessian test now compares the curvature added by the elliptic kernel against the independent NumPy reference for mixed 3/4/6 contacts across two heterogeneous worlds, and the compact solver equivalence test also covers dense Jacobians.

…ptic-fused # Conflicts: # mujoco_warp/_src/types.py

thowell · 2026-11-10T17:22:03Z

Thanks @adenzler-nvidia! Fusing the dense elliptic cone curvature into the per-world JTDAJ tile product via the Householder virtual-row decomposition is a clean design that eliminates atomic Hessian accumulation; please address the items below before merging.

1. `mujoco_warp/_src/solver.py` (lines 168--170, 2855--2932, 2974, 3058)

Issue: ctx.cone_weight is allocated with shape (nworld, njmax), whereas _elliptic_chunk_weights loads TILE_SIZE_K elements with bounds_check=False (reading out of bounds on the final tile when njmax > 16 is not a multiple of 16, unlike d.efc.D and d.efc.state which use d.njmax_pad). In addition, _update_gradient_JTDAJ_dense_tiled and _update_gradient_JTDAJ_dense_tiled_compact do not clamp nefc_in[worldid] to njmax (allowing njmax overflows to index ctx_cone_base_in, ctx_cone_coef_in, and efc_J_in out of bounds), and _is_cone(state) ignores active, triggering _rewrite_cone_rows on stale CONE states past nefc.
Recommendation: Allocate cone_weight with d.njmax_pad, clamp nefc = wp.min(nefc_in[worldid], njmax) in both dense tiled Hessian kernels, and gate _is_cone on active != 0.0.

    # In _create_solver_context:
    cone_coef=wp.empty((nworld, njmax), dtype=types.vec6) if alloc_cone_rows else None,
    cone_weight=wp.empty((nworld, d.njmax_pad), dtype=float) if alloc_cone_rows else None,
    cone_base=wp.empty((nworld, njmax), dtype=int) if alloc_cone_rows else None,

@wp.func
def _is_cone(state: int, active: float) -> int:
  return wp.where(active != 0.0 and state == types.ConstraintState.CONE.value, 1, 0)

    # In _elliptic_chunk_weights:
    ncone = wp.tile_reduce(wp.add, wp.tile_map(_is_cone, state, active))[0]

    # In _update_gradient_JTDAJ_dense_tiled_compact and _update_gradient_JTDAJ_dense_tiled:
    nefc = wp.min(nefc_in[worldid], njmax)

2. `mujoco_warp/_src/solver.py` (lines 1967--2025)

Issue: When an elliptic contact is truncated by njmax or naconmax (conid >= wp.min(nacon_in[0], contact_dim_in.shape[0]), efcid0 < 0, or efcidj < 0), _update_constraint_efc returns early without writing efc_force_out or efc_state_out. Because d.efc.state persists across steps while ctx.cone_coef, ctx.cone_weight, and ctx.cone_base are allocated with wp.empty, a stale CONE state from an earlier step causes the Hessian kernel to read uninitialized ctx.cone_* entries (NaN weights and invalid base indices). When CONE_ROWS is enabled, lines 2018--2020 also re-read ctx_Jaref_in from global memory after uj was already computed in the tangent loop.
Recommendation: Clamp nacon_in[0] to contact_dim_in.shape[0], clear efc_force_out to 0.0 and efc_state_out to SATISFIED before returning early, and populate u[j] = uj directly in the unrolled tangent loop.

    if wp.static(CONE_ROWS):
      dim = int(0)
      friction = types.vec5()
      u = types.vec6()

    if is_elliptic:
      conid = efc_id_in[worldid, efcid]
      if conid >= wp.min(nacon_in[0], contact_dim_in.shape[0]):
        efc_force_out[worldid, efcid] = 0.0
        efc_state_out[worldid, efcid] = types.ConstraintState.SATISFIED.value
        return
      efcid0 = contact_efc_address_in[conid, 0]
      if efcid0 < 0:
        efc_force_out[worldid, efcid] = 0.0
        efc_state_out[worldid, efcid] = types.ConstraintState.SATISFIED.value
        return

      dim = contact_dim_in[conid]
      friction = contact_friction_in[conid]
      mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
      jaref0 = ctx_Jaref_in[worldid, efcid0]
      D0 = efc_D_in[worldid, efcid0]

      for j in range(1, 6):
        if j < dim:
          efcidj = contact_efc_address_in[conid, j]
          if efcidj < 0:
            efc_force_out[worldid, efcid] = 0.0
            efc_state_out[worldid, efcid] = types.ConstraintState.SATISFIED.value
            return
          frictionj = friction[j - 1]
          uj = ctx_Jaref_in[worldid, efcidj] * frictionj
          if wp.static(CONE_ROWS):
            u[j] = uj
          TT += uj * uj
          if efcid == efcidj:
            ufrictionj = uj * frictionj

    if wp.static(CONE_ROWS):
      if new_state == types.ConstraintState.CONE.value:
        coef, weight = _cone_virtual_row(efcid - efcid0, dim, mu, D0, jaref0, u, TT, friction)
        ctx_cone_coef_out[worldid, efcid] = coef
        ctx_cone_weight_out[worldid, efcid] = weight
        ctx_cone_base_out[worldid, efcid] = efcid0

3. `mujoco_warp/_src/types.py`, `mujoco_warp/_src/io.py`, `mujoco_warp/_src/solver.py` (line 5111), `mujoco_warp/_src/solver_test.py` (line 2806)

Issue: _update_gradient_JTCJ_dense was the sole reader of Model.dof_tri_row, Model.dof_tri_col, Data.cdof_tri_row, and Data.cdof_tri_col (ndof_tri and nvmax_pad_sq), leaving unused O(nv^2) and O(nvmax_pad^2) device arrays on Model and compact Data.
Recommendation: Remove dof_tri_row/dof_tri_col and cdof_tri_row/cdof_tri_col (and their size entries/replacements) across types.py, io.py, solver.py, and solver_test.py.

  # In solve_compact (solver.py):
  m2 = dataclasses.replace(m, opt=opt2, nv=nvp, nv_pad=nvp, is_sparse=False)

4. `mujoco_warp/_src/solver_test.py` (lines 2735--2753)

Issue: In _COMPACT_CONTACT_XML, the spheres start with zero tangential velocity on a flat plane (efc.state is QUADRATIC for all rows, so _rewrite_cone_rows in _update_gradient_JTDAJ_dense_tiled_compact is never executed), and d.qacc is not reset before solver.solve_compact(m, d).
Recommendation: Slide the spheres with warmstart disabled so elliptic contacts enter ConstraintState.CONE, assert that at least one row is in CONE, and pre-fill d.qacc.fill_(wp.inf) before solver.solve_compact(m, d).

  @parameterized.product(cone=("pyramidal", "elliptic"), jacobian=("sparse", "dense"))
  def test_constrained_solve_equivalence_all_active(self, cone, jacobian):
    """With every tree active and nvmax=nv, the compacted Newton solve matches baseline qacc."""
    xml = _COMPACT_CONTACT_XML.replace('iterations="20"', f'iterations="20" cone="{cone}"')
    xml = xml.replace('jacobian="sparse"', f'jacobian="{jacobian}"')
    _, _, m, d = _put_compact(xml)
    self.assertEqual(m.is_sparse, jacobian == "sparse")
    m.opt.disableflags |= types.DisableBit.WARMSTART
    qvel = np.zeros((d.nworld, m.nv))
    qvel[:, 1] = 5.0
    qvel[:, 7] = -3.0
    d.qvel.assign(qvel)
    mjw.forward(m, d)  # full baseline solve (also builds efc.J, M)
    self.assertGreater(d.nacon.numpy()[0], 0)  # contacts exist
    if cone == "elliptic":
      nefc = int(d.nefc.numpy()[0])
      self.assertTrue((d.efc.state.numpy()[0, :nefc] == types.ConstraintState.CONE).any())
    baseline_qacc = d.qacc.numpy().copy()

    d.tree_awake = wp.array(np.ones((d.nworld, m.ntree), dtype=int), dtype=int)
    island.update_active_dofs(m, d)
    self.assertEqual(d.ncdof.numpy()[0], m.nv)

    d.qacc.fill_(wp.inf)
    solver.solve_compact(m, d)

    np.testing.assert_allclose(d.qacc.numpy(), baseline_qacc, rtol=1e-3, atol=1e-4)

thowell · 2026-11-10T20:05:38Z

@yuvaltassa investigating porting some of these optimizations to mujoco c

…ptic-fused # Conflicts: # mujoco_warp/_src/determinism_test.py

The constraint update returned early for rows of an elliptic contact that an njmax or naconmax overflow had truncated, without writing their force or state. Both persist across steps, so such a row kept whatever an earlier step had left there. With dense elliptic Newton a stale cone state made the Hessian kernel read cone coefficients and weights that were never written for that row. These rows now get zero force and the satisfied state. The cone weights are tile-loaded in whole row chunks like efc.D, so they are now allocated with the padded row count; with njmax not a multiple of the chunk size the last chunk read past the allocation. Both dense Hessian kernels clamp nefc to njmax, and chunks are only scanned for cone rows below nefc. The tangent residuals of a cone row are kept from the loop that already loads them instead of being read again. The dense JTCJ kernel was the only reader of the upper-triangle index arrays dof_tri_row and dof_tri_col on Model and cdof_tri_row and cdof_tri_col on Data, so they are removed. The compact equivalence test now slides the spheres so that elliptic contacts are in the cone state and, for elliptic cones, takes a single Newton step from the warmstart, where qacc depends on the cone Hessian of the compact kernel.

The two dense Hessian kernels shared a wp.func that returned the chunk's weight tile. Warp allocates a tile returned by a function in shared memory but does not count it in the shared memory the kernel requests, so the kernels overran their shared memory by one weight tile. That went unnoticed at the default chunk size, and is an illegal memory access for some njmax below the chunk size, where the chunk is njmax rows. The weights are now computed in the kernels, which also drops the copy of the returned tile.

The dense cone curvature is built from the virtual rows that the constraint update writes, so rescaling Jaref and only updating the gradient compared two Hessians built from the same rows. The test now updates the constraints after rescaling.

The cone row rewrite indexed a shared-memory copy of the chunk's constraint states, which every row chunk loaded whether or not it had cone rows. The rewrite now reads the state of a row from efc.state, and the chunk's state tile stays in registers.

adenzler-nvidia · 2026-11-11T10:48:08Z

Thanks @thowell, all four are addressed in a7edee1, on top of a merge of current main.

cone_weight is allocated with d.njmax_pad, both dense Hessian kernels clamp nefc to njmax, and _is_cone is gated on the active mask.
Rows of an elliptic contact truncated by an njmax or naconmax overflow now get zero force and the satisfied state, nacon is clamped to the contact capacity, and u is kept from the tangent loop instead of re-reading Jaref. I kept the range(1, dim) loop rather than unrolling it. test_elliptic_truncated_contact covers this: it pre-fills efc.state with the cone state and efc.force with a sentinel, then truncates a contact with njmax=11.
dof_tri_row/dof_tri_col, cdof_tri_row/cdof_tri_col and their size entries are removed.
The compact test slides the spheres, asserts a cone-state row and pre-fills qacc. I deviated in one place. With warmstart disabled the solve converges in one step whether or not the compact kernel adds the cone curvature: with the elliptic term dropped from the compact kernel, qacc still matched to 6e-3, about 0.4% of the tolerance. The test therefore keeps the warmstart and, for elliptic cones, takes a single Newton step, where the same mutation is off by 53.

Two more things came up while testing these:

The two dense kernels shared a wp.func that returned the chunk's weight tile. When a kernel passes a tile returned by a function to an op that needs it in shared memory, Warp also moves the function's copy there without counting it in the shared memory the kernel requests, so both kernels overran their shared memory by one weight tile. That was silent at the default chunk of 16 rows and an illegal memory access for some njmax below 16 (8, 11, 13 and 14 with nv_pad=20). 6f2460f computes the weights in the kernels instead, and test_elliptic_dense_njmax_below_chunk covers it on CUDA.
test_elliptic_hessian_scale_invariance had become vacuous for dense Jacobians: the cone rows are built in the constraint update and the test only updated the gradient. 2442ca1 updates the constraints after rescaling.

While there, a9bcc21 makes the cone row rewrite read a row's state from efc.state instead of from a shared-memory copy of the chunk's states, which every chunk loaded whether or not it had cone rows.

The full suite passes on an L40. Re-measured on an RTX PRO 6000 Blackwell at 8192 worlds (median of 5 interleaved runs, steps/s) against main at d3f27d6 and against this PR before these commits:

model	main	before	now
humanoid, `opt.cone=ELLIPTIC opt.jacobian=DENSE`	4.06M	7.08M	7.14M
unitree_g1 flat (replay), `opt.cone=ELLIPTIC opt.jacobian=DENSE`	1.92M	3.04M	3.07M
aloha_pot	3.33M	3.78M	3.82M
humanoid, `opt.jacobian=DENSE` (pyramidal)	8.46M	8.46M	8.46M

…enzler/dense-elliptic-fused

The dense Hessian kernel gained an nv argument on main, and test_elliptic_dense_hessian launches it directly.

…ptic-fused # Conflicts: # mujoco_warp/_src/io.py # mujoco_warp/_src/solver.py # mujoco_warp/_src/types.py

adenzler-nvidia added 2 commits November 5, 2026 16:25

Trim comments on the dense elliptic Hessian path

32332f3

adenzler-nvidia force-pushed the adenzler/dense-elliptic-fused branch from a9730f0 to 32332f3 Compare November 5, 2026 15:32

Merge remote-tracking branch 'upstream/main' into adenzler/dense-elli…

c010daa

…ptic-fused # Conflicts: # mujoco_warp/_src/types.py

johnnynunez mentioned this pull request Nov 10, 2026

Avoid atomics in the deterministic solver kernels #1933

Closed

adenzler-nvidia added 5 commits November 11, 2026 09:31

Merge remote-tracking branch 'upstream/main' into adenzler/dense-elli…

02f09c8

…ptic-fused # Conflicts: # mujoco_warp/_src/determinism_test.py

adenzler-nvidia added 2 commits November 11, 2026 11:59

Merge branch 'main' of github.com:google-deepmind/mujoco_warp into ad…

7251e10

…enzler/dense-elliptic-fused

Pass nv to the dense Hessian kernel in the elliptic test

249ac2b

The dense Hessian kernel gained an nv argument on main, and test_elliptic_dense_hessian launches it directly.

adenzler-nvidia mentioned this pull request Nov 11, 2026

Kernel requests too little shared memory when it moves a tile returned by a wp.func to shared storage NVIDIA/warp#2294

Open

Merge remote-tracking branch 'upstream/main' into adenzler/dense-elli…

a52b2d0

…ptic-fused # Conflicts: # mujoco_warp/_src/io.py # mujoco_warp/_src/solver.py # mujoco_warp/_src/types.py

adenzler-nvidia mentioned this pull request Nov 13, 2026

Scale the noslip cost-change guard to float32 rounding #1964

Merged

Navigation Menu

Repository navigation

Fuse dense elliptic curvature into the Hessian tile - #1879

Fuse dense elliptic curvature into the Hessian tile#1879
adenzler-nvidia wants to merge 11 commits into
google-deepmind:mainfrom
adenzler-nvidia:adenzler/dense-elliptic-fused

adenzler-nvidia commented Nov 2, 2026 •

edited

Loading

Uh oh!

thowell commented Nov 10, 2026

Uh oh!

thowell commented Nov 10, 2026

Uh oh!

adenzler-nvidia commented Nov 11, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

adenzler-nvidia commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

thowell commented Nov 10, 2026

1. mujoco_warp/_src/solver.py (lines 168--170, 2855--2932, 2974, 3058)

2. mujoco_warp/_src/solver.py (lines 1967--2025)

3. mujoco_warp/_src/types.py, mujoco_warp/_src/io.py, mujoco_warp/_src/solver.py (line 5111), mujoco_warp/_src/solver_test.py (line 2806)

4. mujoco_warp/_src/solver_test.py (lines 2735--2753)

Uh oh!

thowell commented Nov 10, 2026

Uh oh!

adenzler-nvidia commented Nov 11, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

adenzler-nvidia commented Nov 2, 2026 •

edited

Loading

1. `mujoco_warp/_src/solver.py` (lines 168--170, 2855--2932, 2974, 3058)

2. `mujoco_warp/_src/solver.py` (lines 1967--2025)

3. `mujoco_warp/_src/types.py`, `mujoco_warp/_src/io.py`, `mujoco_warp/_src/solver.py` (line 5111), `mujoco_warp/_src/solver_test.py` (line 2806)

4. `mujoco_warp/_src/solver_test.py` (lines 2735--2753)