[ooctipus]

Description

IKSolverTrajectory's CG backend under-converges multi-trajectory solves
(n_problems > 1) on long frame chains: batched solves stall at cg_iterations
budgets that are ample for the identical trajectory solved alone.

Root cause (in warp, not in this repo): warp's internal TiledDot picks its
dot-reduction strategy by batch count, not subproblem length. With
batch_count == 1 it reduces through a tiled tree (a bounded-tree variant above
128 * 512 scalar dofs), but with batch_count > 1 every subproblem is reduced
by a single block whose lanes accumulate their strided share serially in fp32 —
about L/512 serial additions per lane for a subproblem of L scalar dofs. The
O(n) rounding growth this injects into every CG scalar (rho, the step
denominator p·Ap, the stopping-test residual norms) stalls CG on long,
ill-conditioned chains. Filed upstream, with a self-contained repro and a
measured error table, as NVIDIA/warp#1700.

Fix (engine-local): swap the CG state's dot reduction for
_SegmentedTiledDot, a segmented tile tree implementing the same
compute/col/cols surface: one 512-wide tile_sum per 512-entry block of
each subproblem, then one further tile reduction per subproblem over its block
partials. The swap engages only when n_problems > 1, the device is CUDA, and
the per-trajectory length exceeds 128 * 512 scalar dofs — the same threshold
at which warp's own single-batch path abandons serial accumulation. Below the
threshold warp's serial path is numerically healthy and is kept, so shorter
multi-trajectory solves remain bitwise-unchanged (verified below). CPU batched
solves also keep warp's reduction: warp's CPU batched dot is one fully serial
lane per subproblem — the same O(n) issue in principle — but the effect is
unmeasured on CPU and deliberately outside this fix's blast radius.

