Skip to content

Simplify and speed up the GPU projection kernels #318

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

Merged
kshyatt merged 6 commits into from
Oct 30, 2026
Merged

Simplify and speed up the GPU projection kernels #318

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
39c5533
Simplify and speed up the GPU projection kernels
Oct 29, 2026
ea4951b
Make a GPUArrays extension
kshyatt Oct 29, 2026
149058e
more fixes
kshyatt Oct 29, 2026
b440731
Fix dumb errors
kshyatt Oct 29, 2026
d2b2360
Missing qualifier
kshyatt Oct 29, 2026
b9fe75b
One more qualifier
kshyatt Oct 30, 2026
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
3 changes: 3 additions & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
GenericSchur = "c145ed77-6b09-5dd9-b285-bf645a82121e"
GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"

[extensions]
Expand All @@ -26,6 +27,7 @@ MatrixAlgebraKitChainRulesCoreExt = "ChainRulesCore"
MatrixAlgebraKitEnzymeExt = "Enzyme"
MatrixAlgebraKitGenericLinearAlgebraExt = "GenericLinearAlgebra"
MatrixAlgebraKitGenericSchurExt = "GenericSchur"
MatrixAlgebraKitGPUArraysExt = "GPUArrays"
MatrixAlgebraKitMooncakeExt = "Mooncake"

[compat]
Expand All @@ -35,6 +37,7 @@ ChainRulesCore = "1"
Enzyme = "0.13.148"
GenericLinearAlgebra = "0.3.19, 0.4"
GenericSchur = "0.5.6"
GPUArrays = "11, 12"
LinearAlgebra = "1"
PrecompileTools = "1"
Mooncake = "0.5.61"
Expand Down
136 changes: 0 additions & 136 deletions ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,6 @@ end

MatrixAlgebraKit.supports_svd_full(::ROCSOLVER, f::Symbol) = f in (:qr_iteration, :jacobi, :divide_and_conquer, :bisection)

# one kernel for the whole triangle instead of one per column
function MatrixAlgebraKit.uppertriangular!(A::StridedROCMatrix)
m, n = size(A)
A .= ifelse.((1:m) .> transpose(1:n), zero(eltype(A)), A)
return A
end

# rocSOLVER's `gesvd*_batched` functions take the batch as an array
# of device pointers, so a group of equally sized matrices
# doesn't need a copy into a 3D ROCArray.
Expand Down Expand Up @@ -164,117 +157,6 @@ heev!(::ROCSOLVER, A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatri
heevx!(::ROCSOLVER, A::StridedROCMatrix, Dd::StridedROCVector, V::StridedROCMatrix; kwargs...) =
YArocSOLVER.heevx!(A, Dd, V; kwargs...)

function MatrixAlgebraKit.findtruncated_svd(values::StridedROCVector, strategy::TruncationByValue)
return MatrixAlgebraKit.findtruncated(values, strategy)
end

# COV_EXCL_START
function _project_hermitian_offdiag_kernel(Au, Al, Bu, Bl, ::Val{true})
m, n = size(Au)
j = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x
j > n && return
for i in 1:m
@inbounds begin
val = (Au[i, j] - adjoint(Al[j, i])) / 2
Bu[i, j] = val
Bl[j, i] = -adjoint(val)
end
end
return
end

function _project_hermitian_offdiag_kernel(Au, Al, Bu, Bl, ::Val{false})
m, n = size(Au)
j = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x
j > n && return
for i in 1:m
@inbounds begin
val = (Au[i, j] + adjoint(Al[j, i])) / 2
Bu[i, j] = val
Bl[j, i] = adjoint(val)
end
end
return
end

function _project_hermitian_diag_kernel(A, B, ::Val{true})
n = size(A, 1)
j = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x
j > n && return
@inbounds begin
for i in 1:(j - 1)
val = (A[i, j] - adjoint(A[j, i])) / 2
B[i, j] = val
B[j, i] = -adjoint(val)
end
B[j, j] = MatrixAlgebraKit._imimag(A[j, j])
end
return
end

function _project_hermitian_diag_kernel(A, B, ::Val{false})
n = size(A, 1)
j = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x
j > n && return
@inbounds begin
for i in 1:(j - 1)
val = (A[i, j] + adjoint(A[j, i])) / 2
B[i, j] = val
B[j, i] = adjoint(val)
end
B[j, j] = real(A[j, j])
end
return
end
# COV_EXCL_STOP

const SupportedROCMatrix{T} = Union{AnyROCMatrix{T}, SubArray{T, 2, <:AnyROCMatrix{T}}}

function MatrixAlgebraKit._project_hermitian_offdiag!(
Au::SupportedROCMatrix, Al::SupportedROCMatrix, Bu::SupportedROCMatrix, Bl::SupportedROCMatrix, ::Val{anti}
) where {anti}
thread_dim = 512
block_dim = cld(size(Au, 2), thread_dim)
@roc groupsize = thread_dim gridsize = block_dim _project_hermitian_offdiag_kernel(Au, Al, Bu, Bl, Val(anti))
return nothing
end
function MatrixAlgebraKit._project_hermitian_diag!(A::SupportedROCMatrix, B::SupportedROCMatrix, ::Val{anti}) where {anti}
thread_dim = 512
block_dim = cld(size(A, 1), thread_dim)
@roc groupsize = thread_dim gridsize = block_dim _project_hermitian_diag_kernel(A, B, Val(anti))
return nothing
end

# avoids calling the `StridedMatrix` specialization to avoid scalar indexing,
# use (allocating) fallback instead until we write a dedicated kernel
MatrixAlgebraKit.ishermitian_exact(A::StridedROCMatrix) = A == A'
MatrixAlgebraKit.ishermitian_approx(A::StridedROCMatrix; atol, rtol, kwargs...) =
norm(project_antihermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A))
MatrixAlgebraKit.isantihermitian_exact(A::StridedROCMatrix) = A == -A'
MatrixAlgebraKit.isantihermitian_approx(A::StridedROCMatrix; atol, rtol, kwargs...) =
norm(project_hermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A))

