Skip to content

Commit 0f0d856

Browse files
utkarsh530utkarsh530
authored
Merge pull request #322 from jpsamaroo/jps/amdgpu-hip
Updates for HIP-based AMDGPU
2 parents 7971f09 + 579d106 commit 0f0d856

File tree

12 files changed

+60
-44
lines changed

12 files changed

+60
-44
lines changed

‎.buildkite/runtests.yml‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ steps:
44
setup:
55
version:
66
- "1"
7+
- "1.9"
78
env:
89
GROUP: CUDA
910
plugins:
@@ -31,7 +32,7 @@ steps:
3132
setup:
3233
version:
3334
# Known issues on 1.8
34-
- "1.9-nightly"
35+
- "1"
3536
env:
3637
GROUP: AMDGPU
3738
plugins:
@@ -59,7 +60,7 @@ steps:
5960
setup:
6061
version:
6162
- "1"
62-
- "1.9-nightly"
63+
- "1.9"
6364
env:
6465
GROUP: oneAPI
6566
plugins:
@@ -69,8 +70,11 @@ steps:
6970
- |
7071
julia --project=test -e '
7172
import Pkg
72-
Pkg.add(; name = "oneAPI")'
73-
rm test/Manifest.toml
73+
Pkg.add(; name="oneAPI")
74+
75+
println("+++ :julia: Building support library")
76+
include(joinpath(dirname(dirname(Base.find_package("oneAPI"))), "deps", "build_ci.jl"))
77+
Pkg.activate()'
7478
julia --project -e '
7579
import Pkg
7680
println("+++ :julia: Running tests")
@@ -86,8 +90,7 @@ steps:
8690
matrix:
8791
setup:
8892
version:
89-
- "1"
90-
- "1.9-nightly"
93+
- "1.9"
9194
env:
9295
GROUP: Metal
9396
plugins:

‎Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ MetalExt = ["Metal"]
3838
oneAPIExt = ["oneAPI"]
3939

4040
[compat]
41-
AMDGPU = "0.4.9"
41+
AMDGPU = "0.5, 0.6, 0.7, 0.8"
4242
Adapt = "3, 4"
4343
CUDA = "4.1.0, 5"
4444
ChainRulesCore = "1"

‎ext/AMDGPUExt.jl‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@ import DiffEqGPU
55
using .AMDGPU
66
import .AMDGPU: ROCBackend
77

8+
function DiffEqGPU.EnsembleGPUArray(cpu_offload::Float64)
9+
DiffEqGPU.EnsembleGPUArray(ROCBackend(), cpu_offload)
10+
end
811
DiffEqGPU.maxthreads(::ROCBackend) = 256
912
DiffEqGPU.maybe_prefer_blocks(::ROCBackend) = ROCBackend()
1013

‎src/ensemblegpukernel/callbacks.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ struct GPUContinuousCallback{F1, F2, F3, F4, F5, F6, T, T2, T3, I, R} <:
4747
reltol::T2,
4848
repeat_nudge::T3) where {F1, F2, F3, F4, F5, F6, T, T2,
4949
T3, I, R,
50-
}
50+
}
5151
if save_positions != (false, false)
5252
error("Callback `save_positions` are incompatible with kernel-based GPU ODE solvers due requiring static sizing. Please ensure `save_positions = (false,false)` is set in all callback definitions used with such solvers.")
5353
end

‎src/ensemblegpukernel/integrators/integrator_utils.jl‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ end
5454
IIP,
5555
S,
5656
T,
57-
}
57+
}
5858
integrator.retcode = retcode
5959
end
6060

@@ -156,7 +156,7 @@ end
156156
S,
157157
T,
158158
T1,
159-
}
159+
}
160160
# Can get rid of an allocation here with a function
161161
# get_tmp_arr(integrator.cache) which gives a pointer to some
162162
# cache array which can be modified.
@@ -185,7 +185,7 @@ end
185185
S,
186186
T,
187187
T1,
188-
}
188+
}
189189
_change_t_via_interpolation!(integrator, t, modify_save_endpoint)
190190
end
191191

@@ -322,7 +322,7 @@ end
322322
IIP,
323323
S,
324324
T,
325-
}
325+
}
326326
return nothing
327327
end
328328

@@ -333,7 +333,7 @@ end
333333
},
334334
callback,
335335
abst) where {AlgType <: GPUODEAlgorithm, IIP, S, T
336-
}
336+
}
337337
if abst == integrator.t
338338
tmp = integrator.u
339339
elseif abst == integrator.tprev

‎src/ensemblegpukernel/integrators/nonstiff/interpolants.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
IIP,
99
S,
1010
T,
11-
}
11+
}
1212
y₁ = integ.u
1313
k1 = integ.k1
1414
k2 = integ.k2

‎src/ensemblegpukernel/integrators/stiff/types.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,13 @@
33
IIP,
44
S,
55
T,
6-
})(t) where {
6+
})(t) where {
77
AlgType <:
88
GPUODEAlgorithm,
99
IIP,
1010
S,
1111
T,
12-
}
12+
}
1313
Θ = (t - integrator.tprev) / integrator.dt
1414
_ode_interpolant(Θ, integrator.dt, integrator.uprev, integrator)
1515
end

