Backend based AbstractNFFT implementation

AbstractNFFTs/Project.toml

@@ -1,7 +1,7 @@
 name = "AbstractNFFTs"
 uuid = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
 author = ["Tobias Knopp <[email protected]>"]
-version = "0.8.3"
+version = "0.9.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

AbstractNFFTs/src/AbstractNFFTs.jl

@@ -4,6 +4,7 @@ using LinearAlgebra
 using Printf
 
 # interface
+export AbstractNFFTBackend
 export AbstractFTPlan, AbstractRealFTPlan, AbstractComplexFTPlan,
        AbstractNFFTPlan, AbstractNFCTPlan, AbstractNFSTPlan, AbstractNNFFTPlan, 
        plan_nfft, plan_nfct, plan_nfst, mul!, size_in, size_out, nodes!

AbstractNFFTs/src/derived.jl

@@ -10,30 +10,36 @@ planfunc = Symbol("plan_"*"$op")
 
 # The following automatically call the plan_* version for type Array
 
-$(planfunc)(k::AbstractArray, N::Union{Integer,NTuple{D,Int}}, args...; kargs...) where {D} =
-    $(planfunc)(Array, k, N, args...; kargs...)
+$(planfunc)(b::AbstractNFFTBackend, k::AbstractArray, N::Union{Integer,NTuple{D,Int}}, args...; kargs...) where {D} =
+    $(planfunc)(b, Array, k, N, args...; kargs...)
 
-$(planfunc)(k::AbstractArray, y::AbstractArray, args...; kargs...) =
-    $(planfunc)(Array, k, y, args...; kargs...)
+$(planfunc)(b::AbstractNFFTBackend, k::AbstractArray, y::AbstractArray, args...; kargs...) =
+    $(planfunc)(b, Array, k, y, args...; kargs...)
+
+$(planfunc)(k::AbstractArray, args...; kargs...) = $(planfunc)(active_backend(), k, args...; kargs...)
 
 # The follow convert 1D parameters into the format required by the plan
 
-$(planfunc)(Q::Type, k::AbstractVector, N::Integer, rest...; kwargs...)  =
-    $(planfunc)(Q, collect(reshape(k,1,length(k))), (N,), rest...; kwargs...)
+$(planfunc)(b::AbstractNFFTBackend, Q::Type, k::AbstractVector, N::Integer, rest...; kwargs...)  =
+    $(planfunc)(b, Q, collect(reshape(k,1,length(k))), (N,), rest...; kwargs...)
+
+$(planfunc)(b::AbstractNFFTBackend, Q::Type, k::AbstractVector, N::NTuple{D,Int}, rest...; kwargs...) where {D} =
+    $(planfunc)(b, Q, collect(reshape(k,1,length(k))), N, rest...; kwargs...) 
 
-$(planfunc)(Q::Type, k::AbstractVector, N::NTuple{D,Int}, rest...; kwargs...) where {D} =
-    $(planfunc)(Q, collect(reshape(k,1,length(k))), N, rest...; kwargs...) 
+$(planfunc)(b::AbstractNFFTBackend, Q::Type, k::AbstractMatrix, N::NTuple{D,Int}, rest...; kwargs...) where {D}  =
+    $(planfunc)(b, Q, collect(k), N, rest...; kwargs...)
 
-$(planfunc)(Q::Type, k::AbstractMatrix, N::NTuple{D,Int}, rest...; kwargs...) where {D}  =
-    $(planfunc)(Q, collect(k), N, rest...; kwargs...)
+$(planfunc)(Q::Type, args...; kwargs...) = $(planfunc)(active_backend(), Q, args...; kwargs...)
 
+$(planfunc)(::Missing, args...; kwargs...) = no_backend_error()
 end
 end
 
 ## NNFFT constructor
-plan_nnfft(Q::Type, k::AbstractVector, y::AbstractVector, rest...; kwargs...)  =
-    plan_nnfft(Q, collect(reshape(k,1,length(k))), collect(reshape(y,1,length(k))), rest...; kwargs...)
-
+plan_nnfft(Q::Type, args...; kwargs...) = plan_nnfft(active_backend(), Q, args...; kwargs...)
+plan_nnfft(b::AbstractNFFTBackend, Q::Type, k::AbstractVector, y::AbstractVector, rest...; kwargs...)  =
+    plan_nnfft(b, Q, collect(reshape(k,1,length(k))), collect(reshape(y,1,length(k))), rest...; kwargs...)
+plan_nnfft(::Missing, args...; kwargs...) = no_backend_error()
 
 
 ###############################################
@@ -48,26 +54,32 @@ tfunc = Symbol("$(op)_$(trans)")
 
 # TODO fix comments (how?)
 """
-nfft(k, f, rest...; kwargs...)
+nfft(backend, k, f, rest...; kwargs...)
 
 calculates the nfft of the array `f` for the nodes contained in the matrix `k`
 The output is a vector of length M=`size(nodes,2)`
 """
-function $(op)(k, f::AbstractArray; kargs...) 
+$(op)(k, f::AbstractArray; kargs...) = $(op)(active_backend(), k, f::AbstractArray; kargs...) 
+function $(op)(b::AbstractNFFTBackend, k, f::AbstractArray; kargs...) 
   p = $(planfunc)(k, size(f); kargs... )
   return p * f
 end
+$(op)(::Missing, k, f::AbstractArray; kargs...) = no_backend_error()
+
 
 """
-nfft_adjoint(k, N, fHat, rest...; kwargs...)
+nfft_adjoint(backend, k, N, fHat, rest...; kwargs...)
 
 calculates the adjoint nfft of the vector `fHat` for the nodes contained in the matrix `k`.
 The output is an array of size `N`
 """
