From 4f2b4c259ee869d5be4d6d894a66b3005b944532 Mon Sep 17 00:00:00 2001 From: Tim Gymnich Date: Tue, 30 Apr 2024 13:15:38 +0200 Subject: [PATCH] add more tests --- src/host/linalg.jl | 88 ++++++++++++++++---------------- test/testsuite/linalg.jl | 105 +++++++-------------------------------- 2 files changed, 62 insertions(+), 131 deletions(-) diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 54960b9d..9376b46c 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -374,16 +374,16 @@ end function generic_trimatmul!(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("matrix A has dimensions $(size(A)), matrix B has dimensions $(size(B))")) + 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("result C has dimensions $(size(C)), needs $((size(A,1),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' + upper = tfun === identity ? uploc == 'U' : uploc != 'U' unit = isunitc == 'U' function trimatmul(ctx, C, A, B) @@ -414,9 +414,9 @@ function generic_trimatmul!(C::AbstractGPUVecOrMat{R}, uploc, isunitc, tfun::Fun @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 += (unit ? one(Cij) : A[i,i]) * B[i,j] + Cij += (unit ? one(Cij) : transpose(A[i,i])) * B[i,j] for k in (upper ? (i + 1) : 1):(upper ? m : (i - 1)) - Cij += A[k,i] * B[k,j] + Cij += transpose(A[k,i]) * B[k,j] end C[i,j] += Cij end @@ -424,7 +424,7 @@ function generic_trimatmul!(C::AbstractGPUVecOrMat{R}, uploc, isunitc, tfun::Fun return end - function trimatmul_c(ctx, C, A, B) + function trimatmul_a(ctx, C, A, B) idx = @linearidx C assume.(size(C) .> 0) i, j = @inbounds Tuple(CartesianIndices(C)[idx])..., 1 @@ -433,22 +433,22 @@ function generic_trimatmul!(C::AbstractGPUVecOrMat{R}, uploc, isunitc, tfun::Fun @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 += (unit ? one(Cij) : conj(A[i,i])) * B[i,j] + Cij += (unit ? one(Cij) : adjoint(A[i,i])) * B[i,j] for k in (upper ? (i + 1) : 1):(upper ? m : (i - 1)) - Cij += conj(A[k,i]) * B[k,j] + Cij += adjoint(A[k,i]) * B[k,j] end C[i,j] += Cij end return end - + if tfun === identity gpu_call(trimatmul, C, A, B; name="trimatmul") elseif tfun == transpose - gpu_call(trimatmul_t, C, A, B; name="trimatmul") + gpu_call(trimatmul_t, C, A, B; name="trimatmul_t") elseif tfun === adjoint - gpu_call(trimatmul_c, C, A, B; name="trimatmul") + gpu_call(trimatmul_a, C, A, B; name="trimatmul_a") else error("Not supported") end @@ -456,42 +456,42 @@ 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("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("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 +# 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' +# upper = uploc == 'U' +# unit = isunitc == 'U' - # tA = tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C' +# # 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 +# 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 +# @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 +# return +# end - C -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 +507,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 diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl index 71c857da..9149543e 100644 --- a/test/testsuite/linalg.jl +++ b/test/testsuite/linalg.jl @@ -134,93 +134,24 @@ end if VERSION >= v"1.10-" - @testset "trimatmul" begin - n = 128 - b = AT(rand(Float32, n)) - B = AT(rand(Float32, n, n)) - - At = UpperTriangular(AT(rand(Float32, n,n))) - A = AT(At) - Ct = AT(zeros(Float32, n)) - C = zeros(Float32, n) - - mul!(Ct, At, b) - mul!(C, A, b) - @test Ct ≈ C - - Ct = AT(zeros(Float32, n, n)) - C = zeros(Float32, n, n) - - mul!(Ct, At, B) - mul!(C, A, B) - @test Ct ≈ C - - At = UnitUpperTriangular(AT(rand(Float32, n,n))) - A = AT(At) - Ct = AT(zeros(Float32, n)) - C = zeros(Float32, n) - - mul!(Ct, At, b) - mul!(C, A, b) - @test Ct ≈ C - - Ct = AT(zeros(Float32, n, n)) - C = zeros(Float32, n, n) - - mul!(Ct, At, B) - mul!(C, A, B) - @test Ct ≈ C - - - At = LowerTriangular(AT(rand(Float32, n,n))) - A = AT(At) - Ct = AT(zeros(Float32, n)) - C = zeros(Float32, n) - - mul!(Ct, At, b) - mul!(C, A, b) - @test Ct ≈ C - - Ct = AT(zeros(Float32, n, n)) - C = zeros(Float32, n, n) - - mul!(Ct, At, B) - mul!(C, A, B) - @test Ct ≈ C - - At = UnitLowerTriangular(AT(rand(Float32, n,n))) - A = AT(At) - Ct = AT(zeros(Float32, n)) - C = zeros(Float32, n) - - mul!(Ct, At, b) - mul!(C, A, b) - @test Ct ≈ C - - Ct = AT(zeros(Float32, n, n)) - C = zeros(Float32, n, n) - - mul!(Ct, At, B) - mul!(C, A, B) - @test Ct ≈ C - - At = UnitLowerTriangular(AT(rand(ComplexF32, n,n))) - A = AT(At) - b = AT(rand(ComplexF32, n)) - B = AT(rand(ComplexF32, n, n)) - Ct = AT(zeros(ComplexF32, n)) - C = zeros(ComplexF32, n) - - mul!(Ct, At, b) - mul!(C, A, b) - @test Ct ≈ C - - Ct = AT(zeros(ComplexF32, n, n)) - C = zeros(ComplexF32, n, n) - - mul!(Ct, At, B) - mul!(C, A, B) - @test Ct ≈ C + @testset "mul! + Triangular" begin + @testset "trimatmul! ($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)) + Ct = AT(zeros(T, n)) + C = zeros(T, n) + mul!(Ct, f(TR(A)), b) + mul!(C, f(TR(collect(A))), collect(b)) + @test collect(Ct) ≈ C + + B = AT(rand(T, n, n)) + Ct = AT(zeros(T, n, n)) + C = zeros(T, n, n) + mul!(Ct, f(TR(A)), B) + mul!(C, f(TR(collect(A))), collect(B)) + @test collect(Ct) ≈ C + end end end end