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 #1700

Summary

warp.optim.linear.cg with a batched LinearOperator (batch_offsets,
batch_count > 1) computes every CG scalar — rho, the step denominator
p·Ap, and the stopping-test residual norms — through batch_dot_kernel,
which assigns one block per subproblem and lets each lane accumulate its
strided share serially in the payload dtype (subproblem_length / block_dim fp32 additions per lane before a single wp.tile_sum). The code
states its own assumption: "(Assumes all subproblems are small compared to
total length)".

For a single batch, warp already abandons serial accumulation above
128 * tile_size scalars in favor of a bounded tile-tree reduction — but that
switch is gated on batch_count == 1, so batched subproblems stay serial at
every length. The result is a structural accuracy asymmetry that grows with
subproblem length (measured below: up to ~94x the tree-reduction error at
8,192 serial additions per lane), and in our production use it makes batched
CG on long subproblems fail to converge at iteration budgets that are ample
for the identical subproblem solved alone.

The error is deterministic (fixed reduction order, no atomics) — bitwise
repeatable run to run — so it presents downstream as a stable convergence
regression, not as noise.

Environment

• GPU: NVIDIA RTX 6000 Ada Generation (sm_89, 48 GB)
• CUDA: Toolkit 12.9, Driver 13.2
• warp 1.16.0.dev20260720; re-verified on 1.17.0.dev20260725 (nightly
  built from main, fetched 2026-07-25): the code paths cited below are
  unchanged and the measured table is identical on both versions
• python 3.12.13, numpy 2.5.1, Linux x86_64

We found no existing issue mentioning batch_dot_kernel; the closest related
change (GH-1608, inactive trailing-DOF indexing in batch_dot_kernel)
addresses out-of-range access, not accumulation accuracy.

Where in the code

warp/_src/optim/linear.py (line numbers identical at 1.16.0.dev20260720
and 1.17.0.dev20260725):

• Lines ~285-292: use_bounded_tree requires batch_count == 1, so the
  bounded tile-tree reduction is unreachable for batched operators:

  use_bounded_tree = (
      batch_offsets is not None
      and self.batch_count == 1
      and self.device.is_cuda
      and max_length > 128 * self.tile_size
  )

• Lines ~346-373: every batch_count > 1 TiledDot takes the
  batch_dot_kernel branch, commented "(Assumes all subproblems are small
  compared to total length)".

• Lines ~489-508: batch_dot_kernel — per-lane serial accumulation in
  a.dtype:

  acc = a.dtype(0.0)
  for i in range(batch_start, batch_end, wp.block_dim()):
      acc += a[col, i] * b[col, i]
  total = wp.tile_sum(wp.tile(acc))

Serial accumulation has O(n) worst-case rounding growth versus O(log n) for
the tree, and CG recycles these scalars into its direction updates every
iteration, so the per-dot error does not stay put — it steers the Krylov
recurrence.

Measurement (self-contained, warp + numpy only)

warp_dot_repro.py (inline below) computes per-subproblem dot products of
identical fp32 data four ways — the batched serial path under test
(batch_count = 4), warp's own batch_count == 1 policy on each subproblem
alone, the unbatched tiled tree, and an fp64 numpy reference — and reports the
max relative error against fp64 over subproblems and 10 seeded trials. Regime
norm is a == b with N(0,1) entries (the shape of rho = r·r and the
stopping-test norms); widerange is a == b with magnitudes log-uniform in
[1e-4, 1] (residual vectors whose components span scales — the late-CG /
ill-conditioned regime). 10,500 / 43,155 / 111,580 are the 300-, 1,233- and
3,188-frame 35-dof trajectory lengths from the downstream report below
(healthy / healthy / failing).

Measured output (python warp_dot_repro.py --device cuda:0), warp
1.16.0.dev20260720, RTX 6000 Ada (sm_89) — table bitwise-identical on
1.17.0.dev20260725:

What the table shows:

• The single (B=1) column is the smoking gun for the gate. Below
  128 * 512 = 65,536 scalars it matches batched exactly (same serial
  kernel); above, it drops to the tree's error level (the bounded tree
  engages). The identical data loses accuracy purely because sibling
  subproblems share the launch.
• The batched error grows with subproblem length while the tree stays flat
  (~1e-7 at every size). At 8,192 serial adds/lane the gap is 93.6x and still
  widening.
• Consecutive batched compute() calls are bitwise-identical (asserted in the
  script): this is deterministic rounding, not atomics noise.

