jit compile tangent computation in Proximal

[YigitElma]

Still draft!

This seems to be a bottleneck if the tangent computation is chunked. Without any chunking or chunk_size>dim_c, there is no big performance gain (I will check why some benchmarks are actually slower, it is pretty annoying since all I did is just to jit the thing). The performance degradation for chunking is so huge that I wanted to fix it.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -4.20 %    |     3.986e+03      |     3.818e+03      |   -167.49    |       26.31        |       23.10        |
  test_proximal_jac_w7x_with_eq_update   |    0.69 %    |     6.493e+03      |     6.538e+03      |    44.60     |       118.51       |       117.30       |
  test_proximal_freeb_jac                |   -0.04 %    |     1.330e+04      |     1.329e+04      |    -5.57     |       64.65        |       61.94        |
  test_proximal_freeb_jac_blocked        |    0.64 %    |     7.641e+03      |     7.690e+03      |    49.03     |       56.42        |       54.61        |
  test_proximal_freeb_jac_batched        |   -0.10 %    |     7.645e+03      |     7.637e+03      |    -7.96     |       54.91        |       54.43        |
  test_proximal_jac_ripple               |    0.52 %    |     3.391e+03      |     3.408e+03      |    17.67     |       41.59        |       40.47        |
  test_proximal_jac_ripple_bounce1d      |   -0.10 %    |     3.704e+03      |     3.701e+03      |    -3.71     |       52.45        |       50.52        |
  test_eq_solve                          |    1.71 %    |     2.015e+03      |     2.049e+03      |    34.50     |       64.81        |       64.60        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.39 +/- 4.05     | +3.43e-03 +/- 3.52e-02 |  8.72e-01 +/- 2.9e-02  |  8.69e-01 +/- 2.0e-02  |
 test_equilibrium_init_medres            |     +1.20 +/- 4.19     | +8.49e-02 +/- 2.96e-01 |  7.14e+00 +/- 1.9e-01  |  7.06e+00 +/- 2.3e-01  |
 test_equilibrium_init_highres           |     +2.56 +/- 3.78     | +1.97e-01 +/- 2.91e-01 |  7.91e+00 +/- 2.5e-01  |  7.71e+00 +/- 1.5e-01  |
 test_objective_compile_dshape_current   |     -0.40 +/- 1.95     | -1.60e-02 +/- 7.81e-02 |  3.98e+00 +/- 4.8e-02  |  4.00e+00 +/- 6.2e-02  |
 test_objective_compute_dshape_current   |    +12.04 +/- 12.93    | +7.22e-05 +/- 7.75e-05 |  6.72e-04 +/- 7.3e-05  |  6.00e-04 +/- 2.5e-05  |
 test_objective_jac_dshape_current       |     -0.80 +/- 28.77    | -2.03e-04 +/- 7.27e-03 |  2.51e-02 +/- 4.8e-03  |  2.53e-02 +/- 5.5e-03  |
 test_perturb_2                          |     +0.30 +/- 3.61     | +5.76e-02 +/- 6.99e-01 |  1.94e+01 +/- 5.3e-01  |  1.94e+01 +/- 4.6e-01  |
