Use columnwise_dots where applicable (#1127)

Author: abussy

src/DFTK.jl

@@ -44,6 +44,7 @@ include("common/quadrature.jl")
 include("common/hankel.jl")
 include("common/hydrogenic.jl")
 include("common/derivatives.jl")
+include("common/linalg.jl")
 
 export PspHgh
 export PspUpf
@@ -140,7 +141,6 @@ export PreconditionerNone
 export lobpcg_hyper
 export diag_full
 export diagonalize_all_kblocks
-include("eigen/linalg.jl")
 include("eigen/preconditioners.jl")
 include("eigen/diag.jl")
 

src/eigen/linalg.jl src/common/linalg.jlsrc/eigen/linalg.jl renamed to src/common/linalg.jl

@@ -33,8 +33,9 @@
 # other eigenvectors (which is not the case in many - all ? - other
 # implementations)
 
-# - Some functions are reimplemented in a GPU optimized way as part of
-# the DFTK CUDA Extension (ext/DFTKCUDAExt/lobpcg.jl).
+# - Some generic linear algebra functions are used in this implementation. They can be
+#   found in src/common/linalg.jl. GPU optimized versions of these functions are located
+#   in src/gpu/linalg.jl.
 
 
 ## TODO micro-optimization of buffer reuse

src/eigen/lobpcg_hyper_impl.jl

@@ -44,15 +44,12 @@ end
     if isnothing(ψ) || isnothing(occupation)
         return (; E=T(Inf), ops)
     end
-    occupation = [to_cpu(occk) for occk in occupation]
 
     E = zero(T)
     for (ik, ψk) in enumerate(ψ)
-        for iband = 1:size(ψk, 2)
-            ψnk = @views ψk[:, iband]
-            E += (basis.kweights[ik] * occupation[ik][iband]
-                  * real(dot(ψnk, Diagonal(term.kinetic_energies[ik]), ψnk)))
-        end
+        E += basis.kweights[ik] * 
+             sum(occupation[ik] .* 
+                 real(vec(columnwise_dots(ψk, Diagonal(term.kinetic_energies[ik]), ψk))))
     end
     E = mpi_sum(E, basis.comm_kpts)
 

src/terms/kinetic.jl

@@ -52,11 +52,7 @@ function ene_ops(term::TermMagnetic, basis::PlaneWaveBasis{T}, ψ, occupation;
 
     E = zero(T)
     for (ik, k) in enumerate(basis.kpoints)
-        for iband = 1:size(ψ[1], 2)
-            ψnk = @views ψ[ik][:, iband]
-            # TODO optimize this
-            E += basis.kweights[ik] * occupation[ik][iband] * real(dot(ψnk, ops[ik] * ψnk))
-        end
+        E += basis.kweights[ik] * sum(occupation[ik] .* vec(real(columnwise_dots(ψ[ik], ops[ik] * ψ[ik]))))
     end
     E = mpi_sum(E, basis.comm_kpts)
 

src/terms/magnetic.jl

@@ -70,7 +70,6 @@ end
             # We compute the forces from the irreductible BZ; they are symmetrized later.
             G_plus_k_cart = to_cpu(Gplusk_vectors_cart(basis, kpt))
             G_plus_k = Gplusk_vectors(basis, kpt)
-            occupationk = to_cpu(occupation[ik])
             form_factors = to_device(basis.architecture,
                                      build_projector_form_factors(element.psp, G_plus_k_cart))
 
@@ -91,9 +90,7 @@ end
                     map!(p -> -2π*im*p[α], twoπp, G_plus_k)
                     dPdR .= twoπp .* P
                     mul!(δHψk, P, C * (dPdR' * ψ[ik]))
-                    @views -sum(occupationk[iband] * basis.kweights[ik] *
-                                2real(dot(ψ[ik][:, iband], δHψk[:, iband]))
-                            for iband=1:size(ψ[ik], 2))
+                    -basis.kweights[ik]*sum(occupation[ik] .* 2vec(real(columnwise_dots(ψ[ik], δHψk))))
                 end  # α
             end  # r
         end  # kpt
@@ -311,11 +308,8 @@ function compute_dynmat_δH(::TermAtomicNonlocal, basis::PlaneWaveBasis{T}, ψ,
                     δHψk_plus_q = derivative_wrt_αs(model.positions, α, idx) do positions_αs
                         PDPψk(basis, positions_αs, psp_groups, kpt, kpt, ψ[ik])
                     end
-                    -sum(  2occupation[ik][iband] * basis.kweights[ik]
-                               * dot(δψk_plus_q[:, iband], δHψk[:, iband])
-                         + δoccupation[ik][iband]  * basis.kweights[ik]
-                               * 2real(dot(ψk[:, iband], δHψk_plus_q[:, iband]))
-                         for iband=1:size(ψk, 2))
+                    -basis.kweights[ik] * sum(2occupation[ik] .* vec(columnwise_dots(δψk_plus_q, δHψk)) +
+                                              δoccupation[ik] .* 2vec(real(columnwise_dots(ψk, δHψk_plus_q))))
                 end
             end
         end