[leburgel]

uppertriangular! and lowertriangular! zero one column at a time, which on GPU arrays is one kernel launch per column. After a QR (_qr_copyR!) or an LQ this costs as much as a large part of the factorization for small and medium matrices. This adds methods for StridedCuMatrix and StridedROCMatrix that zero the triangle with a single broadcast.

Measured on an H100 NVL, CUDA.jl 6.4.2, median of 20 calls (ms):

	upper/lower triangle, main	PR	qr_compact!, main	PR	lq_compact!, main	PR
Float64 64 × 64	0.15	0.006	0.39	0.25	0.41	0.27
Float64 256 × 256	0.60	0.006	1.43	0.84	1.45	0.86
Float64 1024 × 1024	2.4	0.012	5.9	3.5	6.0	3.6
Float64 4096 × 4096	9.4	0.12	32.5	23.3	33.7	24.5
Float64 8192 × 64	0.15	0.008	0.67	0.53	0.64	0.35
Float64 8192 × 512	1.2	0.030	5.3	4.1	3.7	2.4
ComplexF64 64 × 64	0.18	0.006	0.64	0.45	0.66	0.46
ComplexF64 1024 × 1024	3.0	0.012	8.5	5.5	9.9	5.7
ComplexF64 4096 × 4096	11.4	0.17	56.6	45.6	59.6	48.9

Benchmark

using CUDA, MatrixAlgebraKit, LinearAlgebra, Printf, Statistics
const MAK = MatrixAlgebraKit
function timed(f; n = 20)
    for _ in 1:3; f(); end
    CUDA.synchronize()
    return median(map(_ -> CUDA.@elapsed(f()), 1:n))
end
generic_upper!(A) = invoke(MAK.uppertriangular!, Tuple{AbstractMatrix}, A)
for T in (Float64, ComplexF64), (m, n) in ((64, 64), (256, 256), (1024, 1024), (4096, 4096), (8192, 64), (8192, 512))
    A = CUDA.randn(T, m, n)
    tg, tn = timed(() -> generic_upper!(A)), timed(() -> MAK.uppertriangular!(A))
    tq, tl = timed(() -> qr_compact!(copy(A))), timed(() -> lq_compact!(copy(A)))
    @printf("%-10s %5d × %-5d triangle %8.3f / %8.3f ms, qr_compact! %8.3f ms, lq_compact! %8.3f ms\n", T, m, n, 1e3tg, 1e3tn, 1e3tq, 1e3tl)
end

[kshyatt]

Thanks for this. Some of these duplications make me wonder if we should just make a GPUArrays extension, but that can be done separately.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
...ixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl	75.80% <100.00%> (+0.80%)	⬆️
...MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl	84.21% <100.00%> (+0.87%)	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kshyatt]

Looks like lowertriangular! isn't getting tested, could we cover that too?

[leburgel]

Nothing reaches lowertriangular! on GPU arrays: the GPU LQ goes through the QR of the adjoint, so I dropped those methods and kept uppertriangular! only.