From 67d09be7bcb7bb9a8b59963537e19c7ff0f6030b Mon Sep 17 00:00:00 2001 From: leburgel Date: Mon, 26 Oct 2026 08:08:53 +0100 Subject: [PATCH 1/2] One kernel for uppertriangular!/lowertriangular! on CuArrays and ROCArrays --- .../MatrixAlgebraKitAMDGPUExt.jl | 12 ++++++++++++ .../MatrixAlgebraKitCUDAExt.jl | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 14d849e5e..4508e617f 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -35,6 +35,18 @@ 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 +function MatrixAlgebraKit.lowertriangular!(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. diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl index cbcc9fbe6..c6e2c9948 100644 --- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl +++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl @@ -42,6 +42,18 @@ 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 +function MatrixAlgebraKit.lowertriangular!(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 From 12c5a86d30805550377f99b401392d481667fd44 Mon Sep 17 00:00:00 2001 From: leburgel Date: Mon, 26 Oct 2026 13:58:55 +0100 Subject: [PATCH 2/2] Drop the GPU lowertriangular! methods, which no decomposition reaches --- ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl | 5 ----- ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl | 5 ----- 2 files changed, 10 deletions(-) diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 4508e617f..add91008e 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -41,11 +41,6 @@ function MatrixAlgebraKit.uppertriangular!(A::StridedROCMatrix) A .= ifelse.((1:m) .> transpose(1:n), zero(eltype(A)), A) return A end -function MatrixAlgebraKit.lowertriangular!(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 diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl index c6e2c9948..3da8a7817 100644 --- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl +++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl @@ -48,11 +48,6 @@ function MatrixAlgebraKit.uppertriangular!(A::StridedCuMatrix) A .= ifelse.((1:m) .> transpose(1:n), zero(eltype(A)), A) return A end -function MatrixAlgebraKit.lowertriangular!(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)