Skip to content

Commit 71e0575

Browse files
kshyattkshyatt
committed
Small updates for GPU
1 parent 21db519 commit 71e0575

4 files changed

Lines changed: 72 additions & 69 deletions

File tree

‎Project.toml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
4646
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
4747
ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
4848
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
49+
GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
4950
JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
5051
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
5152
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -56,7 +57,8 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
5657
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
5758

5859
[targets]
59-
test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur", "Random", "Mooncake"]
60+
test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur", "Random", "Mooncake", "GPUArrays"]
6061

6162
[sources]
6263
Mooncake = {url="https://github.com/chalk-lab/Mooncake.jl", rev="ksh/cuarraify"}
64+
GPUArrays = {url="https://github.com/JuliaGPU/GPUArrays.jl", rev="ksh/findlast"}

‎test/testsuite/ad_utils.jl‎

Lines changed: 0 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -160,40 +160,6 @@ function ad_eig_trunc_setup(A, truncalg)
160160
return DV, (Dtrunc, Vtrunc), ΔD2V, (ΔDtrunc, ΔVtrunc)
161161
end
162162

163-
function copy_eigh_full(A; kwargs...)
164-
A = (A + A') / 2
165-
return eigh_full(A; kwargs...)
166-
end
167-
168-
function copy_eigh_full!(A, DV; kwargs...)
169-
A = (A + A') / 2
170-
return eigh_full!(A, DV; kwargs...)
171-
end
172-
173-
function copy_eigh_vals(A; kwargs...)
174-
A = (A + A') / 2
175-
return eigh_vals(A; kwargs...)
176-
end
177-
178-
function copy_eigh_vals!(A, D; kwargs...)
179-
A = (A + A') / 2
180-
return eigh_vals!(A, D; kwargs...)
181-
end
182-
183-
function copy_eigh_trunc(A, alg; kwargs...)
184-
A = (A + A') / 2
185-
return eigh_trunc(A, alg; kwargs...)
186-
end
187-
188-
function copy_eigh_trunc!(A, DV, alg; kwargs...)
189-
A = (A + A') / 2
190-
return eigh_trunc!(A, DV, alg; kwargs...)
191-
end
192-
193-
MatrixAlgebraKit.copy_input(::typeof(copy_eigh_full), A) = MatrixAlgebraKit.copy_input(eigh_full, A)
194-
MatrixAlgebraKit.copy_input(::typeof(copy_eigh_vals), A) = MatrixAlgebraKit.copy_input(eigh_vals, A)
195-
MatrixAlgebraKit.copy_input(::typeof(copy_eigh_trunc), A) = MatrixAlgebraKit.copy_input(eigh_trunc, A)
196-
197163
function ad_eigh_full_setup(A)
198164
m, n = size(A)
199165
T = eltype(A)

‎test/testsuite/chainrules.jl‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ for f in
1010
:svd_compact, :svd_trunc, :svd_vals,
1111
:left_polar, :right_polar,
1212
)
13-
copy_f = Symbol(:copy_, f)
13+
copy_f = Symbol(:cr_copy_, f)
1414
f! = Symbol(f, '!')
1515
_hermitian = startswith(string(f), "eigh")
1616
@eval begin
@@ -62,7 +62,7 @@ function test_chainrules_qr(
6262
QR, ΔQR = ad_qr_compact_setup(A)
6363
ΔQ, ΔR = ΔQR
6464
test_rrule(
65-
copy_qr_compact, A, alg ⊢ NoTangent();
65+
cr_copy_qr_compact, A, alg ⊢ NoTangent();
6666
output_tangent = ΔQR, atol = atol, rtol = rtol
6767
)
6868
test_rrule(
@@ -84,7 +84,7 @@ function test_chainrules_qr(
8484
@testset "qr_null" begin
8585
N, ΔN = ad_qr_null_setup(A)
8686
test_rrule(
87-
copy_qr_null, A, alg ⊢ NoTangent();
87+
cr_copy_qr_null, A, alg ⊢ NoTangent();
8888
output_tangent = ΔN, atol = atol, rtol = rtol
8989
)
9090
test_rrule(
@@ -97,7 +97,7 @@ function test_chainrules_qr(
9797
@testset "qr_full" begin
9898
QR, ΔQR = ad_qr_full_setup(A)
9999
test_rrule(
100-
copy_qr_full, A, alg ⊢ NoTangent();
100+
cr_copy_qr_full, A, alg ⊢ NoTangent();
101101
output_tangent = ΔQR, atol = atol, rtol = rtol
102102
)
103103
test_rrule(
@@ -114,7 +114,7 @@ function test_chainrules_qr(
114114
QR, ΔQR = ad_qr_rd_compact_setup(Ard)
115115
ΔQ, ΔR = ΔQR
116116
test_rrule(
117-
copy_qr_compact, Ard, alg ⊢ NoTangent();
117+
cr_copy_qr_compact, Ard, alg ⊢ NoTangent();
118118
output_tangent = ΔQR, atol = atol, rtol = rtol
119119
)
120120
test_rrule(
@@ -141,7 +141,7 @@ function test_chainrules_lq(
141141
LQ, ΔLQ = ad_lq_compact_setup(A)
142142
ΔL, ΔQ = ΔLQ
143143
test_rrule(
144-
copy_lq_compact, A, alg ⊢ NoTangent();
144+
cr_copy_lq_compact, A, alg ⊢ NoTangent();
145145
output_tangent = ΔLQ, atol = atol, rtol = rtol
146146
)
147147
test_rrule(
@@ -163,7 +163,7 @@ function test_chainrules_lq(
163163
@testset "lq_null" begin
164164
Nᴴ, ΔNᴴ = ad_lq_null_setup(A)
165165
test_rrule(
166-
copy_lq_null, A, alg ⊢ NoTangent();
166+
cr_copy_lq_null, A, alg ⊢ NoTangent();
167167
output_tangent = ΔNᴴ, atol = atol, rtol = rtol
168168
)
169169
test_rrule(
@@ -175,7 +175,7 @@ function test_chainrules_lq(
175175
@testset "lq_full" begin
176176
LQ, ΔLQ = ad_lq_full_setup(A)
177177
test_rrule(
178-
copy_lq_full, A, alg ⊢ NoTangent();
178+
cr_copy_lq_full, A, alg ⊢ NoTangent();
179179
output_tangent = ΔLQ, atol = atol, rtol = rtol
180180
)
181181
test_rrule(
@@ -190,7 +190,7 @@ function test_chainrules_lq(
190190
Ard = instantiate_matrix(T, (m, r)) * instantiate_matrix(T, (r, n))
191191
LQ, ΔLQ = ad_lq_rd_compact_setup(Ard)
192192
test_rrule(
193-
copy_lq_compact, Ard, alg ⊢ NoTangent();
193+
cr_copy_lq_compact, Ard, alg ⊢ NoTangent();
194194
output_tangent = ΔLQ, atol = atol, rtol = rtol
195195
)
196196
test_rrule(
@@ -217,10 +217,10 @@ function test_chainrules_eig(
217217
DV, ΔDV, ΔD2V = ad_eig_full_setup(A)
218218
ΔD, ΔV = ΔDV
219219
test_rrule(
220-
copy_eig_full, A, alg ⊢ NoTangent(); output_tangent = ΔDV, atol, rtol
220+
cr_copy_eig_full, A, alg ⊢ NoTangent(); output_tangent = ΔDV, atol, rtol
221221
)
222222
test_rrule(
223-
copy_eig_full, A, alg ⊢ NoTangent(); output_tangent = ΔD2V, atol, rtol
223+
cr_copy_eig_full, A, alg ⊢ NoTangent(); output_tangent = ΔD2V, atol, rtol
224224
)
225225
test_rrule(
226226
config, eig_full, A, alg ⊢ NoTangent();
@@ -242,7 +242,7 @@ function test_chainrules_eig(
242242
@testset "eig_vals" begin
243243
D, ΔD = ad_eig_vals_setup(A)
244244
test_rrule(
245-
copy_eig_vals, A, alg ⊢ NoTangent(); output_tangent = ΔD, atol, rtol
245+
cr_copy_eig_vals, A, alg ⊢ NoTangent(); output_tangent = ΔD, atol, rtol
246246
)
247247
test_rrule(
248248
config, eig_vals, A, alg ⊢ NoTangent();
@@ -254,7 +254,7 @@ function test_chainrules_eig(
254254
truncalg = TruncatedAlgorithm(alg, truncrank(r; by = abs))
255255
DV, DVtrunc, ΔDV, ΔDVtrunc = ad_eig_trunc_setup(A, truncalg)
256256
test_rrule(
257-
copy_eig_trunc, A, truncalg ⊢ NoTangent();
257+
cr_copy_eig_trunc, A, truncalg ⊢ NoTangent();
258258
output_tangent = (ΔDVtrunc..., zero(real(T))),
259259
atol = atol, rtol = rtol
260260
)
@@ -266,7 +266,7 @@ function test_chainrules_eig(
266266
truncalg = TruncatedAlgorithm(alg, truncrank(5; by = real))
267267
DV, DVtrunc, ΔDV, ΔDVtrunc = ad_eig_trunc_setup(A, truncalg)
268268
test_rrule(
269-
copy_eig_trunc, A, truncalg ⊢ NoTangent();
269+
cr_copy_eig_trunc, A, truncalg ⊢ NoTangent();
270270
output_tangent = (ΔDVtrunc..., zero(real(T))),
271271
atol = atol, rtol = rtol
272272
)
@@ -295,10 +295,10 @@ function test_chainrules_eigh(
295295
DV, ΔDV, ΔD2V = ad_eigh_full_setup(A)
296296
ΔD, ΔV = ΔDV
297297
test_rrule(
298-
copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = ΔDV, atol, rtol
298+
cr_copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = ΔDV, atol, rtol
299299
)
300300
test_rrule(
301-
copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = ΔD2V, atol, rtol
301+
cr_copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = ΔD2V, atol, rtol
302302
)
303303
# eigh_full does not include a projector onto the Hermitian part of the matrix
304304
test_rrule(
@@ -321,7 +321,7 @@ function test_chainrules_eigh(
321321
@testset "eigh_vals" begin
322322
D, ΔD = ad_eigh_vals_setup(A)
323323
test_rrule(
324-
copy_eigh_vals, A, alg ⊢ NoTangent(); output_tangent = ΔD, atol, rtol
324+
cr_copy_eigh_vals, A, alg ⊢ NoTangent(); output_tangent = ΔD, atol, rtol
325325
)
326326
test_rrule(
327327
config, eigh_vals ∘ Matrix ∘ Hermitian, A;
@@ -334,7 +334,7 @@ function test_chainrules_eigh(
334334
truncalg = TruncatedAlgorithm(alg, truncrank(r; by = abs))
335335
DV, DVtrunc, ΔDV, ΔDVtrunc = ad_eigh_trunc_setup(A, truncalg)
336336
test_rrule(
337-
copy_eigh_trunc, A, truncalg ⊢ NoTangent();
337+
cr_copy_eigh_trunc, A, truncalg ⊢ NoTangent();
338338
output_tangent = (ΔDVtrunc..., zero(real(T))),
339339
atol = atol, rtol = rtol
340340
)
@@ -358,7 +358,7 @@ function test_chainrules_eigh(
358358
DV, DVtrunc, ΔDV, ΔDVtrunc = ad_eigh_trunc_setup(A, truncalg)
359359
ind = MatrixAlgebraKit.findtruncated(diagview(DV[1]), truncalg.trunc)
360360
test_rrule(
361-
copy_eigh_trunc, A, truncalg ⊢ NoTangent();
361+
cr_copy_eigh_trunc, A, truncalg ⊢ NoTangent();
362362
output_tangent = (ΔDVtrunc..., zero(real(T))),
363363
atol = atol, rtol = rtol
364364
)
@@ -393,11 +393,11 @@ function test_chainrules_svd(
393393
@testset "svd_compact" begin
394394
USV, ΔUSVᴴ, ΔUS2Vᴴ = ad_svd_compact_setup(A)
395395
test_rrule(
396-
copy_svd_compact, A, alg ⊢ NoTangent();
396+
cr_copy_svd_compact, A, alg ⊢ NoTangent();
397397
output_tangent = ΔUSVᴴ, atol = atol, rtol = rtol
398398
)
399399
test_rrule(
400-
copy_svd_compact, A, alg ⊢ NoTangent();
400+
cr_copy_svd_compact, A, alg ⊢ NoTangent();
401401
output_tangent = ΔUS2Vᴴ, atol = atol, rtol = rtol
402402
)
403403
test_rrule(
@@ -414,7 +414,7 @@ function test_chainrules_svd(
414414
@testset "svd_vals" begin
415415
S, ΔS = ad_svd_vals_setup(A)
416416
test_rrule(
417-
copy_svd_vals, A, alg ⊢ NoTangent();
417+
cr_copy_svd_vals, A, alg ⊢ NoTangent();
418418
output_tangent = ΔS, atol, rtol
419419
)
420420
test_rrule(
@@ -427,7 +427,7 @@ function test_chainrules_svd(
427427
truncalg = TruncatedAlgorithm(alg, truncrank(r))
428428
USVᴴ, ΔUSVᴴ, ΔUSVᴴtrunc = ad_svd_trunc_setup(A, truncalg)
429429
test_rrule(
430-
copy_svd_trunc, A, truncalg ⊢ NoTangent();
430+
cr_copy_svd_trunc, A, truncalg ⊢ NoTangent();
431431
output_tangent = (ΔUSVᴴtrunc..., zero(real(T))),
432432
atol = atol, rtol = rtol
433433
)
@@ -453,7 +453,7 @@ function test_chainrules_svd(
453453
truncalg = TruncatedAlgorithm(alg, trunctol(atol = S[1, 1] / 2))
454454
USVᴴ, ΔUSVᴴ, ΔUSVᴴtrunc = ad_svd_trunc_setup(A, truncalg)
455455
test_rrule(
456-
copy_svd_trunc, A, truncalg ⊢ NoTangent();
456+
cr_copy_svd_trunc, A, truncalg ⊢ NoTangent();
457457
output_tangent = (ΔUSVᴴtrunc..., zero(real(T))),
458458
atol = atol, rtol = rtol
459459
)
@@ -490,7 +490,7 @@ function test_chainrules_polar(
490490
alg = MatrixAlgebraKit.default_polar_algorithm(A)
491491
@testset "left_polar" begin
492492
if m >= n
493-
test_rrule(copy_left_polar, A, alg ⊢ NoTangent(); atol = atol, rtol = rtol)
493+
test_rrule(cr_copy_left_polar, A, alg ⊢ NoTangent(); atol = atol, rtol = rtol)
494494
test_rrule(
495495
config, left_polar, A, alg ⊢ NoTangent();
496496
atol = atol, rtol = rtol, rrule_f = rrule_via_ad, check_inferred = false
@@ -499,7 +499,7 @@ function test_chainrules_polar(
499499
end
500500
@testset "right_polar" begin
501501
if m <= n
502-
test_rrule(copy_right_polar, A, alg ⊢ NoTangent(); atol = atol, rtol = rtol)
502+
test_rrule(cr_copy_right_polar, A, alg ⊢ NoTangent(); atol = atol, rtol = rtol)
503503
test_rrule(
504504
config, right_polar, A, alg ⊢ NoTangent();
505505
atol = atol, rtol = rtol, rrule_f = rrule_via_ad, check_inferred = false

‎test/testsuite/mooncake.jl‎

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,41 @@ using Mooncake, Mooncake.TestUtils
44
using Mooncake: rrule!!
55
using MatrixAlgebraKit: diagview, TruncatedAlgorithm, PolarViaSVD, eigh_trunc
66

7+
function mc_copy_eigh_full(A; kwargs...)
8+
A = (A + A') / 2
9+
return eigh_full(A; kwargs...)
10+
end
11+
12+
function mc_copy_eigh_full!(A, DV; kwargs...)
13+
A = (A + A') / 2
14+
return eigh_full!(A, DV; kwargs...)
15+
end
16+
17+
function mc_copy_eigh_vals(A; kwargs...)
18+
A = (A + A') / 2
19+
return eigh_vals(A; kwargs...)
20+
end
21+
22+
function mc_copy_eigh_vals!(A, D; kwargs...)
23+
A = (A + A') / 2
24+
return eigh_vals!(A, D; kwargs...)
25+
end
26+
27+
function mc_copy_eigh_trunc(A, alg; kwargs...)
28+
A = (A + A') / 2
29+
return eigh_trunc(A, alg; kwargs...)
30+
end
31+
32+
function mc_copy_eigh_trunc!(A, DV, alg; kwargs...)
33+
A = (A + A') / 2
34+
return eigh_trunc!(A, DV, alg; kwargs...)
35+
end
36+
37+
MatrixAlgebraKit.copy_input(::typeof(mc_copy_eigh_full), A) = MatrixAlgebraKit.copy_input(eigh_full, A)
38+
MatrixAlgebraKit.copy_input(::typeof(mc_copy_eigh_vals), A) = MatrixAlgebraKit.copy_input(eigh_vals, A)
39+
MatrixAlgebraKit.copy_input(::typeof(mc_copy_eigh_trunc), A) = MatrixAlgebraKit.copy_input(eigh_trunc, A)
40+
41+
742
make_mooncake_tangent(ΔAelem::T) where {T <: Real} = ΔAelem
843
make_mooncake_tangent(ΔAelem::T) where {T <: Complex} = Mooncake.build_tangent(T, real(ΔAelem), imag(ΔAelem))
944
make_mooncake_tangent(ΔA::AbstractMatrix{<:Real}) = ΔA
@@ -251,31 +286,31 @@ function test_mooncake_eigh(
251286
@testset "eigh_full" begin
252287
DV, ΔDV, ΔD2V = ad_eigh_full_setup(A)
253288
dDV = make_mooncake_tangent(ΔD2V)
254-
Mooncake.TestUtils.test_rule(rng, copy_eigh_full, A; mode = Mooncake.ReverseMode, output_tangent = dDV, is_primitive = false, atol = atol, rtol = rtol)
255-
test_pullbacks_match(copy_eigh_full!, copy_eigh_full, A, DV, ΔD2V)
289+
Mooncake.TestUtils.test_rule(rng, mc_copy_eigh_full, A; mode = Mooncake.ReverseMode, output_tangent = dDV, is_primitive = false, atol = atol, rtol = rtol)
290+
test_pullbacks_match(mc_copy_eigh_full!, mc_copy_eigh_full, A, DV, ΔD2V)
256291
end
257292
@testset "eigh_vals" begin
258293
D, ΔD = ad_eigh_vals_setup(A)
259294
dD = make_mooncake_tangent(ΔD)
260-
Mooncake.TestUtils.test_rule(rng, copy_eigh_vals, A; mode = Mooncake.ReverseMode, output_tangent = dD, is_primitive = false, atol = atol, rtol = rtol)
261-
test_pullbacks_match(copy_eigh_vals!, copy_eigh_vals, A, D, ΔD)
295+
Mooncake.TestUtils.test_rule(rng, mc_copy_eigh_vals, A; mode = Mooncake.ReverseMode, output_tangent = dD, is_primitive = false, atol = atol, rtol = rtol)
296+
test_pullbacks_match(mc_copy_eigh_vals!, mc_copy_eigh_vals, A, D, ΔD)
262297
end
263298
@testset "eigh_trunc" begin
264299
for r in 1:4:m
265300
truncalg = TruncatedAlgorithm(MatrixAlgebraKit.default_eigh_algorithm(A), truncrank(r; by = abs))
266301
DV, _, ΔDV, ΔDVtrunc = ad_eigh_trunc_setup(A, truncalg)
267302
ϵ = zero(real(T))
268303
dDVerr = make_mooncake_tangent((ΔDVtrunc..., ϵ))
269-
Mooncake.TestUtils.test_rule(rng, copy_eigh_trunc, A, truncalg; mode = Mooncake.ReverseMode, output_tangent = dDVerr, atol = atol, rtol = rtol, is_primitive = false)
270-
test_pullbacks_match(copy_eigh_trunc!, copy_eigh_trunc, A, DV, ΔDV, truncalg; rdata = (Mooncake.NoRData(), Mooncake.NoRData(), zero(real(T))))
304+
Mooncake.TestUtils.test_rule(rng, mc_copy_eigh_trunc, A, truncalg; mode = Mooncake.ReverseMode, output_tangent = dDVerr, atol = atol, rtol = rtol, is_primitive = false)
305+
test_pullbacks_match(mc_copy_eigh_trunc!, mc_copy_eigh_trunc, A, DV, ΔDV, truncalg; rdata = (Mooncake.NoRData(), Mooncake.NoRData(), zero(real(T))))
271306
end
272307
D = eigh_vals(A / 2)
273308
truncalg = TruncatedAlgorithm(MatrixAlgebraKit.default_eigh_algorithm(A), trunctol(; atol = maximum(abs, D) / 2))
274309
DV, _, ΔDV, ΔDVtrunc = ad_eigh_trunc_setup(A, truncalg)
275310
ϵ = zero(real(T))
276311
dDVerr = make_mooncake_tangent((ΔDVtrunc..., ϵ))
277-
Mooncake.TestUtils.test_rule(rng, copy_eigh_trunc, A, truncalg; mode = Mooncake.ReverseMode, output_tangent = dDVerr, atol = atol, rtol = rtol, is_primitive = false)
278-
test_pullbacks_match(copy_eigh_trunc!, copy_eigh_trunc, A, DV, ΔDV, truncalg; rdata = (Mooncake.NoRData(), Mooncake.NoRData(), zero(real(T))))
312+
Mooncake.TestUtils.test_rule(rng, mc_copy_eigh_trunc, A, truncalg; mode = Mooncake.ReverseMode, output_tangent = dDVerr, atol = atol, rtol = rtol, is_primitive = false)
313+
test_pullbacks_match(mc_copy_eigh_trunc!, mc_copy_eigh_trunc, A, DV, ΔDV, truncalg; rdata = (Mooncake.NoRData(), Mooncake.NoRData(), zero(real(T))))
279314
end
280315
end
281316
end

0 commit comments

Comments
 (0)