Skip to content

Commit 069ebbb

Browse files
lkdvosclaude
lkdvos
and
claude
committed
test: cover squareroot, logarithm and power on GPU
The domain helpers were rewritten as reductions and broadcasts so they would run on GPU arrays, but nothing exercised that. Add `test_{squareroot,logarithm,power}_gpu`, modelled on the existing `test_exponential_gpu`, comparing device results against the CPU reference for the `DiagonalAlgorithm` fast path and the `MatrixFunctionViaEigh` path. Both domain outcomes are covered, since the clamp and the throw are precisely what the scalar loops were replaced by: roundoff-scale negative eigenvalues must clamp, genuinely negative ones must throw a `DomainError`, and for `logarithm` and negative fractional `power` a zero eigenvalue must throw as well. Scalar indexing anywhere along these paths errors under the GPUArrays guard, so a regression to the previous loops would fail these tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent c7c69a7 commit 069ebbb

3 files changed

Lines changed: 142 additions & 0 deletions

File tree

‎test/logarithm.jl‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ using TestExtras
44
using StableRNGs
55
using MatrixAlgebraKit: diagview
66
using LinearAlgebra
7+
using CUDA, AMDGPU
78
using LinearAlgebra: exp
89

910
BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
@@ -100,3 +101,45 @@ end
100101
end
101102
@test_throws DomainError logarithm(Diagonal(zeros(T, m)))
102103
end
104+
105+
# GPU tests
106+
# ---------
107+
# As for `squareroot`, the `DiagonalAlgorithm` kernel and the shared domain helpers are
108+
# backend-generic, and `MatrixFunctionViaEigh` goes through the device `eigh_full!`. Both
109+
# domain rejections (negative real axis, and the zero eigenvalue where no logarithm exists)
110+
# are exercised, since `_check_nonzero_eigenvalues` is a reduction too. If any step fell back
111+
# to scalar indexing these would error under GPUArrays' scalar-indexing guard.
112+
function test_logarithm_gpu(ArrayT, T)
113+
rng = StableRNG(123)
114+
m = 54
115+
116+
# Diagonal fast path
117+
data = T <: Real ? (abs.(randn(rng, T, m)) .+ one(T)) : (randn(rng, T, m) .+ 4 * one(T))
118+
@test Array(diagview(logarithm(Diagonal(ArrayT(data))))) ≈
119+
diagview(logarithm(Diagonal(data)))
120+
121+
# hermitian positive definite, via the eigenvalue decomposition
122+
X = randn(rng, T, m, m)
123+
A = Matrix(LinearAlgebra.Hermitian(X * X' + one(real(T)) * LinearAlgebra.I))
124+
A_gpu = ArrayT(A)
125+
alg_gpu = MatrixFunctionViaEigh(MatrixAlgebraKit.select_algorithm(eigh_full, A_gpu))
126+
alg_cpu = MatrixFunctionViaEigh(MatrixAlgebraKit.select_algorithm(eigh_full, A))
127+
@test Array(logarithm(A_gpu, alg_gpu)) ≈ logarithm(A, alg_cpu)
128+
129+
# domain handling: negative real axis, and no logarithm at a zero eigenvalue
130+
T <: Real && @test_throws DomainError logarithm(Diagonal(ArrayT(-data)))
131+
@test_throws DomainError logarithm(Diagonal(ArrayT(zeros(T, m))))
132+
return nothing
133+
end
134+
135+
if CUDA.functional()
136+
@testset "logarithm on CUDA for T = $T" for T in BLASFloats
137+
test_logarithm_gpu(CuArray, T)
138+
end
139+
end
140+
141+
if AMDGPU.functional()
142+
@testset "logarithm on AMDGPU for T = $T" for T in BLASFloats
143+
test_logarithm_gpu(ROCArray, T)
144+
end
145+
end

‎test/power.jl‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ using TestExtras
44
using StableRNGs
55
using MatrixAlgebraKit: diagview
66
using LinearAlgebra
7+
using CUDA, AMDGPU
78

89
BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
910
GenericFloats = (Float16, ComplexF16, BigFloat, Complex{BigFloat})
@@ -124,3 +125,53 @@ end
124125
@test_throws SingularException power(Diagonal(zeros(T, m)), -1)
125126
@test_throws DomainError power(Diagonal(zeros(T, m)), -0.5)
126127
end
128+
129+
# GPU tests
130+
# ---------
131+
# As for `squareroot`, the `DiagonalAlgorithm` kernel and the shared domain helpers are
132+
# backend-generic, and `MatrixFunctionViaEigh` goes through the device `eigh_full!`. Both the
133+
# integer branch (which only checks for singularity) and the fractional branch (which clamps
134+
# and rejects out-of-domain eigenvalues) are exercised. If any step fell back to scalar
135+
# indexing these would error under GPUArrays' scalar-indexing guard.
136+
function test_power_gpu(ArrayT, T)
137+
rng = StableRNG(123)
138+
m = 54
139+
140+
# Diagonal fast path, integer and fractional exponents
141+
data = T <: Real ? (abs.(randn(rng, T, m)) .+ one(T)) : (randn(rng, T, m) .+ 4 * one(T))
142+
for p in (2, -1, 0.5, -0.25)
143+
@test Array(diagview(power(Diagonal(ArrayT(data)), p))) ≈
144+
diagview(power(Diagonal(data), p))
145+
end
146+
147+
# hermitian positive definite, via the eigenvalue decomposition
148+
X = randn(rng, T, m, m)
149+
A = Matrix(LinearAlgebra.Hermitian(X * X' + one(real(T)) * LinearAlgebra.I))
150+
A_gpu = ArrayT(A)
151+
alg_gpu = MatrixFunctionViaEigh(MatrixAlgebraKit.select_algorithm(eigh_full, A_gpu))
152+
alg_cpu = MatrixFunctionViaEigh(MatrixAlgebraKit.select_algorithm(eigh_full, A))
153+
for p in (2, -1, 0.5, -0.5)
154+
@test Array(power(A_gpu, p, alg_gpu)) ≈ power(A, p, alg_cpu)
155+
end
156+
157+
# domain handling: integer powers ignore the domain, fractional ones do not
158+
if T <: Real
159+
@test Array(diagview(power(Diagonal(ArrayT(-data)), 2))) ≈ data .^ 2
160+
@test_throws DomainError power(Diagonal(ArrayT(-data)), 0.5)
161+
end
162+
@test_throws SingularException power(Diagonal(ArrayT(zeros(T, m))), -1)
163+
@test_throws DomainError power(Diagonal(ArrayT(zeros(T, m))), -0.5)
164+
return nothing
165+
end
166+
167+
if CUDA.functional()
168+
@testset "power on CUDA for T = $T" for T in BLASFloats
169+
test_power_gpu(CuArray, T)
170+
end
171+
end
172+
173+
if AMDGPU.functional()
174+
@testset "power on AMDGPU for T = $T" for T in BLASFloats
175+
test_power_gpu(ROCArray, T)
176+
end
177+
end

‎test/squareroot.jl‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ using TestExtras
44
using StableRNGs
55
using MatrixAlgebraKit: diagview
66
using LinearAlgebra
7+
using CUDA, AMDGPU
78

89
BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
910
GenericFloats = (Float16, ComplexF16, BigFloat, Complex{BigFloat})
@@ -104,3 +105,50 @@ end
104105
@test squareroot(Diagonal(T[-eps(T), 1])) ≈ Diagonal(T[0, 1])
105106
end
106107
end
108+
109+
# GPU tests
110+
# ---------
111+
# The `DiagonalAlgorithm` kernel and the shared domain helpers are written as reductions and
112+
# broadcasts rather than scalar loops, so the same code runs on GPU; the `MatrixFunctionViaEigh`
113+
# path additionally goes through the device `eigh_full!`. Compare against the CPU reference and
114+
# exercise both domain outcomes, since the clamp and the throw are exactly what the reductions
115+
# replaced. If any step fell back to scalar indexing these would error under GPUArrays'
116+
# scalar-indexing guard.
117+
function test_squareroot_gpu(ArrayT, T)
118+
rng = StableRNG(123)
119+
m = 54
120+
121+
# Diagonal fast path
122+
data = T <: Real ? (abs.(randn(rng, T, m)) .+ one(T)) : randn(rng, T, m)
123+
@test Array(diagview(squareroot(Diagonal(ArrayT(data))))) ≈
124+
diagview(squareroot(Diagonal(data)))
125+
126+
# hermitian positive definite, via the eigenvalue decomposition
127+
X = randn(rng, T, m, m)
128+
A = Matrix(LinearAlgebra.Hermitian(X * X' + one(real(T)) * LinearAlgebra.I))
129+
A_gpu = ArrayT(A)
130+
alg_gpu = MatrixFunctionViaEigh(MatrixAlgebraKit.select_algorithm(eigh_full, A_gpu))
131+
alg_cpu = MatrixFunctionViaEigh(MatrixAlgebraKit.select_algorithm(eigh_full, A))
132+
sqrtA = squareroot(A_gpu, alg_gpu)
133+
@test Array(sqrtA) ≈ squareroot(A, alg_cpu)
134+
@test Array(sqrtA) * Array(sqrtA) ≈ A
135+
136+
# domain handling: roundoff-scale negatives are clamped, genuine ones throw
137+
if T <: Real
138+
@test Array(diagview(squareroot(Diagonal(ArrayT(T[-eps(T), 1]))))) ≈ T[0, 1]
139+
@test_throws DomainError squareroot(Diagonal(ArrayT(-data)))
140+
end
141+
return nothing
142+
end
143+
144+
if CUDA.functional()
145+
@testset "squareroot on CUDA for T = $T" for T in BLASFloats
146+
test_squareroot_gpu(CuArray, T)
147+
end
148+
end
149+
150+
if AMDGPU.functional()
151+
@testset "squareroot on AMDGPU for T = $T" for T in BLASFloats
152+
test_squareroot_gpu(ROCArray, T)
153+
end
154+
end

0 commit comments

Comments
 (0)