Init working GPU_Plan based on package extension

Project.toml

@@ -18,12 +18,14 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
+Adapt = "3, 4"
 AbstractNFFTs = "0.8"
 BasicInterpolators = "0.6.5, 0.7"
 DataFrames = "1.3.1, 1.4.1"
 FFTW = "1.5"
 FINUFFT = "3.0.1"
 FLoops = "0.2"
+GPUArrays = "8, 9, 10"
 Reexport = "1.0"
 PrecompileTools = "1"
 SpecialFunctions = "0.8, 0.10, 1, 2"
@@ -42,7 +44,13 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Ducc0 = "47ec601d-2729-4ac9-bed9-2b3ab5fca9ff"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
 [targets]
 test = ["Test", "BenchmarkTools", "FINUFFT", "NFFT3", "CuNFFT", "Zygote",
         "NFFTTools", "DataFrames", "Ducc0"]  # "NFFTTools"  "CuNFFT"
+
+[extensions]
+NFFTGPUArraysExt = ["Adapt", "GPUArrays"]

ext/NFFTGPUArraysExt/NFFTGPUArraysExt.jl

@@ -0,0 +1,9 @@
+module NFFTGPUArraysExt
+
+using NFFT, NFFT.AbstractNFFTs
+using NFFT.SparseArrays, NFFT.LinearAlgebra, NFFT.FFTW
+using GPUArrays, Adapt
+
+include("implementation.jl")
+
+end

ext/NFFTGPUArraysExt/implementation.jl