Read honestly, the per-dot relative error at our failing production size
(L = 111,580) is small in these synthetic regimes — the point of the table is
the structural, unbounded asymmetry, not any single row. The convergence
consequence needs an ill-conditioned system and CG's feedback loop, which is
what we measured downstream:

Downstream consequence (how we found it) — controlled intervention

In the trajectory-IK solver of a Newton fork (newton-collab, branch
dylanturpin/trajectory-ik @ 6daeb38), CG solves block-Jacobi-preconditioned
banded J^T J systems whose conditioning worsens with horizon length. On an
RTX 6000 Ada with the environment above, solving four bitwise-identical
copies of one 3,188-frame, 35-dof motion clip (111,580 scalar dofs per
subproblem) as one batched operator, 32 outer Levenberg-Marquardt iterations,
CG budget 64 per linearization:

The design isolates the reduction: with four identical subproblems, every
per-batch quantity is identical across copies, so cross-problem coupling
cannot explain a batched-vs-solo difference; the solo control through the same
code path is bitwise-equal to a sequential reference, exonerating padding and
layout. The intervention then confirms the mechanism: swapping only the dot
reduction (a two-stage segmented tile-tree with the same shape as warp's
single-batch bounded tree, everything else untouched) takes the batched solve
from a hard stall at 3.9e8 to bitwise equality with the solo solve. Raising
the CG budget to 128 also recovers the batched solve (with warp's serial
reduction unchanged); batches of shorter clips (300 and 1,233 frames at
35 dofs) are healthy at budget 64, while a 16-clip batch padded to 3,188
frames needed the raised budget too — the failure is thresholded on
subproblem length, not on batch composition or clip content. All of this is
deterministic across reruns and orderings.

These runs use a motion-retargeting pipeline and capture-data corpus that we
cannot attach; we summarize them here as evidence and rely on the
self-contained script above for the mechanism-level reproduction.

Negative result, stated proactively: we tried to build a self-contained
batched-CG convergence repro and failed — synthetic trajectory chains (2-dof
and 24-dof, horizons up to 131k frames) converge fine in batched CG at default
budgets; their systems are too well-conditioned for the extra rounding to
bite. A toy CG demo is therefore not a meaningful test of this defect in
either direction; the dot-level table plus the controlled intervention above
are the reproducible evidence.

Our current workaround (downstream)

The Newton fork now ships an engine-local _SegmentedTiledDot implementing
the compute/col/cols surface of TiledDot: stage 1 reduces every
512-entry block of each subproblem with wp.tile_sum (blocks aligned per
subproblem, never straddling batch boundaries); stage 2 runs one further tile
reduction per subproblem over its block partials. It engages only for
batch_count > 1 on CUDA when the per-subproblem length exceeds
128 * tile_size — the same threshold at which warp's own single-batch path
abandons serial accumulation — and is swapped into the CG solver state after
construction. Because there is no public hook for the dot reduction, the swap
writes the private _tiled_dot attribute, guarded and carrying a TODO to
revert to a plain cg(...) call once warp reduces batched dots with a
multi-block tree. For subproblems up to tile_size^2 scalars the segmented
tree has the same reduction shape as warp's single-batch bounded tree, which
is what makes the batched-equals-solo bitwise parity above possible.

Suggested direction

1. Extend the existing bounded-tree variant to batch_count > 1 — a
   segmented multi-block tree: per-512-block tile_sum partials, then a
   per-subproblem combine (exactly the shape of the workaround above). This
   preserves the reduction shape of the batch_count == 1 path, so downstream
   users who validate batched against solo solves keep bitwise parity where it
   holds today.
2. Cheaper micro-fix: accumulate acc in fp64 (or Kahan-compensated fp32)
   inside batch_dot_kernel for fp32 payloads; the loop is memory-bound, so
   the cost should be negligible. Tradeoff: this changes batched results at
   every subproblem length and cannot be made bitwise-consistent with the
   single-batch tree paths, so batched-vs-solo parity checks downstream break
   even where the current kernel is accurate enough.
3. Either way, choose the strategy by subproblem length — batch_offsets
   deltas are known at TiledDot construction — instead of assuming
   subproblems are small.

Downstream impact (footnote)