Parity envelope: for subproblems up to 512 * 512 = 262,144 scalar dofs the
segmented tree has the same reduction shape as warp's single-batch bounded tree
(block indexing is relative to each subproblem's own offset), so batched solves
match the equivalent single-trajectory solves bitwise. Above that, accuracy
stays tree-like but bitwise parity with the single-batch path ends. The
reduction order is fixed (no atomics), so batched solves are bitwise
deterministic run to run.

Why not fp64 accumulation instead: accumulating the batched dot in fp64
would also fix the accuracy, but it diverges bitwise from warp's single-batch
fp32 tree at every subproblem length. Bitwise parity with the healthy
single-trajectory path is this PR's strongest validation invariant — the
reproduction below lands all four batched copies bitwise-identical (equal
coordinate sha256) to the solo solve — and it is what lets every batched solve
below the engagement threshold ship bitwise-unchanged. fp64 accumulation
remains listed as a warp-side option in NVIDIA/warp#1700, where its
tradeoff is spelled out.

Workaround status: warp exposes no public hook for the CG dot reduction, so
the swap writes the CG state's private _tiled_dot attribute. The attribute is
validated first (present, expected tile size), so a warp internals change fails
loudly at solver construction instead of silently reverting the fix while tests
stay green. A TODO marks the swap for removal once warp reduces batched dots
with a multi-block tree (NVIDIA/warp#1700).

Checklist

• New or existing tests cover these changes
• The documentation is up to date with these changes
• CHANGELOG.md has been updated (if user-facing change)

Test plan

Measurements on RTX 6000 Ada (sm_89), CUDA 12.9, warp 1.16.0.dev20260720,
python 3.12. newton.tests.test_ik_trajectory and
newton.tests.test_ik_trajectory_dynamics pass on cpu + cuda:

uv run --extra dev -m newton.tests -k test_ik_trajectory

New and changed tests (CUDA):

• test_trajectory_cg_batched_long_horizon — above the engagement threshold
  (33,792 frames x 2 dofs, inside the parity envelope) the segmented reduction
  engages, the batched solve is bitwise-equal to the single-trajectory
  solve, and reruns are bitwise deterministic. The bitwise assertion
  discriminates the reduction numerically: under warp's serial batched
  reduction, most coordinates drift by up to ~1e-6.
• test_trajectory_segmented_dot (new) — dot-level accuracy against an fp64
  reference across batch shapes (2-4 segments, partially filled tail block),
  offset independence (identical segment data reduces bitwise-identically to a
  singleton instance's result), repeat-compute determinism, and the two-column
  compute/col interface CG drives (col_offset=1 for the p·Ap column).
• test_trajectory_cg_dot_swap_guard (new) — the swap raises RuntimeError at
  solver construction when warp's CG state stops exposing a compatible
  _tiled_dot (renamed attribute, or unexpected tile size), and installs the
  segmented reduction when compatible.

Scope note: synthetic long-horizon chains do not reproduce the production
convergence failure (2-dof and 24-dof chains up to 131k frames are too well
conditioned and converge fine either way). The bitwise-parity assertion, which
fails under warp's serial reduction, therefore carries the fails-before
evidence at test scale; the reproduction below carries the convergence proof.

Bug fix

Steps to reproduce (without this PR):

1. Build an IKSolverTrajectory problem with a long horizon — above 65,536
   scalar dofs per trajectory, e.g. 3,000+ frames at 35 dofs — whose
   Gauss-Newton system is ill-conditioned enough that CG needs its iteration
   budget (real articulations qualify; toy chains do not).
2. Solve it once with n_problems=1 and once with n_problems > 1 (identical
   copies of the same trajectory), same cg_iterations.
3. Observe the batched solve stalling at a far higher cost than the solo solve,
   deterministically across reruns.

Measured with four bitwise-identical copies of a 3,188-frame, 35-dof humanoid
motion-retargeting trajectory (CMU clip 55_25; 111,580 scalar dofs per
subproblem), ANALYTIC jacobians, 32 LM iterations, default cg_iterations=64:

configuration	final LM cost (fp32)	max landmark error
solo (n_problems=1), before = after fix	328.7804260253906	177.1 mm
batched x4, before fix	3.896e8 (LM stalled at the seed cost; all four copies identical)	425.4 mm
batched x4, after fix	328.7804260253906 (all four)	177.1 mm

With identical copies, every per-batch quantity is identical across copies, so
cross-problem coupling cannot explain the difference; swapping only the dot
reduction takes the batched solve from the stall to bitwise equality with the
solo solve (equal coordinate sha256), confirming the mechanism. Two consecutive
batched runs are bitwise-identical.

On the surrounding retargeting workload (32 humanoid motion-capture clips,
300-7,851 frames at 35 dofs), multi-trajectory solves at the default
cg_iterations=64 now match their sequential single-trajectory references —
the 3,188-frame batch bitwise, the rest within noise — where before the fix,
long batched buckets needed 2-3x the default budget (or single-trajectory
routing) to converge at all; the dominant batch (4 x 3,188 frames) dropped from
4.11 s to 1.91 s wall by converging at the default budget. A 16-trajectory
batch of 300-frame clips, below the engagement threshold, is bitwise-unchanged
by this PR (16/16 clips).

Minimal reproduction:

At toy scale the convergence gap does not manifest (see the scope note above),
but the defective reduction does — without this PR the batched solve silently
diverges from the solo solve's coordinates:

import numpy as np
import warp as wp

import newton.ik as ik

# any articulation with n_frames * n_dofs > 128 * 512, e.g. a 2-dof planar
# arm tracking an end-effector arc over 33,792 frames (objectives as in
# newton/tests/test_ik_trajectory.py)


def solve(model, n_frames, objectives, n_problems):
    solver = ik.IKSolverTrajectory(
        model,
        n_frames,
        objectives,
        n_problems=n_problems,
        jacobian_mode=ik.IKJacobianType.ANALYTIC,
        linear_solver="cg",
    )
    joint_q = wp.zeros((n_problems * n_frames, model.joint_coord_count), dtype=wp.float32)
    solver.step(joint_q, joint_q, iterations=10)
    return joint_q.numpy().reshape(n_problems, n_frames, -1)


q_single = solve(model, n_frames, objectives_x1, 1)
q_batched = solve(model, n_frames, objectives_x2, 2)  # two identical copies

# without this PR: False (most coordinates drift by up to ~1e-6)
# with this PR:    True (bitwise, inside the <= 512*512-dof parity envelope)
print(np.array_equal(q_batched[0], q_single[0]))

test_trajectory_cg_batched_long_horizon is this reproduction, made exact.

Status update (2026-08-05)

• The root cause is now fixed upstream in Warp: optim.linear: batched CG dot products (batch_count > 1) always use serial-lane fp32 accumulation — accuracy degrades with subproblem length and stalls CG on long, ill-conditioned subproblems NVIDIA/warp#1700 was addressed within days (segmented tree reductions for batched dot products, shared by CG/CR/BiCGSTAB/GMRES, in current Warp nightlies). We verified the upstream fix at our production sizes: accuracy equivalent to this PR's workaround (~1e-7 vs an fp64 reference) and batched==singleton bitwise at every size tested.
• This PR remains necessary for every released Warp version. The swap guard plus the in-code TODO make retirement a one-commit follow-up the moment a Warp release ships the fix; we are happy to file that follow-up when the time comes.
• Downstream validation: this branch's code has run continuously in our production retargeting pipeline since 2026-07-26, bitwise-validated against pre-change baselines, and re-validated unchanged after our move to Warp 1.17 nightlies.

[ooctipus]

Status note: NVIDIA landed the root-cause fix for the underlying warp issue (NVIDIA/warp#1700, commit cb9dc55 — batched dots now use segmented tree reductions across CG/CR/BiCGSTAB/GMRES, plus an optional max_batch_length hint). We verified it at our production sizes (accuracy equivalent to this PR's workaround, batched==singleton bitwise at all sizes) and are running warp 1.17 nightlies in production. This PR remains the right fix for anyone on current warp releases; once a warp release ships the fix, the swap can be retired per the in-code TODO — happy to file that follow-up when the time comes, ideally passing max_batch_length through aslinearoperator for the small extra win.

[jaredmejia]

The targeted long-horizon, segmented-dot and guard tests pass on c3307f99 with Warp 1.16 on CUDA. Disabling only the workaround in the numerical probe reproduces the batched-versus-solo mismatch, so it still has a purpose for the older supported versions.

One update to the retirement note: Warp 1.17.0 is now released and includes the newton-physics#1700 reduction fix. Since this branch still allows Warp >=1.15, should we keep this private-kernel workaround version-bounded, or raise the minimum to the fixed release and remove it? I haven't run a fresh 1.17 runtime comparison, so this isn't a claim that the workaround can simply be deleted under the current dependency range.