@@ -0,0 +1,126 @@
+mutable struct GPU_NFFTPlan{T,D, arrTc <: AbstractGPUArray{Complex{T}, D}, vecI <: AbstractGPUVector{Int32}, FP, BP, SM} <: AbstractNFFTPlan{T,D,1} 
+  N::NTuple{D,Int64}
+  NOut::NTuple{1,Int64}
+  J::Int64
+  k::Matrix{T}
+  Ñ::NTuple{D,Int64}
+  dims::UnitRange{Int64}
+  params::NFFTParams{T}
+  forwardFFT::FP
+  backwardFFT::BP
+  tmpVec::arrTc
+  tmpVecHat::arrTc
+  deconvolveIdx::vecI
+  windowLinInterp::Vector{T}
+  windowHatInvLUT::arrTc
+  B::SM
+end
+
+function AbstractNFFTs.plan_nfft(arr::Type{<:AbstractGPUArray}, k::Matrix{T}, N::NTuple{D,Int}, rest...;
+  timing::Union{Nothing,TimingStats} = nothing, kargs...) where {T,D}
+  t = @elapsed begin
+    p = GPU_NFFTPlan(arr, k, N, rest...; kargs...)
+  end
+  if timing != nothing
+    timing.pre = t
+  end
+  return p
+end
+
+function GPU_NFFTPlan(arr, k::Matrix{T}, N::NTuple{D,Int}; dims::Union{Integer,UnitRange{Int64}}=1:D,
+                 fftflags=nothing, kwargs...) where {T,D}
+
+    if dims != 1:D
+      error("GPU NFFT does not work along directions right now!")
+    end
+
+    params, N, NOut, J, Ñ, dims_ = NFFT.initParams(k, N, dims; kwargs...)
+    params.storeDeconvolutionIdx = true # GPU_NFFT only works this way
+    params.precompute = NFFT.FULL # GPU_NFFT only works this way
+
+    tmpVec = adapt(arr, zeros(Complex{T}, Ñ))
+
+    FP = plan_fft!(tmpVec, dims_)
+    BP = plan_bfft!(tmpVec, dims_)
+
+    windowLinInterp, windowPolyInterp, windowHatInvLUT, deconvolveIdx, B = NFFT.precomputation(k, N[dims_], Ñ[dims_], params)
+
+    U = params.storeDeconvolutionIdx ? N : ntuple(d->0,D)
+    tmpVecHat = adapt(arr, zeros(Complex{T}, U))
+
+    deconvIdx = adapt(arr, Int32.(deconvolveIdx))
+    winHatInvLUT = adapt(arr, windowHatInvLUT[1]) 
+    B_ = adapt(arr, Complex{T}.(B)) # Bit hacky
+
+    GPU_NFFTPlan{T,D, typeof(tmpVec), typeof(deconvIdx), typeof(FP), typeof(BP), typeof(B_)}(N, NOut, J, k, Ñ, dims_, params, FP, BP, tmpVec, tmpVecHat, 
+               deconvIdx, windowLinInterp, winHatInvLUT, B_)
+end
+
+AbstractNFFTs.size_in(p::GPU_NFFTPlan) = p.N
+AbstractNFFTs.size_out(p::GPU_NFFTPlan) = p.NOut
+
+
+function AbstractNFFTs.convolve!(p::GPU_NFFTPlan{T,D, arrTc}, g::arrTc, fHat::arr) where {D,T,arr<: AbstractGPUArray, arrTc <: arr}
+  mul!(fHat, transpose(p.B), vec(g)) 
+  return
+end
+
+function AbstractNFFTs.convolve_transpose!(p::GPU_NFFTPlan{T,D, arrTc}, fHat::arr, g::arrTc) where {D,T,arr<: AbstractGPUArray, arrTc <: arr}
+  mul!(vec(g), p.B, fHat)
+  return
+end
+
+function AbstractNFFTs.deconvolve!(p::GPU_NFFTPlan{T,D, arrTc}, f::arr, g::arrTc) where {D,T,arr<: AbstractGPUArray, arrTc <: arr}
+  p.tmpVecHat[:] .= vec(f) .* p.windowHatInvLUT
+  g[p.deconvolveIdx] = p.tmpVecHat
+  return
+end
+
+function AbstractNFFTs.deconvolve_transpose!(p::GPU_NFFTPlan{T,D, arrTc}, g::arrTc, f::arr) where {D,T,arr<: AbstractGPUArray, arrTc <: arr}
+  p.tmpVecHat[:] = g[p.deconvolveIdx]
+  f[:] .= vec(p.tmpVecHat) .* p.windowHatInvLUT
+  return
+end
+
+"""  in-place NFFT on the GPU"""
+function LinearAlgebra.mul!(fHat::arr, p::GPU_NFFTPlan{T,D, arrT}, f::arr; 
+                          verbose=false, timing::Union{Nothing,TimingStats} = nothing) where {T,D,arr<: AbstractGPUArray, arrT <: arr} 
+    NFFT.consistencyCheck(p, f, fHat)
+
+    fill!(p.tmpVec, zero(Complex{T}))
+    t1 = @elapsed @inbounds deconvolve!(p, f, p.tmpVec)
+    t2 = @elapsed p.forwardFFT * p.tmpVec
+    t3 = @elapsed @inbounds convolve!(p, p.tmpVec, fHat)
+    if verbose
+        @info "Timing: deconv=$t1 fft=$t2 conv=$t3"
+    end
+    if timing != nothing
+      timing.conv = t3
+      timing.fft = t2
+      timing.deconv = t1
+    end
+
+    return fHat
+end
+
+"""  in-place adjoint NFFT on the GPU"""
+function LinearAlgebra.mul!(f::arr, pl::Adjoint{Complex{T},<:GPU_NFFTPlan{T,D, arrT}}, fHat::arr;
+                       verbose=false, timing::Union{Nothing,TimingStats} = nothing) where {T,D,arr<: AbstractGPUArray, arrT <: arr}
+    p = pl.parent
+    NFFT.consistencyCheck(p, f, fHat)
+
+    t1 = @elapsed @inbounds convolve_transpose!(p, fHat, p.tmpVec)
+    t2 = @elapsed p.backwardFFT * p.tmpVec
+    t3 = @elapsed @inbounds deconvolve_transpose!(p, p.tmpVec, f)
+    if verbose
+        @info "Timing: conv=$t1 fft=$t2 deconv=$t3"
+    end
+    if timing != nothing
+      timing.conv_adjoint = t1
+      timing.fft_adjoint = t2
+      timing.deconv_adjoint = t3
+    end
+
+    return f
+end
+