Skip to content

Commit 92812f1

Browse files
TendonFFFTendonFFF
authored
Merge pull request #3 from TendonFFF/copilot/add-amdgpu-sparse-matrix-tests
Add AMDGPU extension tests and Buildkite pipeline configuration
2 parents f7be173 + f702e13 commit 92812f1

File tree

5 files changed

+354
-1
lines changed

5 files changed

+354
-1
lines changed

‎.buildkite/AMDGPU_Ext.yml‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
steps:
2+
- label: "AMDGPU Julia {{matrix.version}}"
3+
matrix:
4+
setup:
5+
version:
6+
- "1.10" # oldest
7+
- "1" # latest
8+
plugins:
9+
- JuliaCI/julia#v1:
10+
version: "{{matrix.version}}"
11+
- JuliaCI/julia-test#v1:
12+
test_args: "--quickfail"
13+
- JuliaCI/julia-coverage#v1:
14+
codecov: true
15+
dirs:
16+
- src
17+
- ext
18+
agents:
19+
queue: "juliagpu"
20+
rocm: "*"
21+
rocmgpu: "*"
22+
env:
23+
GROUP: "AMDGPU_Ext"
24+
timeout_in_minutes: 60

‎.buildkite/pipeline.yml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,13 @@ steps:
2323
- "test/ext-test/gpu/**"
2424
- "Project.toml"
2525
target: ".buildkite/CUDA_Ext.yml"
26+
- staticfloat/forerunner: # AMDGPU.jl tests
27+
watch:
28+
- ".buildkite/pipeline.yml"
29+
- ".buildkite/AMDGPU_Ext.yml"
30+
- "src/**"
31+
- "ext/QuantumToolboxAMDGPUExt.jl"
32+
- "test/runtests.jl"
33+
- "test/ext-test/gpu/**"
34+
- "Project.toml"
35+
target: ".buildkite/AMDGPU_Ext.yml"

‎test/ext-test/gpu/Project.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
[deps]
2+
AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
23
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
34
CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
45
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
56
QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
67

78
[compat]
9+
AMDGPU = "1"
810
CUDA = "5"

