diff --git a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl index 14d849e5e..add91008e 100644 --- a/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl +++ b/ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl @@ -35,6 +35,13 @@ 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. diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl index cbcc9fbe6..3da8a7817 100644 --- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl +++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl @@ -42,6 +42,13 @@ 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