Skip to content

ENH: Use cupy #5439

@larsoner

Description

@larsoner

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:

  1. 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).
  2. @jona-sassenhagen and I look at when linalg.svd, linalg.solve are faster for some array sizes and/or common MNE operations.
  3. We test that the outputs are equivalent, to numerical precision (this will be a test, just like we have for CUDA).
  4. 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.
  5. We replace linalg.svd and linalg.cholesky in our codebase with these private functions.
  6. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions