Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@tgymnich
tgymnich committed Apr 30, 2024
1 parent 0335736 commit 74782e1
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 44 deletions.
90 changes: 47 additions & 43 deletions src/host/linalg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -374,18 +374,22 @@ 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

# println("generic_trimatmul!: $(tfun) x $(uploc) x $(isunitc)")

upper = uploc == 'U'
unit = isunitc == 'U'

upper = tfun === identity ? upper : !upper

function trimatmul(ctx, C, A, B)
idx = @linearidx C
assume.(size(C) .> 0)
Expand Down Expand Up @@ -414,17 +418,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 +437,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 +511,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
2 changes: 1 addition & 1 deletion test/testsuite/linalg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@

if VERSION >= v"1.10-"
@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, conj)
@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))
Expand Down

0 comments on commit 74782e1

Please sign in to comment.