add mattrimul

src/host/linalg.jl

@@ -456,42 +456,89 @@ function generic_trimatmul!(C::AbstractGPUVecOrMat{R}, uploc, isunitc, tfun::Fun
     C
 end
 
-# function generic_mattrimul!(C::AbstractGPUVecOrMat{R}, uploc, isunitc, tfun::Function, A::AbstractGPUMatrix{T}, B::AbstractGPUVecOrMat{S}) where {T,S,R}
-#     if size(A,2) != size(B,1)
-#         throw(DimensionMismatch(lazy"matrix A has dimensions $(size(A)), matrix B has dimensions $(size(B))"))
-#     end
-#     if size(C,1) != size(A,1) || size(C,2) != size(B,2)
-#         throw(DimensionMismatch(lazy"result C has dimensions $(size(C)), needs $((size(A,1),size(B,2)))"))
-#     end
-#     if isempty(A) || isempty(B)
-#         return fill!(C, zero(R))
-#     end
-
-#     upper = uploc == 'U'
-#     unit  = isunitc == 'U'
-
-#     # tA = tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C'
-
-#     gpu_call(C, A, B; name="mattrimul") do ctx, C, A, B
-#         idx = @linearidx C
-#         assume.(size(C) .> 0)
-#         i, j = @inbounds Tuple(CartesianIndices(C)[idx])..., 1
-
-#         @inbounds if i <= size(A,1) && j <= size(B,2)
-#             z2 = zero(A[i, 1]*B[1, j] + A[i, 1]*B[1, j])
-#             Cij = convert(promote_type(R, typeof(z2)), z2)
-#             Cij += A[i, j] * (unit ? one(Cij) : B[j, j])
-#             for k in (upper ? 1 : (j + 1)):(upper ? (j - 1) : (size(A,2)))
-#                 Cij += A[i, k]*B[k, j]
-#             end
-#             C[i,j] += Cij
-#         end
-
-#         return
-#     end
-
-#     C
-# end
+function generic_mattrimul!(C::AbstractGPUVecOrMat{R}, uploc, isunitc, tfun::Function, A::AbstractGPUMatrix{T}, B::AbstractGPUVecOrMat{S}) where {T,S,R}
+    if size(A,2) != size(B,1)
+        throw(DimensionMismatch(lazy"matrix A has dimensions $(size(A)), matrix B has dimensions $(size(B))"))
+    end
+    if size(C,1) != size(A,1) || size(C,2) != size(B,2)
+        throw(DimensionMismatch(lazy"result C has dimensions $(size(C)), needs $((size(A,1),size(B,2)))"))
+    end
+    if isempty(A) || isempty(B)
+        return fill!(C, zero(R))
+    end
+
+    upper = tfun === identity ? uploc == 'U' :  uploc != 'U'
+    unit  = isunitc == 'U'
+
+    function mattrimul(ctx, C, A, B)
+        idx = @linearidx C
+        assume.(size(C) .> 0)
+        i, j = @inbounds Tuple(CartesianIndices(C)[idx])..., 1
+        l, m, n = size(A, 1), size(B, 1), size(B, 2)
+
+        @inbounds if i <= l && j <= n
+            z2 = zero(A[i,1] * B[1,j] + A[i,1] * B[1,j])
+            Cij = convert(promote_type(R, typeof(z2)), z2)
+            Cij += A[i,j] * (unit ? one(Cij) : B[j,j])
+            for k in (upper ? 1 : (j + 1)):(upper ? (j - 1) : m)
+                Cij += A[i,k] * B[k,j]
+            end
+            C[i,j] += Cij
+        end
+
+        return
+    end
+
+    function mattrimul_t(ctx, C, A, B)
+        idx = @linearidx C
+        assume.(size(C) .> 0)
+        i, j = @inbounds Tuple(CartesianIndices(C)[idx])..., 1
+        l, m, n = size(A, 1), size(B, 1), size(B, 2)
+
+        @inbounds if i <= l && j <= n
+            z2 = zero(A[i,1] * B[1,j] + A[i,1] * B[1,j])
+            Cij = convert(promote_type(R, typeof(z2)), z2)
+            Cij += A[i,j] * (unit ? one(Cij) : transpose(B[j,j]))
+            for k in (upper ? 1 : (j + 1) ):(upper ? (j - 1) : m)
+                Cij += A[i,k] * transpose(B[j,k])
+            end
+            C[i,j] += Cij
+        end
+
+        return
+    end
+
+    function mattrimul_a(ctx, C, A, B)
+        idx = @linearidx C
+        assume.(size(C) .> 0)
+        i, j = @inbounds Tuple(CartesianIndices(C)[idx])..., 1
+        l, m, n = size(A, 1), size(B, 1), size(B, 2)
+
+        @inbounds if i <= l && j <= n
+            z2 = zero(A[i,1] * B[1,j] + A[i,1] * B[1,j])
+            Cij = convert(promote_type(R, typeof(z2)), z2)
+            Cij += A[i,j] * (unit ? one(Cij) : adjoint(B[j,j]))
+            for k in (upper ? 1 : (j + 1)):(upper ? (j - 1) : m)
+                Cij += A[i,k] * adjoint(B[j,k])
+            end
+            C[i,j] += Cij
+        end
+
+        return
+    end
+
+    if tfun === identity
+        gpu_call(mattrimul, C, A, B; name="mattrimul")
+    elseif tfun == transpose
+        gpu_call(mattrimul_t, C, A, B; name="mattrimul_t")
+    elseif tfun === adjoint
+        gpu_call(mattrimul_a, C, A, B; name="mattrimul_a")
+    else
+        error("Not supported")
+    end
+
+    C
+end
 
 function LinearAlgebra.generic_matvecmul!(C::AbstractGPUVector, tA::AbstractChar, A::AbstractGPUMatrix, B::AbstractGPUVector, _add::MulAddMul = MulAddMul())
     generic_matmatmul!(C, wrap(A, tA), B, _add.alpha, _add.beta)
@@ -507,9 +554,9 @@ function LinearAlgebra.generic_trimatmul!(C::AbstractGPUVecOrMat, uploc, isunitc
     generic_trimatmul!(C, uploc, isunitc, tfun, A, B)
 end
 
-# function LinearAlgebra.generic_mattrimul!(C::AbstractGPUMatrix, uploc, isunitc, tfun::Function, A::AbstractGPUMatrix, B::AbstractGPUMatrix)
-#     generic_mattrimul!(C, uploc, isunitc, tfun, A, B)
-# end
+function LinearAlgebra.generic_mattrimul!(C::AbstractGPUMatrix, uploc, isunitc, tfun::Function, A::AbstractGPUMatrix, B::AbstractGPUMatrix)
+    generic_mattrimul!(C, uploc, isunitc, tfun, A, B)
+end
 
 end
 

test/testsuite/linalg.jl

@@ -152,6 +152,17 @@
                 mul!(C, f(TR(collect(A))), collect(B))
                 @test collect(Ct) ≈ C
             end
+
+            @testset "mattrimul ($TR x $T, $f)" for T in (Float32, ComplexF32), TR in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular), f in (identity, transpose, adjoint)
+                n = 128
+                A = AT(rand(T, n,n))
+                B = AT(rand(T, n, n))
+                Ct = AT(zeros(T, n, n))
+                C = zeros(T, n, n)
+                mul!(Ct, A, f(TR(B)))
+                mul!(C, collect(A), f(TR(collect(B))))
+                @test collect(Ct) ≈ C
+            end
         end
         end
     end