Remove shape-preserving Diagonal conversion constructors. (#2805)

They should always collect the diagonal and return a dense matrix.
For preserving conversions, use Adapt.

Author: maleadt

lib/cublas/CUBLAS.jl

@@ -20,6 +20,8 @@ using LLVM.Interop: assume
 
 using CEnum: @cenum
 
+using Adapt: adapt
+
 
 const cudaDataType_t = cudaDataType
 

lib/cublas/linalg.jl

@@ -328,8 +328,35 @@ for (t, uploc, isunitc) in ((:LowerTriangular, 'L', 'N'),
 end
 
 # Diagonal
-Base.Array(D::Diagonal{T, <:CuArray{T}}) where {T} = Diagonal(Array(D.diag))
-CuArray(D::Diagonal{T, <:Vector{T}}) where {T} = Diagonal(CuArray(D.diag))
+
+# conversions to dense matrices
+Base.Array(D::Diagonal{T, <:CuArray{T}}) where {T} = Array(Diagonal(Array(D.diag)))
+CUDA.CuArray(D::Diagonal{T}) where {T} = CuMatrix(D)
+function CUDA.CuMatrix{T}(D::Diagonal) where {T}
+    n = size(D, 1)
+    B = CUDA.zeros(T, n, n)
+    n == 0 && return B
+
+    gpu_diag = adapt(CuArray, D.diag)
+    ## COV_EXCL_START
+    function fill_diagonal()
+        i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
+        grid_stride = gridDim().x * blockDim().x
+        while i <= n
+            @inbounds B[i,i] = gpu_diag[i]
+            i += grid_stride
+        end
+        return
+    end
+    ## COV_EXCL_STOP
+    kernel = @cuda launch = false fill_diagonal()
+    config = launch_configuration(kernel.fun)
+    threads = min(config.threads, n)
+    blocks = min(config.blocks, cld(n, threads))
+    @cuda threads blocks fill_diagonal()
+
+    return B
+end
 
 function LinearAlgebra.inv(D::Diagonal{T, <:CuArray{T}}) where {T}
     Di = map(inv, D.diag)

test/base/array.jl

@@ -1,5 +1,4 @@
 using LinearAlgebra
-import Adapt
 using ChainRulesCore: add!!, is_inplaceable_destination
 
 @testset "constructors" begin

test/base/kernelabstractions.jl

@@ -1,7 +1,6 @@
 import KernelAbstractions
 import KernelAbstractions as KA
 using SparseArrays
-using Adapt
 
 include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl"))
 

test/core/execution.jl

@@ -1,5 +1,3 @@
-import Adapt
-
 dummy() = return
 
 @testset "@cuda" begin

test/libraries/cublas/level3.jl

@@ -127,7 +127,7 @@ k = 13
                     uplotype in (UpperTriangular, UnitUpperTriangular, LowerTriangular, UnitLowerTriangular)
 
                 A = triu(rand(elty, m, m))
-                dA = CuArray(A) 
+                dA = CuArray(A)
                 Br = rand(elty,m,n)
                 Bl = rand(elty,n,m)
                 d_Br = CuArray(Br)
@@ -136,16 +136,16 @@ k = 13
                 @test Bl/adjtype(uplotype(A)) ≈ Array(d_Bl/adjtype(uplotype(dA)))
             end
             # Check also that scaling parameter works
-            alpha = rand(elty) 
+            alpha = rand(elty)
             A = triu(rand(elty, m, m))
-            dA = CuArray(A) 
+            dA = CuArray(A)
             Br = rand(elty,m,n)
             d_Br = CuArray(Br)
             @test BLAS.trsm('L','U','N','N',alpha,A,Br) ≈ Array(CUBLAS.trsm('L','U','N','N',alpha,dA,d_Br))
         end
 
         @testset "trsm_batched!" begin
-            alpha = rand(elty) 
+            alpha = rand(elty)
             bA = [rand(elty,m,m) for i in 1:10]
             map!((x) -> triu(x), bA, bA)
             bB = [rand(elty,m,n) for i in 1:10]
@@ -288,7 +288,7 @@ k = 13
             B = Diagonal(rand(elty, m))
 
             dA = CuArray(A)
-            dB = CuArray(B)
+            dB = adapt(CuArray, B)
 
             C = A / B
             d_C = dA / dB
@@ -297,7 +297,7 @@ k = 13
             @test C ≈ Array(dA)
             @test C ≈ h_C
 
-            B_bad = Diagonal(rand(elty, m+1))
+            B_bad = Diagonal(CuArray(rand(elty, m+1)))
             @test_throws DimensionMismatch("left hand side has $m columns but D is $(m+1) by $(m+1)") rdiv!(dA, B_bad)
         end
 
@@ -483,12 +483,12 @@ k = 13
 
             @test A ≈ h_A
             @test B ≈ h_B
-            
+
             diagA = diagm(m, m, 0 => d_A)
             diagind_A = diagind(diagA, 0)
             h_A = Array(diagA[diagind_A])
             @test A ≈ h_A
-            
+
             diagA = diagm(m, m, d_A)
             diagind_A = diagind(diagA, 0)
             h_A = Array(diagA[diagind_A])

test/libraries/cusparse/conversions.jl

@@ -1,8 +1,6 @@
 using LinearAlgebra
-using Adapt
 using CUDA.CUSPARSE
 using SparseArrays
-using CUDA
 
 @testset "sparse" begin
     n, m = 4, 4

test/libraries/cusparse/generic.jl

@@ -1,5 +1,3 @@
-using CUDA
-using Adapt
 using CUDA.CUSPARSE
 using SparseArrays
 using LinearAlgebra

test/libraries/cusparse/interfaces.jl

@@ -1,5 +1,3 @@
-using CUDA
-using Adapt
 using CUDA.CUSPARSE
 using LinearAlgebra, SparseArrays
 

test/setup.jl

@@ -14,6 +14,8 @@ testf(f, xs...; kwargs...) = TestSuite.compare(f, CuArray, xs...; kwargs...)
 
 using Random
 
+using Adapt
+
 # detect compute-sanitizer, to disable incompatible tests (e.g. using CUPTI)
 const sanitize = any(contains("NV_SANITIZER"), keys(ENV))