diff --git a/lib/cusparse/linalg.jl b/lib/cusparse/linalg.jl index 4be578e284..e72d1a0841 100644 --- a/lib/cusparse/linalg.jl +++ b/lib/cusparse/linalg.jl @@ -57,79 +57,83 @@ function Base.reshape(A::CuSparseMatrixCOO, dims::Dims) sparse(new_row, new_col, A.nzVal, dims[1], length(dims) == 1 ? 1 : dims[2], fmt = :coo) end -function LinearAlgebra.kron(A::CuSparseMatrixCOO{T, Ti}, B::CuSparseMatrixCOO{T, Ti}) where {Ti, T} - mA,nA = size(A) - mB,nB = size(B) - out_shape = (mA * mB, nA * nB) - Annz = Int64(A.nnz) - Bnnz = Int64(B.nnz) - - if Annz == 0 || Bnnz == 0 - return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape) +_kron_CuSparseMatrixCOO_components(A::CuSparseMatrixCOO) = A.rowInd, A.colInd, A.nzVal, identity, Int(A.nnz) +_kron_CuSparseMatrixCOO_components(At::Transpose{<:Number, <:CuSparseMatrixCOO}) = parent(At).colInd, parent(At).rowInd, parent(At).nzVal, transpose, Int(parent(At).nnz) +_kron_CuSparseMatrixCOO_components(Ah::Adjoint{<:Number, <:CuSparseMatrixCOO}) = parent(Ah).colInd, parent(Ah).rowInd, parent(Ah).nzVal, adjoint, Int(parent(Ah).nnz) + +function LinearAlgebra.kron( + A::Union{CuSparseMatrixCOO{TvA, TiA}, Transpose{TvA, <:CuSparseMatrixCOO{TvA, TiA}}, Adjoint{TvA, <:CuSparseMatrixCOO{TvA, TiA}}}, + B::Union{CuSparseMatrixCOO{TvB, TiB}, Transpose{TvB, <:CuSparseMatrixCOO{TvB, TiB}}, Adjoint{TvB, <:CuSparseMatrixCOO{TvB, TiB}}} + ) where {TvA, TiA, TvB, TiB} + mA, nA = size(A) + mB, nB = size(B) + Ti = promote_type(TiA, TiB) + Tv = typeof(oneunit(TvA)*oneunit(TvB)) + + A_rowInd, A_colInd, A_nzVal, A_nzOp, A_nnz = _kron_CuSparseMatrixCOO_components(A) + B_rowInd, B_colInd, B_nzVal, B_nzOp, B_nnz = _kron_CuSparseMatrixCOO_components(B) + + if A_nnz == 0 || B_nnz == 0 + return CuSparseMatrixCOO{Tv, Ti}(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{Tv}(undef, 0), (mA * mB, nA * nB)) end - row = (A.rowInd .- 1) .* mB - row = repeat(row, inner = Bnnz) - col = (A.colInd .- 1) .* nB - col = repeat(col, inner = Bnnz) - data = repeat(A.nzVal, inner = Bnnz) + C_nnz = A_nnz * B_nnz + C_rowInd = reshape(B_rowInd .+ Ti(mB) .* (reshape(A_rowInd, (1, A_nnz)) .- one(Ti)), C_nnz) + C_colInd = reshape(B_colInd .+ Ti(nB) .* (reshape(A_colInd, (1, A_nnz)) .- one(Ti)), C_nnz) + C_nzVal = reshape(B_nzOp.(B_nzVal) .* A_nzOp.(reshape(A_nzVal, (1, A_nnz))), C_nnz) - row .+= repeat(B.rowInd .- 1, outer = Annz) .+ 1 - col .+= repeat(B.colInd .- 1, outer = Annz) .+ 1 - - data .*= repeat(B.nzVal, outer = Annz) - - sparse(row, col, data, out_shape..., fmt = :coo) + C = CuSparseMatrixCOO{Tv, Ti}(C_rowInd, C_colInd, C_nzVal, (mA * mB, nA * nB), C_nnz) + return sort_coo(C) end -function LinearAlgebra.kron(A::CuSparseMatrixCOO{T, Ti}, B::Diagonal) where {Ti, T} - mA,nA = size(A) - mB,nB = size(B) - out_shape = (mA * mB, nA * nB) - Annz = Int64(A.nnz) - Bnnz = nB - - if Annz == 0 || Bnnz == 0 - return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape) - end +function LinearAlgebra.kron( + A::Union{CuSparseMatrixCOO{TvA, TiA}, Transpose{TvA, <:CuSparseMatrixCOO{TvA, TiA}}, Adjoint{TvA, <:CuSparseMatrixCOO{TvA, TiA}}}, + B::Diagonal{TvB, <:Union{CuVector{TvB}, Base.ReshapedArray{TvB, 1, <:Adjoint{TvB, <:CuVector{TvB}}}}} + ) where {TvA, TiA, TvB} + mA, nA = size(A) + mB, nB = size(B) + Ti = TiA + Tv = typeof(oneunit(TvA)*oneunit(TvB)) - row = (A.rowInd .- 1) .* mB - row = repeat(row, inner = Bnnz) - col = (A.colInd .- 1) .* nB - col = repeat(col, inner = Bnnz) - data = repeat(A.nzVal, inner = Bnnz) + A_rowInd, A_colInd, A_nzVal, A_nzOp, A_nnz = _kron_CuSparseMatrixCOO_components(A) + B_rowInd, B_colInd, B_nzVal, B_nnz = one(Ti):Ti(nB), one(Ti):Ti(nB), B.diag, Int(nB) - row .+= CuVector(repeat(0:nB-1, outer = Annz)) .+ 1 - col .+= CuVector(repeat(0:nB-1, outer = Annz)) .+ 1 + if A_nnz == 0 || B_nnz == 0 + return CuSparseMatrixCOO{Tv, Ti}(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{Tv}(undef, 0), (mA * mB, nA * nB)) + end - data .*= repeat(CUDA.ones(T, nB), outer = Annz) + C_nnz = A_nnz * B_nnz + C_rowInd = reshape(B_rowInd .+ Ti(mB) .* (reshape(A_rowInd, (1, A_nnz)) .- one(Ti)), C_nnz) + C_colInd = reshape(B_colInd .+ Ti(nB) .* (reshape(A_colInd, (1, A_nnz)) .- one(Ti)), C_nnz) + C_nzVal = reshape(B_nzVal .* A_nzOp.(reshape(A_nzVal, (1, A_nnz))), C_nnz) - sparse(row, col, data, out_shape..., fmt = :coo) + C = CuSparseMatrixCOO{Tv, Ti}(C_rowInd, C_colInd, C_nzVal, (mA * mB, nA * nB), C_nnz) + return sort_coo(C) end -function LinearAlgebra.kron(A::Diagonal, B::CuSparseMatrixCOO{T, Ti}) where {Ti, T} - mA,nA = size(A) - mB,nB = size(B) - out_shape = (mA * mB, nA * nB) - Annz = nA - Bnnz = Int64(B.nnz) - - if Annz == 0 || Bnnz == 0 - return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape) - end +function LinearAlgebra.kron( + A::Diagonal{TvA, <:Union{CuVector{TvA}, Base.ReshapedArray{TvA, 1, <:Adjoint{TvA, <:CuVector{TvA}}}}}, + B::Union{CuSparseMatrixCOO{TvB, TiB}, Transpose{TvB, <:CuSparseMatrixCOO{TvB, TiB}}, Adjoint{TvB, <:CuSparseMatrixCOO{TvB, TiB}}} + ) where {TvA, TvB, TiB} + mA, nA = size(A) + mB, nB = size(B) + Ti = TiB + Tv = typeof(oneunit(TvA)*oneunit(TvB)) - row = (0:nA-1) .* mB - row = CuVector(repeat(row, inner = Bnnz)) - col = (0:nA-1) .* nB - col = CuVector(repeat(col, inner = Bnnz)) - data = repeat(CUDA.ones(T, nA), inner = Bnnz) + A_rowInd, A_colInd, A_nzVal, A_nnz = one(Ti):Ti(nA), one(Ti):Ti(nA), A.diag, Int(nA) + B_rowInd, B_colInd, B_nzVal, B_nzOp, B_nnz = _kron_CuSparseMatrixCOO_components(B) - row .+= repeat(B.rowInd .- 1, outer = Annz) .+ 1 - col .+= repeat(B.colInd .- 1, outer = Annz) .+ 1 + if A_nnz == 0 || B_nnz == 0 + return CuSparseMatrixCOO{Tv, Ti}(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{Tv}(undef, 0), (mA * mB, nA * nB)) + end - data .*= repeat(B.nzVal, outer = Annz) + C_nnz = A_nnz * B_nnz + C_rowInd = reshape(B_rowInd .+ Ti(mB) .* (reshape(A_rowInd, (1, A_nnz)) .- one(Ti)), C_nnz) + C_colInd = reshape(B_colInd .+ Ti(nB) .* (reshape(A_colInd, (1, A_nnz)) .- one(Ti)), C_nnz) + C_nzVal = reshape(B_nzOp.(B_nzVal) .* reshape(A_nzVal, (1, A_nnz)), C_nnz) - sparse(row, col, data, out_shape..., fmt = :coo) + C = CuSparseMatrixCOO{Tv, Ti}(C_rowInd, C_colInd, C_nzVal, (mA * mB, nA * nB), C_nnz) + return sort_coo(C) end function LinearAlgebra.dot(y::CuVector{T}, A::CuSparseMatrixCSC{T}, x::CuVector{T}) where {T<:Union{BlasInt, BlasFloat}} diff --git a/test/libraries/cusparse/linalg.jl b/test/libraries/cusparse/linalg.jl index 1f865e558c..30c529821c 100644 --- a/test/libraries/cusparse/linalg.jl +++ b/test/libraries/cusparse/linalg.jl @@ -1,5 +1,5 @@ using CUDA.CUSPARSE -using LinearAlgebra, SparseArrays +using LinearAlgebra, SparseArrays, Adapt m = 10 @testset "T = $T" for T in [Float32, Float64, ComplexF32, ComplexF64] @@ -7,10 +7,13 @@ m = 10 B = sprand(T, m, m, 0.3) ZA = spzeros(T, m, m) C = I(div(m, 2)) + D = Diagonal(rand(T, m)) @testset "type = $typ" for typ in [CuSparseMatrixCSR, CuSparseMatrixCSC] dA = typ(A) dB = typ(B) dZA = typ(ZA) + dD = adapt(CuArray, D) + dC = adapt(CuArray, C) @testset "opnorm and norm" begin @test opnorm(A, Inf) ≈ opnorm(dA, Inf) @test opnorm(A, 1) ≈ opnorm(dA, 1) @@ -37,11 +40,60 @@ m = 10 end end @testset "kronecker product with I opa = $opa" for opa in (identity, transpose, adjoint) - @test collect(kron(opa(dA), C)) ≈ kron(opa(A), C) - @test collect(kron(C, opa(dA))) ≈ kron(C, opa(A)) - @test collect(kron(opa(dZA), C)) ≈ kron(opa(ZA), C) - @test collect(kron(C, opa(dZA))) ≈ kron(C, opa(ZA)) + @test collect(kron(opa(dA), dC)) ≈ kron(opa(A), C) + @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 + @testset "kronecker product with Diagonal opa = $opa" for opa in (identity, transpose, adjoint) + @test collect(kron(opa(dA), dD)) ≈ kron(opa(A), D) + @test collect(kron(dD, opa(dA))) ≈ kron(D, opa(A)) + @test collect(kron(opa(dZA), dD)) ≈ kron(opa(ZA), D) + @test collect(kron(dD, opa(dZA))) ≈ kron(D, opa(ZA)) + end + end +end + +@testset "T = $T" for T in [Float32, Float64, ComplexF32, ComplexF64] + mat_sizes = [(2, 3), (2, 0)] + @testset "size(A) = ($(mA), $(nA)), size(B) = ($(mB), $(nB))" for (mA, nA) in mat_sizes, (mB, nB) in mat_sizes + A = sprand(T, mA, nA, 0.5) + B = sprand(T, mB, nB, 0.5) + + A_I, A_J, A_V = findnz(A) + dA = CuSparseMatrixCOO{T, Cint}(adapt(CuVector{Cint}, A_I), adapt(CuVector{Cint}, A_J), adapt(CuVector{T}, A_V), size(A)) + B_I, B_J, B_V = findnz(B) + dB = CuSparseMatrixCOO{T, Cint}(adapt(CuVector{Cint}, B_I), adapt(CuVector{Cint}, B_J), adapt(CuVector{T}, B_V), size(B)) + + @testset "kronecker (COO ⊗ COO) opa = $opa, opb = $opb" for opa in (identity, transpose, adjoint), opb in (identity, transpose, adjoint) + dC = kron(opa(dA), opb(dB)) + @test collect(dC) ≈ kron(opa(A), opb(B)) + @test eltype(dC) == typeof(oneunit(T) * oneunit(T)) + @test dC isa CuSparseMatrixCOO + end + end +end + +@testset "TA = $TA, TvB = $TvB" for TvB in [Float32, Float64, ComplexF32, ComplexF64], TA in [Bool, TvB] + A = Diagonal(rand(TA, 2)) + B = sprand(TvB, 3, 4, 0.5) + dA = adapt(CuArray, A) + + B_I, B_J, B_V = findnz(B) + dB = CuSparseMatrixCOO{TvB, Cint}(adapt(CuVector{Cint}, B_I), adapt(CuVector{Cint}, B_J), adapt(CuVector{TvB}, B_V), size(B)) + + @testset "kronecker (diagonal ⊗ COO) opa = $opa, opb = $opb" for opa in (identity, adjoint), opb in (identity, transpose, adjoint) + dC = kron(opa(dA), opb(dB)) + @test collect(dC) ≈ kron(opa(A), opb(B)) + @test eltype(dC) == typeof(oneunit(TA) * oneunit(TvB)) + @test dC isa CuSparseMatrixCOO + end + + @testset "kronecker (COO ⊗ diagonal) opa = $opa, opb = $opb" for opa in (identity, adjoint), opb in (identity, transpose, adjoint) + dC = kron(opb(dB), opa(dA)) + @test collect(dC) ≈ kron(opb(B), opa(A)) + @test eltype(dC) == typeof(oneunit(TvB) * oneunit(TA)) + @test dC isa CuSparseMatrixCOO end end