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

Infer arraytype for plan_nfft from k #142

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Infer arraytype for plan_nfft from k #142

wants to merge 4 commits into from

Conversation

nHackel
Copy link
Collaborator

@nHackel nHackel commented Sep 7, 2024

This PR (hopefully) closes #100. The goal of this PR is to seamlessly create either CPU or GPU plans depending on the given argument(s). At the bottom of this comment I'll list some limitations of the current situation of NFFT.jl in regards to this switching.

k = ...

p_cpu = plan_nfft(k, ...; kwargs...)

using CUDA # or AMDGPU, ...

p_gpu = plan_nfft(CuArray(k), ...; kwargs...)

The basic problem of correctly inferring the type is that we need the "base" type of a given array. For example if given a Array{ComplexF32, 2} we want Array and if given CuArray{Float32, 2, ...} we want CuArray. We require this base type, since it is needed for Adapt.jl to correctly move temporary plan fields to the GPU. At the moment there is no stable API in Julia to achieve this, so we have to use some workaround.

A first option would be to introduce a new function. Since there is no suitable base function we could extend for this case (as far as I know) we have to define our own, for example plan_type(arr::AbstractArray) = Array. Then with package extensions we can implement plan_type(::CuArray{...} = CuArray and so on. Any array type we don't define, would have to be defined by a user and/or a new PR needs to be made.

The second option is to try to determine the "base" type of a given array using a (not stable) API as described in here. The benefit is that this also works without us specifying types for the different GPU backends, but it can break in future Julia versions. Potentially, we would need to support different variants of this function for different Julia versions.

Both approaches are just "best-effort" approaches and there will likely be cases where they fall. For example if k is not a dense array but some lazy-wrapped array, for example adjoint(k). In those edge-cases a user still needs to be able to create a fitting plan. For the first option a user could implement their plan_type, but since the plan_nfft(Array, k, ...; ...) interface is still available, I feel like this is superfluous and opted to implement the second option.

For convenience I also implemented the inferrence for adjoint and transpose. If there are any more common wrappers we want to support we can add them as package extensions.


With this PR we can chose the correct plan_type based on k. However, at the moment the initParams functions in which k is used only accepts dense CPU matrices. This means that for the GPU backends, we first infer the type correctly from k, then move k to the CPU and construct our GPU plan. I think this limitation exceeds the scope of this PR, but it is a little bit counter-intuitive from a user perspective.

Similarly neither our initParams nor the inferrence will work with a matrix-free construct like from LinearOperators.jl or LinearMappings.jl.

@nHackel nHackel marked this pull request as draft September 7, 2024 03:45
@tknopp
Copy link
Member

tknopp commented Sep 7, 2024

Thanks for pushing this forward!

I think the second approach is fine. Reading the PR in Julia base it looks like there is consensus that this is a needed feature. So if this will change at some point then only because there will be an official API and we can switch to that.

Regarding the limitations you basically nailed the issue. Originally my thinking was that the operator is always fully constructed on the CPU and only the final arrays are moved to the GPU. In that thinking it makes sense that initParams remains on the CPU and furthermore k would actually never need to get on the GPU which makes it strange that we move it there just to infer the type. I elaborated about this concern in this comment #100 (comment). However, the initialization will definitely also profit from GPU computing and thus I simply did not think it to the end (which you did now). And I agree that this would be another issue and not in the scope of this PR.

@JeffFessler
Copy link
Collaborator

Exciting! When I first started #100 I did not realize that something like the magic of strip_type_parameters would be needed. Neat solution.
Would the test suite benefit from some @inferred macros?

@nHackel nHackel marked this pull request as ready for review September 11, 2024 20:26
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.

Suggestion: infer plan type from x type
3 participants