[CUSPARSE] Fix wrong method for kron with Diagonal and add support to kron with FillArrays

[albertomercurio]

The previous implementation of kron between a CUSPARSE matrix and a Diagonal was wrong. It was assuming the diagonal to be the identity, which is not always the case. I have extended it to any general Diagonal matrix embedding a CuVector.

In order to support the kron with identity like matrices, I have added the extension to FillArrays.jl. Here the diagonal is filled by an unique value.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic master) to apply these changes.

Click here to view the suggested changes.

diff --git a/ext/FillArraysExt.jl b/ext/FillArraysExt.jl
index 8470d47e6..9db0459ed 100644
--- a/ext/FillArraysExt.jl
+++ b/ext/FillArraysExt.jl
@@ -29,13 +29,13 @@ function LinearAlgebra.kron(A::CUSPARSE.CuSparseMatrixCOO{T1, Ti}, B::Diagonal{T
     col = repeat(col, inner = Bnnz)
     data = repeat(convert(CUDA.CuVector{T}, A.nzVal), inner = Bnnz)
 
-    row .+= CUDA.CuVector(repeat(0:nB-1, outer = Annz)) .+ 1
-    col .+= CUDA.CuVector(repeat(0:nB-1, outer = Annz)) .+ 1
+    row .+= CUDA.CuVector(repeat(0:(nB - 1), outer = Annz)) .+ 1
+    col .+= CUDA.CuVector(repeat(0:(nB - 1), outer = Annz)) .+ 1
 
     # Multiply by the fill value (already promoted type T)
     data .*= fill_value
 
-    CUSPARSE.sparse(row, col, data, out_shape..., fmt = :coo)
+    return CUSPARSE.sparse(row, col, data, out_shape..., fmt = :coo)
 end
 
 # kron between Diagonal{T, AbstractFill} and CuSparseMatrixCOO
@@ -54,9 +54,9 @@ function LinearAlgebra.kron(A::Diagonal{T1, <:FillArrays.AbstractFill{T1}}, B::C
     # Get the fill value from the diagonal
     fill_value = FillArrays.getindex_value(A.diag)
 
-    row = (0:nA-1) .* mB
+    row = (0:(nA - 1)) .* mB
     row = CUDA.CuVector(repeat(row, inner = Bnnz))
-    col = (0:nA-1) .* nB
+    col = (0:(nA - 1)) .* nB
     col = CUDA.CuVector(repeat(col, inner = Bnnz))
     data = CUDA.fill(T(fill_value), nA * Bnnz)
 
@@ -65,7 +65,7 @@ function LinearAlgebra.kron(A::Diagonal{T1, <:FillArrays.AbstractFill{T1}}, B::C
 
     data .*= repeat(B.nzVal, outer = Annz)
 
-    CUSPARSE.sparse(row, col, data, out_shape..., fmt = :coo)
+    return CUSPARSE.sparse(row, col, data, out_shape..., fmt = :coo)
 end
 
 end  # extension module
diff --git a/test/libraries/cusparse/linalg.jl b/test/libraries/cusparse/linalg.jl
index 2fb4008c7..07cd40f9e 100644
--- a/test/libraries/cusparse/linalg.jl
+++ b/test/libraries/cusparse/linalg.jl
@@ -6,7 +6,7 @@ m = 10
     A  = sprand(T, m, m, 0.2)
     B  = sprand(T, m, m, 0.3)
     ZA = spzeros(T, m, m)
-    C  = Diagonal(rand(T, div(m, 2)))
+    C = Diagonal(rand(T, div(m, 2)))
 
     dC = adapt(CuArray, C)
     @testset "type = $typ" for typ in [CuSparseMatrixCSR, CuSparseMatrixCSC]
@@ -38,25 +38,25 @@ m = 10
         end
         @testset "kronecker product with I opa = $opa" for opa in (identity, transpose, adjoint)
             @test collect(kron(opa(dA), dC)) ≈ kron(opa(A), C)
-            @test collect(kron(dC, opa(dA))) ≈ kron(C, opa(A)) 
+            @test collect(kron(dC, opa(dA))) ≈ kron(C, opa(A))
             @test collect(kron(opa(dZA), dC)) ≈ kron(opa(ZA), C)
             @test collect(kron(dC, opa(dZA))) ≈ kron(C, opa(ZA))
         end
     end
-    
+
     # Test type promotion for kron with Diagonal
     @testset "Type promotion - T1 = $T1, T2 = $T2" for (T1, T2) in [(Float32, Float64), (Float64, ComplexF64)]
         A_T1 = sprand(T1, m, m, 0.2)
         dA_T1 = CuSparseMatrixCSR(A_T1)
         dC_T2 = Diagonal(CUDA.rand(T2, div(m, 2)))
         C_T2 = collect(dC_T2)
-        
+
         # Test kron(sparse_T1, diagonal_T2)
         result_gpu = kron(dA_T1, dC_T2)
         result_cpu = kron(A_T1, C_T2)
         @test eltype(result_gpu) == promote_type(T1, T2)
         @test collect(result_gpu) ≈ result_cpu
-        
+
         # Test kron(diagonal_T2, sparse_T1)
         result_gpu = kron(dC_T2, dA_T1)
         result_cpu = kron(C_T2, A_T1)
diff --git a/test/libraries/fillarrays.jl b/test/libraries/fillarrays.jl
index fa108e197..bf2b203aa 100644
--- a/test/libraries/fillarrays.jl
+++ b/test/libraries/fillarrays.jl
@@ -6,91 +6,91 @@ using FillArrays
     @testset "T = $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
         m, n = 100, 80
         A = sprand(T, m, n, 0.2)
-        
+
         # Test with Ones
         @testset "Ones - size $diag_size" for diag_size in [5, 10]
             C_ones = Diagonal(Ones{T}(diag_size))
             dA = CuSparseMatrixCSR(A)
-            
+
             # Test kron(sparse, diagonal)
             result_gpu = collect(kron(dA, C_ones))
             result_cpu = kron(A, collect(C_ones))
             @test result_gpu ≈ result_cpu
-            
+
             # Test kron(diagonal, sparse)
             result_gpu = collect(kron(C_ones, dA))
             result_cpu = kron(collect(C_ones), A)
             @test result_gpu ≈ result_cpu
         end
-        
+
         # Test with Fill (constant value)
         @testset "Fill - value $val, size $diag_size" for val in [T(2.5), T(-1.0)], diag_size in [5, 10]
             C_fill = Diagonal(Fill(val, diag_size))
             dA = CuSparseMatrixCSR(A)
-            
+
             # Test kron(sparse, diagonal)
             result_gpu = collect(kron(dA, C_fill))
             result_cpu = kron(A, collect(C_fill))
             @test result_gpu ≈ result_cpu
-            
+
             # Test kron(diagonal, sparse)
             result_gpu = collect(kron(C_fill, dA))
             result_cpu = kron(collect(C_fill), A)
             @test result_gpu ≈ result_cpu
         end
-        
+
         # Test with Zeros
         @testset "Zeros - size $diag_size" for diag_size in [5, 10]
             C_zeros = Diagonal(Zeros{T}(diag_size))
             dA = CuSparseMatrixCSR(A)
-            
+
             # Test kron(sparse, diagonal)
             result_gpu = collect(kron(dA, C_zeros))
             result_cpu = kron(A, collect(C_zeros))
             @test result_gpu ≈ result_cpu
             @test nnz(kron(dA, C_zeros)) == 0
-            
+
             # Test kron(diagonal, sparse)
             result_gpu = collect(kron(C_zeros, dA))
             result_cpu = kron(collect(C_zeros), A)
             @test result_gpu ≈ result_cpu
             @test nnz(kron(C_zeros, dA)) == 0
         end
-        
+
         # Test with transpose and adjoint wrappers
         @testset "Transpose/Adjoint with Fill" begin
             diag_size = 5
             val = T(3.0)
             C_fill = Diagonal(Fill(val, diag_size))
-            
+
             @testset "opa = $opa" for opa in (identity, transpose, adjoint)
                 dA = CuSparseMatrixCSR(A)
-                
+
                 # Test kron(opa(sparse), diagonal)
                 result_gpu = collect(kron(opa(dA), C_fill))
                 result_cpu = kron(opa(A), collect(C_fill))
                 @test result_gpu ≈ result_cpu
-                
+
                 # Test kron(diagonal, opa(sparse))
                 result_gpu = collect(kron(C_fill, opa(dA)))
                 result_cpu = kron(collect(C_fill), opa(A))
                 @test result_gpu ≈ result_cpu
             end
         end
-        
+
         # Test type promotion
         @testset "Type promotion" begin
             @testset "T1 = $T1, T2 = $T2" for (T1, T2) in [(Float32, Float64), (Float64, ComplexF64), (Float32, ComplexF32)]
                 A_T1 = sprand(T1, m, n, 0.2)
                 dA_T1 = CuSparseMatrixCSR(A_T1)
                 C_fill_T2 = Diagonal(Fill(T2(2.5), 5))
-                
+
                 # Test kron(sparse_T1, diagonal_T2)
                 result_gpu = kron(dA_T1, C_fill_T2)
                 result_cpu = kron(A_T1, collect(C_fill_T2))
                 @test eltype(result_gpu) == promote_type(T1, T2)
                 @test collect(result_gpu) ≈ result_cpu
-                
+
                 # Test kron(diagonal_T2, sparse_T1)
                 result_gpu = kron(C_fill_T2, dA_T1)
                 result_cpu = kron(collect(C_fill_T2), A_T1)