From 39c55334455a583830575ade23ab13c3ecb6036a Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 14:15:17 +0100 Subject: [PATCH 1/6] Simplify and speed up the GPU projection kernels --- .../MatrixAlgebraKitAMDGPUExt.jl | 91 +++++-------------- .../MatrixAlgebraKitCUDAExt.jl | 89 +++++------------- 2 files changed, 42 insertions(+), 138 deletions(-) diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index add91008e..3c2a48083 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -168,81 +168,34 @@ function MatrixAlgebraKit.findtruncated_svd(values::StridedROCVector, strategy:: 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}) +function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} 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]) + i = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x + j = workitemIdx().y + (workgroupIdx().y - 1) * workgroupDim().y + (i > n || j > n || i > j) && return + @inbounds if i == j + B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) + else + aij = A[i, j] + aji = A[j, i] + val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 + B[i, j] = val + B[j, i] = anti ? -adjoint(val) : adjoint(val) 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 +function MatrixAlgebraKit.project_hermitian_native!( + A::SupportedROCMatrix, B::SupportedROCMatrix, anti::Val; blocksize = 32 + ) + n = size(A, 1) + n == 0 && return B + thread_dim = (16, 16) + block_dim = (cld(n, 16), cld(n, 16)) + @roc groupsize = thread_dim gridsize = block_dim _project_hermitian_kernel(A, B, anti) + return B end # avoids calling the `StridedMatrix` specialization to avoid scalar indexing, @@ -256,7 +209,6 @@ MatrixAlgebraKit.isantihermitian_approx(A::StridedROCMatrix; atol, rtol, kwargs. 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 @@ -268,7 +220,6 @@ function MatrixAlgebraKit._avgdiff!(A::StridedROCMatrix, B::StridedROCMatrix) 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) diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl index 3da8a7817..94dcbccec 100644 --- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl +++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl @@ -91,81 +91,34 @@ function MatrixAlgebraKit.findtruncated_svd(values::StridedCuVector, strategy::T 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}) +function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} 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]) + i = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x + j = workitemIdx().y + (workgroupIdx().y - 1) * workgroupDim().y + (i > n || j > n || i > j) && return + @inbounds if i == j + B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) + else + aij = A[i, j] + aji = A[j, i] + val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 + B[i, j] = val + B[j, i] = anti ? -adjoint(val) : adjoint(val) 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 +function MatrixAlgebraKit.project_hermitian_native!( + A::SupportedCuMatrix, B::SupportedCuMatrix, anti::Val; blocksize = 32 + ) + n = size(A, 1) + n == 0 && return B + thread_dim = (16, 16) + block_dim = (cld(n, 16), cld(n, 16)) + @roc groupsize = thread_dim gridsize = block_dim _project_hermitian_kernel(A, B, anti) + return B end # avoids calling the `StridedMatrix` specialization to avoid scalar indexing, From ea4951b0cef652ac2cb808cef463490a77c78053 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 15:42:19 +0100 Subject: [PATCH 2/6] Make a GPUArrays extension --- Project.toml | 3 + .../MatrixAlgebraKitAMDGPUExt.jl | 87 --------------- .../MatrixAlgebraKitCUDAExt.jl | 101 +---------------- ext/MatrixAlgebraKitGPUArraysExt.jl | 102 ++++++++++++++++++ 4 files changed, 106 insertions(+), 187 deletions(-) create mode 100644 ext/MatrixAlgebraKitGPUArraysExt.jl diff --git a/Project.toml b/Project.toml index 553a21f04..39ee0620c 100644 --- a/Project.toml +++ b/Project.toml @@ -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] @@ -26,6 +27,7 @@ MatrixAlgebraKitChainRulesCoreExt = "ChainRulesCore" MatrixAlgebraKitEnzymeExt = "Enzyme" MatrixAlgebraKitGenericLinearAlgebraExt = "GenericLinearAlgebra" MatrixAlgebraKitGenericSchurExt = "GenericSchur" +MatrixAlgebraKitGPUArraysExt = "GPUArrays" MatrixAlgebraKitMooncakeExt = "Mooncake" [compat] @@ -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" diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 3c2a48083..e20b81da5 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -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. @@ -164,68 +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 - -function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} - n = size(A, 1) - i = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x - j = workitemIdx().y + (workgroupIdx().y - 1) * workgroupDim().y - (i > n || j > n || i > j) && return - @inbounds if i == j - B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) - else - aij = A[i, j] - aji = A[j, i] - val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 - B[i, j] = val - B[j, i] = anti ? -adjoint(val) : adjoint(val) - end - return -end - -const SupportedROCMatrix{T} = Union{AnyROCMatrix{T}, SubArray{T, 2, <:AnyROCMatrix{T}}} - -function MatrixAlgebraKit.project_hermitian_native!( - A::SupportedROCMatrix, B::SupportedROCMatrix, anti::Val; blocksize = 32 - ) - n = size(A, 1) - n == 0 && return B - thread_dim = (16, 16) - block_dim = (cld(n, 16), cld(n, 16)) - @roc groupsize = thread_dim gridsize = block_dim _project_hermitian_kernel(A, B, anti) - return B -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()) - 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 - 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} @@ -234,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 diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl index 94dcbccec..0c04bb235 100644 --- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl +++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl @@ -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 @@ -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 @@ -87,70 +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 - -function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} - n = size(A, 1) - i = workitemIdx().x + (workgroupIdx().x - 1) * workgroupDim().x - j = workitemIdx().y + (workgroupIdx().y - 1) * workgroupDim().y - (i > n || j > n || i > j) && return - @inbounds if i == j - B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) - else - aij = A[i, j] - aji = A[j, i] - val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 - B[i, j] = val - B[j, i] = anti ? -adjoint(val) : adjoint(val) - end - return -end - -const SupportedCuMatrix{T} = Union{AnyCuMatrix{T}, SubArray{T, 2, <:AnyCuMatrix{T}}} - -function MatrixAlgebraKit.project_hermitian_native!( - A::SupportedCuMatrix, B::SupportedCuMatrix, anti::Val; blocksize = 32 - ) - n = size(A, 1) - n == 0 && return B - thread_dim = (16, 16) - block_dim = (cld(n, 16), cld(n, 16)) - @roc groupsize = thread_dim gridsize = block_dim _project_hermitian_kernel(A, B, anti) - return B -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} @@ -159,16 +88,6 @@ 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 @@ -176,22 +95,4 @@ function _sylvester(A::AnyCuMatrix, B::AnyCuMatrix, C::AnyCuMatrix) 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 diff --git a/ext/MatrixAlgebraKitGPUArraysExt.jl b/ext/MatrixAlgebraKitGPUArraysExt.jl new file mode 100644 index 000000000..52486fadb --- /dev/null +++ b/ext/MatrixAlgebraKitGPUArraysExt.jl @@ -0,0 +1,102 @@ +module MatrixAlgebraKitGPUArraysExt + +using MatrixAlgebraKit, GPUArrays + +# one kernel for the whole triangle instead of one per column +function MatrixAlgebraKit.uppertriangular!(A::AnyGPUMatrix) + m, n = size(A) + A .= ifelse.((1:m) .> transpose(1:n), zero(eltype(A)), A) + return A +end + +@kernel function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} + n = size(A, 1) + i, j = @index(Global) + (i > n || j > n || i > j) && return + @inbounds if i == j + B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) + else + aij = A[i, j] + aji = A[j, i] + val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 + B[i, j] = val + B[j, i] = anti ? -adjoint(val) : adjoint(val) + end + return +end + +const SupportedGPUMatrix{T} = Union{AnyGPUMatrix{T}, SubArray{T, 2, <:AnyGPUMatrix{T}}} + +function MatrixAlgebraKit.project_hermitian_native!( + A::SupportedGPUMatrix, B::SupportedGPUMatrix, anti::Val; blocksize = 32 + ) + n = size(A, 1) + n == 0 && return B + thread_dim = (16, 16) + block_dim = (cld(n, 16), cld(n, 16)) + _project_hermitian_kernel!(backend, thread_dim, block_dim)(A, B, anti) + return B +end + +function MatrixAlgebraKit._avgdiff!(A::StridedGPUMatrix, B::StrideGPUMatrix) + axes(A) == axes(B) || throw(DimensionMismatch()) + @kernel function _avgdiff_kernel(A, B) + j = @index(Global) + j > length(A) && return + @inbounds begin + a = A[j] + b = B[j] + A[j] = (a + b) / 2 + B[j] = b - a + end + return + end + thread_dim = 512 + block_dim = cld(length(A), thread_dim) + backend = backend(A) + _avgdiff_kernel(backend, thread_dim, length(A))(A, B) + return A, B +end + +function MatrixAlgebraKit.findtruncated_svd(values::StridedGPUVector, strategy::TruncationByValue) + return MatrixAlgebraKit.findtruncated(values, strategy) +end + +# avoids calling the `StridedMatrix` specialization to avoid scalar indexing, +# use (allocating) fallback instead until we write a dedicated kernel +MatrixAlgebraKit.ishermitian_exact(A::StridedGPUMatrix) = A == A' +MatrixAlgebraKit.ishermitian_approx(A::StridedGPUMatrix; atol, rtol, kwargs...) = + norm(project_antihermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A)) +MatrixAlgebraKit.isantihermitian_exact(A::StridedGPUMatrix) = A == -A' +MatrixAlgebraKit.isantihermitian_approx(A::StridedGPUMatrix; atol, rtol, kwargs...) = + norm(project_hermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A)) + +# TODO: intersect/union don't work on GPU +MatrixAlgebraKit._ind_intersect(A::StridedGPUVector{Int}, B::StridedGPUVector{Int}) = + MatrixAlgebraKit._ind_intersect(collect(A), collect(B)) +MatrixAlgebraKit._ind_union(A::AbstractVector{<:Integer}, B::StridedGPUVector{Int}) = + MatrixAlgebraKit._ind_union(A, collect(B)) +MatrixAlgebraKit._ind_union(A::StridedGPUVector{Int}, B::AbstractVector{<:Integer}) = + MatrixAlgebraKit._ind_union(collect(A), B) +MatrixAlgebraKit._ind_union(A::StridedGPUVector{Int}, B::StridedGPUVector{Int}) = + MatrixAlgebraKit._ind_union(collect(A), collect(B)) + +function MatrixAlgebraKit.svd_rank(S::AnyGPUVector; rank_atol = MatrixAlgebraKit.default_pullback_rank_atol(S)) + return something(findlast(≥(rank_atol), S), 0) +end + +function MatrixAlgebraKit.svd_pullback!(ΔA::AnyGPUMatrix, A, USVᴴ, ΔUSVᴴ, ind::AnyGPUVector; 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 MatrixAlgebraKit.svd_pushforward!( + ΔA::Diagonal{T, <:AnyGPUVector{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 From 149058efb914e8f47d895220fc4631fb170f5f11 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 12:13:07 -0400 Subject: [PATCH 3/6] more fixes --- ext/MatrixAlgebraKitGPUArraysExt.jl | 63 +++++++++++++++-------------- 1 file changed, 32 insertions(+), 31 deletions(-) diff --git a/ext/MatrixAlgebraKitGPUArraysExt.jl b/ext/MatrixAlgebraKitGPUArraysExt.jl index 52486fadb..e2dc1944f 100644 --- a/ext/MatrixAlgebraKitGPUArraysExt.jl +++ b/ext/MatrixAlgebraKitGPUArraysExt.jl @@ -1,6 +1,7 @@ module MatrixAlgebraKitGPUArraysExt -using MatrixAlgebraKit, GPUArrays +using MatrixAlgebraKit, GPUArrays, LinearAlgebra +using GPUArrays.KernelAbstractions: @kernel, @index, Global, get_backend # one kernel for the whole triangle instead of one per column function MatrixAlgebraKit.uppertriangular!(A::AnyGPUMatrix) @@ -11,18 +12,18 @@ end @kernel function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} n = size(A, 1) - i, j = @index(Global) - (i > n || j > n || i > j) && return - @inbounds if i == j - B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) - else - aij = A[i, j] - aji = A[j, i] - val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 - B[i, j] = val - B[j, i] = anti ? -adjoint(val) : adjoint(val) + i, j = @index(Global, NTuple) + if !(i > n || j > n || i > j) + @inbounds if i == j + B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) + else + aij = A[i, j] + aji = A[j, i] + val = anti ? (aij - adjoint(aji)) / 2 : (aij + adjoint(aji)) / 2 + B[i, j] = val + B[j, i] = anti ? -adjoint(val) : adjoint(val) + end end - return end const SupportedGPUMatrix{T} = Union{AnyGPUMatrix{T}, SubArray{T, 2, <:AnyGPUMatrix{T}}} @@ -34,51 +35,51 @@ function MatrixAlgebraKit.project_hermitian_native!( n == 0 && return B thread_dim = (16, 16) block_dim = (cld(n, 16), cld(n, 16)) - _project_hermitian_kernel!(backend, thread_dim, block_dim)(A, B, anti) + _project_hermitian_kernel(get_backend(A), thread_dim, block_dim)(A, B, anti) return B end -function MatrixAlgebraKit._avgdiff!(A::StridedGPUMatrix, B::StrideGPUMatrix) +function MatrixAlgebraKit._avgdiff!(A::AnyGPUMatrix, B::AnyGPUMatrix) axes(A) == axes(B) || throw(DimensionMismatch()) @kernel function _avgdiff_kernel(A, B) j = @index(Global) - j > length(A) && return - @inbounds begin - a = A[j] - b = B[j] - A[j] = (a + b) / 2 - B[j] = b - a + if j ≤ length(A) + @inbounds begin + a = A[j] + b = B[j] + A[j] = (a + b) / 2 + B[j] = b - a + end end - return end thread_dim = 512 block_dim = cld(length(A), thread_dim) - backend = backend(A) + backend = get_backend(A) _avgdiff_kernel(backend, thread_dim, length(A))(A, B) return A, B end -function MatrixAlgebraKit.findtruncated_svd(values::StridedGPUVector, strategy::TruncationByValue) +function MatrixAlgebraKit.findtruncated_svd(values::AnyGPUVector, strategy::MatrixAlgebraKit.TruncationByValue) return MatrixAlgebraKit.findtruncated(values, strategy) end # avoids calling the `StridedMatrix` specialization to avoid scalar indexing, # use (allocating) fallback instead until we write a dedicated kernel -MatrixAlgebraKit.ishermitian_exact(A::StridedGPUMatrix) = A == A' -MatrixAlgebraKit.ishermitian_approx(A::StridedGPUMatrix; atol, rtol, kwargs...) = +MatrixAlgebraKit.ishermitian_exact(A::AnyGPUMatrix) = A == A' +MatrixAlgebraKit.ishermitian_approx(A::AnyGPUMatrix; atol, rtol, kwargs...) = norm(project_antihermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A)) -MatrixAlgebraKit.isantihermitian_exact(A::StridedGPUMatrix) = A == -A' -MatrixAlgebraKit.isantihermitian_approx(A::StridedGPUMatrix; atol, rtol, kwargs...) = +MatrixAlgebraKit.isantihermitian_exact(A::AnyGPUMatrix) = A == -A' +MatrixAlgebraKit.isantihermitian_approx(A::AnyGPUMatrix; atol, rtol, kwargs...) = norm(project_hermitian(A; kwargs...)) ≤ max(atol, rtol * norm(A)) # TODO: intersect/union don't work on GPU -MatrixAlgebraKit._ind_intersect(A::StridedGPUVector{Int}, B::StridedGPUVector{Int}) = +MatrixAlgebraKit._ind_intersect(A::AnyGPUVector{Int}, B::AnyGPUVector{Int}) = MatrixAlgebraKit._ind_intersect(collect(A), collect(B)) -MatrixAlgebraKit._ind_union(A::AbstractVector{<:Integer}, B::StridedGPUVector{Int}) = +MatrixAlgebraKit._ind_union(A::AbstractVector{<:Integer}, B::AnyGPUVector{Int}) = MatrixAlgebraKit._ind_union(A, collect(B)) -MatrixAlgebraKit._ind_union(A::StridedGPUVector{Int}, B::AbstractVector{<:Integer}) = +MatrixAlgebraKit._ind_union(A::AnyGPUVector{Int}, B::AbstractVector{<:Integer}) = MatrixAlgebraKit._ind_union(collect(A), B) -MatrixAlgebraKit._ind_union(A::StridedGPUVector{Int}, B::StridedGPUVector{Int}) = +MatrixAlgebraKit._ind_union(A::AnyGPUVector{Int}, B::AnyGPUVector{Int}) = MatrixAlgebraKit._ind_union(collect(A), collect(B)) function MatrixAlgebraKit.svd_rank(S::AnyGPUVector; rank_atol = MatrixAlgebraKit.default_pullback_rank_atol(S)) From b440731d151124e39231c9ba93c8e0959293e00d Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 13:16:41 -0400 Subject: [PATCH 4/6] Fix dumb errors --- ext/MatrixAlgebraKitGPUArraysExt.jl | 36 +++++++++++++---------------- 1 file changed, 16 insertions(+), 20 deletions(-) diff --git a/ext/MatrixAlgebraKitGPUArraysExt.jl b/ext/MatrixAlgebraKitGPUArraysExt.jl index e2dc1944f..ad7092026 100644 --- a/ext/MatrixAlgebraKitGPUArraysExt.jl +++ b/ext/MatrixAlgebraKitGPUArraysExt.jl @@ -1,7 +1,7 @@ module MatrixAlgebraKitGPUArraysExt using MatrixAlgebraKit, GPUArrays, LinearAlgebra -using GPUArrays.KernelAbstractions: @kernel, @index, Global, get_backend +using GPUArrays.KernelAbstractions: @kernel, @index, get_backend # one kernel for the whole triangle instead of one per column function MatrixAlgebraKit.uppertriangular!(A::AnyGPUMatrix) @@ -12,7 +12,7 @@ end @kernel function _project_hermitian_kernel(A, B, ::Val{anti}) where {anti} n = size(A, 1) - i, j = @index(Global, NTuple) + i, j = @index(Global, NTuple) if !(i > n || j > n || i > j) @inbounds if i == j B[j, j] = anti ? MatrixAlgebraKit._imimag(A[j, j]) : real(A[j, j]) @@ -33,29 +33,25 @@ function MatrixAlgebraKit.project_hermitian_native!( ) n = size(A, 1) n == 0 && return B - thread_dim = (16, 16) - block_dim = (cld(n, 16), cld(n, 16)) - _project_hermitian_kernel(get_backend(A), thread_dim, block_dim)(A, B, anti) + _project_hermitian_kernel(get_backend(A), (16, 16), (n, n))(A, B, anti) return B end -function MatrixAlgebraKit._avgdiff!(A::AnyGPUMatrix, B::AnyGPUMatrix) - axes(A) == axes(B) || throw(DimensionMismatch()) - @kernel function _avgdiff_kernel(A, B) - j = @index(Global) - if j ≤ length(A) - @inbounds begin - a = A[j] - b = B[j] - A[j] = (a + b) / 2 - B[j] = b - a - end +@kernel function _avgdiff_kernel(A, B) + j = @index(Global) + if j ≤ length(A) + @inbounds begin + a = A[j] + b = B[j] + A[j] = (a + b) / 2 + B[j] = b - a end end - thread_dim = 512 - block_dim = cld(length(A), thread_dim) - backend = get_backend(A) - _avgdiff_kernel(backend, thread_dim, length(A))(A, B) +end + +function MatrixAlgebraKit._avgdiff!(A::AnyGPUMatrix, B::AnyGPUMatrix) + axes(A) == axes(B) || throw(DimensionMismatch()) + _avgdiff_kernel(get_backend(A), 512, length(A))(A, B) return A, B end From d2b2360b89f8fc989ca6defef2f3df169a4b335b Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 15:33:35 -0400 Subject: [PATCH 5/6] Missing qualifier --- ext/MatrixAlgebraKitGPUArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/MatrixAlgebraKitGPUArraysExt.jl b/ext/MatrixAlgebraKitGPUArraysExt.jl index ad7092026..d7fcf9289 100644 --- a/ext/MatrixAlgebraKitGPUArraysExt.jl +++ b/ext/MatrixAlgebraKitGPUArraysExt.jl @@ -83,7 +83,7 @@ function MatrixAlgebraKit.svd_rank(S::AnyGPUVector; rank_atol = MatrixAlgebraKit end function MatrixAlgebraKit.svd_pullback!(ΔA::AnyGPUMatrix, A, USVᴴ, ΔUSVᴴ, ind::AnyGPUVector; kwargs...) - return svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, collect(ind); kwargs...) + return MatrixAlgebraKit.svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, collect(ind); kwargs...) end # have to override this as methods are missing in GPUArrays for the various From b9fe75ba05add98699181b1b9a7acd06e9392e35 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Oct 2026 02:48:35 -0400 Subject: [PATCH 6/6] One more qualifier --- ext/MatrixAlgebraKitGPUArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/MatrixAlgebraKitGPUArraysExt.jl b/ext/MatrixAlgebraKitGPUArraysExt.jl index d7fcf9289..8058ec3be 100644 --- a/ext/MatrixAlgebraKitGPUArraysExt.jl +++ b/ext/MatrixAlgebraKitGPUArraysExt.jl @@ -93,7 +93,7 @@ function MatrixAlgebraKit.svd_pushforward!( 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) + return MatrixAlgebraKit.svd_pushforward!(diagm(MatrixAlgebraKit.diagview(ΔA)), A, USVᴴ, ΔUSVᴴ, ind; rank_atol, degeneracy_atol) end end