Make some computations in DFTK GPU-compatible

src/PlaneWaveBasis.jl

@@ -15,7 +15,7 @@ Discretization information for ``k``-point-dependent quantities such as orbitals
 More generally, a ``k``-point is a block of the Hamiltonian;
 eg collinear spin is treated by doubling the number of kpoints.
 """
-struct Kpoint{T <: Real, GT <: AbstractArray}
+struct Kpoint{T <: Real, GT <: AbstractVector{Vec3{Int}}}
     spin::Int                     # Spin component can be 1 or 2 as index into what is
     #                             # returned by the `spin_components` function
     coordinate::Vec3{T}           # Fractional coordinate of k-point
@@ -39,10 +39,9 @@ Normalization conventions:
 
 `ifft` and `fft` convert between these representations.
 """
-struct PlaneWaveBasis{T, VT, T_G_vectors, T_r_vectors, T_kpt_G_vecs} <: AbstractBasis{
-    T
-} where {VT <: Real, T_G_vectors <: AbstractArray{Vec3{Int}}, T_r_vectors <: AbstractArray,
-         T_kpt_G_vecs <: AbstractVector{Vec3{Int}}}
+struct PlaneWaveBasis{T, VT, T_G_vectors, T_r_vectors, T_kpt_G_vecs} <: AbstractBasis{T} where
+    {VT <: Real, T_G_vectors <: AbstractArray{Vec3{Int}},
+    T_r_vectors <: AbstractArray{Vec3}, T_kpt_G_vecs <: AbstractVector{Vec3{Int}}}
     # T is the default type to express data, VT the corresponding bare value type (i.e. not dual)
     model::Model{T, VT}
 
@@ -141,7 +140,7 @@ Base.eltype(::PlaneWaveBasis{T}) where {T} = T
         mapping_inv = Dict(ifull => iball for (iball, ifull) in enumerate(mapping))
         for iσ = 1:model.n_spin_components
             push!(kpoints_per_spin[iσ],
-                  Kpoint{T,typeof(Gvecs_k)}(iσ, k, mapping, mapping_inv, Gvecs_k))
+                  Kpoint(iσ, k, mapping, mapping_inv, Gvecs_k))
         end
     end
     vcat(kpoints_per_spin...)  # put all spin up first, then all spin down
@@ -399,7 +398,7 @@ end
 The list of ``G + k`` vectors, in reduced coordinates.
 """
 function Gplusk_vectors(basis::PlaneWaveBasis, kpt::Kpoint)
-    coordinate = kpt.coordinate  # Avoid closure on kpt (not isbits)
+    coordinate = kpt.coordinate  # Accelerator: avoid closure on kpt (not isbits)
     map(G -> G + coordinate, G_vectors(basis, kpt))
 end
 
@@ -450,7 +449,7 @@ Returns nothing if outside the range of valid wave vectors.
     end
 end
 
-@inline function index_G_vectors(basis::PlaneWaveBasis, G::AbstractVector{<:Integer})
+function index_G_vectors(basis::PlaneWaveBasis, G::AbstractVector{<:Integer})
     index_G_vectors(basis.fft_size, G)
 end
 

src/scf/chi0models.jl

@@ -55,7 +55,7 @@ function (χ0::DielectricModel)(basis; kwargs...)
     C0  = 1 - εr
     iszero(C0) && return nothing  # Will yield no contribution
 
-    Gsq = [norm2(G) for G in G_vectors_cart(basis)]
+    Gsq = norm2.(G_vectors_cart(basis))
     apply_sqrtL = identity
     if χ0.localization != identity
         sqrtL = sqrt.(χ0.localization.(r_vectors(basis)))

src/scf/mixing.jl

@@ -50,7 +50,7 @@ end
 @timing "KerkerMixing" function mix_density(mixing::KerkerMixing, basis::PlaneWaveBasis,
                                             δF; kwargs...)
     T      = eltype(δF)
-    G²     = map(norm2, G_vectors_cart(basis))
+    G²     = norm2.(G_vectors_cart(basis))
     kTF    = T.(mixing.kTF)
     ΔDOS_Ω = T.(mixing.ΔDOS_Ω)
 

src/terms/hartree.jl

@@ -30,10 +30,9 @@ function TermHartree(basis::PlaneWaveBasis{T}, scaling_factor) where {T}
     # Solving the Poisson equation ΔV = -4π ρ in Fourier space
     # is multiplying elementwise by 4π / |G|^2.
 
-    GPUArraysCore.allowscalar() do
-        poisson_green_coeffs = 4T(π) ./ [norm2(G) for G in G_vectors_cart(basis)]
-        poisson_green_coeffs[1] = 0  # Compensating charge background => Zero DC
-    end
+    poisson_green_coeffs = 4T(π) ./ norm2.(G_vectors_cart(basis))
+    GPUArraysCore.@allowscalar poisson_green_coeffs[1] = 0  # Compensating charge background => Zero DC
+
     enforce_real!(basis, poisson_green_coeffs)  # Symmetrize Fourier coeffs to have real iFFT
     poisson_green_coeffs = to_device(basis.architecture, poisson_green_coeffs)
 

src/terms/xc.jl

@@ -97,7 +97,7 @@ end
         end
         if haskey(terms, :Vl) && any(x -> abs(x) > term.potential_threshold, terms.Vl)
             @warn "Meta-GGAs with a Δρ term have not yet been thoroughly tested." maxlog=1
-            mG² = [-norm2(G) for G in G_vectors_cart(basis)]
+            mG² = -norm2.(G_vectors_cart(basis))
             Vl  = reshape(terms.Vl, n_spin, basis.fft_size...)
             Vl_fourier = fft(basis, Vl[s, :, :, :])
             # TODO: forcing real-valued ifft; should be enforced at creation of array
@@ -255,7 +255,7 @@ function LibxcDensities(basis, max_derivative::Integer, ρ, τ)
     # Compute Δρ
     if max_derivative > 1
         Δρ_real = similar(ρ_real, n_spin, basis.fft_size...)
-        mG² = [-norm2(G) for G in G_vectors_cart(basis)]
+        mG² = -norm2.(G_vectors_cart(basis))
         for σ = 1:n_spin
             # TODO: forcing real-valued ifft; should be enforced at creation of array
             Δρ_real[σ, :, :, :] .= irfft(basis, mG² .* @view ρ_fourier[σ, :, :, :];