function MatrixAlgebraKit._avgdiff!(A::StridedROCMatrix, B::StridedROCMatrix)
axes(A) == axes(B) || throw(DimensionMismatch())
# COV_EXCL_START
function _avgdiff_kernel(A, B)
j = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x
j > length(A) && return
@inbounds begin
a = A[j]
b = B[j]
A[j] = (a + b) / 2
B[j] = b - a
end
return
end
# COV_EXCL_STOP
thread_dim = 512
block_dim = cld(length(A), thread_dim)
@roc groupsize = thread_dim gridsize = block_dim _avgdiff_kernel(A, B)
return A, B
end

# avoids calling the BlasMat specialization that assumes syrk! or herk! is called
# TODO: remove once syrk! or herk! is defined
function MatrixAlgebraKit._mul_herm!(C::StridedROCMatrix{T}, A::StridedROCMatrix{T}) where {T <: BlasFloat}
Expand All @@ -283,27 +165,9 @@ function MatrixAlgebraKit._mul_herm!(C::StridedROCMatrix{T}, A::StridedROCMatrix
return C
end

# TODO: intersect/union don't work on GPU
MatrixAlgebraKit._ind_intersect(A::ROCVector{Int}, B::ROCVector{Int}) =
MatrixAlgebraKit._ind_intersect(collect(A), collect(B))
MatrixAlgebraKit._ind_union(A::AbstractVector{<:Integer}, B::ROCVector{Int}) =
MatrixAlgebraKit._ind_union(A, collect(B))
MatrixAlgebraKit._ind_union(A::ROCVector{Int}, B::AbstractVector{<:Integer}) =
MatrixAlgebraKit._ind_union(collect(A), B)
MatrixAlgebraKit._ind_union(A::ROCVector{Int}, B::ROCVector{Int}) =
MatrixAlgebraKit._ind_union(collect(A), collect(B))

function _sylvester(A::AnyROCMatrix, B::AnyROCMatrix, C::AnyROCMatrix)
hX = sylvester(collect(A), collect(B), collect(C))
return ROCArray(hX)
end

function svd_rank(S::AnyROCVector; rank_atol = MatrixAlgebraKit.default_pullback_rank_atol(S))
return something(findlast(≥(rank_atol), S), 0)
end

function svd_pullback!(ΔA::AnyROCMatrix, A, USVᴴ, ΔUSVᴴ, ind::AnyROCVector; kwargs...)
return svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, collect(ind); kwargs...)
end

