Skip to content

[WIP] Allow turning off allowscalar copies for yacusolver #270

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
kshyatt wants to merge 1 commit into
base: main
Choose a base branch
from
+79 −48
Draft

[WIP] Allow turning off allowscalar copies for yacusolver #270

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 79 additions & 48 deletions ext/MatrixAlgebraKitCUDAExt/yacusolver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@ for (bname, fname, elty, relty) in
A::StridedCuMatrix{$elty},
S::StridedCuVector{$relty} = similar(A, $relty, min(size(A)...)),
U::StridedCuMatrix{$elty} = similar(A, $elty, size(A, 1), min(size(A)...)),
Vᴴ::StridedCuMatrix{$elty} = similar(A, $elty, min(size(A)...), size(A, 2))
Vᴴ::StridedCuMatrix{$elty} = similar(A, $elty, min(size(A)...), size(A, 2));
check::Bool = false,
)
chkstride1(A, U, Vᴴ, S)
m, n = size(A)
Expand Down Expand Up @@ -90,8 +91,10 @@ for (bname, fname, elty, relty) in
end
CUDA.unsafe_free!(rwork)

info = @allowscalar dh.info[1]
cuSOLVER.chkargsok(BlasInt(info))
if check
info = @allowscalar dh.info[1]
cuSOLVER.chkargsok(BlasInt(info))
end

return (S, U, Vᴴ)
end
Expand All @@ -103,7 +106,8 @@ function gesvdp!(
S::StridedCuVector = similar(A, real(T), min(size(A)...)),
U::StridedCuMatrix{T} = similar(A, T, size(A, 1), min(size(A)...)),
Vᴴ::StridedCuMatrix{T} = similar(A, T, min(size(A)...), size(A, 2));
tol = norm(A) * eps(real(T))
tol = norm(A) * eps(real(T)),
check::Bool = false,
) where {T <: BlasFloat}
chkstride1(A, U, S, Vᴴ)
m, n = size(A)
Expand Down Expand Up @@ -166,8 +170,10 @@ function gesvdp!(
err = h_err_sigma[]
err > tol && @warn "gesvdp! did not attain the requested tolerance: error = $err > tolerance = $tol"

flag = @allowscalar dh.info[1]
cuSOLVER.chklapackerror(BlasInt(flag))
if check
flag = @allowscalar dh.info[1]
cuSOLVER.chklapackerror(BlasInt(flag))
end
if Ũ !== U && length(U) > 0
U .= view(Ũ, 1:m, 1:size(U, 2))
end
Expand Down Expand Up @@ -196,6 +202,7 @@ for (bname, fname, elty, relty) in
U::StridedCuMatrix{$elty} = similar(A, $elty, size(A, 1), min(size(A)...)),
Vᴴ::StridedCuMatrix{$elty} = similar(A, $elty, min(size(A)...), size(A, 2));
tol::$relty = eps($relty),
check::Bool = false,
max_sweeps::Int = 100,
kwargs...
)
Expand Down Expand Up @@ -253,8 +260,10 @@ for (bname, fname, elty, relty) in
)
end

info = @allowscalar dh.info[1]
cuSOLVER.chkargsok(BlasInt(info))
if check
info = @allowscalar dh.info[1]
cuSOLVER.chkargsok(BlasInt(info))
end

cuSOLVER.cusolverDnDestroyGesvdjInfo(params[])

