Merge pull request #116 from EHTJulia/ptiede-copyto

Switch to copyto! for better performance with Enzyme

Author: ptiede

Project.toml

@@ -1,6 +1,6 @@
 name = "VLBISkyModels"
 uuid = "d6343c73-7174-4e0f-bb64-562643efbeca"
-version = "0.6.20"
+version = "0.6.21"
 authors = ["Paul Tiede <ptiede91@gmail.com> and contributors"]
 
 [deps]
@@ -36,18 +36,21 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 FINUFFT = "d8beea63-0952-562e-9c6a-8e8ef7364055"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 NonuniformFFTs = "cd96f58b-6017-4a02-bb9e-f4d81626177f"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+
 
 [extensions]
 VLBISkyModelsFINUFFT = ["FINUFFT"]
 VLBISkyModelsMakieExt = ["Makie", "DimensionalData"]
 VLBISkyModelsNonuniformFFTs = ["NonuniformFFTs"]
+VLBISkyModelsReactantExt = ["Reactant"]
 
 [compat]
 AbstractFFTs = "1"
 Accessors = "0.1"
 ArgCheck = "2"
 ChainRulesCore = "1"
-ComradeBase = "^0.9.6"
+ComradeBase = "^0.9.8"
 DelimitedFiles = "1"
 DimensionalData = "0.29 - 0.29.24, ^0.29.26"
 DocStringExtensions = "0.6,0.7,0.8,0.9"
@@ -66,6 +69,7 @@ NonuniformFFTs = "0.9"
 PaddedViews = "0.5"
 PolarizedTypes = "^0.1.1"
 Printf = "1.8"
+Reactant = "0.2"
 RecipesBase = "1"
 Reexport = "1"
 Serialization = "1.8"
@@ -78,7 +82,8 @@ julia = "1.9"
 FINUFFT = "d8beea63-0952-562e-9c6a-8e8ef7364055"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 NonuniformFFTs = "cd96f58b-6017-4a02-bb9e-f4d81626177f"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "FINUFFT", "Makie", "NonuniformFFTs"]
+test = ["Test", "FINUFFT", "Makie", "NonuniformFFTs", "Reactant"]

docs/src/base_api.md

@@ -89,6 +89,7 @@ ComradeBase.phasecenter
 ComradeBase.executor
 ComradeBase.Serial
 ComradeBase.ThreadsEx
+ComradeBase.ReactantEx
 ComradeBase.header
 ComradeBase.NoHeader
 ComradeBase.MinimalHeader

ext/VLBISkyModelsFINUFFT.jl

@@ -70,7 +70,7 @@ EnzymeRules.inactive_type(::Type{<:FINUFFT.finufft_plan}) = true
 
 function VLBISkyModels._jlnuft!(out, A::FINUFFTPlan, b::AbstractArray{<:Real})
     bc = getcache(A)
-    bc .= b
+    copyto!(bc, b)
     FINUFFT.finufft_exec!(A.forward, bc, out)
     return nothing
 end

ext/VLBISkyModelsReactantExt.jl

