diff --git a/src/host/linalg.jl b/src/host/linalg.jl index a198c3de..3c6788e2 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -326,6 +326,14 @@ end Base.copy(D::Diagonal{T, <:AbstractGPUArray{T, N}}) where {T, N} = Diagonal(copy(D.diag)) +# copy only the diagonal; the generic `AnyGPUArray` method would launch a kernel over all +# n^2 entries, calling `setindex!(::Diagonal, ...)` whose error path isn't GPU-compatible +function Base.copyto!(D1::Diagonal{T, <:AbstractGPUArray{T, N}}, + D2::Diagonal{S, <:AbstractGPUArray{S, M}}) where {T, N, S, M} + copyto!(D1.diag, D2.diag) + return D1 +end + _isrealandpositive(x) = isreal(x) && real(x) > 0 function LinearAlgebra.cholesky!(D::Diagonal{<:Any, <:AbstractGPUArray}, diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl index 0df896c7..6a3835b1 100644 --- a/test/testsuite/linalg.jl +++ b/test/testsuite/linalg.jl @@ -512,6 +512,18 @@ end @test collect(D) == collect(C) end + @testset "copyto! diagonal" begin + D = Diagonal(AT(rand(Float32, 10))) + C = Diagonal(AT(zeros(Float32, 10))) + @test copyto!(C, D) === C + @test C.diag isa AT + @test collect(C) == collect(D) + + D = Diagonal(AT(Int32.(1:10))) + copyto!(C, D) + @test collect(C) == Diagonal(Float32.(1:10)) + end + @testset "cholesky + Diagonal" begin n = 128 # Add one in order prevent failures due to random numbers being zero