Reorganize macros

Author: abussy

src/elements.jl

@@ -39,14 +39,18 @@ eval_psp_energy_correction(T, ::Element) = zero(T)
 eval_psp_energy_correction(psp::Element) = eval_psp_energy_correction(Float64, psp)
 
 # Fall back to the Gaussian table for Elements without pseudopotentials
-function valence_charge_density_fourier(el::Element, p::T)::T where {T <: Real}
+function valence_charge_density_fourier(el::Element, p)
     gaussian_valence_charge_density_fourier(el, p)
 end
 
 """Gaussian valence charge density using Abinit's coefficient table, in Fourier space."""
 function gaussian_valence_charge_density_fourier(el::Element, p::T)::T where {T <: Real}
     charge_ionic(el) * exp(-(p * atom_decay_length(el))^2)
 end
+function gaussian_valence_charge_density_fourier(el::Element, ps::AbstractVector{T}) where {T <: Real}
+    arch = architecture(ps)
+    to_device(arch, map(p -> gaussian_valence_charge_density_fourier(el, p), to_cpu(ps)))
+end
 
 function core_charge_density_fourier(::Element, ::T)::T where {T <: Real}
     error("Abstract elements do not necesesarily provide core charge density.")
@@ -160,40 +164,17 @@ charge_ionic(el::ElementPsp)      = charge_ionic(el.psp)
 has_core_density(el::ElementPsp)  = has_core_density(el.psp)
 eval_psp_energy_correction(T, el::ElementPsp) = eval_psp_energy_correction(T, el.psp)
 
-function local_potential_fourier(el::ElementPsp, p::T) where {T <: Real}
-    eval_psp_local_fourier(el.psp, p)
-end
-local_potential_real(el::ElementPsp, r::Real) = eval_psp_local_real(el.psp, r)
+local_potential_fourier(el::ElementPsp, p) = eval_psp_local_fourier(el.psp, p)
+local_potential_real(el::ElementPsp, r) = eval_psp_local_real(el.psp, r)
 
-function valence_charge_density_fourier(el::ElementPsp, p::T) where {T <: Real}
+function valence_charge_density_fourier(el::ElementPsp, p)
     if has_valence_density(el.psp)
         eval_psp_density_valence_fourier(el.psp, p)
     else
         gaussian_valence_charge_density_fourier(el, p)
     end
 end
-function core_charge_density_fourier(el::ElementPsp, p::T) where {T <: Real}
-    eval_psp_density_core_fourier(el.psp, p)
-end
-
-# Vectorized versions of the above, specific implementation depending on the Psp type
-function local_potential_fourier(el::ElementPsp, ps::AbstractVector{T}) where {T <: Real}
-    eval_psp_local_fourier(el.psp, ps)
-end
-function local_potential_real(el::ElementPsp, rs::AbstractVector{T}) where {T <: Real}
-    eval_psp_local_real(el.psp, rs)
-end
-function valence_charge_density_fourier(el::ElementPsp, ps::AbstractVector{T}) where {T <: Real}
-    if has_valence_density(el.psp)
-        eval_psp_density_valence_fourier(el.psp, ps)
-    else
-        gaussian_valence_charge_density_fourier(el, ps)
-    end
-end
-function core_charge_density_fourier(el::ElementPsp, ps::AbstractVector{T}) where {T <: Real}
-    eval_psp_density_core_fourier(el.psp, ps)
-end
-
+core_charge_density_fourier(el::ElementPsp, p) = eval_psp_density_core_fourier(el.psp, p)
 
 #
 # ElementCohenBergstresser
