enable selecting gpu

[brokenhammer]

Add a parameter "gpuid" in "set_device" to manually assign a GPU.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +6.45 +/- 6.93     | +3.44e-02 +/- 3.69e-02 |  5.67e-01 +/- 1.4e-02  |  5.33e-01 +/- 3.4e-02  |
 test_equilibrium_init_medres            |     -0.01 +/- 1.77     | -2.77e-04 +/- 7.29e-02 |  4.12e+00 +/- 5.5e-02  |  4.12e+00 +/- 4.8e-02  |
 test_equilibrium_init_highres           |     +0.68 +/- 2.06     | +3.65e-02 +/- 1.11e-01 |  5.43e+00 +/- 8.0e-02  |  5.39e+00 +/- 7.7e-02  |
 test_objective_compile_dshape_current   |     +3.00 +/- 2.48     | +1.17e-01 +/- 9.74e-02 |  4.04e+00 +/- 8.4e-02  |  3.92e+00 +/- 4.9e-02  |
 test_objective_compute_dshape_current   |     +4.09 +/- 3.34     | +2.09e-04 +/- 1.71e-04 |  5.33e-03 +/- 1.6e-04  |  5.12e-03 +/- 6.3e-05  |
 test_objective_jac_dshape_current       |     +2.61 +/- 7.51     | +1.10e-03 +/- 3.17e-03 |  4.33e-02 +/- 2.3e-03  |  4.22e-02 +/- 2.1e-03  |
 test_perturb_2                          |     -2.31 +/- 4.33     | -4.63e-01 +/- 8.68e-01 |  1.96e+01 +/- 2.6e-01  |  2.00e+01 +/- 8.3e-01  |
 test_proximal_freeb_jac                 |     -3.96 +/- 1.88     | -3.06e-01 +/- 1.45e-01 |  7.43e+00 +/- 5.0e-02  |  7.73e+00 +/- 1.4e-01  |
+test_solve_fixed_iter                   |     -5.62 +/- 1.73     | -1.90e+00 +/- 5.85e-01 |  3.19e+01 +/- 3.9e-01  |  3.38e+01 +/- 4.3e-01  |
 test_LinearConstraintProjection_build   |     -3.52 +/- 7.06     | -3.73e-01 +/- 7.49e-01 |  1.02e+01 +/- 1.9e-01  |  1.06e+01 +/- 7.3e-01  |
 test_build_transform_fft_midres         |     -0.51 +/- 4.20     | -3.14e-03 +/- 2.57e-02 |  6.08e-01 +/- 2.3e-02  |  6.11e-01 +/- 1.2e-02  |
 test_build_transform_fft_highres        |     -0.35 +/- 2.24     | -3.36e-03 +/- 2.18e-02 |  9.69e-01 +/- 1.6e-02  |  9.72e-01 +/- 1.5e-02  |
 test_equilibrium_init_lowres            |     -1.37 +/- 3.05     | -5.38e-02 +/- 1.19e-01 |  3.86e+00 +/- 8.4e-02  |  3.92e+00 +/- 8.5e-02  |
 test_objective_compile_atf              |     -0.48 +/- 0.84     | -3.91e-02 +/- 6.82e-02 |  8.07e+00 +/- 4.3e-02  |  8.11e+00 +/- 5.3e-02  |
 test_objective_compute_atf              |     +1.44 +/- 2.32     | +2.30e-04 +/- 3.70e-04 |  1.62e-02 +/- 2.2e-04  |  1.59e-02 +/- 3.0e-04  |
 test_objective_jac_atf                  |     -1.72 +/- 3.38     | -3.41e-02 +/- 6.69e-02 |  1.95e+00 +/- 5.3e-02  |  1.98e+00 +/- 4.1e-02  |
 test_perturb_1                          |     +0.02 +/- 1.14     | +3.02e-03 +/- 1.67e-01 |  1.47e+01 +/- 9.1e-02  |  1.46e+01 +/- 1.4e-01  |
 test_proximal_jac_atf                   |     +0.05 +/- 0.78     | +4.27e-03 +/- 6.43e-02 |  8.26e+00 +/- 4.7e-02  |  8.26e+00 +/- 4.4e-02  |
 test_proximal_freeb_compute             |     +0.09 +/- 1.03     | +1.86e-04 +/- 2.03e-03 |  1.96e-01 +/- 1.1e-03  |  1.96e-01 +/- 1.7e-03  |
 test_solve_fixed_iter_compiled          |     -0.45 +/- 0.58     | -9.24e-02 +/- 1.18e-01 |  2.03e+01 +/- 7.1e-02  |  2.04e+01 +/- 9.3e-02  |

[dpanici]

Thanks for the PR! Approving but I'd like @f0uriest to also look over before merging

[f0uriest]

My main question is what benefit this has over setting CUDA_VISIBLE_DEVICES=... which is what we currently recommend. I could also see there being possible issues in cases where a user tries to do both.

We'll also need to test this a bit manually since our CI doesn't do any GPU stuff.

[YigitElma]

To @f0uriest's point, I believe

os.environ["CUDA_VISIBLE_DEVICES"] = str(gpuid)

is also fine, but having this functionality in set_device will help printing correct device type and available memory at the initialization. Also, probably better for the default values of jac_chunk_size which checks the config["avail_mem"], iirc.

[codecov]

Codecov Report

Attention: Patch coverage is 10.00000% with 9 lines in your changes missing coverage. Please review.

Project coverage is 95.63%. Comparing base (f8c66c7) to head (2fb7584).

Files with missing lines	Patch %	Lines
desc/__init__.py	10.00%	9 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1490      +/-   ##
==========================================
- Coverage   95.63%   95.63%   -0.01%     
==========================================
  Files         101      101              
  Lines       25542    25546       +4     
==========================================
+ Hits        24428    24431       +3     
- Misses       1114     1115       +1     

Files with missing lines	Coverage Δ
desc/__init__.py	44.26% <10.00%> (-3.11%)	⬇️

... and 2 files with indirect coverage changes

[brokenhammer]

To @f0uriest's point, I believe

os.environ["CUDA_VISIBLE_DEVICES"] = str(gpuid)

is also fine, but having this functionality in set_device will help printing correct device type and available memory at the initialization. Also, probably better for the default values of jac_chunk_size which checks the config["avail_mem"], iirc.

Besides, adding the parameter makes it convenient to assign different GPUs to different processes in a multi-process run.

[YigitElma]

The black formatting test seems to fail. I believe if you fix that this should be good to go.