‎test/ext-test/gpu/amdgpu_ext.jl‎

Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
@testset "AMDGPU Extension" verbose = true begin
2+
# Test that scalar indexing is disallowed
3+
@test_throws ErrorException AMDGPU.rand(1)[1]
4+
5+
ψdi = Qobj(Int64[1, 0])
6+
ψdf = Qobj(Float64[1, 0])
7+
ψdc = Qobj(ComplexF64[1, 0])
8+
ψsi = to_sparse(ψdi)
9+
ψsf = to_sparse(ψdf)
10+
ψsc = to_sparse(ψdc)
11+
12+
Xdi = Qobj(Int64[0 1; 1 0])
13+
Xdf = Qobj(Float64[0 1; 1 0])
14+
Xdc = Qobj(ComplexF64[0 1; 1 0])
15+
Xsi = to_sparse(Xdi)
16+
Xsf = to_sparse(Xdf)
17+
Xsc = to_sparse(Xdc)
18+
19+
@test_throws DomainError roc(ψdi; word_size = 16)
20+
21+
# type conversion of AMDGPU dense arrays
22+
@test typeof(roc(ψdi; word_size = 64).data) == typeof(ROCArray(ψdi).data) <: ROCArray{Int64,1}
23+
@test typeof(roc(ψdi; word_size = 32).data) == typeof(ROCArray{Int32}(ψdi).data) <: ROCArray{Int32,1}
24+
@test typeof(roc(ψdf; word_size = 64).data) == typeof(ROCArray(ψdf).data) <: ROCArray{Float64,1}
25+
@test typeof(roc(ψdf; word_size = 32).data) == typeof(ROCArray{Float32}(ψdf).data) <: ROCArray{Float32,1}
26+
@test typeof(roc(ψdc; word_size = 64).data) == typeof(ROCArray(ψdc).data) <: ROCArray{ComplexF64,1}
27+
@test typeof(roc(ψdc; word_size = 32).data) == typeof(ROCArray{ComplexF32}(ψdc).data) <: ROCArray{ComplexF32,1}
28+
@test typeof(roc(Xdi; word_size = 64).data) == typeof(ROCArray(Xdi).data) <: ROCArray{Int64,2}
29+
@test typeof(roc(Xdi; word_size = 32).data) == typeof(ROCArray{Int32}(Xdi).data) <: ROCArray{Int32,2}
30+
@test typeof(roc(Xdf; word_size = 64).data) == typeof(ROCArray(Xdf).data) <: ROCArray{Float64,2}
31+
@test typeof(roc(Xdf; word_size = 32).data) == typeof(ROCArray{Float32}(Xdf).data) <: ROCArray{Float32,2}
32+
@test typeof(roc(Xdc; word_size = 64).data) == typeof(ROCArray(Xdc).data) <: ROCArray{ComplexF64,2}
33+
@test typeof(roc(Xdc; word_size = 32).data) == typeof(ROCArray{ComplexF32}(Xdc).data) <: ROCArray{ComplexF32,2}
34+
35+
# type conversion of AMDGPU sparse arrays
36+
@test typeof(roc(ψsi; word_size = 64).data) == typeof(ROCSparseVector(ψsi).data) == ROCSparseVector{Int64,Int32}
37+
@test typeof(roc(ψsi; word_size = 32).data) == typeof(ROCSparseVector{Int32}(ψsi).data) == ROCSparseVector{Int32,Int32}
38+
@test typeof(roc(ψsf; word_size = 64).data) == typeof(ROCSparseVector(ψsf).data) == ROCSparseVector{Float64,Int32}
39+
@test typeof(roc(ψsf; word_size = 32).data) ==
40+
typeof(ROCSparseVector{Float32}(ψsf).data) ==
41+
ROCSparseVector{Float32,Int32}
42+
@test typeof(roc(ψsc; word_size = 64).data) == typeof(ROCSparseVector(ψsc).data) == ROCSparseVector{ComplexF64,Int32}
43+
@test typeof(roc(ψsc; word_size = 32).data) ==
44+
typeof(ROCSparseVector{ComplexF32}(ψsc).data) ==
45+
ROCSparseVector{ComplexF32,Int32}
46+
@test typeof(roc(Xsi; word_size = 64).data) == typeof(ROCSparseMatrixCSC(Xsi).data) == ROCSparseMatrixCSC{Int64,Int32}
47+
@test typeof(roc(Xsi; word_size = 32).data) ==
48+
typeof(ROCSparseMatrixCSC{Int32}(Xsi).data) ==
49+
ROCSparseMatrixCSC{Int32,Int32}
50+
@test typeof(roc(Xsf; word_size = 64).data) ==
51+
typeof(ROCSparseMatrixCSC(Xsf).data) ==
52+
ROCSparseMatrixCSC{Float64,Int32}
53+
@test typeof(roc(Xsf; word_size = 32).data) ==
54+
typeof(ROCSparseMatrixCSC{Float32}(Xsf).data) ==
55+
ROCSparseMatrixCSC{Float32,Int32}
56+
@test typeof(roc(Xsc; word_size = 64).data) ==
57+
typeof(ROCSparseMatrixCSC(Xsc).data) ==
58+
ROCSparseMatrixCSC{ComplexF64,Int32}
59+
@test typeof(roc(Xsc; word_size = 32).data) ==
60+
typeof(ROCSparseMatrixCSC{ComplexF32}(Xsc).data) ==
61+
ROCSparseMatrixCSC{ComplexF32,Int32}
62+
@test typeof(ROCSparseMatrixCSR(Xsi).data) == ROCSparseMatrixCSR{Int64,Int32}
63+
@test typeof(ROCSparseMatrixCSR{Int32}(Xsi).data) == ROCSparseMatrixCSR{Int32,Int32}
64+
@test typeof(ROCSparseMatrixCSR(Xsf).data) == ROCSparseMatrixCSR{Float64,Int32}
65+
@test typeof(ROCSparseMatrixCSR{Float32}(Xsf).data) == ROCSparseMatrixCSR{Float32,Int32}
66+
@test typeof(ROCSparseMatrixCSR(Xsc).data) == ROCSparseMatrixCSR{ComplexF64,Int32}
67+
@test typeof(ROCSparseMatrixCSR{ComplexF32}(Xsc).data) == ROCSparseMatrixCSR{ComplexF32,Int32}
68+
69+
# type conversion of AMDGPU Diagonal arrays
70+
@test roc(qeye(10), word_size = Val(32)).data isa Diagonal{ComplexF32,<:ROCVector{ComplexF32}}
71+
@test roc(qeye(10), word_size = Val(64)).data isa Diagonal{ComplexF64,<:ROCVector{ComplexF64}}
72+
73+
# Sparse To Dense
74+
# @test to_dense(roc(ψsi; word_size = 64)).data isa ROCVector{Int64} # TODO: Fix this in AMDGPU.jl
75+
@test to_dense(roc(ψsf; word_size = 64)).data isa ROCVector{Float64}
76+
@test to_dense(roc(ψsc; word_size = 64)).data isa ROCVector{ComplexF64}
77+
# @test to_dense(roc(Xsi; word_size = 64)).data isa ROCMatrix{Int64} # TODO: Fix this in AMDGPU.jl
78+
@test to_dense(roc(Xsf; word_size = 64)).data isa ROCMatrix{Float64}
79+
@test to_dense(roc(Xsc; word_size = 64)).data isa ROCMatrix{ComplexF64}
80+
81+
# @test to_dense(Int32, roc(ψsf; word_size = 64)).data isa ROCVector{Int32} # TODO: Fix this in AMDGPU.jl
82+
# @test to_dense(Float32, roc(ψsf; word_size = 64)).data isa ROCVector{Float32} # TODO: Fix this in AMDGPU.jl
83+
# @test to_dense(ComplexF32, roc(ψsf; word_size = 64)).data isa ROCVector{ComplexF32} # TODO: Fix this in AMDGPU.jl
84+
# @test to_dense(Int64, roc(Xsf; word_size = 32)).data isa ROCMatrix{Int64} # TODO: Fix this in AMDGPU.jl
85+
# @test to_dense(Float64, roc(Xsf; word_size = 32)).data isa ROCMatrix{Float64} # TODO: Fix this in AMDGPU.jl
86+
# @test to_dense(ComplexF64, roc(Xsf; word_size = 32)).data isa ROCMatrix{ComplexF64} # TODO: Fix this in AMDGPU.jl
87+
88+
# brief example in README and documentation
89+
N = 20
90+
ω64 = 1.0 # Float64
91+
ω32 = 1.0f0 # Float32
92+
γ64 = 0.1 # Float64
93+
γ32 = 0.1f0 # Float32
94+
tlist = range(0, 10, 100)
95+
96+
## calculate by CPU
97+
a_cpu = destroy(N)
98+
ψ0_cpu = fock(N, 3)
99+
H_cpu = ω64 * a_cpu' * a_cpu
100+
sol_cpu = mesolve(H_cpu, ψ0_cpu, tlist, [sqrt(γ64) * a_cpu], e_ops = [a_cpu' * a_cpu], progress_bar = Val(false))
101+
102+
## calculate by GPU (with 64-bit)
103+
a_gpu64 = roc(destroy(N))
104+
ψ0_gpu64 = roc(fock(N, 3))
105+
H_gpu64 = ω64 * a_gpu64' * a_gpu64
106+
sol_gpu64 = mesolve(
107+
H_gpu64,
108+
ψ0_gpu64,
109+
tlist,
110+
[sqrt(γ64) * a_gpu64],
111+
e_ops = [a_gpu64' * a_gpu64],
112+
progress_bar = Val(false),
113+
)
114+
115+
## calculate by GPU (with 32-bit)
116+
a_gpu32 = roc(destroy(N), word_size = 32)
117+
ψ0_gpu32 = roc(fock(N, 3), word_size = 32)
118+
H_gpu32 = ω32 * a_gpu32' * a_gpu32
119+
sol_gpu32 = mesolve(
120+
H_gpu32,
121+
ψ0_gpu32,
122+
tlist,
123+
[sqrt(γ32) * a_gpu32],
124+
e_ops = [a_gpu32' * a_gpu32],
125+
progress_bar = Val(false),
126+
)
127+
128+
@test all([isapprox(sol_cpu.expect[i], sol_gpu64.expect[i]) for i in 1:length(tlist)])
129+
@test all([isapprox(sol_cpu.expect[i], sol_gpu32.expect[i]; atol = 1e-6) for i in 1:length(tlist)])
130+
end
131+
132+
@testset "AMDGPU steadystate" begin
133+
N = 50
134+
Δ = 0.01
135+
F = 0.1
136+
γ = 0.1
137+
nth = 2
138+
139+
a = destroy(N)
140+
H = Δ * a' * a + F * (a + a')
141+
c_ops = [sqrt(γ * (nth + 1)) * a, sqrt(γ * nth) * a']
142+
143+
ρ_ss_cpu = steadystate(H, c_ops)
144+
145+
H_gpu_csc = roc(H)
146+
c_ops_gpu_csc = [roc(c_op) for c_op in c_ops]
147+
ρ_ss_gpu_csc = steadystate(H_gpu_csc, c_ops_gpu_csc, solver = SteadyStateLinearSolver())
148+
149+
H_gpu_csr = ROCSparseMatrixCSR(H_gpu_csc)
150+
c_ops_gpu_csr = [ROCSparseMatrixCSR(c_op) for c_op in c_ops_gpu_csc]
151+
ρ_ss_gpu_csr = steadystate(H_gpu_csr, c_ops_gpu_csr, solver = SteadyStateLinearSolver())
152+
153+
@test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csc.data) atol = 1e-8 * length(ρ_ss_cpu)
154+
@test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csr.data) atol = 1e-8 * length(ρ_ss_cpu)
155+
end
156+
157+
@testset "AMDGPU spectrum" begin
158+
N = 10
159+
a = roc(destroy(N))
160+
H = a' * a
161+
c_ops = [sqrt(0.1 * (0.01 + 1)) * a, sqrt(0.1 * (0.01)) * a']
162+
solver = Lanczos(steadystate_solver = SteadyStateLinearSolver())
163+
164+
ω_l = range(0, 3, length = 1000)
165+
spec = spectrum(H, ω_l, c_ops, a', a; solver = solver)
166+
167+
spec = collect(spec)
168+
spec = spec ./ maximum(spec)
169+
170+
test_func = maximum(real.(spec)) * (0.1 / 2)^2 ./ ((ω_l .- 1) .^ 2 .+ (0.1 / 2)^2)
171+
idxs = test_func .> 0.05
172+
@test sum(abs2.(spec[idxs] .- test_func[idxs])) / sum(abs2.(test_func[idxs])) < 0.01
173+
174+
# TODO: Fix this
175+
# @testset "Type Inference spectrum" begin
176+
# @inferred spectrum(H, ω_l, c_ops, a', a; solver = solver)
177+
# end
178+
end
179+
180+
@testset "AMDGPU ptrace" begin
181+
g = fock(2, 1)
182+
e = fock(2, 0)
183+
α = sqrt(0.7)
184+
β = sqrt(0.3) * 1im
185+
ψ = α * kron(g, e) + β * kron(e, g) |> roc
186+
187+
ρ1 = ptrace(ψ, 1)
188+
ρ2 = ptrace(ψ, 2)
189+
@test ρ1.data isa ROCArray
190+
@test ρ2.data isa ROCArray
191+
@test Array(ρ1.data) ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
192+
@test Array(ρ2.data) ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
193+
194+
ψ_d = ψ'
195+
196+
ρ1 = ptrace(ψ_d, 1)
197+
ρ2 = ptrace(ψ_d, 2)
198+
@test ρ1.data isa ROCArray
199+
@test ρ2.data isa ROCArray
200+
@test Array(ρ1.data) ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
201+
@test Array(ρ2.data) ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
202+
203+
ρ = ket2dm(ψ)
204+
ρ1 = ptrace(ρ, 1)
205+
ρ2 = ptrace(ρ, 2)
206+
@test ρ.data isa ROCArray
207+
@test ρ1.data isa ROCArray
208+
@test ρ2.data isa ROCArray
209+
@test Array(ρ1.data) ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
210+
@test Array(ρ2.data) ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
211+
212+
ψ1 = normalize(g + 1im * e)
213+
ψ2 = normalize(g + e)
214+
ρ1 = ket2dm(ψ1)
215+
ρ2 = ket2dm(ψ2)
216+
ρ = kron(ρ1, ρ2) |> roc
217+
ρ1_ptr = ptrace(ρ, 1)
218+
ρ2_ptr = ptrace(ρ, 2)
219+
@test ρ1_ptr.data isa ROCArray
220+
@test ρ2_ptr.data isa ROCArray
221+
@test ρ1.data ≈ Array(ρ1_ptr.data) atol = 1e-10
222+
@test ρ2.data ≈ Array(ρ2_ptr.data) atol = 1e-10
223+
224+
ψlist = [rand_ket(2), rand_ket(3), rand_ket(4), rand_ket(5)]
225+
ρlist = [rand_dm(2), rand_dm(3), rand_dm(4), rand_dm(5)]
226+
ψtotal = tensor(ψlist...) |> roc
227+
ρtotal = tensor(ρlist...) |> roc
228+
sel_tests = [
229+
SVector{0,Int}(),
230+
1,
231+
2,
232+
3,
233+
4,
234+
(1, 2),
235+
(1, 3),
236+
(1, 4),
237+
(2, 3),
238+
(2, 4),
239+
(3, 4),
240+
(1, 2, 3),
241+
(1, 2, 4),
242+
(1, 3, 4),
243+
(2, 3, 4),
244+
(1, 2, 3, 4),
245+
]
246+
for sel in sel_tests
247+
if length(sel) == 0
248+
@test ptrace(ψtotal, sel) ≈ 1.0
249+
@test ptrace(ρtotal, sel) ≈ 1.0
250+
else
251+
@test ptrace(ψtotal, sel) ≈ roc(tensor([ket2dm(ψlist[i]) for i in sel]...))
252+
@test ptrace(ρtotal, sel) ≈ roc(tensor([ρlist[i] for i in sel]...))
253+
end
254+
end
255+
@test ptrace(ψtotal, (1, 3, 4)) ≈ ptrace(ψtotal, (4, 3, 1)) # check sort of sel
256+
@test ptrace(ρtotal, (1, 3, 4)) ≈ ptrace(ρtotal, (3, 1, 4)) # check sort of sel
257+
258+
@testset "Type Inference (ptrace)" begin
259+
@inferred ptrace(ρ, 1)
260+
@inferred ptrace(ρ, 2)
261+
@inferred ptrace(ψ_d, 1)
262+
@inferred ptrace(ψ_d, 2)
263+
@inferred ptrace(ψ, 1)
264+
@inferred ptrace(ψ, 2)
265+
end
266+
end
267+
268+
@testset "AMDGPU eigsolve" begin
269+
N = 30
270+
Δ = 0.5
271+
U = 0.1
272+
κ = 0.1
273+
F = 0.5
274+
275+
a = destroy(N)
276+
H = Δ * a' * a + U / 2 * a' * a' * a * a + F * (a + a')
277+
278+
c_ops = [sqrt(κ) * a]
279+
280+
L = liouvillian(H, c_ops)
281+
L_gpu = ROCSparseMatrixCSR(L)
282+
283+
vals_cpu, vecs_cpu = eigenstates(L; sparse = true, sigma = 0.01, eigvals = 4, krylovdim = 30)
284+
vals_gpu, vecs_gpu = eigenstates(
285+
L_gpu;
286+
sparse = true,
287+
sigma = 0.01,
288+
eigvals = 4,
289+
krylovdim = 30,
290+
solver = LUFactorization(),
291+
v0 = AMDGPU.rand(ComplexF64, size(L_gpu, 1)),
292+
)
293+
294+
@test vals_cpu ≈ vals_gpu atol = 1e-8
295+
@test all(zip(vecs_cpu, vecs_gpu)) do (v_cpu, v_gpu)
296+
return isapprox(abs(dot(v_cpu.data, Array(v_gpu.data))), 1; atol = 1e-8)
297+
end
298+
end

‎test/runtests.jl‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ using Test
22
using TestItemRunner
33
using Pkg
44

5-
const GROUP_LIST = String["All", "Core", "Code-Quality", "AutoDiff_Ext", "Makie_Ext", "CUDA_Ext", "Arbitrary-Precision"]
5+
const GROUP_LIST = String["All", "Core", "Code-Quality", "AutoDiff_Ext", "Makie_Ext", "CUDA_Ext", "AMDGPU_Ext", "Arbitrary-Precision"]
66

77
const GROUP = get(ENV, "GROUP", "All")
88
(GROUP in GROUP_LIST) || throw(ArgumentError("Unknown GROUP = $GROUP\nThe allowed groups are: $GROUP_LIST\n"))
@@ -83,6 +83,25 @@ if (GROUP == "CUDA_Ext")
8383
include(joinpath(testdir, "ext-test", "gpu", "cuda_ext.jl"))
8484
end
8585

86+
if (GROUP == "AMDGPU_Ext")
87+
Pkg.activate("ext-test/gpu")
88+
Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
89+
Pkg.update()
90+
91+
using QuantumToolbox
92+
import LinearAlgebra: Diagonal, dot
93+
import StaticArraysCore: SVector
94+
using AMDGPU
95+
using AMDGPU: ROCArray, ROCVector, ROCMatrix
96+
using AMDGPU.rocSPARSE: ROCSparseVector, ROCSparseMatrixCSC, ROCSparseMatrixCSR
97+
using LinearSolve
98+
99+
QuantumToolbox.about()
100+
AMDGPU.versioninfo()
101+
102+
include(joinpath(testdir, "ext-test", "gpu", "amdgpu_ext.jl"))
103+
end
104+
86105
if (GROUP == "Arbitrary-Precision")
87106
Pkg.activate("ext-test/cpu/arbitrary_precision")
88107
Pkg.develop(PackageSpec(path = dirname(@__DIR__)))

0 commit comments

Comments
 (0)