Skip to content

Misc - #2293

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

Misc#2293
YigitElma wants to merge 22 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.51 %    |     4.231e+03      |     4.083e+03      |   -148.32    |       34.76        |       30.34        |
  test_proximal_jac_w7x_with_eq_update   |   -2.14 %    |     6.844e+03      |     6.698e+03      |   -146.19    |       156.28       |       161.30       |
  test_proximal_freeb_jac                |   -0.32 %    |     1.355e+04      |     1.351e+04      |    -43.79    |       81.03        |       83.75        |
  test_proximal_freeb_jac_blocked        |   -0.52 %    |     7.878e+03      |     7.838e+03      |    -40.89    |       71.30        |       74.15        |
  test_proximal_freeb_jac_batched        |   -0.98 %    |     7.885e+03      |     7.808e+03      |    -77.32    |       72.23        |       73.04        |
  test_proximal_jac_ripple               |   -0.42 %    |     3.759e+03      |     3.743e+03      |    -15.89    |       55.59        |       57.35        |
  test_proximal_jac_ripple_bounce1d      |   -0.43 %    |     3.962e+03      |     3.945e+03      |    -17.05    |       70.32        |       71.22        |
  test_eq_solve                          |   -6.20 %    |     1.835e+03      |     1.721e+03      |   -113.75    |       51.74        |       54.16        |
  test_objective_quadratic_flux_jac      |   -0.33 %    |     1.899e+03      |     1.893e+03      |    -6.27     |       35.98        |       35.82        |

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         |     +1.36 +/- 4.32     | +1.16e-02 +/- 3.68e-02 |  8.65e-01 +/- 2.9e-02  |  8.54e-01 +/- 2.2e-02  |
 test_equilibrium_init_lowres            |     -0.72 +/- 2.53     | -4.72e-02 +/- 1.65e-01 |  6.50e+00 +/- 1.1e-01  |  6.54e+00 +/- 1.2e-01  |
 test_objective_compile_atf              |     -1.14 +/- 3.18     | -6.98e-02 +/- 1.94e-01 |  6.04e+00 +/- 1.1e-01  |  6.11e+00 +/- 1.6e-01  |
 test_objective_compute_atf              |     +1.22 +/- 13.17    | +2.54e-05 +/- 2.74e-04 |  2.10e-03 +/- 6.7e-05  |  2.08e-03 +/- 2.7e-04  |
 test_objective_jac_atf                  |     -0.22 +/- 5.43     | -3.38e-03 +/- 8.47e-02 |  1.55e+00 +/- 4.0e-02  |  1.56e+00 +/- 7.4e-02  |
 test_perturb_1                          |     -1.73 +/- 2.64     | -2.07e-01 +/- 3.17e-01 |  1.18e+01 +/- 2.6e-01  |  1.20e+01 +/- 1.8e-01  |
 test_proximal_jac_atf                   |     -0.50 +/- 1.64     | -2.61e-02 +/- 8.56e-02 |  5.18e+00 +/- 5.7e-02  |  5.20e+00 +/- 6.4e-02  |
 test_proximal_freeb_compute             |     -0.34 +/- 3.17     | -5.76e-04 +/- 5.31e-03 |  1.67e-01 +/- 4.4e-03  |  1.68e-01 +/- 3.0e-03  |
 test_solve_fixed_iter                   |     +0.10 +/- 2.48     | +2.48e-02 +/- 6.01e-01 |  2.42e+01 +/- 3.9e-01  |  2.42e+01 +/- 4.6e-01  |
 test_LinearConstraintProjection_build   |    -10.91 +/- 3.93     | -7.46e-01 +/- 2.68e-01 |  6.09e+00 +/- 1.7e-01  |  6.84e+00 +/- 2.1e-01  |
 test_objective_compute_ripple           |     -1.00 +/- 5.11     | -2.37e-03 +/- 1.21e-02 |  2.34e-01 +/- 9.3e-03  |  2.36e-01 +/- 7.7e-03  |
 test_objective_grad_ripple              |     +0.90 +/- 3.03     | +7.89e-03 +/- 2.66e-02 |  8.86e-01 +/- 1.9e-02  |  8.78e-01 +/- 1.8e-02  |
 test_objective_quadratic_flux_compute   |     +0.20 +/- 1.48     | +4.24e-05 +/- 3.18e-04 |  2.15e-02 +/- 2.4e-04  |  2.14e-02 +/- 2.1e-04  |
 test_build_transform_fft_midres         |     -2.33 +/- 2.83     | -2.16e-02 +/- 2.62e-02 |  9.07e-01 +/- 1.5e-02  |  9.28e-01 +/- 2.1e-02  |
 test_build_transform_fft_highres        |     -2.72 +/- 3.06     | -3.31e-02 +/- 3.73e-02 |  1.18e+00 +/- 2.1e-02  |  1.22e+00 +/- 3.0e-02  |
 test_equilibrium_init_medres            |     -2.36 +/- 2.63     | -1.70e-01 +/- 1.89e-01 |  7.02e+00 +/- 1.4e-01  |  7.19e+00 +/- 1.2e-01  |
 test_objective_compile_dshape_current   |     -1.07 +/- 5.08     | -4.64e-02 +/- 2.19e-01 |  4.27e+00 +/- 6.0e-02  |  4.32e+00 +/- 2.1e-01  |
 test_objective_compute_dshape_current   |     -8.19 +/- 10.17    | -6.07e-05 +/- 7.53e-05 |  6.80e-04 +/- 3.9e-05  |  7.41e-04 +/- 6.4e-05  |
 test_objective_jac_dshape_current       |     +2.36 +/- 25.13    | +5.61e-04 +/- 5.97e-03 |  2.43e-02 +/- 4.0e-03  |  2.38e-02 +/- 4.4e-03  |
 test_perturb_2                          |     -0.07 +/- 1.57     | -1.09e-02 +/- 2.47e-01 |  1.57e+01 +/- 1.7e-01  |  1.57e+01 +/- 1.7e-01  |
