JuliaMath / NFFT.jl

Julia implementation of the Non-equidistant Fast Fourier Transform (NFFT)
Other
153 stars 28 forks source link

Infer arraytype for `plan_nfft` from `k` #142

Open nHackel opened 2 months ago

nHackel commented 2 months ago

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.

tknopp commented 2 months ago

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 https://github.com/JuliaMath/NFFT.jl/issues/100#issuecomment-1162676968. 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 commented 1 month ago

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?