end
148 changes: 1 addition & 147 deletions ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ using MatrixAlgebraKit: default_qr_algorithm, default_lq_algorithm, default_svd_
import MatrixAlgebraKit: geqrf!, ungqr!, unmqr!, gesvd!, gesvdp!, gesvdr!, gesvdj!
import MatrixAlgebraKit: gesvdj_batched!
import MatrixAlgebraKit: heevj!, heevd!, geev!
import MatrixAlgebraKit: _gpu_Xgesvdr!, _sylvester, svd_rank, svd_pullback!, eigh_pullback!, eig_pullback!, svd_pushforward!
import MatrixAlgebraKit: _gpu_Xgesvdr!, _sylvester
using CUDA, CUDA.cuBLAS
using CUDA: i32
using LinearAlgebra
Expand Down Expand Up @@ -42,13 +42,6 @@ end

MatrixAlgebraKit.prefers_ungqr(::CUSOLVER) = true

# one kernel for the whole triangle instead of one per column
function MatrixAlgebraKit.uppertriangular!(A::StridedCuMatrix)
m, n = size(A)
A .= ifelse.((1:m) .> transpose(1:n), zero(eltype(A)), A)
return A
end

MatrixAlgebraKit.supports_svd_full(::CUSOLVER, f::Symbol) = f in (:qr_iteration, :jacobi, :svd_polar)

# `cusolverDnXgesvdjBatched` only accepts matrices up to 32x32
Expand Down Expand Up @@ -87,117 +80,6 @@ heevj!(::CUSOLVER, A::StridedCuMatrix, Dd::StridedCuVector, V::StridedCuMatrix;
heevd!(::CUSOLVER, A::StridedCuMatrix, Dd::StridedCuVector, V::StridedCuMatrix; kwargs...) =
YACUSOLVER.heevd!(A, Dd, V; kwargs...)

function MatrixAlgebraKit.findtruncated_svd(values::StridedCuVector, strategy::TruncationByValue)
return MatrixAlgebraKit.findtruncated(values, strategy)
end

# COV_EXCL_START
function _project_hermitian_offdiag_kernel(Au, Al, Bu, Bl, ::Val{true})
m, n = size(Au)
j = threadIdx().x + (blockIdx().x - 1i32) * blockDim().x
j > n && return
for i in 1:m
@inbounds begin
val = (Au[i, j] - adjoint(Al[j, i])) / 2
Bu[i, j] = val
Bl[j, i] = -adjoint(val)
end
end
return
end

function _project_hermitian_offdiag_kernel(Au, Al, Bu, Bl, ::Val{false})
m, n = size(Au)
j = threadIdx().x + (blockIdx().x - 1) * blockDim().x
j > n && return
for i in 1:m
@inbounds begin
val = (Au[i, j] + adjoint(Al[j, i])) / 2
Bu[i, j] = val
Bl[j, i] = adjoint(val)
end
end
return
end

function _project_hermitian_diag_kernel(A, B, ::Val{true})
n = size(A, 1)
j = threadIdx().x + (blockIdx().x - 1) * blockDim().x
j > n && return
@inbounds begin
for i in 1i32:(j - 1i32)
val = (A[i, j] - adjoint(A[j, i])) / 2
B[i, j] = val
B[j, i] = -adjoint(val)
end
B[j, j] = MatrixAlgebraKit._imimag(A[j, j])
end
return
end

function _project_hermitian_diag_kernel(A, B, ::Val{false})
n = size(A, 1)
j = threadIdx().x + (blockIdx().x - 1i32) * blockDim().x
j > n && return
@inbounds begin
for i in 1i32:(j - 1i32)
val = (A[i, j] + adjoint(A[j, i])) / 2
B[i, j] = val
B[j, i] = adjoint(val)
end
B[j, j] = real(A[j, j])
end
return
end
# COV_EXCL_STOP

const SupportedCuMatrix{T} = Union{AnyCuMatrix{T}, SubArray{T, 2, <:AnyCuMatrix{T}}}

function MatrixAlgebraKit._project_hermitian_offdiag!(
Au::SupportedCuMatrix, Al::SupportedCuMatrix, Bu::SupportedCuMatrix, Bl::SupportedCuMatrix, ::Val{anti}
) where {anti}
thread_dim = 512
block_dim = cld(size(Au, 2), thread_dim)
@cuda threads = thread_dim blocks = block_dim _project_hermitian_offdiag_kernel(Au, Al, Bu, Bl, Val(anti))
return nothing
end
function MatrixAlgebraKit._project_hermitian_diag!(A::SupportedCuMatrix, B::SupportedCuMatrix, ::Val{anti}) where {anti}
thread_dim = 512
block_dim = cld(size(A, 1), thread_dim)
@cuda threads = thread_dim blocks = block_dim _project_hermitian_diag_kernel(A, B, Val(anti))
return nothing
end

# avoids calling the `StridedMatrix` specialization to avoid scalar indexing,
# use (allocating) fallback instead until we write a dedicated kernel
MatrixAlgebraKit.ishermitian_exact(A::StridedCuMatrix) = A == A'
MatrixAlgebraKit.ishermitian_approx(A::StridedCuMatrix; atol, rtol, kwargs...) =
norm(project_antihermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A))
MatrixAlgebraKit.isantihermitian_exact(A::StridedCuMatrix) = A == -A'
MatrixAlgebraKit.isantihermitian_approx(A::StridedCuMatrix; atol, rtol, kwargs...) =
norm(project_hermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A))

