Optimizing XC instantiation for GPUs (#1262)

Co-authored-by: Michael F. Herbst <info@michael-herbst.com>

Author: abussy

src/density_methods.jl

@@ -163,9 +163,9 @@ function atomic_density_superposition(basis::PlaneWaveBasis{T},
 
     # Pre-allocation of large arrays for GPU efficiency
     Gs = G_vectors(basis)
-    ρ = to_device(basis.architecture, zeros(Complex{T}, length(Gs)))
-    ρ_tmp = similar(ρ)
     indices = to_device(basis.architecture, collect(1:length(Gs)))
+    ρ = zeros_like(indices, Complex{T})
+    ρ_tmp = similar(ρ)
 
     for (igroup, group) in enumerate(basis.model.atom_groups)
         for iatom in group
@@ -184,8 +184,7 @@ end
 Returns the form factors at unique values of |G + q| (in Cartesian coordinates).
 Additionally, returns a mapping from any G index to the corresponding entry in the form_factors array.
 """
-function atomic_density_form_factors(basis::PlaneWaveBasis{T},
-                                     method::AtomicDensity ) where {T<:Real}
+function atomic_density_form_factors(basis::PlaneWaveBasis{T}, method::AtomicDensity) where {T<:Real}
     G_cart = to_cpu(G_vectors_cart(basis))
 
     iG2ifnorm_cpu = zeros(Int, length(G_cart))
@@ -194,49 +193,94 @@ function atomic_density_form_factors(basis::PlaneWaveBasis{T},
         p = norm(G)
         iG2ifnorm_cpu[iG] = get!(norm_indices, p, length(norm_indices) + 1)
     end
+    iG2ifnorm = to_device(basis.architecture, iG2ifnorm_cpu)
 
-    form_factors_cpu = zeros(T, length(norm_indices), length(basis.model.atom_groups))
-    for (p, ifnorm) in norm_indices
-        for (igroup, group) in enumerate(basis.model.atom_groups)
-            element = basis.model.atoms[first(group)]
-            form_factors_cpu[ifnorm, igroup] = atomic_density(element, p, method)
-        end
+    ni_pairs = collect(pairs(norm_indices))
+    ps = to_device(basis.architecture, first.(ni_pairs))
+    indices = to_device(basis.architecture, last.(ni_pairs))
+
+    form_factors = similar(ps, length(norm_indices), length(basis.model.atom_groups))
+    for (igroup, group) in enumerate(basis.model.atom_groups)
+        element = basis.model.atoms[first(group)]
+        @inbounds form_factors[indices, igroup] .= atomic_density(element, ps, method)
     end
 
-    form_factors = to_device(basis.architecture, form_factors_cpu)
-    iG2ifnorm = to_device(basis.architecture, iG2ifnorm_cpu)
     (; form_factors, iG2ifnorm)
 end
 
-function atomic_density(element::Element, Gnorm::T,
-                        ::ValenceDensityGaussian)::T where {T <: Real}
-    gaussian_valence_charge_density_fourier(element, Gnorm)
-end
+#
+# Check for presence of certain densities in elements
+#
+
+"""Check presence of model valence density (as an initial guess)."""
+has_valence_density(::Element) = false
+
+"""Check presence of model core charge density (non-linear core correction)."""
+has_core_density(::Element) = false
+
+"""Check presence of model core kinetic energy density (non-linear core correction for τ)."""
+has_core_kinetic_energy_density(::Element) = false
+
+has_core_density(el::ElementPsp)  = has_core_density(el.psp)
+has_core_kinetic_energy_density(el::ElementPsp) = has_core_kinetic_energy_density(el.psp)
+has_valence_density(el::ElementPsp) = has_valence_density(el.psp)
 
-function atomic_density(element::Element, Gnorm::T,
-                        ::ValenceDensityPseudo)::T where {T <: Real}
-    eval_psp_valence_density_fourier(element.psp, Gnorm)
+#
+# Atomic density dispatches
+#
+
+function atomic_density(element::Union{Element,<:ElementPsp}, Gnorm::Real, dens::AtomicDensity)
+    first(atomic_density(element, [Gnorm], dens))  # Dispatch to vector version
 end
 
-function atomic_density(element::Element, Gnorm::T,
-                        ::ValenceDensityAuto)::T where {T <: Real}
-    valence_charge_density_fourier(element, Gnorm)
+"""Gaussian valence charge density using Abinit's coefficient table, in Fourier space."""
+function atomic_density(element::Element, Gnorms::AbstractVector, ::ValenceDensityGaussian)
+    charge = charge_ionic(element)
+    decay_length = atom_decay_length(element)
+    map(Gnorms) do Gnorm
+        # Note: cannot pass element to GPU kernel, because not isbit
+        charge * exp(-(Gnorm * decay_length)^2)
+    end
+end
+function atomic_density(element::Element, Gnorms::AbstractVector, ::ValenceDensityAuto)
+    if has_valence_density(element)
+        atomic_density(element, Gnorms, ValenceDensityPseudo())
+    else
+        atomic_density(element, Gnorms, ValenceDensityGaussian())
+    end
 end
 
-function atomic_density(element::Element, Gnorm::T,
-                        ::CoreDensity)::T where {T <: Real}
-    has_core_density(element) ? core_charge_density_fourier(element, Gnorm) : zero(T)
+function atomic_density(element::Element, Gnorms::AbstractVector, ::CoreDensity)
+    @assert !has_core_density(element)
+    zeros_like(Gnorms)
+end
+function atomic_density(element::Element, Gnorms::AbstractVector, ::CoreKineticEnergyDensity)
+    @assert !has_core_kinetic_energy_density(element)
+    zeros_like(Gnorms)
 end
 
-function atomic_density(element::Element, Gnorm::T,
-                        ::CoreKineticEnergyDensity)::T where {T <: Real}
+function atomic_density(element::ElementPsp, Gnorms::AbstractVector, ::ValenceDensityPseudo)
+    eval_psp_valence_density_fourier(element.psp, Gnorms)
+end
+function atomic_density(element::ElementPsp, Gnorms::AbstractVector, ::CoreDensity)
+    if has_core_density(element)
+        eval_psp_core_density_fourier(element.psp, Gnorms)
+    else
+        zeros_like(Gnorms)
+    end
+end
+function atomic_density(element::ElementPsp, Gnorms::AbstractVector, ::CoreKineticEnergyDensity)
     if has_core_kinetic_energy_density(element)
-        eval_psp_core_kinetic_energy_density_fourier(element.psp, Gnorm)
+        eval_psp_core_kinetic_energy_density_fourier(element.psp, Gnorms)
     else
-        zero(T)
+        zeros_like(Gnorms)
     end
 end
 
+#
+# Some data we need for the Gaussian densities
+#
+
 # Get the lengthscale of the valence density for an atom with `n_elec_core` core
 # and `n_elec_valence` valence electrons.
 function atom_decay_length(n_elec_core, n_elec_valence)

src/elements.jl

@@ -7,6 +7,13 @@ using AtomsBase
 # Very likely `species`, and `charge_ionic` need to be defined as well.
 abstract type Element end
 
+"""Return the local potential for a passed coordinate or vector of coordinates in frequency space"""
+function local_potential_fourier end
+
+"""Return the local potential for a passed coordinate or vector of coordinates in real space"""
+function local_potential_real end
+
+
 """Return the chemical species corresponding to an element"""
 AtomsBase.species(::Element) = ChemicalSpecies(0)  # dummy atom
 
@@ -31,40 +38,23 @@ n_elec_core(el::Element) = charge_nuclear(el) - charge_ionic(el)
 """Return the pseudopotential family for the element if this is known, else `nothing`"""
 pseudofamily(::Element) = nothing
 
-"""Check presence of model core charge density (non-linear core correction)."""
-has_core_density(::Element) = false
-
-"""Check presence of model core kinetic energy density (non-linear core correction for τ)."""
-has_core_kinetic_energy_density(::Element) = false
-
 # The preceding functions are fallback implementations that should be altered as needed.
 
 eval_psp_energy_correction(T, ::Element) = zero(T)
-eval_psp_energy_correction(psp::Element) = eval_psp_energy_correction(Float64, psp)
+eval_psp_energy_correction(el::Element) = eval_psp_energy_correction(Float64, el)
 
 # Fall back to the Gaussian table for Elements without pseudopotentials
-function valence_charge_density_fourier(el::Element, p::T)::T where {T <: Real}
-    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)
+# Vector to single-argument fallbacks in the Elements
+for fn in [:local_potential_fourier, :local_potential_real]
+    @eval begin
+        function DFTK.$fn(el::Element, arg::AbstractVector{<:Real})
+            arch = architecture(arg)
+            to_device(arch, map(p -> $fn(el, p), to_cpu(arg)))
+        end
+    end
 end
 
-function core_charge_density_fourier(::Element, ::T)::T where {T <: Real}
-    error("Abstract elements do not necesesarily provide core charge density.")
-end
-
-function core_kinetic_energy_density_fourier(::Element, ::T)::T where {T <: Real}
-    error("Abstract elements do not necesesarily provide core kinetic energy density.")
-end
-
-# Generic vectorized version of local_potential_fourier, GPU-safe
-function local_potential_fourier(el::Element, ps::AbstractVector{T}) where {T <: Real}
-    arch = architecture(ps)
-    to_device(arch, map(p -> local_potential_fourier(el, p), to_cpu(ps)))
-end
 
 Base.show(io::IO, el::Element) = print(io, "$(typeof(el))(:$(species(el)))")
 
@@ -171,30 +161,13 @@ end
 AtomsBase.mass(el::ElementPsp)    = el.mass
 AtomsBase.species(el::ElementPsp) = el.species
 charge_ionic(el::ElementPsp)      = charge_ionic(el.psp)
-has_core_density(el::ElementPsp)  = has_core_density(el.psp)
-has_core_kinetic_energy_density(el::ElementPsp) = has_core_kinetic_energy_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
-# Vectorized version of the above
-function local_potential_fourier(el::ElementPsp, ps::AbstractVector{T}) where {T <: Real}
-    eval_psp_local_fourier(el.psp, ps)
-end
+# Function forwarding for ElementPsp (for each case both versions are needed to resolve ambiguities)
+local_potential_fourier(el::ElementPsp, p::Real) = eval_psp_local_fourier(el.psp, p)
+local_potential_fourier(el::ElementPsp, p::AbstractVector{<:Real}) = eval_psp_local_fourier(el.psp, p)
 local_potential_real(el::ElementPsp, r::Real) = eval_psp_local_real(el.psp, r)
-
-function valence_charge_density_fourier(el::ElementPsp, p::T) where {T <: Real}
-    if has_valence_density(el.psp)
-        eval_psp_valence_density_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_core_density_fourier(el.psp, p)
-end
-
+local_potential_real(el::ElementPsp, r::AbstractVector{<:Real}) = eval_psp_local_real(el.psp, r)
 
 #
 # ElementCohenBergstresser
@@ -264,7 +237,6 @@ function local_potential_fourier(el::ElementCohenBergstresser, p::T) where {T <:
 end
 # TODO Strictly speaking needs a eval_psp_energy_correction
 
-
 #
 # ElementGaussian
 #
@@ -289,7 +261,7 @@ function ElementGaussian(α, L; symbol=:X, mass=nothing)
     T = promote_type(typeof(α), typeof(L))
     ElementGaussian{T}(α, L, symbol, mass)
 end
-function local_potential_real(el::ElementGaussian, r)
+function local_potential_real(el::ElementGaussian, r::Real)
     -el.α / (√(2π) * el.L) * exp(- (r / el.L)^2 / 2)
 end
 function local_potential_fourier(el::ElementGaussian, p::Real)