@@ -263,7 +244,6 @@ function local_potential_fourier(el::ElementCohenBergstresser, p::T) where {T <:
 end
 # TODO Strictly speaking needs a eval_psp_energy_correction
 
-
 #
 # ElementGaussian
 #
@@ -298,25 +278,22 @@ end
 
 # Generic API expectations: all element functions taking a real space scalar |r| or a
 # reciprocal space scalar |p| should have a vectorized version accepting vectors of |r| or |p|.
-# This macro vectorizes element functions by calling existing single-value version elementwise.
-# This is GPU safe and generic. Performance critical functions should have their own
-# GPU-optimized implementation.
-macro vectorize_element_function(fn)
-    quote
-        function $fn(el::Element, arg::AbstractVector{T}) where {T <: Real}
-            arch = architecture(arg)
-            to_device(arch, map(p -> $fn(el, p), to_cpu(arg)))
+# The following loop vectorizes element functions by calling the single-value version
+# elementwise. This is GPU safe and generic. Performance critical functions should have their
+# own GPU-optimized implementation. Note: explicit loop over Element types in order to avoid
+# ambiguities with the specific ElementPsp functions.
+for fn in [:gaussian_valence_charge_density_fourier, :core_charge_density_fourier,
+           :local_potential_fourier, :local_potential_real]
+    for el_type in [ElementCoulomb, ElementCohenBergstresser, ElementGaussian]
+        @eval begin
+            function DFTK.$fn(el::$el_type, arg::AbstractVector{T}) where {T <: Real}
+                arch = architecture(arg)
+                to_device(arch, map(p -> $fn(el, p), to_cpu(arg)))
+            end
         end
     end
 end
 
-# Generic vectorized element functions
-@vectorize_element_function DFTK.valence_charge_density_fourier
-@vectorize_element_function DFTK.gaussian_valence_charge_density_fourier
-@vectorize_element_function DFTK.core_charge_density_fourier
-@vectorize_element_function DFTK.local_potential_fourier
-@vectorize_element_function DFTK.local_potential_real
-
 #
 # Helper functions
 #

src/pseudo/NormConservingPsp.jl

@@ -48,15 +48,15 @@ abstract type NormConservingPsp end
 # have their own GPU-optimized implementation instead of relying on this macro. The
 # different norm-conserving pseudopotential types are responsible for the implementation
 # of vectorized functions, whether by using these macros or not.
-macro vectorize_psp_function(fn, PspType)
+macro vectorize_psp_function(PspType, fn)
     quote
         function $fn(psp::$PspType, vec::AbstractVector{T}) where {T <: Real}
             arch = architecture(vec)
             to_device(arch, map(p -> $fn(psp, p), to_cpu(vec)))
         end
     end
 end
-macro vectorize_psp_projector_function(fn, PspType)
+macro vectorize_psp_projector_function(PspType, fn)
     quote
         function $fn(psp::$PspType, i, l, vec::AbstractVector{T}) where {T <: Real}
             arch = architecture(vec)

src/pseudo/PspHgh.jl

@@ -121,7 +121,7 @@ function eval_psp_local_fourier(psp::PspHgh, p::T) where {T <: Real}
 
     4T(π) * rloc^2 * (-Zion + sqrt(T(π) / 2) * rloc * t^2 * P) * exp(-t^2 / 2) / t^2
 end
-@vectorize_psp_function DFTK.eval_psp_local_fourier PspHgh
+@vectorize_psp_function PspHgh DFTK.eval_psp_local_fourier
 
 # [GTH98] (1)
 function eval_psp_local_real(psp::PspHgh, r::T) where {T <: Real}
@@ -133,7 +133,7 @@ function eval_psp_local_real(psp::PspHgh, r::T) where {T <: Real}
         + exp(-rr^2 / 2) * (cloc[1] + cloc[2] * rr^2 + cloc[3] * rr^4 + cloc[4] * rr^6)
     )
 end
-@vectorize_psp_function DFTK.eval_psp_local_real PspHgh
+@vectorize_psp_function PspHgh DFTK.eval_psp_local_real
 
 
 # [HGH98] (7-15) except they do it with plane waves normalized by 1/sqrt(Ω).
@@ -162,15 +162,15 @@ function eval_psp_projector_fourier(psp::PspHgh, i, l, p::T) where {T <: Real}
 
     error("Not implemented for l=$l and i=$i")
 end
-@vectorize_psp_projector_function DFTK.eval_psp_projector_fourier PspHgh
+@vectorize_psp_projector_function PspHgh DFTK.eval_psp_projector_fourier
 
 # [HGH98] (3)
 function eval_psp_projector_real(psp::PspHgh, i, l, r::T) where {T <: Real}
     rp = T(psp.rp[l + 1])
     ired = (4i - 1) / T(2)
     sqrt(T(2)) * r^(l + 2(i - 1)) * exp(-r^2 / 2rp^2) / rp^(l + ired) / sqrt(gamma(l + ired))
 end
-@vectorize_psp_projector_function DFTK.eval_psp_projector_real PspHgh
+@vectorize_psp_projector_function PspHgh DFTK.eval_psp_projector_real
 
 function eval_psp_energy_correction(T, psp::PspHgh)
     # By construction we need to compute the DC component of the difference

src/pseudo/PspLinComb.jl

@@ -82,11 +82,11 @@ end
 @make_psplincomb_call DFTK.eval_psp_density_core_real
 @make_psplincomb_call DFTK.eval_psp_density_core_fourier
 
-@vectorize_psp_function DFTK.eval_psp_local_real PspLinComb
-@vectorize_psp_function DFTK.eval_psp_local_fourier PspLinComb
-@vectorize_psp_function DFTK.eval_psp_density_valence_real PspLinComb
-@vectorize_psp_function DFTK.eval_psp_density_valence_fourier PspLinComb
-@vectorize_psp_function DFTK.eval_psp_density_core_real PspLinComb
-@vectorize_psp_function DFTK.eval_psp_density_core_fourier PspLinComb
-@vectorize_psp_projector_function DFTK.eval_psp_projector_real PspLinComb
-@vectorize_psp_projector_function DFTK.eval_psp_projector_fourier PspLinComb
+@vectorize_psp_function PspLinComb DFTK.eval_psp_local_real
+@vectorize_psp_function PspLinComb DFTK.eval_psp_local_fourier
+@vectorize_psp_function PspLinComb DFTK.eval_psp_density_valence_real
+@vectorize_psp_function PspLinComb DFTK.eval_psp_density_valence_fourier
+@vectorize_psp_function PspLinComb DFTK.eval_psp_density_core_real
+@vectorize_psp_function PspLinComb DFTK.eval_psp_density_core_fourier
+@vectorize_psp_projector_function PspLinComb DFTK.eval_psp_projector_real
+@vectorize_psp_projector_function PspLinComb DFTK.eval_psp_projector_fourier

