Working pullback and tests for EIGH + CUDA (#235)

Author: kshyatt

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

@@ -8,7 +8,7 @@ using MatrixAlgebraKit: CUSOLVER, LQViaTransposedQR, TruncationByValue, Abstract
 using MatrixAlgebraKit: default_qr_algorithm, default_lq_algorithm, default_svd_algorithm, default_eig_algorithm, default_eigh_algorithm
 import MatrixAlgebraKit: geqrf!, ungqr!, unmqr!, gesvd!, gesvdp!, gesvdr!, gesvdj!
 import MatrixAlgebraKit: heevj!, heevd!, geev!
-import MatrixAlgebraKit: _gpu_Xgesvdr!, _sylvester, svd_rank, svd_pullback!
+import MatrixAlgebraKit: _gpu_Xgesvdr!, _sylvester, svd_rank, svd_pullback!, eigh_pullback!
 using CUDA, CUDA.cuBLAS
 using CUDA: i32
 using LinearAlgebra
@@ -205,4 +205,8 @@ function svd_pullback!(ΔA::AnyCuMatrix, A, USVᴴ, ΔUSVᴴ, ind::AnyCuVector;
     return svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, collect(ind); kwargs...)
 end
 
+function eigh_pullback!(ΔA::AnyCuMatrix, A, DV, ΔDV, ind::AnyCuVector; kwargs...)
+    return eigh_pullback!(ΔA, A, DV, ΔDV, collect(ind); kwargs...)
+end
+
 end

src/pullbacks/eigh.jl

@@ -31,7 +31,7 @@ function check_and_prepare_eigh_cotangents(
     bc = Base.broadcasted(transpose(D), D, aVᴴΔV₁) do d₁, d₂, v
         return abs(d₁ - d₂) < degeneracy_atol ? v : zero(v)
     end
-    Δgauge = norm(bc, Inf)
+    Δgauge = maximum(abs, bc; init = abs(zero(eltype(D))))
 
     Δgauge ≤ gauge_atol ||
         @warn "`eigh` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
@@ -42,7 +42,8 @@ function check_and_prepare_eigh_cotangents(
     if !iszerotangent(ΔDmat)
         ΔD = diagview(ΔDmat)
         length(indD) == length(ΔD) || throw(DimensionMismatch())
-        view(diagview(VᴴAΔV), indD) .+= real.(ΔD)
+        # needed to avoid GPUCompiler errors
+        VᴴAΔV[diagind(VᴴAΔV)[indD]] .+= real.(ΔD)
     else
         ΔD = nothing
     end

test/mooncake/eigh.jl

@@ -18,4 +18,9 @@ for T in (BLASFloats..., GenericFloats...)
         AT = Diagonal{T, Vector{T}}
         TestSuite.test_mooncake_eigh(AT, m; atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
+    if T ∈ BLASFloats && CUDA.functional()
+        TestSuite.test_mooncake_eigh(CuMatrix{T}, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        AT = Diagonal{T, CuVector{T}}
+        TestSuite.test_mooncake_eigh(AT, m; atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+    end
 end