@@ -0,0 +1,187 @@
+module VLBISkyModelsReactantExt
+
+using VLBISkyModels
+using AbstractFFTs
+using Reactant
+using NFFT
+using NFFT: AbstractNFFTs
+using VLBISkyModels: ReactantAlg
+using LinearAlgebra
+
+
+struct ReactantNFFTPlan{T, D, K <: AbstractArray, arrTc, vecI, vecII, FP, BP, INV, SM} <:
+    AbstractNFFTPlan{T, D, 1}
+    N::NTuple{D, Int}
+    NOut::NTuple{1, Int}
+    J::Int
+    k::K
+    Ñ::NTuple{D, Int}
+    dims::UnitRange{Int}
+    forwardFFT::FP
+    backwardFFT::BP
+    tmpVec::arrTc
+    tmpVecHat::arrTc
+    deconvolveIdx::vecI
+    windowLinInterp::vecII
+    windowHatInvLUT::INV
+    B::SM
+end
+
+
+function VLBISkyModels.plan_nuft_spatial(
+        ::ReactantAlg,
+        imgdomain::ComradeBase.AbstractRectiGrid,
+        visdomain::ComradeBase.UnstructuredDomain,
+    )
+    visp = domainpoints(visdomain)
+    uv2 = similar(visp.U, (2, length(visdomain)))
+    dpx = pixelsizes(imgdomain)
+    dx = dpx.X
+    dy = dpx.Y
+    rm = ComradeBase.rotmat(imgdomain)'
+    # Here we flip the sign because the NFFT uses the -2pi convention
+    uv2[1, :] .= -VLBISkyModels._rotatex.(visp.U, visp.V, Ref(rm)) .* dx
+    uv2[2, :] .= -VLBISkyModels._rotatey.(visp.U, visp.V, Ref(rm)) .* dy
+    return ReactantNFFTPlan(uv2, size(imgdomain))
+end
+
+function VLBISkyModels.make_phases(
+        ::ReactantAlg,
+        imgdomain::ComradeBase.AbstractRectiGrid,
+        visdomain::ComradeBase.UnstructuredDomain,
+    )
+    return Reactant.to_rarray(VLBISkyModels.make_phases(NFFTAlg(), imgdomain, visdomain))
+end
+
+function VLBISkyModels._jlnuft!(out, A::ReactantNFFTPlan, inp::Reactant.AnyTracedRArray)
+    LinearAlgebra.mul!(out, A, inp)
+    return nothing
+end
+
+
+Base.adjoint(p::ReactantNFFTPlan) = p
+
+
+function AbstractNFFTs.plan_nfft(
+        arr::Type{<:Reactant.AnyTracedRArray},
+        k::AbstractMatrix,
+        N::NTuple{D, Int},
+        rest...;
+        kargs...,
+    ) where {D}
+    p = ReactantNFFTPlan(arr, k, N; kargs...)
+    return p
+end
+
+function ReactantNFFTPlan(
+        k::AbstractArray{T}, N::NTuple{D, Int}; fftflags = nothing, kwargs...
+    ) where {T, D}
+
+
+    dims = 1:D
+    CT = complex(T)
+    params, N, NOut, J, Ñ, dims_ = NFFT.initParams(k, N, dims; kwargs...)
+
+    # Get the correct type
+    FP = @jit plan_fft!(zeros(ComplexF64, 2, 2))
+    BP = @jit plan_bfft!(zeros(ComplexF64, 2, 2))
+
+    params.storeDeconvolutionIdx = true # GPU_NFFT only works this way
+    params.precompute = NFFT.FULL # GPU_NFFT only works this way
+
+    windowLinInterp, windowPolyInterp, windowHatInvLUT, deconvolveIdx, B = NFFT.precomputation(
+        k, N[dims_], Ñ[dims_], params
+    )
+
+    U = params.storeDeconvolutionIdx ? N : ntuple(d -> 0, Val(D))
+
+    tmpVec = Reactant.to_rarray(zeros(CT, Ñ))
+    tmpVecHat = Reactant.to_rarray(zeros(CT, U))
+    deconvIdx = Reactant.to_rarray(Int.(deconvolveIdx))
+    winHatInvLUT = Reactant.to_rarray(complex(windowHatInvLUT[1]))
+    B_ = (Reactant.to_rarray(complex.(Array(B))))
+
+    return ReactantNFFTPlan{
+        T,
+        D,
+        typeof(k),
+        typeof(tmpVec),
+        typeof(deconvIdx),
+        typeof(windowLinInterp),
+        typeof(FP),
+        typeof(BP),
+        typeof(winHatInvLUT),
+        typeof(B_),
+    }(
+        N,
+        NOut,
+        J,
+        k,
+        Ñ,
+        dims_,
+        FP,
+        BP,
+        tmpVec,
+        tmpVecHat,
+        deconvIdx,
+        windowLinInterp,
+        winHatInvLUT,
+        B_,
+    )
+end
+
+AbstractNFFTs.size_in(p::ReactantNFFTPlan) = p.N
+AbstractNFFTs.size_out(p::ReactantNFFTPlan) = p.NOut
+
+function AbstractNFFTs.convolve!(
+        p::ReactantNFFTPlan{T, D}, g::Reactant.AnyTracedRArray, fHat::Reactant.AnyTracedRArray
+    ) where {D, T}
+    mul!(fHat, transpose(p.B), vec(g))
+    return nothing
+end
+
+function AbstractNFFTs.convolve_transpose!(
+        p::ReactantNFFTPlan{T, D}, fHat::Reactant.AnyTracedRArray, g::Reactant.AnyTracedRArray
+    ) where {D, T}
+    mul!(vec(g), p.B, fHat)
+    return nothing
+end
+
+function Base.:*(p::ReactantNFFTPlan{T}, f::Reactant.AnyTracedRArray; kargs...) where {T}
+    fHat = similar(f, complex(T), size_out(p))
+    mul!(fHat, p, f; kargs...)
+    return fHat
+end
+
+function AbstractNFFTs.deconvolve!(
+        p::ReactantNFFTPlan{T, D}, f::AbstractArray, g::AbstractArray
+    ) where {D, T}
+    tmp = f .* reshape(p.windowHatInvLUT, size(f))
+    @allowscalar g[p.deconvolveIdx] = reshape(tmp, :)
+    return nothing
+end
+
+"""  in-place NFFT on the GPU"""
+function LinearAlgebra.mul!(
+        fHat::Reactant.AnyTracedRArray,
+        p::ReactantNFFTPlan{T, D},
+        f::Reactant.AnyTracedRArray;
+        verbose = false,
+        timing::Union{Nothing, TimingStats} = nothing,
+    ) where {T, D}
+    NFFT.consistencyCheck(p, f, fHat)
+
+    fill!(p.tmpVec, zero(Complex{T}))
+    t1 = @elapsed @inbounds deconvolve!(p, f, p.tmpVec)
+    fHat .= p.tmpVec[1:length(fHat)]
+    p.forwardFFT * p.tmpVec
+    return t3 = @elapsed @inbounds NFFT.convolve!(p, p.tmpVec, fHat)
+end
+
+function NFFT.nfft(k::AbstractMatrix, f::Reactant.AnyTracedRArray, args...; kwargs...)
+    p = ReactantNFFTPlan(typeof(f), k, size(f))
+    return p * f
+end
+
+
+end