Before any workaround, long batched subproblems in our retargeting workload
needed 2-3x the default CG iteration budget (or routing to batch_count == 1
solves) to converge at all; with the segmented-tree workaround they converge
at the default budget, and the dominant batch of our capacity benchmark
(4 x 3,188 frames) dropped from 4.11 s to 1.91 s at unchanged results. A warp-side
fix lets us delete the private-attribute swap entirely.

Reproduction script (warp_dot_repro.py)

#!/usr/bin/env python
"""Accuracy of warp's batched TiledDot (serial-lane accumulation) vs its own tree reductions.

``warp.optim.linear`` reduces CG dot products with ``TiledDot``. With
``batch_offsets`` describing more than one subproblem (``batch_count > 1``),
every dot product goes through ``batch_dot_kernel``: one block per subproblem,
each lane serially striding its subproblem range and accumulating in the
payload dtype (``subproblem_length / block_dim`` fp32 additions per lane
before a single ``wp.tile_sum``). The unbatched path instead reduces through a
tiled tree (one 512-wide ``tile_sum`` per block, then log-rounds of block
sums), and for a *single* batch longer than ``128 * tile_size`` warp itself
abandons serial accumulation for a bounded-tree variant — but that switch is
gated on ``batch_count == 1``, so batched subproblems stay serial at every
length.

This script measures what that asymmetry costs. For a batch of B fp32
subproblems of equal length L it computes every per-subproblem dot product
four ways on identical data:

  batched : one ``TiledDot(batch_offsets=...)`` with ``batch_count == B``
            (the serial-lane ``batch_dot_kernel`` under test)
  single  : per-subproblem ``TiledDot(batch_offsets=[0, L])`` with
            ``batch_count == 1`` — warp's own policy for the same data when
            it is not batched with siblings (serial up to ``128 * 512``,
            bounded tree above)
  tree    : per-subproblem unbatched ``TiledDot`` (plain tiled tree)
  fp64    : numpy float64 reference on the same fp32 inputs

and reports the maximum relative error of each warp reduction against fp64
across subproblems and trials. It also checks that consecutive batched
``compute()`` calls are bitwise identical (the error is deterministic, not
atomics noise).

Two CG-relevant input regimes:

  norm      : a == b, entries ~ N(0, 1) — the shape of rho = r.r and the
              stopping-test residual norms (positive summands)
  widerange : a == b, |entries| log-uniform in [1e-4, 1] — residual vectors
              whose components span scales, the late-CG / ill-conditioned
              regime in which CG direction updates are most sensitive to
              noise in rho / alpha

Requires only warp and numpy. Usage:

  python warp_dot_repro.py [--device cuda:0] [--batch 4] [--trials 10]
"""

from __future__ import annotations

import argparse

import numpy as np
import warp as wp

try:  # warp >= 1.10 module layout (the public shim does not re-export TiledDot)
    from warp._src.optim.linear import TiledDot
except ImportError:  # older layouts exposed the implementation module directly
    from warp.optim.linear import TiledDot

# Per-subproblem scalar lengths. 65,536 = 128 * 512 is the threshold above
# which warp's own single-batch path abandons serial accumulation for its
# bounded tree; 10,500 / 43,155 / 111,580 are 300-, 1,233- and 3,188-frame
# 35-dof trajectory chains from the downstream report (healthy / healthy /
# failing in batched CG at the default iteration budget).
DEFAULT_LENGTHS = (10_500, 43_155, 65_536, 111_580, 262_144, 1_048_576, 4_194_304)


def make_inputs(rng: np.random.Generator, length: int, regime: str) -> tuple[np.ndarray, np.ndarray]:
    if regime == "norm":
        a = rng.standard_normal(length).astype(np.float32)
        return a, a
    if regime == "widerange":
        mag = 10.0 ** rng.uniform(-4.0, 0.0, length)
        a = (rng.choice((-1.0, 1.0), length) * mag).astype(np.float32)
        return a, a
    raise ValueError(regime)


def batched_dot(a_np: np.ndarray, b_np: np.ndarray, offsets_np: np.ndarray, device) -> np.ndarray:
    """All subproblem dots through one batched TiledDot (serial batch_dot_kernel)."""
    offsets = wp.array(offsets_np, dtype=wp.int32, device=device)
    dot = TiledDot(int(offsets_np[-1]), wp.float32, device=device, batch_offsets=offsets)
    a = wp.array(a_np, dtype=wp.float32, device=device)
    b = wp.array(b_np, dtype=wp.float32, device=device)
    dot.compute(a, b)
    first = dot.col(0).numpy().copy()
    dot.compute(a, b)  # determinism check: bitwise-equal on recompute
    assert np.array_equal(first, dot.col(0).numpy()), "batched dot not deterministic"
    return first


