Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add GPU support based on (exchangable) structs for solver state #81

Merged
merged 71 commits into from
Jul 4, 2024

Conversation

nHackel
Copy link
Member

@nHackel nHackel commented Apr 2, 2024

This PR adds GPU support by changing the solver structs to the following pattern:

mutable struct Solver{matT, ...}
  A::matT
  # Other "static" fields
  state::AbstractSolverState{<:Solver}
end

mutable struct SolverState{T, tempT} <: AbstractSolverState{Solver}
  x::tempT
  rho::T
  # ...
  iteration::Int64
end

The trick here is to not fully constrain the state in the solver. This allows us to adapt the state based on the measurements b given to the solve/init! functions.

In this PR we use this to detect if the array type of our preallocated variables are the same as the one of b. If not we adapt our variables with similar, store the new state and initialize as usual. This also doesn't require any dependency on CUDA.

First tests have shown that the actual iteration method is as type-stable as it was before. I am seeing small performance regressions, however those seem to be in the init! and solve method and seem to stem from my setup having to handle both "old" and "new" variants.

ToDos:

  • CGNR
  • FISTA
  • POGM
  • OptISTA
  • ADMM
  • (Kaczmarz)
  • Direct Solvers
  • Update Callbacks
  • Update Proximal Maps
  • Update Tests for GPU Support

Future options for this setup:

We could also use different states per solver to encode variantes of an algorithm. Randomized Kaczmarz or FISTA with restart might get different types. A danger here is to add a lot of maintenance cost if the number of variants explodes.

We can use the States to implement "multiframe" measurements where b is a matrix. Then we need to allocate a State with a matrix of preallocated variables (if a solver supports this directly. Otherwise we can just loop over each measurement as a fallback. For this I would also introduce a trait for Single-/Multiframe Solvers).

We can still use package extensions to write specialized versions of the iteration functions which are then conditionally loaded. We could also wrap common array operations of our solvers into "our" own function and then just specialize those for combinations of CUDA, Sparse and Dense as we did it for Kaczmarz

@nHackel
Copy link
Member Author

nHackel commented Apr 3, 2024

This will close #54 and #64

@nHackel nHackel marked this pull request as ready for review June 28, 2024 07:34
@nHackel nHackel merged commit 1741afd into master Jul 4, 2024
2 of 8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants