Implement \ and ldiv! for Diagonal (#391)

Author: simsurace

src/host/linalg.jl

@@ -127,6 +127,33 @@ else
     end
 end
 
+function Base.:\(D::Diagonal{<:Any, <:AbstractGPUArray}, B::AbstractGPUVecOrMat)
+    z = D.diag .== 0
+    if any(z)
+        i = findfirst(collect(z))
+        throw(SingularException(i))
+    else
+        return D.diag .\ B
+    end
+end
+
+function LinearAlgebra.ldiv!(D::Diagonal{<:Any, <:AbstractGPUArray}, B::StridedVecOrMat) 
+    m, n = size(B, 1), size(B, 2)
+    if m != length(D.diag)
+        throw(DimensionMismatch("diagonal matrix is $(length(D.diag)) by $(length(D.diag)) but right hand side has $m rows"))
+    end
+    (m == 0 || n == 0) && return B
+    z = D.diag .== 0
+    if any(z)
+        i = findfirst(collect(z))
+        throw(SingularException(i))
+    else
+        B .= D.diag .\ B
+    end
+    return B
+end
+
+
 ## matrix multiplication
 
 function generic_matmatmul!(C::AbstractArray{R}, A::AbstractArray{T}, B::AbstractArray{S}, a::Number, b::Number) where {T,S,R}

test/testsuite/linalg.jl

@@ -130,6 +130,51 @@
             @test cholesky(D, check = false).info == 3
         end
 
+        @testset "\\ + Diagonal" begin
+            n = 128
+            d = AT(rand(Float32, n))
+            D = Diagonal(d)
+            b = AT(rand(Float32, n))
+            B = AT(rand(Float32, n, n))
+            @test collect(D \ b) ≈ Diagonal(collect(d)) \ collect(b)
+            @test collect(D \ B) ≈ Diagonal(collect(d)) \ collect(B)
+
+            d = ones(Float32, n)
+            d[rand(1:n)] = 0
+            d = AT(d)
+            D = Diagonal(d)
+            @test_throws SingularException D \ B
+        end
+
+        @testset "ldiv! + Diagonal" begin
+            n = 128
+            d = AT(rand(Float32, n))
+            D = Diagonal(d)
+            b = AT(rand(Float32, n))
+            B = AT(rand(Float32, n, n))
+            X = AT(zeros(Float32, n, n))
+            Y = zeros(Float32, n, n)
+            ldiv!(X, D, B)
+            ldiv!(Y, Diagonal(collect(d)), collect(B))
+            @test collect(X) ≈ Y
+            ldiv!(D, B)
+            @test collect(B) ≈ collect(X)
+
+            d = ones(Float32, n)
+            d[rand(1:n)] = 0
+            d = AT(d)
+            D = Diagonal(d)
+            B = AT(rand(Float32, n, n))
+
+            # three-argument version does not throw SingularException
+            ldiv!(X, D, B)
+            ldiv!(Y, Diagonal(collect(d)), collect(B))
+            @test collect(X) ≈ Y
+
+            # two-argument version throws SingularException
+            @test_throws SingularException ldiv!(D, B)
+        end
+
         @testset "$f! with diagonal $d" for (f, f!) in ((triu, triu!), (tril, tril!)),
                                             d in -2:2
             A = randn(Float32, 10, 10)