set constraint jacobian nonzeros njmax_nnz

[thowell]

introduces a new Data field nJefc njmax_nnz, the number of non-zeros in the constraint Jacobian (per-world).

enables significant memory savings:

---

humanoid with sparse constraint Jacobian

total memory (MiB): 432.00 -> 368.00

setup

• SPARSE_CONSTRAINT_JACOBIAN=True
• fixes from updates for constraint.py #1220

mjwarp-testspeed benchmarks/humanoid/humanoid.xml --nworld=8192 --nconmax=24 --njmax=64 --memory

Loading model from: benchmarks/humanoid/humanoid.xml...

Model
  nq: 28 nv: 27 nu: 21 nbody: 17 ngeom: 20
Option
  integrator: EULER
  cone: PYRAMIDAL
  solver: NEWTON iterations: 100 ls_iterations: 50
  is_sparse: False
  ls_parallel: False
  broadphase: NXN broadphase_filter: PLANE|SPHERE|OBB
Data
  nworld: 8192 naconmax: 196608 njmax: 64
Rolling out 1000 steps at dt = 0.005...

Summary for 8192 parallel rollouts

Total JIT time: 0.29 s
Total simulation time: 2.71 s
Total steps per second: 3,025,771
Total realtime factor: 15,128.85 x
Total time per step: 330.49 ns
Total converged worlds: 8192 / 8192

Model memory 0.03 MiB (0.01% of used memory):
 (no field >= 1% of used memory)
Data memory 298.03 MiB (68.99% of used memory):
 xmat: 4.78 MiB (1.11%)
 ximat: 4.78 MiB (1.11%)
 geom_xmat: 5.62 MiB (1.30%)
 cdof: 5.06 MiB (1.17%)
 cinert: 5.31 MiB (1.23%)
 crb: 5.31 MiB (1.23%)
 qM: 24.50 MiB (5.67%)
 qLD: 22.78 MiB (5.27%)
 cdof_dot: 5.06 MiB (1.17%)
 contact.frame: 6.75 MiB (1.56%)
 efc.J_colind: 54.00 MiB (12.50%)
 efc.J: 54.00 MiB (12.50%)
Other memory: 133.94 MiB (31.00% of used memory)
Total memory: 432.00 MiB (0.89% of total device memory)

reducing from the default njmax * nv = 64 * 27 = 1728 -> nJefc=700

mjwarp-testspeed benchmarks/humanoid/humanoid.xml --nworld=8192 --nconmax=24 --njmax=64 --nJefc=700 --memory

Loading model from: benchmarks/humanoid/humanoid.xml...

Model
  nq: 28 nv: 27 nu: 21 nbody: 17 ngeom: 20
Option
  integrator: EULER
  cone: PYRAMIDAL
  solver: NEWTON iterations: 100 ls_iterations: 50
  is_sparse: False
  ls_parallel: False
  broadphase: NXN broadphase_filter: PLANE|SPHERE|OBB
Data
  nworld: 8192 naconmax: 196608 njmax: 64
Rolling out 1000 steps at dt = 0.005...

Summary for 8192 parallel rollouts

Total JIT time: 0.30 s
Total simulation time: 2.71 s
Total steps per second: 3,025,884
Total realtime factor: 15,129.42 x
Total time per step: 330.48 ns
Total converged worlds: 8192 / 8192

Model memory 0.03 MiB (0.01% of used memory):
 (no field >= 1% of used memory)
Data memory 233.78 MiB (63.53% of used memory):
 xmat: 4.78 MiB (1.30%)
 ximat: 4.78 MiB (1.30%)
 geom_xmat: 5.62 MiB (1.53%)
 cdof: 5.06 MiB (1.38%)
 cinert: 5.31 MiB (1.44%)
 crb: 5.31 MiB (1.44%)
 qM: 24.50 MiB (6.66%)
 qLD: 22.78 MiB (6.19%)
 cdof_dot: 5.06 MiB (1.38%)
 contact.frame: 6.75 MiB (1.83%)
 contact.friction: 3.75 MiB (1.02%)
 contact.solimp: 3.75 MiB (1.02%)
 efc.J_colind: 21.88 MiB (5.94%)
 efc.J: 21.88 MiB (5.94%)
Other memory: 134.19 MiB (36.46% of used memory)
Total memory: 368.00 MiB (0.76% of total device memory)

[erikfrey]

It pains me to add another configurable param but after discussing with @thowell I understand why this one is important. LGTM! Just two nits.