Skip to content

Commit

Permalink
add more tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@tgymnich
tgymnich committed May 11, 2024
1 parent a796f89 commit 4f2b4c2
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 131 deletions.
88 changes: 44 additions & 44 deletions src/host/linalg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -414,17 +414,17 @@ 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

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
Expand All @@ -433,65 +433,65 @@ 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

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)
Expand All @@ -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

Expand Down
105 changes: 18 additions & 87 deletions test/testsuite/linalg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4f2b4c2

Please sign in to comment.