src/pseudo/PspUpf.jl

@@ -174,7 +174,7 @@ has_core_density(psp::PspUpf) = !all(iszero, psp.r2_ρcore)
 function eval_psp_projector_real(psp::PspUpf, i, l, r::T)::T where {T<:Real}
     psp.r2_projs_interp[l+1][i](r) / r^2  # TODO if r is below a threshold, return zero
 end
-@vectorize_psp_projector_function DFTK.eval_psp_projector_real PspUpf
+@vectorize_psp_projector_function PspUpf DFTK.eval_psp_projector_real
 
 function eval_psp_projector_fourier(psp::PspUpf, i, l, p::T)::T where {T<:Real}
     # The projectors may have been cut off before the end of the radial mesh
@@ -185,7 +185,7 @@ function eval_psp_projector_fourier(psp::PspUpf, i, l, p::T)::T where {T<:Real}
     r2_proj = @view psp.r2_projs[l+1][i][1:ircut_proj]
     hankel(rgrid, r2_proj, l, p)
 end
-@vectorize_psp_projector_function DFTK.eval_psp_projector_fourier PspUpf
+@vectorize_psp_projector_function PspUpf DFTK.eval_psp_projector_fourier
 
 count_n_pswfc_radial(psp::PspUpf, l) = length(psp.r2_pswfcs[l+1])
 
@@ -194,7 +194,7 @@ pswfc_label(psp::PspUpf, i, l) = psp.pswfc_labels[l+1][i]
 function eval_psp_pswfc_real(psp::PspUpf, i, l, r::T)::T where {T<:Real}
     psp.r2_pswfcs_interp[l+1][i](r) / r^2  # TODO if r is below a threshold, return zero
 end
-@vectorize_psp_projector_function DFTK.eval_psp_pswfc_real PspUpf
+@vectorize_psp_projector_function PspUpf DFTK.eval_psp_pswfc_real
 
 function eval_psp_pswfc_fourier(psp::PspUpf, i, l, p::T)::T where {T<:Real}
     # Pseudo-atomic wavefunctions are _not_ currently cut off like the other
@@ -203,10 +203,10 @@ function eval_psp_pswfc_fourier(psp::PspUpf, i, l, p::T)::T where {T<:Real}
     # If issues arise, try cutting them off too.
     return hankel(psp.rgrid, psp.r2_pswfcs[l+1][i], l, p)
 end
-@vectorize_psp_projector_function DFTK.eval_psp_pswfc_fourier PspUpf
+@vectorize_psp_projector_function PspUpf DFTK.eval_psp_pswfc_fourier
 
 eval_psp_local_real(psp::PspUpf, r::T) where {T<:Real} = psp.vloc_interp(r)
-@vectorize_psp_function DFTK.eval_psp_local_real PspUpf
+@vectorize_psp_function PspUpf DFTK.eval_psp_local_real
 
 # Low-level function for the local part of the pseudopotential in reciprocal space
 function _eval_psp_local_fourier(quadrature, rgrid, vloc, Zion, p::T)::T where {T<:Real}
@@ -247,19 +247,19 @@ end
 function eval_psp_density_valence_real(psp::PspUpf, r::T) where {T<:Real}
     psp.r2_ρion_interp(r) / r^2  # TODO if r is below a threshold, return zero
 end
-@vectorize_psp_function DFTK.eval_psp_density_valence_real PspUpf
+@vectorize_psp_function PspUpf DFTK.eval_psp_density_valence_real
 
 function eval_psp_density_valence_fourier(psp::PspUpf, p::T) where {T<:Real}
     rgrid = @view psp.rgrid[1:psp.ircut]
     r2_ρion = @view psp.r2_ρion[1:psp.ircut]
     return hankel(rgrid, r2_ρion, 0, p)
 end
-@vectorize_psp_function DFTK.eval_psp_density_valence_fourier PspUpf
+@vectorize_psp_function PspUpf DFTK.eval_psp_density_valence_fourier
 
 function eval_psp_density_core_real(psp::PspUpf, r::T) where {T<:Real}
     psp.r2_ρcore_interp(r) / r^2  # TODO if r is below a threshold, return zero
 end
-@vectorize_psp_function DFTK.eval_psp_density_core_real PspUpf
+@vectorize_psp_function PspUpf DFTK.eval_psp_density_core_real
 
 function eval_psp_density_core_fourier(psp::PspUpf, p::T) where {T<:Real}
     rgrid = @view psp.rgrid[1:psp.ircut]