Skip to content

Misc - #2293

Draft
YigitElma wants to merge 24 commits into
masterfrom
yge/lcp-init-perf
Draft

Misc#2293
YigitElma wants to merge 24 commits into
masterfrom
yge/lcp-init-perf

Conversation

@YigitElma

@YigitElma YigitElma commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator
  • Improve the LCP build method when dealing with large arrays. It also removes some of the big arrays completely, like versions of _Z to cut down the base memory. Same spirit as Remove large arrays before big memory operations #2281, but this one also improves the optimization initialization time by roughly half on GPU.
  • Improves compilation time of qr_multiply by using a jax.lax.scan instead of the unrolled for loop.
  • Updates the initial value for alpha. The previous value of 0 was basically recalculating the same p_newton
  • Includes changes from Fix tests for older dependencies #2287 for some tests

Just a rough eq.solve memory benchmark:
memory-01_eq_solve-gpu

@YigitElma YigitElma added the run_benchmarks Run timing benchmarks on this PR against current master branch label Aug 14, 2026
@github-actions

github-actions Bot commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -3.31 %    |     4.231e+03      |     4.091e+03      |   -140.05    |       24.73        |       21.87        |
  test_proximal_jac_w7x_with_eq_update   |   -2.46 %    |     6.858e+03      |     6.689e+03      |   -169.03    |       109.49       |       114.28       |
  test_proximal_freeb_jac                |   -0.13 %    |     1.353e+04      |     1.351e+04      |    -17.23    |       70.24        |       68.02        |
  test_proximal_freeb_jac_blocked        |   -0.33 %    |     7.854e+03      |     7.829e+03      |    -25.54    |       70.79        |       59.57        |
  test_proximal_freeb_jac_batched        |   -0.54 %    |     7.862e+03      |     7.820e+03      |    -42.53    |       59.94        |       60.36        |
  test_proximal_jac_ripple               |   -0.82 %    |     3.723e+03      |     3.693e+03      |    -30.53    |       36.53        |       36.73        |
  test_proximal_jac_ripple_bounce1d      |   -1.64 %    |     3.950e+03      |     3.885e+03      |    -64.92    |       46.21        |       46.37        |
  test_eq_solve                          |   -3.81 %    |     1.828e+03      |     1.758e+03      |    -69.68    |       32.77        |       34.24        |
  test_objective_quadratic_flux_jac      |   -0.56 %    |     1.895e+03      |     1.885e+03      |    -10.54    |       23.67        |       24.46        |

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

@github-actions

github-actions Bot commented Aug 14, 2026