-function $(tfunc)(k, N, fHat;  kargs...) 
+$(tfunc)(k, N, fHat;  kargs...) = $(tfunc)(active_backend(), k, N, fHat;  kargs...)
+function $(tfunc)(b::AbstractNFFTBackend, k, N, fHat;  kargs...) 
   p = $(planfunc)(k, N;  kargs...)
   return $(trans)(p) * fHat
 end
+$(tfunc)(::Missing, k, N, fHat;  kargs...) = no_backend_error()
+
 
 end
 end

AbstractNFFTs/src/interface.jl

@@ -1,3 +1,25 @@
+abstract type AbstractNFFTBackend end
+const ACTIVE_BACKEND = Ref{Union{Missing, AbstractNFFTBackend}}(missing)
+
+"""
+    set_active_backend!(back::Union{Missing, Module, AbstractNFFTBackend})
+
+Set the default NFFT plan backend. A module `back` must implement `back.backend()`.
+"""
+set_active_backend!(back::Module) = set_active_backend!(back.backend())
+function set_active_backend!(back::Union{Missing, AbstractNFFTBackend})
+  ACTIVE_BACKEND[] = back
+end
+active_backend() = ACTIVE_BACKEND[]
+function no_backend_error() 
+  error(
+    """
+    No default backend available!
+    Make sure to also "import/using" an NFFT backend such as NFFT or NonuniformFFTs.
+    """
+  )
+end
+
 """
   AbstractFTPlan{T,D,R}
 

NFFTTools/Project.toml

@@ -12,5 +12,5 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 [compat]
 julia = "1.6"
 AbstractFFTs = "1.0"
-AbstractNFFTs = "0.6, 0.7, 0.8"
+AbstractNFFTs = "0.6, 0.7, 0.8, 0.9"
 FFTW = "1"

Project.toml

@@ -1,7 +1,7 @@
 name = "NFFT"
 uuid = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 authors = ["Tobias Knopp <[email protected]>"]
-version = "0.13.7"
+version = "0.14"
 
 [deps]
 AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
@@ -19,7 +19,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
 Adapt = "3, 4"
-AbstractNFFTs = "0.8"
+AbstractNFFTs = "0.9"
 BasicInterpolators = "0.6.5, 0.7"
 DataFrames = "1.3.1, 1.4.1"
 FFTW = "1.5"

ext/NFFTGPUArraysExt/implementation.jl

@@ -16,7 +16,7 @@ mutable struct GPU_NFFTPlan{T,D, arrTc <: AbstractGPUArray{Complex{T}, D}, vecI
   B::SM
 end
 
-function AbstractNFFTs.plan_nfft(arr::Type{<:AbstractGPUArray}, k::Matrix{T}, N::NTuple{D,Int}, rest...;
+function AbstractNFFTs.plan_nfft(::NFFTBackend, 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...)

src/NFFT.jl

@@ -16,7 +16,7 @@ using BasicInterpolators
 
 @reexport using AbstractNFFTs
 
-export NDFTPlan, NDCTPlan, NDSTPlan, NNDFTPlan, 
+export NFFTBackend, NDFTPlan, NDCTPlan, NDSTPlan, NNDFTPlan, 
        NFFTPlan, NFFTParams
 
 
@@ -37,13 +37,16 @@ include("precomputation.jl")
 #################
 # factory methods
 #################
+struct NFFTBackend <: AbstractNFFTBackend end
+activate!() = AbstractNFFTs.set_active_backend!(NFFT)
+backend() = NFFTBackend()
 
 """
-        plan_nfft(k::Matrix{T}, N::NTuple{D,Int}, rest...;  kargs...)
+    NFFT.plan_nfft(k::Matrix{T}, N::NTuple{D,Int}, rest...;  kargs...)
 
 compute a plan for the NFFT of a size-`N` array at the nodes contained in `k`.
 """
-function AbstractNFFTs.plan_nfft(::Type{<:Array}, k::Matrix{T}, N::NTuple{D,Int}, rest...;
+function AbstractNFFTs.plan_nfft(::NFFTBackend, ::Type{<:Array}, k::Matrix{T}, N::NTuple{D,Int}, rest...;
                    timing::Union{Nothing,TimingStats} = nothing, kargs...) where {T,D}
   t = @elapsed begin
     p = NFFTPlan(k, N, rest...; kargs...)
@@ -60,6 +63,7 @@ include("convolution.jl")
 
 function __init__()
   NFFT._use_threads[] = (Threads.nthreads() > 1)
+  activate!()
 end
 
 include("precompile.jl")

src/precompile.jl

@@ -6,7 +6,7 @@ using PrecompileTools
       J, N = 8, 16
       k = range(-0.4, stop=0.4, length=J)  # nodes at which the NFFT is evaluated
       f = randn(ComplexF64, J)             # data to be transformed
-      p = plan_nfft(k, N, reltol=1e-9)     # create plan
+      p = plan_nfft(NFFTBackend(), k, N, reltol=1e-9)     # create plan
       fHat = adjoint(p) * f                # calculate adjoint NFFT
       y = p * fHat                         # calculate forward NFFT
 

test/constructors.jl

@@ -1,5 +1,13 @@
 @testset "Constructors" begin
 
+  @testset "Backend" begin
+    @test NFFT.backend() isa NFFTBackend
+    AbstractNFFTs.set_active_backend!(missing)
+    @test ismissing(AbstractNFFTs.active_backend())
+    NFFT.activate!()
+    @test AbstractNFFTs.active_backend() isa NFFTBackend
+  end
+
   @test_throws ArgumentError NFFTPlan(zeros(1,4), (2,2))
 
   p = NFFTPlan(zeros(2,4), (2,2))