Repository: JuliaMolSim/DFTK.jl
Pull request #1068: GPU optimization of LOBPCG
Author: login: abussy
Created: 2025-02-24T14:50:22Z | Merged: 2025-05-15T10:42:04Z

This PR is the result of a detailed profiling of the LOBPCG solver with NVIDIA's Nsight Systems. It allowed for the identification of various hot spots, where code is very slow during GPU runs.

In particular, there are many instances of explicit loops over matrix columns. This access pattern is not ideal, as the massive parallelism of the GPU is not fully exploited. Array operations on the whole matrix are far more efficient.

I measured speed-ups of the order of 30% on the whole LOBPCG iterative solver. Excluding the cost of the H x Psi product (not modified in this PR), the speed-ups reach 50%.

Unfortunately, this comes at the cost of some code readability. I left comments describing what is calculated when necessary.

Finally, I scrapped 2 loops using DFTK's custom threading (in `ortho! X vs Y` and `ldiv!` for the preconditioner). I made sure the effect is negligible on CPU runs ( tested with the default `n_DFTK` = `n_blas` thread option). It seems that simple BLAS threading on large array operations is quite efficient by itself.


Diff:
diff --git a/src/DFTK.jl b/src/DFTK.jl
index 1ec5bea8f1..ab698b2691 100644
--- a/src/DFTK.jl
+++ b/src/DFTK.jl
@@ -138,6 +138,7 @@ export PreconditionerNone
 export lobpcg_hyper
 export diag_full
 export diagonalize_all_kblocks
+include("eigen/linalg.jl")
 include("eigen/preconditioners.jl")
 include("eigen/diag.jl")
 
@@ -231,7 +232,10 @@ include("postprocess/refine.jl")
 # Workarounds
 include("workarounds/dummy_inplace_fft.jl")
 include("workarounds/forwarddiff_rules.jl")
-include("workarounds/gpu_arrays.jl")
+
+# Optimized generic GPU functions and GPU workarounds
+include("gpu/linalg.jl")
+include("gpu/gpu_arrays.jl")
 
 # Precompilation block with a basic workflow
 
diff --git a/src/eigen/linalg.jl b/src/eigen/linalg.jl
new file mode 100644
index 0000000000..0036cf69ec
--- /dev/null
+++ b/src/eigen/linalg.jl
@@ -0,0 +1,15 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function columnwise_dots(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of real(dot(A[:, i], M, B[:, i])), for all columns of
+# A, B, and matrix M
+@views function columnwise_dots(A::AbstractArray{T}, M, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], M, B[:, i])) for i = 1:size(A, 2)]
+end
diff --git a/src/eigen/lobpcg_hyper_impl.jl b/src/eigen/lobpcg_hyper_impl.jl
index d1771e9e5e..e74a7e0c8c 100644
--- a/src/eigen/lobpcg_hyper_impl.jl
+++ b/src/eigen/lobpcg_hyper_impl.jl
@@ -33,6 +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).
+
 
 ## TODO micro-optimization of buffer reuse
 ## TODO write a version that doesn't assume that B is well-conditioned, and doesn't reuse B applications at all
@@ -170,7 +173,7 @@ function B_ortho!(X, BX)
     rdiv!(BX, U)
 end
 
-normest(M) = maximum(abs.(diag(M))) + norm(M - Diagonal(diag(M)))
+normest(M) = maximum(abs, diag(M)) + norm(M - Diagonal(diag(M)))
 # Orthogonalizes X to tol
 # Returns the new X, the number of Cholesky factorizations algorithm, and the
 # growth factor by which small perturbations of X can have been magnified
@@ -252,24 +255,15 @@ end
 
 # Randomize the columns of X if the norm is below tol
 function drop_small!(X::AbstractArray{T}; tol=2eps(real(T))) where {T}
-    dropped = Int[]
-    for i=1:size(X,2)
-        n = norm(@views X[:,i])
-        if n <= tol
-            X[:,i] = randn(T, size(X,1))
-            push!(dropped, i)
-        end
-    end
+    dropped = findall(n -> n <= tol, columnwise_norms(X))
+    @views randn!(TaskLocalRNG(), X[:, dropped])
     dropped
 end
 
 # Find X that is orthogonal, and B-orthogonal to Y, up to a tolerance tol.
 @timing "ortho! X vs Y" function ortho!(X::AbstractArray{T}, Y, BY; tol=2eps(real(T))) where {T}
     # normalize to try to cheaply improve conditioning
-    parallel_loop_over_range(1:size(X, 2)) do i
-        n = norm(@views X[:,i])
-        @views X[:,i] ./= n
-    end
+    X ./= columnwise_norms(X)'
 
     niter = 1
     ninners = zeros(Int, 0)
@@ -322,7 +316,7 @@ end
 end
 
 function final_retval(X, AX, BX, λ, resid_history, niter, n_matvec)
-    λ_host = oftype(ones(eltype(λ), 1), λ)  # Copy to CPU for element-wise access
+    λ_host = to_cpu(λ)  # Copy to CPU for element-wise access
     if !issorted(λ_host)
         p = sortperm(λ_host)
         λ_host = λ_host[p]
@@ -336,6 +330,12 @@ function final_retval(X, AX, BX, λ, resid_history, niter, n_matvec)
      residual_history=resid_history[:, 1:niter+1], n_matvec)
 end
 
+# Computes λ = real((X' * AX) / (X' *BX)), for each column of X
+function compute_λ(X, AX, BX)
+    λs = @views [real((X[:, n]'*AX[:, n]) / (X[:, n]'BX[:, n])) for n=1:size(X, 2)]
+    oftype(real(X[:, 1]), λs)  # Offload to GPU if needed
+end
+
 ### The algorithm is Xn+1 = rayleigh_ritz(hcat(Xn, A*Xn, Xn-Xn-1))
 ### We follow the strategy of Hetmaniuk and Lehoucq, and maintain a B-orthonormal basis Y = (X,R,P)
 ### After each rayleigh_ritz step, the B-orthonormal X and P are deduced by an orthogonal rotation from Y
@@ -389,8 +389,7 @@ end
     end
     nlocked = 0
     niter = 0  # the first iteration is fake