+test_proximal_jac_atf_chunked           |    -66.85 +/- 1.23     | -1.16e+01 +/- 2.15e-01 |  5.77e+00 +/- 4.9e-02  |  1.74e+01 +/- 2.1e-01  |
 test_proximal_freeb_jac                 |     +1.61 +/- 3.46     | +7.59e-02 +/- 1.62e-01 |  4.77e+00 +/- 6.7e-02  |  4.70e+00 +/- 1.5e-01  |
 test_solve_fixed_iter                   |     +1.03 +/- 2.81     | +2.74e-01 +/- 7.51e-01 |  2.70e+01 +/- 6.7e-01  |  2.67e+01 +/- 3.4e-01  |
 test_LinearConstraintProjection_build   |     +3.07 +/- 3.98     | +2.06e-01 +/- 2.67e-01 |  6.93e+00 +/- 1.6e-01  |  6.72e+00 +/- 2.1e-01  |
 test_objective_compute_ripple_bounce1d  |     -0.98 +/- 5.79     | -2.77e-03 +/- 1.64e-02 |  2.80e-01 +/- 1.2e-02  |  2.83e-01 +/- 1.1e-02  |
 test_objective_grad_ripple_bounce1d     |     +5.06 +/- 2.58     | +5.07e-02 +/- 2.59e-02 |  1.05e+00 +/- 2.2e-02  |  1.00e+00 +/- 1.4e-02  |
 test_build_transform_fft_midres         |     -7.79 +/- 3.19     | -7.50e-02 +/- 3.08e-02 |  8.88e-01 +/- 2.3e-02  |  9.63e-01 +/- 2.1e-02  |
 test_build_transform_fft_highres        |     -3.87 +/- 3.36     | -4.83e-02 +/- 4.20e-02 |  1.20e+00 +/- 3.6e-02  |  1.25e+00 +/- 2.2e-02  |
 test_equilibrium_init_lowres            |    -10.44 +/- 3.54     | -7.65e-01 +/- 2.59e-01 |  6.56e+00 +/- 2.1e-01  |  7.32e+00 +/- 1.6e-01  |
 test_objective_compile_atf              |     -8.67 +/- 3.17     | -5.86e-01 +/- 2.14e-01 |  6.18e+00 +/- 1.5e-01  |  6.76e+00 +/- 1.6e-01  |
 test_objective_compute_atf              |     -8.70 +/- 13.62    | -1.90e-04 +/- 2.97e-04 |  1.99e-03 +/- 1.4e-04  |  2.18e-03 +/- 2.6e-04  |
 test_objective_jac_atf                  |     -5.32 +/- 2.72     | -8.94e-02 +/- 4.57e-02 |  1.59e+00 +/- 3.9e-02  |  1.68e+00 +/- 2.4e-02  |
 test_perturb_1                          |     -3.12 +/- 2.92     | -5.38e-01 +/- 5.04e-01 |  1.67e+01 +/- 1.8e-01  |  1.73e+01 +/- 4.7e-01  |
 test_proximal_jac_atf                   |     +3.44 +/- 1.59     | +1.84e-01 +/- 8.48e-02 |  5.51e+00 +/- 5.0e-02  |  5.33e+00 +/- 6.8e-02  |
 test_proximal_jac_atf_with_eq_update    |     +4.46 +/- 1.74     | +5.41e-01 +/- 2.11e-01 |  1.27e+01 +/- 1.9e-01  |  1.21e+01 +/- 9.0e-02  |
 test_proximal_freeb_compute             |     +2.34 +/- 2.81     | +3.96e-03 +/- 4.75e-03 |  1.73e-01 +/- 3.1e-03  |  1.69e-01 +/- 3.6e-03  |
 test_solve_fixed_iter_compiled          |     +1.49 +/- 2.28     | +1.20e-01 +/- 1.84e-01 |  8.19e+00 +/- 1.5e-01  |  8.07e+00 +/- 1.0e-01  |
 test_objective_compute_ripple           |     +0.85 +/- 4.81     | +2.04e-03 +/- 1.16e-02 |  2.43e-01 +/- 6.6e-03  |  2.41e-01 +/- 9.5e-03  |
 test_objective_grad_ripple              |     +8.20 +/- 4.22     | +7.29e-02 +/- 3.76e-02 |  9.62e-01 +/- 2.0e-02  |  8.89e-01 +/- 3.2e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.33%. Comparing base (1c7178d) to head (dd9f654).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2239   +/-   ##
=======================================
  Coverage   94.33%   94.33%           
=======================================
  Files         101      101           
  Lines       28855    28859    +4     
=======================================
+ Hits        27220    27224    +4     
  Misses       1635     1635           

Files with missing lines	Coverage Δ
desc/optimize/_constraint_wrappers.py	96.93% <100.00%> (+0.03%)	⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[YigitElma]

Motivation

For below code,

eq = get("precise_QA")
grid = LinearGrid(M=eq.M, N=eq.N, NFP=eq.NFP, rho=np.array([0.6, 0.8, 1.0]), sym=True)

obj1 = ObjectiveFunction(
    QuasisymmetryTwoTerm(
        eq=eq, helicity=(1, 0), grid=grid, normalize=False
    )
)
k = 7
R_modes = np.vstack(
    (
        [0, 0, 0],
        eq.surface.R_basis.modes[
            np.max(np.abs(eq.surface.R_basis.modes), 1) > k, :
        ],
    )
)
Z_modes = eq.surface.Z_basis.modes[
    np.max(np.abs(eq.surface.Z_basis.modes), 1) > k, :
]
constraints = (
    FixBoundaryR(eq=eq, modes=R_modes),
    FixBoundaryZ(eq=eq, modes=Z_modes),
    FixPressure(eq=eq),
    FixCurrent(eq=eq),
    FixPsi(eq=eq),
)
cons = ObjectiveFunction(constraints)
con = ObjectiveFunction(ForceBalance(eq=eq), jac_chunk_size=300)
prox = ProximalProjection(objective=obj1, constraint=con, eq=eq)
lcp = LinearConstraintProjection(objective=prox, constraint=cons)
lcp.build()
xr = lcp.x()
J = lcp.jac_scaled_error(xr)
print(len(xr))   # prints 224

This PR on my laptop GPU,

Master branch,

[YigitElma]

@f0uriest do you see any obvious reason why jitting the tangents part can make the vmaped route slightly slower?