‎test/distributed_multi_gpu.jl‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ addprocs(2)
99
du[2] = u[1] * (p[2] - u[3]) - u[2]
1010
du[3] = u[1] * u[2] - p[3] * u[3]
1111
end
12-
CUDA.allowscalar(false)
1312
u0 = Float32[1.0; 0.0; 0.0]
1413
tspan = (0.0f0, 100.0f0)
1514
p = (10.0f0, 28.0f0, 8 / 3.0f0)

‎test/ensemblegpuarray.jl‎

Lines changed: 26 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -74,15 +74,18 @@ monteprob_jac = EnsembleProblem(prob_jac, prob_func = prob_func)
7474
@time solve(monteprob_jac, Rodas5(), EnsembleCPUArray(), dt = 0.1,
7575
trajectories = 10,
7676
saveat = 1.0f0)
77-
@time solve(monteprob_jac, Rodas5(), EnsembleGPUArray(backend), dt = 0.1,
78-
trajectories = 10,
79-
saveat = 1.0f0)
8077
@time solve(monteprob_jac, TRBDF2(), EnsembleCPUArray(), dt = 0.1,
8178
trajectories = 10,
8279
saveat = 1.0f0)
83-
@time solve(monteprob_jac, TRBDF2(), EnsembleGPUArray(backend), dt = 0.1,
84-
trajectories = 10,
85-
saveat = 1.0f0)
80+
81+
if GROUP == "CUDA"
82+
@time solve(monteprob_jac, Rodas5(), EnsembleGPUArray(backend), dt = 0.1,
83+
trajectories = 10,
84+
saveat = 1.0f0)
85+
@time solve(monteprob_jac, TRBDF2(), EnsembleGPUArray(backend), dt = 0.1,
86+
trajectories = 10,
87+
saveat = 1.0f0)
88+
end
8689

8790
@info "Callbacks"
8891

@@ -183,18 +186,19 @@ sol = solve(rober_prob, Rodas5(), abstol = 1.0f-8, reltol = 1.0f-8)
183186
sol = solve(rober_prob, TRBDF2(), abstol = 1.0f-4, reltol = 1.0f-1)
184187
rober_monteprob = EnsembleProblem(rober_prob, prob_func = prob_func)
185188

186-
# TODO: Does not work with Linearsolve.jl v1.35.0 https://github.com/SciML/DiffEqGPU.jl/pull/229
189+
if GROUP == "CUDA"
190+
@time sol = solve(rober_monteprob, Rodas5(),
191+
EnsembleGPUArray(backend), trajectories = 10,
192+
saveat = 1.0f0,
193+
abstol = 1.0f-8,
194+
reltol = 1.0f-8)
195+
@time sol = solve(rober_monteprob, TRBDF2(),
196+
EnsembleGPUArray(backend), trajectories = 10,
197+
saveat = 1.0f0,
198+
abstol = 1.0f-4,
199+
reltol = 1.0f-1)
200+
end
187201

188-
@time sol = solve(rober_monteprob, Rodas5(),
189-
EnsembleGPUArray(backend), trajectories = 10,
190-
saveat = 1.0f0,
191-
abstol = 1.0f-8,
192-
reltol = 1.0f-8)
193-
@time sol = solve(rober_monteprob, TRBDF2(),
194-
EnsembleGPUArray(backend), trajectories = 10,
195-
saveat = 1.0f0,
196-
abstol = 1.0f-4,
197-
reltol = 1.0f-1)
198202
@time sol = solve(rober_monteprob, TRBDF2(), EnsembleThreads(),
199203
trajectories = 10,
200204
abstol = 1e-4, reltol = 1e-1, saveat = 1.0f0)
@@ -241,5 +245,8 @@ monteprob = EnsembleProblem(prob_jac,
241245
sol = solve(monteprob, Tsit5(), EnsembleGPUArray(backend, 0.0), trajectories = 10,
242246
adaptive = false, dt = 0.01f0, save_everystep = false)
243247

244-
sol = solve(monteprob, Rosenbrock23(), EnsembleGPUArray(backend, 0.0), trajectories = 10,
245-
adaptive = false, dt = 0.01f0, save_everystep = false)
248+
if GROUP == "CUDA"
249+
sol = solve(monteprob, Rosenbrock23(), EnsembleGPUArray(backend, 0.0),
250+
trajectories = 10,
251+
adaptive = false, dt = 0.01f0, save_everystep = false)
252+
end

‎test/ensemblegpuarray_oop.jl‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,12 @@ prob_func = (prob, i, repeat) -> remake(prob, p = rand(Float32, 3) .* p)
3434
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
3535
@time sol = solve(monteprob, Tsit5(), EnsembleGPUArray(backend), trajectories = 10_000,
3636
saveat = 1.0f0)
37-
@time sol = solve(monteprob, Rosenbrock23(), EnsembleGPUArray(backend),
38-
trajectories = 10_000,
39-
saveat = 1.0f0)
40-
@time sol = solve(monteprob, TRBDF2(), EnsembleGPUArray(backend),
41-
trajectories = 10_000,
42-
saveat = 1.0f0)
37+
38+
if GROUP == "CUDA"
39+
@time sol = solve(monteprob, Rosenbrock23(), EnsembleGPUArray(backend),
40+
trajectories = 10_000,
41+
saveat = 1.0f0)
42+
@time sol = solve(monteprob, TRBDF2(), EnsembleGPUArray(backend),
43+
trajectories = 10_000,
44+
saveat = 1.0f0)
45+
end

0 commit comments

Comments
 (0)