From 909b31d834b958f05964806721df7f85ce6fceba Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sun, 15 Dec 2024 15:03:40 +0100 Subject: [PATCH 1/4] Add vector-vector Kronecker product --- .gitignore | 2 ++ src/host/linalg.jl | 27 +++++++++++++++++++++++++++ test/testsuite/linalg.jl | 6 ++++++ 3 files changed, 35 insertions(+) diff --git a/.gitignore b/.gitignore index 9b73c974..349b1c4a 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,5 @@ Manifest.toml # MacOS generated files *.DS_Store + +/.vscode diff --git a/src/host/linalg.jl b/src/host/linalg.jl index b4acc8e7..6af06f27 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -736,3 +736,30 @@ function Base.isone(x::AbstractGPUMatrix{T}) where {T} Array(y)[] end + +## Kronecker product + +function LinearAlgebra.kron!(z::AbstractGPUVector{T1}, x::AbstractGPUVector{T2}, y::AbstractGPUVector{T3}) where {T1,T2,T3} + @assert length(z) == length(x) * length(y) + + @kernel function kron_kernel!(z, @Const(x), @Const(y)) + i, j = @index(Global, NTuple) + + ly = length(y) + + @inbounds z[(i - 1) * ly + j] = x[i] * y[j] + end + + backend = KernelAbstractions.get_backend(z) + kernel = kron_kernel!(backend) + + kernel(z, x, y, ndrange=(length(x), length(y))) + + return z +end + +function LinearAlgebra.kron(x::AbstractGPUVector{T1}, y::AbstractGPUVector{T2}) where {T1,T2} + T = promote_type(T1, T2) + z = similar(x, T, length(x) * length(y)) + return LinearAlgebra.kron!(z, x, y) +end diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl index 8de31854..d9b90adc 100644 --- a/test/testsuite/linalg.jl +++ b/test/testsuite/linalg.jl @@ -310,6 +310,12 @@ @test iszero(A) @test isone(A) == false end + + @testset "kron" begin + for T in eltypes + @test compare(kron, AT, rand(T, 32), rand(T, 64)) + end + end end @testsuite "linalg/mul!/vector-matrix" (AT, eltypes)->begin From e83f5827dcdc998fbe17b9cad93588c311dd414c Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sun, 15 Dec 2024 16:01:31 +0100 Subject: [PATCH 2/4] Implement also the matrix-matrix kronecker product --- src/host/linalg.jl | 33 +++++++++++++++++++++++++++++++++ test/testsuite/linalg.jl | 1 + 2 files changed, 34 insertions(+) diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 6af06f27..5abe94d0 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -763,3 +763,36 @@ function LinearAlgebra.kron(x::AbstractGPUVector{T1}, y::AbstractGPUVector{T2}) z = similar(x, T, length(x) * length(y)) return LinearAlgebra.kron!(z, x, y) end + +function LinearAlgebra.kron!(Z::AbstractGPUMatrix{T1}, A::AbstractGPUMatrix{T2}, B::AbstractGPUMatrix{T3}) where {T1, T2, T3} + @assert size(Z, 1) == size(A, 1) * size(B, 1) + @assert size(Z, 2) == size(A, 2) * size(B, 2) + + @kernel function kron_kernel!(Z, @Const(A), @Const(B)) + ai, aj = @index(Global, NTuple) # Indices in the result matrix + + lb1, lb2 = size(B) # Dimensions of B + + # Map global indices (ai, aj) to submatrices of the Kronecker product + i_a = (ai - 1) ÷ lb1 + 1 # Corresponding row index in A + i_b = (ai - 1) % lb1 + 1 # Corresponding row index in B + j_a = (aj - 1) ÷ lb2 + 1 # Corresponding col index in A + j_b = (aj - 1) % lb2 + 1 # Corresponding col index in B + + @inbounds Z[ai, aj] = A[i_a, j_a] * B[i_b, j_b] + end + + backend = KernelAbstractions.get_backend(Z) + kernel = kron_kernel!(backend) + + kernel(Z, A, B, ndrange=(size(Z, 1), size(Z, 2))) + + return Z +end + +function LinearAlgebra.kron(A::AbstractGPUMatrix{T1}, B::AbstractGPUMatrix{T2}) where {T1, T2} + T = promote_type(T1, T2) + size_Z = (size(A, 1) * size(B, 1), size(A, 2) * size(B, 2)) + Z = similar(A, T, size_Z...) + return kron!(Z, A, B) +end diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl index d9b90adc..60394a20 100644 --- a/test/testsuite/linalg.jl +++ b/test/testsuite/linalg.jl @@ -314,6 +314,7 @@ @testset "kron" begin for T in eltypes @test compare(kron, AT, rand(T, 32), rand(T, 64)) + @test compare(kron, AT, rand(T, 32, 64), rand(T, 128, 16)) end end end From 5bdfbed280f714a4db3d5093e289be8706555810 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sun, 15 Dec 2024 17:49:13 +0100 Subject: [PATCH 3/4] Add methods also for Transpose and Adjoint --- src/host/linalg.jl | 69 +++++++++++++++++++++++++--------------- test/testsuite/linalg.jl | 4 ++- 2 files changed, 47 insertions(+), 26 deletions(-) diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 5abe94d0..6c5114f7 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -764,35 +764,54 @@ function LinearAlgebra.kron(x::AbstractGPUVector{T1}, y::AbstractGPUVector{T2}) return LinearAlgebra.kron!(z, x, y) end -function LinearAlgebra.kron!(Z::AbstractGPUMatrix{T1}, A::AbstractGPUMatrix{T2}, B::AbstractGPUMatrix{T3}) where {T1, T2, T3} - @assert size(Z, 1) == size(A, 1) * size(B, 1) - @assert size(Z, 2) == size(A, 2) * size(B, 2) +trans_adj_wrappers = ((T -> :(AbstractGPUMatrix{$T}), T -> 'N', identity), + (T -> :(Transpose{$T, <:AbstractGPUMatrix{$T}}), T -> 'T', A -> :(parent($A))), + (T -> :(Adjoint{$T, <:AbstractGPUMatrix{$T}}), T -> T <: Real ? 'T' : 'C', A -> :(parent($A)))) - @kernel function kron_kernel!(Z, @Const(A), @Const(B)) - ai, aj = @index(Global, NTuple) # Indices in the result matrix +for (wrapa, transa, unwrapa) in trans_adj_wrappers, (wrapb, transb, unwrapb) in trans_adj_wrappers + TypeA = wrapa(:(T1)) + TypeB = wrapb(:(T2)) + TypeC = :(AbstractGPUMatrix{T3}) - lb1, lb2 = size(B) # Dimensions of B + @eval function LinearAlgebra.kron!(C::$TypeC, A::$TypeA, B::$TypeB) where {T1,T2,T3} + @assert size(C, 1) == size(A, 1) * size(B, 1) + @assert size(C, 2) == size(A, 2) * size(B, 2) - # Map global indices (ai, aj) to submatrices of the Kronecker product - i_a = (ai - 1) ÷ lb1 + 1 # Corresponding row index in A - i_b = (ai - 1) % lb1 + 1 # Corresponding row index in B - j_a = (aj - 1) ÷ lb2 + 1 # Corresponding col index in A - j_b = (aj - 1) % lb2 + 1 # Corresponding col index in B - - @inbounds Z[ai, aj] = A[i_a, j_a] * B[i_b, j_b] - end - - backend = KernelAbstractions.get_backend(Z) - kernel = kron_kernel!(backend) + ta = $transa(T1) + tb = $transb(T2) + + @kernel function kron_kernel!(C, @Const(A), @Const(B)) + ai, aj = @index(Global, NTuple) # Indices in the result matrix + + # lb1, lb2 = size(B) # Dimensions of B + lb1, lb2 = tb == 'N' ? size(B) : reverse(size(B)) + + # Map global indices (ai, aj) to submatrices of the Kronecker product + i_a = (ai - 1) ÷ lb1 + 1 # Corresponding row index in A + i_b = (ai - 1) % lb1 + 1 # Corresponding row index in B + j_a = (aj - 1) ÷ lb2 + 1 # Corresponding col index in A + j_b = (aj - 1) % lb2 + 1 # Corresponding col index in B - kernel(Z, A, B, ndrange=(size(Z, 1), size(Z, 2))) + @inbounds begin + a_ij = ta == 'N' ? A[i_a, j_a] : (ta == 'T' ? A[j_a, i_a] : conj(A[j_a, i_a])) + b_ij = tb == 'N' ? B[i_b, j_b] : (tb == 'T' ? B[j_b, i_b] : conj(B[j_b, i_b])) - return Z -end + C[ai, aj] = a_ij * b_ij + end + end + + backend = KernelAbstractions.get_backend(C) + kernel = kron_kernel!(backend) + + kernel(C, $(unwrapa(:A)), $(unwrapb(:B)), ndrange=(size(C, 1), size(C, 2))) + + return C + end -function LinearAlgebra.kron(A::AbstractGPUMatrix{T1}, B::AbstractGPUMatrix{T2}) where {T1, T2} - T = promote_type(T1, T2) - size_Z = (size(A, 1) * size(B, 1), size(A, 2) * size(B, 2)) - Z = similar(A, T, size_Z...) - return kron!(Z, A, B) + @eval function LinearAlgebra.kron(A::$TypeA, B::$TypeB) where {T1, T2} + T = promote_type(T1, T2) + size_C = (size(A, 1) * size(B, 1), size(A, 2) * size(B, 2)) + C = similar(A, T, size_C...) + return kron!(C, A, B) + end end diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl index 60394a20..d4833d3e 100644 --- a/test/testsuite/linalg.jl +++ b/test/testsuite/linalg.jl @@ -314,7 +314,9 @@ @testset "kron" begin for T in eltypes @test compare(kron, AT, rand(T, 32), rand(T, 64)) - @test compare(kron, AT, rand(T, 32, 64), rand(T, 128, 16)) + for opa in (identity, transpose, adjoint), opb in (identity, transpose, adjoint) + @test compare(kron, AT, opa(rand(T, 32, 64)), opb(rand(T, 128, 16))) + end end end end From fc94e651b1c4d4ace9ba7875b24d9d2bf750f870 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Mon, 16 Dec 2024 09:45:52 +0100 Subject: [PATCH 4/4] Simplify kernel --- src/host/linalg.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 6c5114f7..2c928747 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -745,9 +745,7 @@ function LinearAlgebra.kron!(z::AbstractGPUVector{T1}, x::AbstractGPUVector{T2}, @kernel function kron_kernel!(z, @Const(x), @Const(y)) i, j = @index(Global, NTuple) - ly = length(y) - - @inbounds z[(i - 1) * ly + j] = x[i] * y[j] + @inbounds z[(i - 1) * length(y) + j] = x[i] * y[j] end backend = KernelAbstractions.get_backend(z)