Specialize cholesky for Diagonal (#384)

Author: simsurace

src/host/linalg.jl

@@ -101,6 +101,20 @@ end
 
 Base.copy(D::Diagonal{T, <:AbstractGPUArray{T, N}}) where {T, N} = Diagonal(copy(D.diag))
 
+# prevent scalar indexing
+function LinearAlgebra.cholesky!(D::Diagonal{T, <:AbstractGPUArray{T, N}}, 
+    ::Val{false} = Val(false); check::Bool = true
+) where {T, N}
+    info = 0
+    if mapreduce(x -> isreal(x) && isposdef(x), &, D.diag)
+        D.diag .= sqrt.(D.diag)
+    else
+        info = findfirst(x -> !isreal(x) || !isposdef(x), collect(D.diag))
+        check && throw(PosDefException(info))
+    end
+    Cholesky(D, 'U', convert(LinearAlgebra.BlasInt, info))
+end
+
 
 ## matrix multiplication
 

test/testsuite/linalg.jl

@@ -21,6 +21,7 @@
         end
     end
 
+
     @testset "symmetric" begin
         @testset "Hermitian" begin
             A    = rand(Float32,2,2)
@@ -103,6 +104,19 @@
             @test collect(D) == collect(C)
         end
 
+        @testset "cholesky + Diagonal" begin
+            n = 128
+            d = AT(rand(Float32, n))
+            D = Diagonal(d)
+            F = collect(D)
+            @test collect(cholesky(D).U) ≈ collect(cholesky(F).U)
+            @test collect(cholesky(D).L) ≈ collect(cholesky(F).L)
+
+            d = AT([1f0, 2f0, -1f0, 0f0])
+            D = Diagonal(d)
+            @test cholesky(D, check = false).info == 3
+        end
+
         @testset "$f! with diagonal $d" for (f, f!) in ((triu, triu!), (tril, tril!)),
                                             d in -2:2
             A = randn(Float32, 10, 10)