Open nHackel opened 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.
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?
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.
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 wantArray
and if givenCuArray{Float32, 2, ...}
we wantCuArray
. We require this base type, since it is needed forAdapt.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 implementplan_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 exampleadjoint(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 theirplan_type
, but since theplan_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
andtranspose
. 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 whichk
is used only accepts dense CPU matrices. This means that for the GPU backends, we first infer the type correctly fromk
, then movek
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.