+test_proximal_jac_atf_with_eq_update    |     -3.90 +/- 1.03     | -4.67e-01 +/- 1.23e-01 |  1.15e+01 +/- 5.9e-02  |  1.20e+01 +/- 1.1e-01  |
 test_proximal_freeb_jac                 |     +1.10 +/- 9.26     | +5.19e-02 +/- 4.36e-01 |  4.76e+00 +/- 4.3e-01  |  4.71e+00 +/- 7.1e-02  |
 test_solve_fixed_iter_compiled          |     -1.76 +/- 3.19     | -1.04e-01 +/- 1.88e-01 |  5.78e+00 +/- 1.6e-01  |  5.88e+00 +/- 9.9e-02  |
 test_objective_compute_ripple_bounce1d  |     +0.97 +/- 3.39     | +2.91e-03 +/- 1.02e-02 |  3.05e-01 +/- 3.8e-03  |  3.02e-01 +/- 9.5e-03  |
 test_objective_grad_ripple_bounce1d     |     +0.17 +/- 2.91     | +1.63e-03 +/- 2.74e-02 |  9.41e-01 +/- 2.3e-02  |  9.39e-01 +/- 1.5e-02  |
 test_objective_quadratic_flux_jac       |     +2.17 +/- 1.64     | +4.76e-02 +/- 3.61e-02 |  2.24e+00 +/- 3.6e-02  |  2.20e+00 +/- 6.0e-03  |

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.90110% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 94.35%. Comparing base (7582fb3) to head (9a31c64).

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.35%           
=======================================
  Files         101      101           
  Lines       29036    29064   +28     
=======================================
+ Hits        27396    27424   +28     
  Misses       1640     1640           
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% <ø> (ø)
desc/optimize/least_squares.py 99.45% <ø> (ø)
desc/optimize/optimizer.py 96.23% <100.00%> (ø)
desc/optimize/tr_subproblems.py 99.43% <ø> (ø)
desc/optimize/utils.py 95.63% <100.00%> (+0.19%) ⬆️
desc/perturbations.py 89.21% <100.00%> (+0.05%) ⬆️
desc/backend.py 90.87% <95.65%> (+0.22%) ⬆️
🚀 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
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.

1 participant