Optimization: Single kernel iterative linesearch

[Kenny-Vilella]

Merged all linesearch iterative pipeline into a single kernel for perf optimization.

Perf on the asv benchmark:

Branch	Scene	Step time
this branch	AlohaPot	1900
main	AlohaPot	4087
-------------------	------------	-------
this branch	ApolloFlat	910
main	ApolloFlat	1139
-------------------	------------	-------
this branch	ApolloTerrain	4228
main	ApolloTerrain	4296
-------------------	------------	-------
this branch	FrankaEmikaPanda	80
main	FrankaEmikaPanda	101
-------------------	------------	-------
this branch	Humanoid	642
main	Humanoid	781
-------------------	------------	-------
this branch	ThreeHumanoids	12861
main	ThreeHumanoids	11489

[Kenny-Vilella]

For some of these benchmarks, we now get closer performance between iterative and parallel linesearch.
8% slower on the Humanoid, 30% slower on the Apollo, 15% slower on the AlohaPot.

[thowell]

can we benchmark the aloha + pot scene with keyframes: 0, 1, 2? thanks!

[Kenny-Vilella]

Voila !

Keyframe	main	this branch
0	508.66 x	1,073.98 x
1	869.31 x	1,192.79 x
2	491.00 x	1,050.33 x

The command:

python mujoco_warp/testspeed.py --function=step benchmark/aloha_pot/scene.xml --keyframe=1 --nworld=8192 --nstep=1000 --njmax=80 --nconmax=500000 -o opt.is_sparse=True -o opt.ls_parallel=False

[Kenny-Vilella]

Two more notes:

• With this approach, the linesearch terminates early when it is fully converged, so we can choose a fairly high value of ls_iterations without affecting too much perf.
• I am not yet satisfied with the code duplication in _eval_elliptic and _eval_init, but I am not sure how to improve the situation without affecting perf.

[erikfrey]

@Kenny-Vilella this PR is like opening a Christmas present. Better perf, fewer kernels, AND could delete a bunch of temp arrays off data!?

Will do a proper review soon!