Expand All @@ -272,6 +281,7 @@ function gesvdr!(
S::StridedCuVector = similar(A, real(T), min(size(A)...)),
U::StridedCuMatrix{T} = similar(A, T, size(A, 1), min(size(A)...)),
Vᴴ::StridedCuMatrix{T} = similar(A, T, min(size(A)...), size(A, 2));
check::Bool = false,
k::Int = length(S),
p::Int = min(size(A)...) - k - 1,
niters::Int = 1
Expand Down Expand Up @@ -319,8 +329,10 @@ function gesvdr!(
)
end

flag = @allowscalar dh.info[1]
cuSOLVER.chklapackerror(BlasInt(flag))
if check
flag = @allowscalar dh.info[1]
cuSOLVER.chklapackerror(BlasInt(flag))
end
if Ũ !== U && length(U) > 0
U .= view(Ũ, 1:m, 1:size(U, 2))
end
Expand All @@ -336,7 +348,7 @@ end
# Wrapper for general eigensolver
for (celty, elty) in ((:ComplexF32, :Float32), (:ComplexF64, :Float64), (:ComplexF32, :ComplexF32), (:ComplexF64, :ComplexF64))
@eval begin
function Xgeev!(A::StridedCuMatrix{$elty}, D::StridedCuVector{$celty}, V::StridedCuMatrix{$celty})
function Xgeev!(A::StridedCuMatrix{$elty}, D::StridedCuVector{$celty}, V::StridedCuMatrix{$celty}; check::Bool = false)
require_one_based_indexing(A, V, D)
chkstride1(A, V, D)
n = checksquare(A)
Expand Down Expand Up @@ -391,8 +403,10 @@ for (celty, elty) in ((:ComplexF32, :Float32), (:ComplexF64, :Float64), (:Comple
sizeof(buffer_gpu), buffer_cpu, sizeof(buffer_cpu), dh.info
)
end
flag = @allowscalar dh.info[1]
cuSOLVER.chkargsok(BlasInt(flag))
if check
flag = @allowscalar dh.info[1]
cuSOLVER.chkargsok(BlasInt(flag))
end
if eltype(A) <: Real
work = CuVector{$elty}(undef, n)
DR = view(D2, 1:n)
Expand All @@ -414,7 +428,8 @@ end
# jobz::Char,
# uplo::Char,
# A::StridedCuMatrix{$elty},
# B::StridedCuMatrix{$elty})
# B::StridedCuMatrix{$elty};
# check::Bool = false)
# chkuplo(uplo)
# nA, nB = checksquare(A, B)
# if nB != nA
Expand All @@ -436,9 +451,11 @@ end
# return $fname(dh, itype, jobz, uplo, n, A, lda, B, ldb, W,
# buffer, sizeof(buffer) ÷ sizeof($elty), dh.info)
# end

# info = @allowscalar dh.info[1]
# chkargsok(BlasInt(info))
#
# if check
# info = @allowscalar dh.info[1]
# chkargsok(BlasInt(info))
# end

# if jobz == 'N'
# return W
Expand All @@ -460,6 +477,7 @@ end
# uplo::Char,
# A::StridedCuMatrix{$elty},
# B::StridedCuMatrix{$elty};
# check::Bool = false,
# tol::$relty=eps($relty),
# max_sweeps::Int=100)
# chkuplo(uplo)
Expand Down Expand Up @@ -488,10 +506,10 @@ end
# return $fname(dh, itype, jobz, uplo, n, A, lda, B, ldb, W,
# buffer, sizeof(buffer) ÷ sizeof($elty), dh.info, params[])
# end

# info = @allowscalar dh.info[1]
# chkargsok(BlasInt(info))

# if check
# info = @allowscalar dh.info[1]
# chkargsok(BlasInt(info))
# end
# cusolverDnDestroySyevjInfo(params[])

# if jobz == 'N'
Expand All @@ -516,6 +534,7 @@ end
# function $jname(jobz::Char,
# uplo::Char,
# A::StridedCuArray{$elty};
# check::Bool = false,
# tol::$relty=eps($relty),
# max_sweeps::Int=100)

Expand Down Expand Up @@ -548,14 +567,15 @@ end
# sizeof(buffer) ÷ sizeof($elty), dh.info, params[], batchSize)
# end

# # Copy the solver info and delete the device memory
# info = @allowscalar collect(dh.info)
# if check
# # Copy the solver info and delete the device memory
# info = collect(dh.info)

# # Double check the solver's exit status
# for i in 1:batchSize
# chkargsok(BlasInt(info[i]))
# # Double check the solver's exit status
# for i in 1:batchSize
# chkargsok(BlasInt(info[i]))
# end
# end

# cusolverDnDestroySyevjInfo(params[])

# # Return eigenvalues (in W) and possibly eigenvectors (in A)
Expand All @@ -575,7 +595,8 @@ end
# @eval begin
# function potrsBatched!(uplo::Char,
# A::Vector{<:StridedCuMatrix{$elty}},
# B::Vector{<:StridedCuVecOrMat{$elty}})
# B::Vector{<:StridedCuVecOrMat{$elty}};
# check::Bool = false,)
# if length(A) != length(B)
# throw(DimensionMismatch(""))
# end
Expand All @@ -602,10 +623,11 @@ end
# # Run the solver
# $fname(dh, uplo, n, nrhs, Aptrs, lda, Bptrs, ldb, dh.info, batchSize)

# # Copy the solver info and delete the device memory
# info = @allowscalar dh.info[1]
# chklapackerror(BlasInt(info))

# if check
# # Copy the solver info and delete the device memory
# info = @allowscalar dh.info[1]
# chklapackerror(BlasInt(info))
# end
# return B
# end
# end
Expand All @@ -616,7 +638,7 @@ end
# (:cusolverDnCpotrfBatched, :ComplexF32),
# (:cusolverDnZpotrfBatched, :ComplexF64))
# @eval begin
# function potrfBatched!(uplo::Char, A::Vector{<:StridedCuMatrix{$elty}})
# function potrfBatched!(uplo::Char, A::Vector{<:StridedCuMatrix{$elty}}; check::Bool = false)

# # Set up information for the solver arguments
# chkuplo(uplo)
Expand All @@ -632,14 +654,15 @@ end
# # Run the solver
# $fname(dh, uplo, n, Aptrs, lda, dh.info, batchSize)

# # Copy the solver info and delete the device memory
# info = @allowscalar collect(dh.info)
# if check
# # Copy the solver info and delete the device memory
# info = collect(dh.info)

# # Double check the solver's exit status
# for i in 1:batchSize
# chkargsok(BlasInt(info[i]))
# # Double check the solver's exit status
# for i in 1:batchSize
# chkargsok(BlasInt(info[i]))
# end
# end

# # info[i] > 0 means the leading minor of order info[i] is not positive definite
# # LinearAlgebra.LAPACK does not throw Exception here
# # to simplify calls to isposdef! and factorize
Expand All @@ -649,7 +672,8 @@ end
# end

# # gesv
# function gesv!(X::CuVecOrMat{T}, A::CuMatrix{T}, B::CuVecOrMat{T}; fallback::Bool=true,
# function gesv!(X::CuVecOrMat{T}, A::CuMatrix{T}, B::CuVecOrMat{T};
# fallback::Bool=true, check::Bool = false,
# residual_history::Bool=false, irs_precision::String="AUTO",
# refinement_solver::String="CLASSICAL",
# maxiters::Int=0, maxiters_inner::Int=0, tol::Float64=0.0,
Expand Down Expand Up @@ -703,10 +727,11 @@ end
# X, ldx, buffer, sizeof(buffer), niters, dh.info)
# end

# # Copy the solver flag and delete the device memory
# flag = @allowscalar dh.info[1]
# chklapackerror(BlasInt(flag))

# if check
# # Copy the solver flag and delete the device memory
# flag = @allowscalar dh.info[1]
# chklapackerror(BlasInt(flag))
# end
# return X, info
# end

Expand All @@ -721,6 +746,7 @@ for (bname, fname, elty, relty) in (
A::StridedCuMatrix{$elty},
W::StridedCuVector{$relty},
V::StridedCuMatrix{$elty};
check::Bool = false,
uplo::Char = 'U',
tol::$relty = eps($relty),
max_sweeps::Int = 100
Expand Down Expand Up @@ -752,8 +778,10 @@ for (bname, fname, elty, relty) in (
)
end

info = @allowscalar dh.info[1]
chkargsok(BlasInt(info))
if check
info = @allowscalar dh.info[1]
chkargsok(BlasInt(info))
end

if jobz == 'V' && V !== A
copy!(V, A)
Expand All @@ -767,6 +795,7 @@ function heevd!(
A::StridedCuMatrix{T},
W::StridedCuVector{Tr},
V::StridedCuMatrix{T};
check::Bool = false,
uplo::Char = 'U'
) where {T <: BlasFloat, Tr <: BlasReal}
chkuplo(uplo)
Expand Down Expand Up @@ -800,8 +829,10 @@ function heevd!(
)
end

info = @allowscalar dh.info[1]
chkargsok(BlasInt(info))
if check
info = @allowscalar dh.info[1]
chkargsok(BlasInt(info))
end

if jobz == 'V' && V !== A
copy!(V, A)
Expand Down
Loading