Apparently cupy has improved the linalg situation for CUDA-capable machines to the point we could benefit from using it.
I propose adding a MNE_USE_CUPY config var that people can set that automagically uses cupy when it's installed and useful to do so for linalg functions. Concretely this would mean:
- I create
_safe_svd and a new _cholesky (or others if people think of some we use) to use cupy or scipy.linalg based on MNE_USE_CUPY (and cupy presence, just like cuda).
- @jona-sassenhagen and I look at when
linalg.svd, linalg.solve are faster for some array sizes and/or common MNE operations.
- We test that the outputs are equivalent, to numerical precision (this will be a test, just like we have for CUDA).
- We test whether or not these actually work in parallel code. I suspect that they will not, in which case we can decide if
n_jobs > 1 that we should not use cupy for linalg functions.
- We replace
linalg.svd and linalg.cholesky in our codebase with these private functions.
- We extend to other functions like
linalg.qr, etc. as necessary (maybe separate PR).
CC @rkmaddox I know you suffer at the hands of slow scipy.linalg sometimes, this might make your life easier.
EDIT: From @jona-sassenhagen in #5253.
Apparently
cupyhas improved thelinalgsituation for CUDA-capable machines to the point we could benefit from using it.I propose adding a
MNE_USE_CUPYconfig var that people can set that automagically usescupywhen it's installed and useful to do so forlinalgfunctions. Concretely this would mean:_safe_svdand a new_cholesky(or others if people think of some we use) to usecupyorscipy.linalgbased on MNE_USE_CUPY (andcupypresence, just likecuda).linalg.svd,linalg.solveare faster for some array sizes and/or common MNE operations.n_jobs > 1that we should not usecupyforlinalgfunctions.linalg.svdandlinalg.choleskyin our codebase with these private functions.linalg.qr, etc. as necessary (maybe separate PR).CC @rkmaddox I know you suffer at the hands of slow
scipy.linalgsometimes, this might make your life easier.EDIT: From @jona-sassenhagen in #5253.