-    λs = @views [real((X[:, n]'*AX[:, n]) / (X[:, n]'BX[:, n])) for n=1:M]
-    λs = oftype(real(X[:, 1]), λs)  # Offload to GPU if needed
+    λs = compute_λ(X, AX, BX)
     new_X  = X
     new_AX = AX
     new_BX = BX
@@ -431,9 +430,8 @@ end
         ### Compute new residuals
         @timing "Update residuals" begin
             new_R = new_AX .- new_BX .* λs'
-            @views for i = 1:size(X, 2)
-                resid_history[i + nlocked, niter+1] = norm(new_R[:, i])
-            end
+            norms = to_cpu(columnwise_norms(new_R))
+            @views resid_history[1 + nlocked: size(new_R, 2) + nlocked, niter+1] .= norms[:]
         end
         @debug niter resid_history[:, niter+1]
 
@@ -512,10 +510,9 @@ end
         end
 
         # Quick sanity check
-        for i = 1:size(X, 2)
-            @views if abs(BX[:, i]'X[:, i] - 1) >= sqrt(eps(real(eltype(X))))
-                error("LOBPCG is badly failing to keep the vectors normalized; this should never happen")
-            end
+        diffs = abs.(columnwise_dots(BX, X) .-1)
+        if any(diffs .>= sqrt(eps(real(eltype(X)))))
+           error("LOBPCG is badly failing to keep the vectors normalized; this should never happen")
         end
 
         # Restrict all views to active
diff --git a/src/eigen/preconditioners.jl b/src/eigen/preconditioners.jl
index a9c213d9a8..dd0335a0b0 100644
--- a/src/eigen/preconditioners.jl
+++ b/src/eigen/preconditioners.jl
@@ -24,11 +24,11 @@ PreconditionerNone(::HamiltonianBlock) = I
 (Simplified version of)
 [Tetter-Payne-Allan preconditioning](https://doi.org/10.1103/physrevb.40.12255).
 """
-mutable struct PreconditionerTPA{T <: Real}
+mutable struct PreconditionerTPA{T <: Real, Tkin <: AbstractVector{T}}
     basis::PlaneWaveBasis
     kpt::Kpoint
-    kin::AbstractVector{T}  # kinetic energy of every G
-    mean_kin::Union{Nothing, Vector{T}}  # mean kinetic energy of every band
+    kin::Tkin  # kinetic energy of every G
+    mean_kin::Union{Nothing, Tkin}  # mean kinetic energy of every band
     default_shift::T  # if mean_kin is not set by `precondprep!`, this will be used for the shift
 end
 
@@ -40,7 +40,7 @@ function PreconditionerTPA(basis::PlaneWaveBasis{T}, kpt::Kpoint; default_shift=
     #      it's better to pass a HamiltonianBlock directly and read the computed values.
     kinetic_term = only(kinetic_term)
     kin = kinetic_energy(kinetic_term, basis.Ecut, Gplusk_vectors_cart(basis, kpt))
-    PreconditionerTPA{T}(basis, kpt, kin, nothing, default_shift)
+    PreconditionerTPA{T, typeof(kin)}(basis, kpt, kin, nothing, default_shift)
 end
 function PreconditionerTPA(ham::HamiltonianBlock; kwargs...)
     PreconditionerTPA(ham.basis, ham.kpoint; kwargs...)
@@ -50,7 +50,7 @@ end
     if P.mean_kin === nothing
         ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
-    parallel_loop_over_range(1:size(Y, 2)) do n
+        parallel_loop_over_range(1:size(Y, 2)) do n
             Y[:, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[:, n]
         end
     end
@@ -73,7 +73,7 @@ end
 (Base.:*)(P::PreconditionerTPA, R) = mul!(copy(R), P, R)
 
 function precondprep!(P::PreconditionerTPA, X::AbstractArray)
-    P.mean_kin = [real(dot(x, Diagonal(P.kin), x)) for x in eachcol(X)]
+    P.mean_kin = vec(real(columnwise_dots(X, Diagonal(P.kin), X)))
 end
 precondprep!(P::PreconditionerTPA, ::Nothing) = 1  # fallback for edge cases
 
diff --git a/src/workarounds/gpu_arrays.jl b/src/gpu/gpu_arrays.jl
similarity index 100%
rename from src/workarounds/gpu_arrays.jl
rename to src/gpu/gpu_arrays.jl
diff --git a/src/gpu/linalg.jl b/src/gpu/linalg.jl
new file mode 100644
index 0000000000..989d78407a
--- /dev/null
+++ b/src/gpu/linalg.jl
@@ -0,0 +1,45 @@
+### GPU-specific implementations of functions called during LOBPCG
+# The massive parallelism of the GPU can only be fully exploited when
+# operating on whole arrays. For performance reasons, one should avoid
+# explicitly looping over columns or elements. This approach is not
+# necessarily the most performant on CPU, as the allocation of large
+# temporary arrays hurts cache locality. It is also harder to read.
+
+using LinearAlgebra
+using GPUArraysCore
+
+function compute_λ(X::AbstractGPUArray{T}, AX::AbstractGPUArray{T}, BX::AbstractGPUArray{T}) where {T}
+    num = sum(conj(X) .* AX, dims=1)
+    den = sum(conj(X) .* BX, dims=1)
+    vec(real.(num ./ den))
+end
+
+function columnwise_dots(A::AbstractGPUArray{T}, B::AbstractGPUArray{T}) where {T}
+    sum(conj(A) .* B; dims=1)
+end
+
+function columnwise_dots(A::AbstractGPUArray{T}, M, B::AbstractGPUArray{T}) where {T}
+    sum(conj(A) .* (M * B); dims=1)
+end
+
+function columnwise_dots(A::AbstractGPUArray{T}, D::Diagonal, B::AbstractGPUArray{T}) where {T}
+    sum(conj(A) .* (D.diag .* B); dims=1)
+end
+
+function ldiv!(Y::AbstractGPUArray{T}, P::PreconditionerTPA, R::AbstractGPUArray{T}) where {T}
+    if P.mean_kin === nothing
+        ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
+    else
+        Y .= (P.mean_kin' ./ (P.mean_kin' .+ P.kin)) .* R
+    end
+    Y
+end
+
+function mul!(Y::AbstractGPUArray{T}, P::PreconditionerTPA, R::AbstractGPUArray{T}) where {T}
+    if P.mean_kin === nothing
+        mul!(Y, Diagonal(P.kin .+ P.default_shift), R)
+    else
+        Y .= ((P.mean_kin' .+ P.kin) ./ P.mean_kin') .* R
+    end
+    Y
+end
diff --git a/src/interpolation.jl b/src/interpolation.jl
index b0f456dc5f..2bc96f5185 100644
--- a/src/interpolation.jl
+++ b/src/interpolation.jl
@@ -105,12 +105,11 @@ function interpolate_kpoint(data_in::AbstractVecOrMat,
     n_bands  = size(data_in, 2)
     n_Gk_out = length(G_vectors(basis_out, kpoint_out))
     data_out = similar(data_in, n_Gk_out, n_bands) .= 0
-    # TODO: use a map, or this will not be GPU compatible (scalar indexing)
-    for iin = 1:size(data_in, 1)
-        idx_fft = kpoint_in.mapping[iin]
-        idx_fft in keys(kpoint_out.mapping_inv) || continue
-        iout = kpoint_out.mapping_inv[idx_fft]
-        data_out[iout, :] = data_in[iin, :]
-    end
+
+    max_nG = max(length(G_vectors(basis_in)), length(G_vectors(basis_out)))
+    tmp = similar(data_in, max_nG, n_bands) .= 0
+
+    tmp[kpoint_in.mapping, :] .= data_in
+    data_out .= @view tmp[kpoint_out.mapping, :]
     ortho_qr(data_out)  # Re-orthogonalize and renormalize
 end


Reviews:
user: login: antoine-levitt
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-24T16:08:43Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-24T16:38:51Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-24T16:43:44Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-24T16:44:10Z
user: login: antoine-levitt
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-24T16:51:26Z
user: login: mfherbst
body: Not great, but within what we discussed looks like a good solution to me.

@antoine-levitt may also have some thoughts.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T08:58:14Z
user: login: antoine-levitt
body: Thanks for the PR. Not very satisfactory, but pluses outweigh minuses, unfortunately...
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T12:21:33Z
user: login: mfherbst
body: Some small comments, but mostly fine.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-12T09:41:24Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T11:31:23Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T11:32:25Z
user: login: mfherbst
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-12T13:34:23Z
user: login: mfherbst
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-12T13:34:55Z
user: login: mfherbst
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T06:40:57Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T08:35:38Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T08:41:43Z
user: login: mfherbst
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T11:13:40Z
user: login: antoine-levitt
body: LGTM, thanks!
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T14:03:16Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T13:23:12Z
user: login: antoine-levitt
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T06:34:00Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-15T07:29:08Z
user: login: Technici4n
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-07-09T13:59:23Z
user: login: Technici4n
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-07-09T14:01:18Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-07-09T14:08:47Z
user: login: Technici4n
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-07-09T15:20:08Z
user: login: abussy
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-07-09T16:05:38Z

Review comments:
diff_hunk: @@ -249,24 +249,19 @@ end
 
 # Randomize the columns of X if the norm is below tol
 function drop!(X::AbstractArray{T}, tol=2eps(real(T))) where {T}
-    dropped = Int[]
-    for i=1:size(X,2)
-        n = norm(@views X[:,i])
-        if n <= tol
-            X[:,i] = randn(T, size(X,1))
-            push!(dropped, i)
-        end
-    end
+    # Using array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, X; dims=1)))
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: `norm.(eachcol(X))`?
created_at: 2025-02-24T15:51:41Z
author_association: MEMBER
diff_hunk: @@ -249,24 +249,19 @@ end
 
 # Randomize the columns of X if the norm is below tol
 function drop!(X::AbstractArray{T}, tol=2eps(real(T))) where {T}
-    dropped = Int[]
-    for i=1:size(X,2)
-        n = norm(@views X[:,i])
-        if n <= tol
-            X[:,i] = randn(T, size(X,1))
-            push!(dropped, i)
-        end
-    end
+    # Using array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, X; dims=1)))
+    dropped = findall(n -> n <= tol, norms)
+    @views randn!(TaskLocalRNG(), X[:, dropped])
     dropped
 end
 
 # Find X that is orthogonal, and B-orthogonal to Y, up to a tolerance tol.
 @timing "ortho! X vs Y" function ortho!(X::AbstractArray{T}, Y, BY; tol=2eps(real(T))) where {T}
     # normalize to try to cheaply improve conditioning
-    parallel_loop_over_range(1:size(X, 2)) do i
-        n = norm(@views X[:,i])
-        @views X[:,i] ./= n
-    end
+    # using array operations for GPU efficiency
+    norms = sqrt.(sum(abs2, X; dims=1))
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: same as above
created_at: 2025-02-24T15:53:28Z
author_association: MEMBER
diff_hunk: @@ -310,10 +305,17 @@ end
     X
 end
 
+# Computes λ = real((X' * AX) / (X' *BX)), for each column of X
+function compute_λ(X, AX, BX)
+    # using array operations for GPU performance
+    num = sum(conj(X) .* AX, dims=1)
+    den = sum(conj(X) .* BX, dims=1)
+    vec(real.(num ./ den))
+end
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: Is it clear that this way of doing is faster? Eg on CPU you're possibly allocating big arrays and losing out on cache locality?
created_at: 2025-02-24T15:57:51Z
author_association: MEMBER
diff_hunk: @@ -310,10 +305,17 @@ end
     X
 end
 
+# Computes λ = real((X' * AX) / (X' *BX)), for each column of X
+function compute_λ(X, AX, BX)
+    # using array operations for GPU performance
+    num = sum(conj(X) .* AX, dims=1)
+    den = sum(conj(X) .* BX, dims=1)
+    vec(real.(num ./ den))
+end
 
 function final_retval(X, AX, BX, resid_history, niter, n_matvec)
-    λ = @views [real((X[:, n]'*AX[:, n]) / (X[:, n]'BX[:, n])) for n=1:size(X, 2)]
-    λ_device = oftype(X[:, 1], λ)  # Offload to GPU if needed
+    λ_device = compute_λ(X, AX, BX)
+    λ = oftype(ones(eltype(λ_device), 1), λ_device)
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: This will need to be merged with master
created_at: 2025-02-24T16:01:20Z
author_association: MEMBER
diff_hunk: @@ -324,8 +326,10 @@ function final_retval(X, AX, BX, resid_history, niter, n_matvec)
         BX = BX[:, p]
         resid_history = resid_history[p, :]
     end
+    # compute norms of column vectors with array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, residuals; dims=1)))
     (; λ=λ_device, X, AX, BX,
-     residual_norms=norm.(eachcol(residuals)),
+     residual_norms=oftype(ones(eltype(norms), 1), norms),
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: Add a `to_cpu` method?
created_at: 2025-02-24T16:01:39Z
author_association: MEMBER
diff_hunk: @@ -249,24 +249,19 @@ end
 
 # Randomize the columns of X if the norm is below tol
 function drop!(X::AbstractArray{T}, tol=2eps(real(T))) where {T}
-    dropped = Int[]
-    for i=1:size(X,2)
-        n = norm(@views X[:,i])
-        if n <= tol
-            X[:,i] = randn(T, size(X,1))
-            push!(dropped, i)
-        end
-    end
+    # Using array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, X; dims=1)))
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: Wait, does that not work? I can't test I don't have a GPU here.
created_at: 2025-02-24T16:08:39Z
author_association: MEMBER
diff_hunk: @@ -249,24 +249,19 @@ end
 
 # Randomize the columns of X if the norm is below tol
 function drop!(X::AbstractArray{T}, tol=2eps(real(T))) where {T}
-    dropped = Int[]
-    for i=1:size(X,2)
-        n = norm(@views X[:,i])
-        if n <= tol
-            X[:,i] = randn(T, size(X,1))
-            push!(dropped, i)
-        end
-    end
+    # Using array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, X; dims=1)))
path: src/eigen/lobpcg_hyper_impl.jl
user: login: abussy
body: There is a factor of 10-20x performance boost with the array operation over `norm.(eachcol(X))`. Essentially, it goes from being the main bottleneck of `ortho! X vs Y` to being almost negligible.
created_at: 2025-02-24T16:38:50Z
author_association: COLLABORATOR
diff_hunk: @@ -310,10 +305,17 @@ end
     X
 end
 
+# Computes λ = real((X' * AX) / (X' *BX)), for each column of X
+function compute_λ(X, AX, BX)
+    # using array operations for GPU performance
+    num = sum(conj(X) .* AX, dims=1)
+    den = sum(conj(X) .* BX, dims=1)
+    vec(real.(num ./ den))
+end
path: src/eigen/lobpcg_hyper_impl.jl
user: login: abussy
body: It did not show up as a an issue when I was testing the CPU performance of this PR. This might be the case sometimes, maybe. As for the norms above, the performance gain on the GPU is massive.

For code that could be problematic on the CPU, one option would be to have a GPU definition of it in `DFTKCUDAExt.jl` and dispatch to it when needed. The main issue with it is code inflation though.
created_at: 2025-02-24T16:43:44Z
author_association: COLLABORATOR
diff_hunk: @@ -324,8 +326,10 @@ function final_retval(X, AX, BX, resid_history, niter, n_matvec)
         BX = BX[:, p]
         resid_history = resid_history[p, :]
     end
+    # compute norms of column vectors with array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, residuals; dims=1)))
     (; λ=λ_device, X, AX, BX,
-     residual_norms=norm.(eachcol(residuals)),
+     residual_norms=oftype(ones(eltype(norms), 1), norms),
path: src/eigen/lobpcg_hyper_impl.jl
user: login: abussy
body: Valid point, it would make a lot of sense here
created_at: 2025-02-24T16:44:10Z
author_association: COLLABORATOR
diff_hunk: @@ -249,24 +249,19 @@ end
 
 # Randomize the columns of X if the norm is below tol
 function drop!(X::AbstractArray{T}, tol=2eps(real(T))) where {T}
-    dropped = Int[]
-    for i=1:size(X,2)
-        n = norm(@views X[:,i])
-        if n <= tol
-            X[:,i] = randn(T, size(X,1))
-            push!(dropped, i)
-        end
-    end
+    # Using array operations for GPU performance
+    norms = vec(sqrt.(sum(abs2, X; dims=1)))
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: Is this an issue with GPUArrays (that it doesn't implement the right overloads) or something more intrinsic?
created_at: 2025-02-24T16:51:26Z
author_association: MEMBER
diff_hunk: @@ -33,6 +33,14 @@
 # other eigenvectors (which is not the case in many - all ? - other
 # implementations)
 
+# - The massive parallelism of the GPU can only be fully expoloited when
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: typo
created_at: 2025-03-25T12:10:13Z
author_association: MEMBER
diff_hunk: @@ -33,6 +33,14 @@
 # other eigenvectors (which is not the case in many - all ? - other
 # implementations)
 
+# - The massive parallelism of the GPU can only be fully expoloited when
+# operating on whole arrays. For performance reasons, one should avoid
+# explicitly looping over columns or element. For example, computing
+# the norms of the columns of an array A should be dones as:
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: typo
created_at: 2025-03-25T12:10:20Z
author_association: MEMBER
diff_hunk: @@ -33,6 +33,14 @@
 # other eigenvectors (which is not the case in many - all ? - other
 # implementations)
 
+# - The massive parallelism of the GPU can only be fully expoloited when
+# operating on whole arrays. For performance reasons, one should avoid
+# explicitly looping over columns or element. For example, computing
+# the norms of the columns of an array A should be dones as:
+# norms = vec(sqrt.(sum(abs2, A; dims=1)))
+# Rather than the more eloquant:
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: typo
created_at: 2025-03-25T12:10:27Z
author_association: MEMBER
diff_hunk: @@ -247,26 +255,22 @@ normest(M) = maximum(abs.(diag(M))) + norm(M - Diagonal(diag(M)))
     (; X, nchol, growth_factor)
 end
 
+# Calculate the norms of the columns of an array X
+function colnorms(X::AbstractArray{T}) where{T}
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: columnwise_norm
created_at: 2025-03-25T12:10:42Z
author_association: MEMBER
diff_hunk: @@ -247,26 +255,22 @@ normest(M) = maximum(abs.(diag(M))) + norm(M - Diagonal(diag(M)))
     (; X, nchol, growth_factor)
 end
 
+# Calculate the norms of the columns of an array X
+function colnorms(X::AbstractArray{T}) where{T}
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: you don't need to type X
created_at: 2025-03-25T12:11:04Z
author_association: MEMBER
diff_hunk: @@ -501,10 +504,9 @@ end
         end
 
         # Quick sanity check
-        for i = 1:size(X, 2)
-            @views if abs(BX[:, i]'X[:, i] - 1) >= sqrt(eps(real(eltype(X))))
-                error("LOBPCG is badly failing to keep the vectors normalized; this should never happen")
-            end
+        diffs = abs.(sum(conj(BX) .* X, dims=1) .-1)
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: spaces before and after .-
created_at: 2025-03-25T12:14:48Z
author_association: MEMBER
diff_hunk: @@ -501,10 +504,9 @@ end
         end
 
         # Quick sanity check
-        for i = 1:size(X, 2)
-            @views if abs(BX[:, i]'X[:, i] - 1) >= sqrt(eps(real(eltype(X))))
-                error("LOBPCG is badly failing to keep the vectors normalized; this should never happen")
-            end
+        diffs = abs.(sum(conj(BX) .* X, dims=1) .-1)
path: src/eigen/lobpcg_hyper_impl.jl
user: login: antoine-levitt
body: Should we have a `diag_prod` function?
created_at: 2025-03-25T12:15:26Z
author_association: MEMBER
diff_hunk: @@ -28,7 +28,7 @@ mutable struct PreconditionerTPA{T <: Real}
     basis::PlaneWaveBasis
     kpt::Kpoint
     kin::AbstractVector{T}  # kinetic energy of every G
-    mean_kin::Union{Nothing, Vector{T}}  # mean kinetic energy of every band
+    mean_kin::Union{Nothing, AbstractVector{T}}  # mean kinetic energy of every band
path: src/eigen/preconditioners.jl
user: login: antoine-levitt
body: This should probably be specialized, ie parametrize the type with TA and have kin::TA, mean_kin::Union{Nothing, TA}. But it doesn't matter much.
created_at: 2025-03-25T12:16:51Z
author_association: MEMBER
diff_hunk: @@ -50,9 +50,8 @@ end
     if P.mean_kin === nothing
         ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
-    parallel_loop_over_range(1:size(Y, 2)) do n
-            Y[:, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[:, n]
-        end
+        # same as loop Y[:, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[:, n]
+        Y .= (P.mean_kin' ./ (P.mean_kin' .+ P.kin)) .* R
path: src/eigen/preconditioners.jl
user: login: antoine-levitt
body: That's the kind of matlabism/numpyisms that I was glad to leave behind when going to julia... sigh... I don't have a good solution here.
created_at: 2025-03-25T12:19:37Z
author_association: MEMBER
diff_hunk: @@ -64,16 +63,16 @@ ldiv!(P::PreconditionerTPA, R) = ldiv!(R, P, R)
     if P.mean_kin === nothing
         mul!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
-        parallel_loop_over_range(1:size(Y, 2)) do n
-            Y[:, n] .= (P.mean_kin[n] .+ P.kin) ./ P.mean_kin[n] .* R[:, n]
-        end
+        # same as loop Y[:, n] .= (P.mean_kin[n] .+ P.kin) ./ P.mean_kin[n] .* R[:, n]
+        Y .= ((P.mean_kin' .+ P.kin) ./ P.mean_kin') .* R
     end
     Y
 end
 (Base.:*)(P::PreconditionerTPA, R) = mul!(copy(R), P, R)
 
 function precondprep!(P::PreconditionerTPA, X::AbstractArray)
-    P.mean_kin = [real(dot(x, Diagonal(P.kin), x)) for x in eachcol(X)]
+    # real(dot(x, Diagonal(P.kin), x)), for all columns x of X (efficient on the GPU)
+    P.mean_kin = vec(real(sum(conj(X) .* (P.kin .* X); dims=1)))
path: src/eigen/preconditioners.jl
user: login: antoine-levitt
body: Maybe we should have a `diagprod(X, P.kin, X)` method
created_at: 2025-03-25T12:20:46Z
author_association: MEMBER
diff_hunk: @@ -11,6 +11,19 @@
 import LinearAlgebra: ldiv!
 import LinearAlgebra: mul!
 
+# Returns dot(A[:, i], Diagonal(P.kin), B[:, i])), for all columns of A, B
path: src/eigen/preconditioners.jl
user: login: mfherbst
body: Mismatch between comment and implementation.

Also I would actually do this like the `dot` function. I.e. have a two-argument version (where you just call the two-argument `dot`) and a three-argument version, where you call the three-argument `dot`. I think the three-argument `dot` with an `I` (uniform scaling) compiles away, so you can just have that as your implementation and then dispatch `diag_prod(A, B) = diag_prod(A, I, B)`.
created_at: 2025-05-12T09:29:28Z
author_association: MEMBER
diff_hunk: @@ -11,6 +11,19 @@
 import LinearAlgebra: ldiv!
 import LinearAlgebra: mul!
 
+# Returns dot(A[:, i], Diagonal(P.kin), B[:, i])), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}; diag=nothing) where {T}
path: src/eigen/preconditioners.jl
user: login: mfherbst
body: Maybe we should move all of these helper functions to a single file in `src/common`. Because some being in `preconditioners.jl` and some in LOBPCG does not really help finding them in the future.
created_at: 2025-05-12T09:31:19Z
author_association: MEMBER
diff_hunk: @@ -28,7 +28,7 @@ mutable struct PreconditionerTPA{T <: Real}
     basis::PlaneWaveBasis
     kpt::Kpoint
     kin::AbstractVector{T}  # kinetic energy of every G
-    mean_kin::Union{Nothing, Vector{T}}  # mean kinetic energy of every band
+    mean_kin::Union{Nothing, AbstractVector{T}}  # mean kinetic energy of every band
path: src/eigen/preconditioners.jl
user: login: mfherbst
body: Agree. I'd have this:

```suggestion
mutable struct PreconditionerTPA{T <: Real, Tkin <: AbstractVector{T}}
    basis::PlaneWaveBasis
    kpt::Kpoint
    kin::Tkin  # kinetic energy of every G
    mean_kin::Tkin  # mean kinetic energy of every band or arbitrary shift
end
```
and I would pre-set the `mean_kin` with a single elements set to the default shift of 1 or something like that.
created_at: 2025-05-12T09:39:04Z
author_association: MEMBER
diff_hunk: @@ -50,7 +63,7 @@ end
     if P.mean_kin === nothing
         ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
-    parallel_loop_over_range(1:size(Y, 2)) do n
path: src/eigen/preconditioners.jl
user: login: mfherbst
body: Removing parallelisation ok ?
created_at: 2025-05-12T09:39:19Z
author_association: MEMBER
diff_hunk: @@ -11,6 +11,19 @@
 import LinearAlgebra: ldiv!
 import LinearAlgebra: mul!
 
+# Returns dot(A[:, i], Diagonal(P.kin), B[:, i])), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}; diag=nothing) where {T}
path: src/eigen/preconditioners.jl
user: login: abussy
body: Maybe in a new file `src/eigen/common.jl` ? This would keep the locality of functions needed by LOBPCG for the moment when it branches out of DFTK.
created_at: 2025-05-12T11:31:23Z
author_association: COLLABORATOR
diff_hunk: @@ -50,7 +63,7 @@ end
     if P.mean_kin === nothing
         ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
-    parallel_loop_over_range(1:size(Y, 2)) do n
path: src/eigen/preconditioners.jl
user: login: abussy
body: I need to double check performance, but this also goes in the direction of eventually branching out LOBPCG  from DFTK.
created_at: 2025-05-12T11:32:25Z
author_association: COLLABORATOR
diff_hunk: @@ -11,6 +11,19 @@
 import LinearAlgebra: ldiv!
 import LinearAlgebra: mul!
 
+# Returns dot(A[:, i], Diagonal(P.kin), B[:, i])), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}; diag=nothing) where {T}
path: src/eigen/preconditioners.jl
user: login: mfherbst
body: Hmm, maybe. I was more thinking `src/common/linalg.jl` or something like that. But `src/eigen/linalg.jl` is also fine, I'd say. I don't like using "common" and "utils" for too many things as these words don't really have a meaning per se.
created_at: 2025-05-12T13:34:23Z
author_association: MEMBER
diff_hunk: @@ -50,7 +63,7 @@ end
     if P.mean_kin === nothing
         ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
-    parallel_loop_over_range(1:size(Y, 2)) do n
path: src/eigen/preconditioners.jl
user: login: mfherbst
body: I see, but I suppose the preconditioner would remain in DFTK in any case.
created_at: 2025-05-12T13:34:55Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    @assert size(A) == size(B)
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of dot(A[:, i], Diagonal(diag), B[:, i]), for all columns of
+# A, B, and a vector diag
+@views function diag_prod(A::AbstractArray{T}, diag::AbstractVector,
path: src/eigen/linalg.jl
user: login: mfherbst
body: Why not explicitly have `diag::Diagonal` and than pass it through to `dot` ? Or rather here it could even be D::AbstractMatrix without much changes, right ? (May not be so easy for the GPU implementation, though)
created_at: 2025-05-13T06:37:08Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    @assert size(A) == size(B)
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of dot(A[:, i], Diagonal(diag), B[:, i]), for all columns of
+# A, B, and a vector diag
+@views function diag_prod(A::AbstractArray{T}, diag::AbstractVector,
path: src/eigen/linalg.jl
user: login: abussy
body: Yes, that's for the GPU implementation, it's is assumed that an array is passed:

```
function DFTK.diag_prod(A::CUDA.CuArray{T}, diag::CUDA.CuArray{T}, B::CUDA.CuArray{T}) where {T}
    @assert size(A) == size(B)
    @assert length(diag) == size(B, 1)
    sum(conj(A) .* (diag .* B); dims=1)
end
```
created_at: 2025-05-13T08:35:38Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    @assert size(A) == size(B)
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of dot(A[:, i], Diagonal(diag), B[:, i]), for all columns of
+# A, B, and a vector diag
+@views function diag_prod(A::AbstractArray{T}, diag::AbstractVector,
path: src/eigen/linalg.jl
user: login: abussy
body: I could probably pass an `AbstractMatrix` to the function, and have a special case with a `isdiag()` test on the GPU (and otherwise do the full MM). Not sure if it is worth the extra code.
created_at: 2025-05-13T08:41:43Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    @assert size(A) == size(B)
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of dot(A[:, i], Diagonal(diag), B[:, i]), for all columns of
+# A, B, and a vector diag
+@views function diag_prod(A::AbstractArray{T}, diag::AbstractVector,
path: src/eigen/linalg.jl
user: login: mfherbst
body: Just specialise on the type ::Diagonal for GPU ? From the diagonal object you can extract the actual vector as a GPU array
created_at: 2025-05-13T11:13:40Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,47 @@
+### GPU-specific implementations of functions called during LOBPCG
path: ext/DFTKCUDAExt/lobpcg.jl
user: login: antoine-levitt
body: Hm that's very annoying. We should really take lobpcg out of DFTK... Maybe add a line in lobpcg saying it's overloaded here?
created_at: 2025-05-13T13:57:54Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
path: src/eigen/linalg.jl
user: login: antoine-levitt
body: add real()
created_at: 2025-05-13T13:58:20Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
path: src/eigen/linalg.jl
user: login: antoine-levitt
body: Maybe columnwise_dot?
created_at: 2025-05-13T13:59:40Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function diag_prod(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    @assert size(A) == size(B)
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of dot(A[:, i], M, B[:, i]), for all columns of
+# A, B, and matrix M
+@views function diag_prod(A::AbstractArray{T}, M::AbstractMatrix,
+                          B::AbstractArray{T}) where {T}
+    @assert size(A) == size(B)
+    @assert size(M, 2) == size(B, 1)
path: src/eigen/linalg.jl
user: login: antoine-levitt
body: I wouldn't do this (and I would not type M) because M can be a bit more abstract (eg I or something)
created_at: 2025-05-13T14:00:56Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,47 @@
+### GPU-specific implementations of functions called during LOBPCG
path: ext/DFTKCUDAExt/lobpcg.jl
user: login: antoine-levitt
body: Sorry I see you've done this already.
created_at: 2025-05-13T14:01:15Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
path: src/eigen/linalg.jl
user: login: abussy
body: mmh, I personally feel like `columnwise_norms` is more descriptive
created_at: 2025-05-14T13:23:12Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
path: src/eigen/linalg.jl
user: login: antoine-levitt
body: This is not about that function it's another one
created_at: 2025-05-15T06:34:00Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,19 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
path: src/eigen/linalg.jl
user: login: abussy
body: Oh, my bad. Right, that makes sense, especially since `diag_prod` was generalized to non-diagonal matrices.
created_at: 2025-05-15T07:29:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,15 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function columnwise_dots(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of real(dot(A[:, i], M, B[:, i])), for all columns of
+# A, B, and matrix M
+@views function columnwise_dots(A::AbstractArray{T}, M, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], M, B[:, i])) for i = 1:size(A, 2)]
+end
path: src/eigen/linalg.jl
user: login: Technici4n
body: What is the reasoning for the `real` calls here?
created_at: 2025-07-09T13:59:23Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,15 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function columnwise_dots(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of real(dot(A[:, i], M, B[:, i])), for all columns of
+# A, B, and matrix M
+@views function columnwise_dots(A::AbstractArray{T}, M, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], M, B[:, i])) for i = 1:size(A, 2)]
+end
path: src/eigen/linalg.jl
user: login: Technici4n
body: (asking because I don't see them in the GPU version - did I miss something?)
created_at: 2025-07-09T14:01:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,15 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function columnwise_dots(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of real(dot(A[:, i], M, B[:, i])), for all columns of
+# A, B, and matrix M
+@views function columnwise_dots(A::AbstractArray{T}, M, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], M, B[:, i])) for i = 1:size(A, 2)]
+end
path: src/eigen/linalg.jl
user: login: abussy
body: I think that's most likely a mistake. I assume it comes from the original preconditioner code: https://github.com/JuliaMolSim/DFTK.jl/blob/3c34d3855f6fec995edc3681caa075a4bd35ed1e/src/eigen/preconditioners.jl#L75-L77
But `real` is still taken there now, so I think it can go.
created_at: 2025-07-09T14:08:47Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,15 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function columnwise_dots(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of real(dot(A[:, i], M, B[:, i])), for all columns of
+# A, B, and matrix M
+@views function columnwise_dots(A::AbstractArray{T}, M, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], M, B[:, i])) for i = 1:size(A, 2)]
+end
path: src/eigen/linalg.jl
user: login: Technici4n
body: Do you want me to open an issue to remember?
created_at: 2025-07-09T15:20:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,15 @@
+# Calculate the norms of the columns of an array
+function columnwise_norms(X::AbstractArray{T}) where{T}
+    vec(sqrt.(sum(abs2, X; dims=1)))
+end
+
+# Returns a vector of dot(A[:, i], B[:, i]), for all columns of A, B
+@views function columnwise_dots(A::AbstractArray{T}, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], B[:, i])) for i = 1:size(A, 2)]
+end
+
+# Returns a vector of real(dot(A[:, i], M, B[:, i])), for all columns of
+# A, B, and matrix M
+@views function columnwise_dots(A::AbstractArray{T}, M, B::AbstractArray{T}) where {T}
+    [real(dot(A[:, i], M, B[:, i])) for i = 1:size(A, 2)]
+end
path: src/eigen/linalg.jl
user: login: abussy
body: I'll take care of it immediately

created_at: 2025-07-09T16:05:38Z
author_association: COLLABORATOR

PR comments:
user: login: antoine-levitt
created_at: 2025-02-24T16:16:52Z
body: Hm, I'm kind of split on this.
- Looks like GPU style clashes with threading in a number of places. Threading has been mostly an annoyance, and for very limited results. I think I'm fine with just scrapping threading from DFTK altogether, and pointing people to GPU if they want to do large-scale.
- It's pretty unfortunate if GPU forces us towards a matlab/numpy style. Those dot products are not very nice... I'd rather use things like eachcol, map, etc.

Also don't add comments at each place we use vector-style operations: we should converge on one uniform style and use it without comments.
author_association: MEMBER
user: login: mfherbst
created_at: 2025-02-26T15:22:00Z
body: Yeah same here. It's a shame this has such an impact on readability. From my point of view the key points are:

- Avoid code duplication: I think it's better to have a bit of ugly code at selected places (with nice comments and a common convention of course) rather than a GPU and a CPU version with different code. Maybe for some primitives we use often (e.g. column-wise norms or outer products) we can invent functions that are easier to understand and hide the array syntax a little.

- Avoid performance regressions on CPU: For small arrays this may not matter but we should test on non-trivial calculations if the extra temporaries don't end up killing us
author_association: MEMBER
user: login: abussy
created_at: 2025-03-24T16:55:15Z
body: I went back and worked on this. Here are the changes:
- Single comment on GPU performance style at the beginning of `lobpcg_hyper_impl.jl` rather than inconsistently scattered GPU performance comments
- Introduced a function to calculate norms of column vectors, since used multiple times
- Modified the SCF guess interpolation in `interpolation.jl` to make it GPU efficient (following example from `fft.jl`). This significantly reduces the time spent in the first SCF step.
- used the `to_cpu` function where possible

> Is it clear that this way of doing is faster? Eg on CPU you're possibly allocating big arrays and losing out on cache locality?

I carefully tested this in the CPU, and indeed, there is a performance hit on CPU with large systems. I reverted this change back to the original, since the CPU hit is greater than the GPU gain. This is not a massive issue, since this operation is only done once per `LOBPCG` call (compared to norms, that can be calculated 1000s of times).
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-03-25T09:02:59Z
body: If there really turns out to be a bad CPU perf issue can we specialise for CPU maybe, but certainly I'd try to avoid that.
author_association: MEMBER
user: login: antoine-levitt
created_at: 2025-03-25T12:22:54Z
body: Also careful in refactoring that lobpcg should be self-contained (not use too much stuff from other parts of DFTK) so that users can steal it for their own projects. We really should pull it off to its own project and depend on it. We had a plan to integrate it with either iterativesolvers or krylovkit but I think we should just give up on that and split it off.
author_association: MEMBER
user: login: mfherbst
created_at: 2025-03-25T14:53:52Z
body: Agree with trying to keep LOBPCG self-contained (and split off hopefully soonish) and I think a `diagprod` function would be good
author_association: MEMBER
user: login: abussy
created_at: 2025-03-25T15:09:10Z
body: Sure, then I'll define a local `to_cpu` function. 

I've run additional tests, and I think we cannot escape some level of CPU/GPU code separation. Out of all the changes I have proposed, only the calculation of the norms is consistently faster on both GPU and CPU. 

I think the best course of action is to define alternative GPU optimized code in `DFTCUDAExt.jl` (which should maybe become a directory), while keeping the easy to read Julian code elsewhere. When the time comes to split `LOBPCG` into its own project, it would inherit this CUDA extension as well. As much as possible, I would try to do this for low level functions.
 
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-03-26T12:59:13Z
body: Sounds very reasonable, thanks @abussy .
author_association: MEMBER
user: login: abussy
created_at: 2025-04-04T07:42:00Z
body: With this latest commit, I did the following:
- moved performance critical code into low-level functions (e.g. `compute_λ`, `diag_prod`), written in a CPU and Human friendly manner
- Changed the `DFTKCUDAExt` to a directory, in which the above function are re-implemented in a GPU optimized way. They will be called appropriately with multiple dispatch.

In the spirit of splitting LOBPCG from DFTK in the future, I also removed usage of DFTK's explicit threading. Note that I kept using DFTK's `to_cpu` function in order to avoid redefinition warnings. Since this is a one line function, it will be easy to move it to a LOBPCG module.

Finally, I kept the calculation of norms in the function `columnwise_norms` as an array operation, as it performs better than `norm.(eachcol(X))`, even on the CPU (marginally so, but still).
author_association: COLLABORATOR
user: login: abussy
created_at: 2025-05-12T09:38:37Z
body: Can I get some feedback on this PR? I have a couple of other GPU optimization contributions that would benefit from this being merged first.
author_association: COLLABORATOR
user: login: abussy
created_at: 2025-05-12T16:27:29Z
body: Implemented most of the points brought up by @mfherbst:
- Moved help linear algebra functions into their own file in `src/eigen/linalg.jl`
- Reintroduced DFTK thread parallelization in the preconditioner: my measurements show that this is faster than BLAS threading alone, and if it stays in DFTK, it is not a problem. Overall, this has a close to negligible impact on timings.
- Wrote separate functions for 2- and 3-arguments `diag_prod`

I however did not change the `PreconditionerTPA` struct yet. Forcing the `mean_kin` field to be an `AbstracVector{T}` complicates the handling of the preconditioner application in the special case when there is only a default shift. Currently, we have
```
if P.mean_kin === nothing
    ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
else
```
but to get the same with the above constraint, I do not find better than
```
if length(P.mean_kin) == 1 && P.mean_kin[1] == P.default_shift
    ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
else
```
which looks very clunky to me.
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-05-13T06:41:29Z
body: > I however did not change the PreconditionerTPA struct yet. Forcing the mean_kin field to be an AbstracVector{T} complicates the handling of the preconditioner application in the special case when there is only a default shift. Currently, we have

Yeah ok, but I'd still use type templates instead of the `AbstractArray` objects as antoine suggested.
author_association: MEMBER
user: login: abussy
created_at: 2025-05-13T13:53:22Z
body: - Made `diag_prod` signature more like `dot` by taking a matrix rather than a vector
- Parametrized type of `kin` and `mean_kin` in `PreconditionerTPA`
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-05-14T06:20:50Z
body: A general point that came to me: These GPU specific optimisations should be helpful for AMD as well, right? So perhaps we should not make the types CuArrays but GpuArrays (and thus not put the code into the extmodule)?
author_association: MEMBER
user: login: abussy
created_at: 2025-05-14T08:39:47Z
body: Indeed, these optimizations should be vendor agnostic. I have access to some Mi200 GPUs, I'll try to run there and confirm it works/performs. If so, I suggest creating a new subdirectory in `src/common` to put the generic GPU optimizations (maybe `src/common/gpu_opts`). Then the DFTKCUDAExt could also be reverted to a single file.
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-05-14T09:48:32Z
body: Yeah or src/common/linalg.jl or src/eigen/linalg.jl with both versions: CPU and GPU.
author_association: MEMBER
user: login: abussy
created_at: 2025-05-14T13:27:32Z
body: I switched to a more generic implementation using `AbstractGPUArray`, and moved the code back to DFTK's `src`. I also checked that it works both on NVIDIA and AMD GPUs.

Regarding the location, I put it in `src/common/gpu_opts/lobpcg.jl`, rather than in `src/eigen/linalg.jl` with the rest of the helper functions. My reasoning is that there will be more of these coming (e.g. symmetry in the next PR), and I think it is easier to navigate if all GPU specific implementations are in the same place.
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-05-14T13:56:15Z
body: > Regarding the location, I put it in src/common/gpu_opts/lobpcg.jl, rather than in src/eigen/linalg.jl with the rest of the helper functions. My reasoning is that there will be more of these coming (e.g. symmetry in the next PR), and I think it is easier to navigate if all GPU specific implementations are in the same place.

Ok, I see. Then I'd not have this in common and rather use a separate subdirectory `src/gpu` with all these files inside. I would also move `src/workarounds/gpu_arrays.jl` in there. Maybe good to keep the naming between cpu and gpu consistent, i.e. if there is `src/eigen/linalg.jl` there should be `src/gpu/linalg.jl` with the respective GPU things ?
author_association: MEMBER
user: login: abussy
created_at: 2025-05-14T15:10:12Z
body: > Maybe good to keep the naming between cpu and gpu consistent, i.e. if there is src/eigen/linalg.jl there should be src/gpu/linalg.jl with the respective GPU things ?

That would probably be the clearest solution. Only caveat is possible file proliferation:  `linalg.jl`, `lobpcg.jl` and `preconditioner.jl` for this PR alone. I think it is ok, since writing GPU specific code is only done when truly necessary.
author_association: COLLABORATOR
user: login: mfherbst
created_at: 2025-05-14T16:26:19Z
body: > Only caveat is possible file proliferation

True, but to me the preconditioner and lobpcg GPU things make sense to just be in a file `src/gpu/linalg.jl` ... but whatever, probably not too important after all. I think what matters is that the GPU functions should not start to scatter around multiple locations in the code base, but are in one clearly identified file or folder.
author_association: MEMBER
user: login: mfherbst
created_at: 2025-05-15T08:33:11Z
body: I just removed the assertions for the size checks. Such functions are pretty low-level and either the called functions anyway perform such size checks or the size agreement is ensured by the surrounding algorithm. Such `@assert` insertions will in general only discourage the compiler from inlineing.

@abussy Thanks very much for the good work, laying the foundation for more GPU improvements in the future.
author_association: MEMBER