Copy link
Copy Markdown
Contributor
|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +5.29 +/- 7.26     | +4.48e-02 +/- 6.15e-02 |  8.92e-01 +/- 5.5e-02  |  8.48e-01 +/- 2.7e-02  |
 test_equilibrium_init_lowres            |     +2.81 +/- 5.27     | +1.80e-01 +/- 3.37e-01 |  6.57e+00 +/- 2.7e-01  |  6.39e+00 +/- 2.1e-01  |
 test_objective_compile_atf              |     -0.71 +/- 6.96     | -4.37e-02 +/- 4.29e-01 |  6.12e+00 +/- 7.4e-02  |  6.16e+00 +/- 4.2e-01  |
 test_objective_compute_atf              |    +12.78 +/- 15.53    | +2.74e-04 +/- 3.33e-04 |  2.42e-03 +/- 2.9e-04  |  2.14e-03 +/- 1.6e-04  |
 test_objective_jac_atf                  |     +2.55 +/- 5.02     | +4.06e-02 +/- 8.00e-02 |  1.63e+00 +/- 7.0e-02  |  1.59e+00 +/- 3.9e-02  |
 test_perturb_1                          |     -3.29 +/- 3.39     | -3.96e-01 +/- 4.08e-01 |  1.16e+01 +/- 2.6e-01  |  1.20e+01 +/- 3.2e-01  |
 test_proximal_jac_atf                   |     -1.56 +/- 4.06     | -8.31e-02 +/- 2.16e-01 |  5.24e+00 +/- 9.9e-02  |  5.32e+00 +/- 1.9e-01  |
 test_proximal_freeb_compute             |     -2.49 +/- 3.21     | -4.15e-03 +/- 5.36e-03 |  1.63e-01 +/- 3.3e-03  |  1.67e-01 +/- 4.2e-03  |
 test_solve_fixed_iter                   |     -2.71 +/- 3.58     | -6.81e-01 +/- 9.01e-01 |  2.45e+01 +/- 5.7e-01  |  2.51e+01 +/- 7.0e-01  |
 test_LinearConstraintProjection_build   |    -12.79 +/- 4.65     | -8.88e-01 +/- 3.23e-01 |  6.05e+00 +/- 2.0e-01  |  6.94e+00 +/- 2.5e-01  |
 test_objective_compute_ripple           |     -1.42 +/- 5.31     | -3.37e-03 +/- 1.26e-02 |  2.34e-01 +/- 9.1e-03  |  2.37e-01 +/- 8.7e-03  |
 test_objective_grad_ripple              |     -3.33 +/- 5.45     | -3.02e-02 +/- 4.94e-02 |  8.76e-01 +/- 2.1e-02  |  9.07e-01 +/- 4.5e-02  |
 test_objective_quadratic_flux_compute   |     -1.51 +/- 3.47     | -3.26e-04 +/- 7.52e-04 |  2.13e-02 +/- 1.4e-04  |  2.16e-02 +/- 7.4e-04  |
 test_build_transform_fft_midres         |    +15.55 +/- 10.30    | +8.91e-02 +/- 5.90e-02 |  6.62e-01 +/- 2.8e-02  |  5.73e-01 +/- 5.2e-02  |
 test_build_transform_fft_highres        |    +12.96 +/- 6.05     | +9.57e-02 +/- 4.47e-02 |  8.34e-01 +/- 4.0e-02  |  7.39e-01 +/- 1.9e-02  |
 test_equilibrium_init_medres            |    +14.53 +/- 5.67     | +6.75e-01 +/- 2.63e-01 |  5.32e+00 +/- 2.3e-01  |  4.65e+00 +/- 1.2e-01  |
 test_objective_compile_dshape_current   |    +10.29 +/- 8.34     | +2.51e-01 +/- 2.03e-01 |  2.69e+00 +/- 1.9e-01  |  2.44e+00 +/- 7.9e-02  |
 test_objective_compute_dshape_current   |     +3.34 +/- 7.44     | +1.09e-05 +/- 2.42e-05 |  3.36e-04 +/- 2.1e-05  |  3.26e-04 +/- 1.3e-05  |
 test_objective_jac_dshape_current       |    +18.60 +/- 39.89    | +2.62e-03 +/- 5.61e-03 |  1.67e-02 +/- 4.5e-03  |  1.41e-02 +/- 3.3e-03  |
 test_perturb_2                          |     +9.08 +/- 4.45     | +8.65e-01 +/- 4.25e-01 |  1.04e+01 +/- 3.8e-01  |  9.53e+00 +/- 1.8e-01  |
 test_proximal_jac_atf_with_eq_update    |     +1.30 +/- 2.44     | +8.37e-02 +/- 1.58e-01 |  6.55e+00 +/- 9.9e-02  |  6.46e+00 +/- 1.2e-01  |
 test_proximal_freeb_jac                 |     +2.55 +/- 2.21     | +7.88e-02 +/- 6.81e-02 |  3.17e+00 +/- 5.6e-02  |  3.09e+00 +/- 3.9e-02  |
 test_solve_fixed_iter_compiled          |     +1.42 +/- 4.75     | +5.54e-02 +/- 1.85e-01 |  3.95e+00 +/- 1.8e-01  |  3.89e+00 +/- 3.9e-02  |
 test_objective_compute_ripple_bounce1d  |     -5.07 +/- 6.91     | -8.64e-03 +/- 1.18e-02 |  1.62e-01 +/- 7.1e-03  |  1.70e-01 +/- 9.4e-03  |
 test_objective_grad_ripple_bounce1d     |     +0.89 +/- 4.20     | +5.33e-03 +/- 2.52e-02 |  6.06e-01 +/- 2.3e-02  |  6.01e-01 +/- 1.1e-02  |
 test_objective_quadratic_flux_jac       |     -1.92 +/- 1.88     | -2.18e-02 +/- 2.14e-02 |  1.11e+00 +/- 1.6e-02  |  1.14e+00 +/- 1.4e-02  |

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

@codecov

codecov Bot commented Aug 14, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.92473% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 94.36%. Comparing base (3d01c2b) to head (2d6f3ff).

Files with missing lines Patch % Lines
desc/backend.py 95.65% 1 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##           master    #2293   +/-   ##
=======================================
  Coverage   94.35%   94.36%           
=======================================
  Files         101      101           
  Lines       29046    29074   +28     
=======================================
+ Hits        27407    27436   +29     
+ Misses       1639     1638    -1     
Files with missing lines Coverage Δ
desc/objectives/utils.py 100.00% <100.00%> (ø)
desc/optimize/_constraint_wrappers.py 96.89% <100.00%> (-0.02%) ⬇️
desc/optimize/_desc_wrappers.py 92.68% <100.00%> (+1.50%) ⬆️
desc/optimize/aug_lagrangian_ls.py 96.00% <100.00%> (ø)
desc/optimize/least_squares.py 99.45% <100.00%> (ø)
desc/optimize/optimizer.py 96.23% <100.00%> (ø)
desc/optimize/tr_subproblems.py 99.43% <ø> (ø)
desc/optimize/utils.py 95.73% <100.00%> (+0.18%) ⬆️
desc/perturbations.py 89.21% <100.00%> (+0.05%) ⬆️
desc/backend.py 90.87% <95.65%> (+0.22%) ⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@YigitElma
YigitElma changed the base branch from master to yge/fix-compute-everything-jax-finufft August 16, 2026 17:14
@YigitElma
YigitElma changed the base branch from yge/fix-compute-everything-jax-finufft to master August 24, 2026 00:14
Comment thread desc/optimize/least_squares.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run_benchmarks Run timing benchmarks on this PR against current master branch waiting for other PRs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants