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         |     -3.65 +/- 8.75     | -2.01e-02 +/- 4.83e-02 |  5.32e-01 +/- 3.2e-02  |  5.52e-01 +/- 3.6e-02  |
 test_equilibrium_init_medres            |     -5.63 +/- 1.96     | -2.47e-01 +/- 8.62e-02 |  4.15e+00 +/- 5.6e-02  |  4.40e+00 +/- 6.6e-02  |
 test_equilibrium_init_highres           |     +4.37 +/- 2.12     | +2.37e-01 +/- 1.15e-01 |  5.67e+00 +/- 6.9e-02  |  5.43e+00 +/- 9.2e-02  |
 test_objective_compile_dshape_current   |     +3.81 +/- 1.42     | +1.49e-01 +/- 5.58e-02 |  4.07e+00 +/- 5.0e-02  |  3.92e+00 +/- 2.4e-02  |
 test_objective_compute_dshape_current   |     +0.06 +/- 1.88     | +2.89e-06 +/- 9.75e-05 |  5.19e-03 +/- 6.3e-05  |  5.19e-03 +/- 7.4e-05  |
 test_objective_jac_dshape_current       |     +3.86 +/- 9.41     | +1.68e-03 +/- 4.09e-03 |  4.51e-02 +/- 3.1e-03  |  4.35e-02 +/- 2.6e-03  |
 test_perturb_2                          |     +5.19 +/- 4.06     | +1.04e+00 +/- 8.12e-01 |  2.10e+01 +/- 2.9e-01  |  2.00e+01 +/- 7.6e-01  |
 test_proximal_freeb_jac                 |     +2.72 +/- 1.17     | +2.03e-01 +/- 8.75e-02 |  7.68e+00 +/- 6.8e-02  |  7.47e+00 +/- 5.5e-02  |
 test_solve_fixed_iter                   |     +1.95 +/- 2.71     | +6.36e-01 +/- 8.84e-01 |  3.33e+01 +/- 5.7e-01  |  3.26e+01 +/- 6.7e-01  |
 test_LinearConstraintProjection_build   |     +3.09 +/- 2.34     | +3.23e-01 +/- 2.45e-01 |  1.08e+01 +/- 2.3e-01  |  1.04e+01 +/- 8.9e-02  |
 test_build_transform_fft_midres         |     +0.44 +/- 4.47     | +2.66e-03 +/- 2.72e-02 |  6.11e-01 +/- 2.6e-02  |  6.08e-01 +/- 7.0e-03  |
 test_build_transform_fft_highres        |     +3.32 +/- 3.19     | +3.24e-02 +/- 3.12e-02 |  1.01e+00 +/- 1.8e-02  |  9.77e-01 +/- 2.6e-02  |
 test_equilibrium_init_lowres            |     +6.26 +/- 4.13     | +2.42e-01 +/- 1.60e-01 |  4.11e+00 +/- 1.5e-01  |  3.87e+00 +/- 6.4e-02  |
 test_objective_compile_atf              |     +0.33 +/- 1.99     | +2.70e-02 +/- 1.63e-01 |  8.21e+00 +/- 5.7e-02  |  8.18e+00 +/- 1.5e-01  |
 test_objective_compute_atf              |     +7.74 +/- 3.90     | +1.23e-03 +/- 6.19e-04 |  1.71e-02 +/- 5.8e-04  |  1.59e-02 +/- 2.1e-04  |
 test_objective_jac_atf                  |     -0.15 +/- 2.67     | -3.03e-03 +/- 5.24e-02 |  1.96e+00 +/- 3.5e-02  |  1.97e+00 +/- 3.9e-02  |
 test_perturb_1                          |     +2.05 +/- 3.04     | +3.04e-01 +/- 4.50e-01 |  1.51e+01 +/- 4.1e-01  |  1.48e+01 +/- 1.9e-01  |
 test_proximal_jac_atf                   |     -0.82 +/- 0.77     | -6.78e-02 +/- 6.38e-02 |  8.20e+00 +/- 4.8e-02  |  8.27e+00 +/- 4.2e-02  |
 test_proximal_freeb_compute             |     -0.13 +/- 0.78     | -2.66e-04 +/- 1.57e-03 |  2.00e-01 +/- 1.2e-03  |  2.00e-01 +/- 1.0e-03  |
 test_solve_fixed_iter_compiled          |     -0.76 +/- 0.98     | -1.57e-01 +/- 2.01e-01 |  2.04e+01 +/- 1.4e-01  |  2.06e+01 +/- 1.5e-01  |

[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.62%. Comparing base (2fb7584) to head (c2eee70).
Report is 17 commits behind head on master.

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.62%   -0.02%     
==========================================
  Files         101      101              
  Lines       25546    25546              
==========================================
- Hits        24431    24428       -3     
- Misses       1115     1118       +3     

Files with missing lines	Coverage Δ
desc/__init__.py	44.26% <10.00%> (ø)

... 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.