function MatrixAlgebraKit._avgdiff!(A::StridedCuMatrix, B::StridedCuMatrix)
axes(A) == axes(B) || throw(DimensionMismatch())
# COV_EXCL_START
function _avgdiff_kernel(A, B)
j = threadIdx().x + (blockIdx().x - 1i32) * blockDim().x
j > length(A) && return
@inbounds begin
a = A[j]
b = B[j]
A[j] = (a + b) / 2
B[j] = b - a
end
return
end
# COV_EXCL_STOP
thread_dim = 512
block_dim = cld(length(A), thread_dim)
@cuda threads = thread_dim blocks = block_dim _avgdiff_kernel(A, B)
return A, B
end

# avoids calling the BlasMat specialization that assumes syrk! or herk! is called
# TODO: remove once syrk! or herk! is defined
function MatrixAlgebraKit._mul_herm!(C::StridedCuMatrix{T}, A::StridedCuMatrix{T}) where {T <: BlasFloat}
Expand All @@ -206,39 +88,11 @@ function MatrixAlgebraKit._mul_herm!(C::StridedCuMatrix{T}, A::StridedCuMatrix{T
return C
end

# TODO: intersect/union don't work on GPU
MatrixAlgebraKit._ind_intersect(A::CuVector{Int}, B::CuVector{Int}) =
MatrixAlgebraKit._ind_intersect(collect(A), collect(B))
MatrixAlgebraKit._ind_union(A::AbstractVector{<:Integer}, B::CuVector{Int}) =
MatrixAlgebraKit._ind_union(A, collect(B))
MatrixAlgebraKit._ind_union(A::CuVector{Int}, B::AbstractVector{<:Integer}) =
MatrixAlgebraKit._ind_union(collect(A), B)
MatrixAlgebraKit._ind_union(A::CuVector{Int}, B::CuVector{Int}) =
MatrixAlgebraKit._ind_union(collect(A), collect(B))

function _sylvester(A::AnyCuMatrix, B::AnyCuMatrix, C::AnyCuMatrix)
# https://github.com/JuliaGPU/CUDA.jl/issues/3021
# to add native sylvester to CUDA
hX = sylvester(collect(A), collect(B), collect(C))
return CuArray(hX)
end

function svd_rank(S::AnyCuVector; rank_atol = MatrixAlgebraKit.default_pullback_rank_atol(S))
return something(findlast(≥(rank_atol), S), 0)
end

function svd_pullback!(ΔA::AnyCuMatrix, A, USVᴴ, ΔUSVᴴ, ind::AnyCuVector; kwargs...)
return svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, collect(ind); kwargs...)
end

# have to override this as methods are missing in GPUArrays for the various
# views of Diagonal of ΔA
function svd_pushforward!(
ΔA::Diagonal{T, <:CuVector{T}}, A, USVᴴ, ΔUSVᴴ, ind = Colon();
rank_atol::Real = MatrixAlgebraKit.default_pullback_rank_atol(USVᴴ[2]),
degeneracy_atol::Real = MatrixAlgebraKit.default_pullback_rank_atol(USVᴴ[2])
) where {T}
return MatrixAlgebraKit.svd_pushforward!(diagm(diagview(ΔA)), A, USVᴴ, ΔUSVᴴ, ind; rank_atol, degeneracy_atol)
end

end
Loading
Loading