src/fourierdomain/nuft/nfft_alg.jl

@@ -83,7 +83,7 @@ end
 @inline function _nuft!(out::AbstractArray, A, b::AbstractArray)
     tmp = similar(out)
     _jlnuft!(tmp, A, b)
-    out .= tmp
+    copyto!(out, tmp)
     return nothing
 end
 
@@ -183,7 +183,7 @@ function EnzymeRules.reverse(
     for (db, dout) in zip(dbs, douts)
         # TODO open PR on NFFT so we can do this in place.
         _jlnuft_adjointadd!(db, A.val, dout)
-        dout .= 0
+        fill!(dout, 0)
     end
     return (nothing, nothing, nothing)
 end

src/fourierdomain/nuft/nfft_reactant.jl

@@ -0,0 +1 @@
+struct ReactantAlg <: NUFT end

src/fourierdomain/nuft/nuft.jl

@@ -206,3 +206,5 @@ include(joinpath(@__DIR__, "dft_alg.jl"))
 include(joinpath(@__DIR__, "finufft.jl"))
 
 include(joinpath(@__DIR__, "nonuniformffts.jl"))
+
+include(joinpath(@__DIR__, "nfft_reactant.jl"))

src/models/combinators.jl

@@ -144,9 +144,9 @@ end
 function intensitymap_numeric!(sim::IntensityMap, m::AddModel)
     csim = copy(sim)
     intensitymap!(csim, m.m1)
-    sim .= csim
+    copyto!(sim, csim)
     intensitymap!(csim, m.m2)
-    sim .= sim .+ csim
+    sim .+= csim
     return nothing
 end
 
@@ -296,7 +296,7 @@ end
     ) where {M1, M2}
     cvis = similar(vis)
     visibilitymap!(cvis, m.m1)
-    vis .= cvis
+    copyto!(vis, cvis)
     visibilitymap!(cvis, m.m2)
     vis .*= cvis
     return nothing

src/models/continuous_image.jl

@@ -194,9 +194,10 @@ function applypulse!(vis, pulse, gfour::AbstractFourierDualDomain)
     # through the broadcast
     pvis = parent(vis)
     dp = domainpoints(guv)
-    for i in eachindex(pvis, dp)
-        pvis[i] *= visibility_point(mp, dp[i])
-    end
+    pvis .*= visibility_point.(Ref(mp), dp)
+    # for i in eachindex(pvis, dp)
+    #     pvis[i] *= visibility_point(mp, dp[i])
+    # end
     # pvis .*= visibility_point.(Ref(mp), dp)
     return vis
 end

src/models/geometric_models.jl

@@ -61,7 +61,6 @@ end
     return exp(-2 * T(π)^2 * (u^2 + v^2)) + zero(T)im
 end
 
-
 """
     $(TYPEDEF)