def single_dot(a_np: np.ndarray, b_np: np.ndarray, offsets_np: np.ndarray, device) -> np.ndarray:
    """Each subproblem through warp's batch_count == 1 policy.

    Serial up to 128 * tile_size scalars, bounded tree above — i.e. what warp
    itself does with the identical data when it is not batched with siblings.
    """
    out = np.empty(len(offsets_np) - 1, dtype=np.float32)
    a = wp.array(a_np, dtype=wp.float32, device=device)
    b = wp.array(b_np, dtype=wp.float32, device=device)
    for k in range(len(offsets_np) - 1):
        lo, hi = int(offsets_np[k]), int(offsets_np[k + 1])
        offsets = wp.array(np.array([0, hi - lo], dtype=np.int32), device=device)
        dot = TiledDot(hi - lo, wp.float32, device=device, batch_offsets=offsets)
        dot.compute(a[lo:hi], b[lo:hi])
        out[k] = dot.col(0).numpy()[0]
    return out


def tree_dot(a_np: np.ndarray, b_np: np.ndarray, offsets_np: np.ndarray, device) -> np.ndarray:
    """Each subproblem dot through warp's unbatched tiled tree reduction."""
    out = np.empty(len(offsets_np) - 1, dtype=np.float32)
    a = wp.array(a_np, dtype=wp.float32, device=device)
    b = wp.array(b_np, dtype=wp.float32, device=device)
    for k in range(len(offsets_np) - 1):
        lo, hi = int(offsets_np[k]), int(offsets_np[k + 1])
        dot = TiledDot(hi - lo, wp.float32, device=device)
        dot.compute(a[lo:hi], b[lo:hi])
        out[k] = dot.col(0).numpy()[0]
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--batch", type=int, default=4, help="subproblems per batched dot")
    parser.add_argument("--trials", type=int, default=10, help="random seeds per row")
    parser.add_argument("--lengths", type=int, nargs="+", default=list(DEFAULT_LENGTHS))
    args = parser.parse_args()

    wp.init()
    device = wp.get_device(args.device)
    print(f"warp {wp.config.version}, device {device.name} (sm_{device.arch}), "
          f"batch={args.batch}, trials={args.trials}")
    print()
    print("| regime | length L | serial adds/lane | batched (B>1) | single (B=1) | tree | batched/tree |")
    print("|---|---|---|---|---|---|---|")

    for regime in ("norm", "widerange"):
        for length in args.lengths:
            offsets_np = np.arange(args.batch + 1, dtype=np.int32) * length
            err = {"batched": 0.0, "single": 0.0, "tree": 0.0}
            for trial in range(args.trials):
                rng = np.random.default_rng(1000 * trial + length)
                a_np, b_np = make_inputs(rng, args.batch * length, regime)
                ref = np.array(
                    [
                        np.dot(
                            a_np[offsets_np[k] : offsets_np[k + 1]].astype(np.float64),
                            b_np[offsets_np[k] : offsets_np[k + 1]].astype(np.float64),
                        )
                        for k in range(args.batch)
                    ]
                )
                for name, fn in (("batched", batched_dot), ("single", single_dot), ("tree", tree_dot)):
                    got = fn(a_np, b_np, offsets_np, device)
                    err[name] = max(err[name], float(np.max(np.abs(got - ref) / np.abs(ref))))
            ratio = err["batched"] / err["tree"] if err["tree"] > 0 else float("inf")
            print(
                f"| {regime} | {length:>9,} | {length // 512:>4} "
                f"| {err['batched']:.3e} | {err['single']:.3e} | {err['tree']:.3e} | {ratio:>7.1f}x |"
            )

    print()
    print("batched = TiledDot(batch_offsets=..., batch_count=B) serial-lane accumulation (batch_dot_kernel)")
    print("single  = TiledDot(batch_offsets=[0, L], batch_count=1) on each subproblem alone:")
    print("          warp's own switch — serial <= 128 * 512 scalars, bounded tree above")
    print("tree    = unbatched TiledDot tiled tree reduction on the same subproblem data")
    print("rel err = |result - fp64 reference| / |fp64 reference|, max over subproblems and trials")
    print("determinism: consecutive batched compute() calls bitwise-equal at every row (asserted)")


if __name__ == "__main__":
    main()