Improve kron implementation (#600)

Significantly simplify indexing by not applying transposes and adjoints
manually and by using `fld1` and `mod1`. Also add some combinations
involving mixed vectors and matrices for generality

Author: simeonschaub

src/host/linalg.jl

@@ -739,17 +739,17 @@ end
 
 ## Kronecker product
 
+@kernel function kron_kernel_vec!(z, @Const(x), @Const(y))
+    i, j = @index(Global, NTuple)
+
+    @inbounds z[(i - 1) * length(y) + j] = x[i] * y[j]
+end
+
 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)
-    
-        @inbounds z[(i - 1) * length(y) + j] = x[i] * y[j]
-    end
-
     backend = KernelAbstractions.get_backend(z)
-    kernel = kron_kernel!(backend)
+    kernel = kron_kernel_vec!(backend)
 
     kernel(z, x, y, ndrange=(length(x), length(y)))
 
@@ -759,57 +759,51 @@ 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)
+    return kron!(z, x, y)
+end
+
+@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 = size(B, 1)
+    lb2 = size(B, 2)
+
+    # Map global indices (ai, aj) to submatrices of the Kronecker product
+    i_a = fld1(ai, lb1)  # Corresponding row index in A
+    i_b = mod1(ai, lb1)  # Corresponding row index in B
+    j_a = fld1(aj, lb2)  # Corresponding col index in A
+    j_b = mod1(aj, lb2)  # Corresponding col index in B
+
+    @inbounds C[ai, aj] = A[i_a, j_a] * B[i_b, j_b]
 end
 
-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))))
+trans_adj_wrappers = (
+    T -> :(AbstractGPUVecOrMat{$T}),
+    T -> :(Transpose{$T, <:AbstractGPUVecOrMat{$T}}),
+    T -> :(Adjoint{$T, <:AbstractGPUVecOrMat{$T}}),
+)
 
-for (wrapa, transa, unwrapa) in trans_adj_wrappers, (wrapb, transb, unwrapb) in trans_adj_wrappers
-    TypeA = wrapa(:(T1))
-    TypeB = wrapb(:(T2))
-    TypeC = :(AbstractGPUMatrix{T3})
+for wrapa in trans_adj_wrappers, wrapb in trans_adj_wrappers
+    TypeA = wrapa(:T1)
+    TypeB = wrapb(:T2)
+    TypeC = :(AbstractGPUVecOrMat{T3})
 
-    @eval function LinearAlgebra.kron!(C::$TypeC, A::$TypeA, B::$TypeB) where {T1,T2,T3}
+    @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)
 
-        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
-
-            @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]))
-
-                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)))
-    
+
+        kernel(C, A, B, ndrange=(size(C, 1), size(C, 2)))
+
         return C
     end
 
     @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...)
+        C = similar(A, T, size(A, 1) * size(B, 1), size(A, 2) * size(B, 2))
         return kron!(C, A, B)
     end
 end

test/testsuite/linalg.jl

@@ -313,8 +313,7 @@
 
     @testset "kron" begin
         for T in eltypes
-            @test compare(kron, AT, rand(T, 32), rand(T, 64))
-            for opa in (identity, transpose, adjoint), opb in (identity, transpose, adjoint)
+            for opa in (vec, identity, transpose, adjoint), opb in (vec, identity, transpose, adjoint)
                 @test compare(kron, AT, opa(rand(T, 32, 64)), opb(rand(T, 128, 16)))
             end
         end