Skip to content

Commit 41d4d1f

Browse files
Merge pull request #1794 from CliMA/ck/reduce_cuda_dep
Reduce CUDA dependency in tests
2 parents c99c172 + 6a0ba7b commit 41d4d1f

15 files changed

Lines changed: 29 additions & 66 deletions

File tree

.buildkite/pipeline.yml

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1145,39 +1145,39 @@ steps:
11451145
key: "perf_gpu_spectral_ops_cuda_float32"
11461146
command:
11471147
- "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
1148-
- "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --device CUDA --float-type Float32"
1148+
- "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --float-type Float32"
11491149
env:
11501150
CLIMACOMMS_DEVICE: "CUDA"
11511151
agents:
11521152
slurm_gpus: 1
11531153

11541154
- label: "Perf: SEM operator benchmarks (CPU Float32)"
11551155
key: "perf_gpu_spectral_ops_cpu_float32"
1156-
command: "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --device CPU --float-type Float32"
1156+
command: "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --float-type Float32"
11571157

11581158
- label: "Perf: SEM operator benchmarks (cuda Float64)"
11591159
key: "perf_gpu_spectral_ops_cuda_float64"
11601160
command:
11611161
- "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
1162-
- "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --device CUDA --float-type Float64"
1162+
- "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --float-type Float64"
11631163
env:
11641164
CLIMACOMMS_DEVICE: "CUDA"
11651165
agents:
11661166
slurm_gpus: 1
11671167

11681168
- label: "Perf: SEM operator benchmarks (CPU Float64)"
11691169
key: "perf_gpu_spectral_ops_cpu_float64"
1170-
command: "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --device CPU --float-type Float64"
1170+
command: "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --float-type Float64"
11711171

11721172
- label: "Perf: SEM operator benchmarks (extruded CPU Float64)"
11731173
key: "perf_gpu_spectral_ops_extruded_cpu_float64"
1174-
command: "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --device CPU --float-type Float64 --space-type ExtrudedFiniteDifferenceSpace"
1174+
command: "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --float-type Float64 --space-type ExtrudedFiniteDifferenceSpace"
11751175

11761176
- label: "Perf: SEM operator benchmarks"
11771177
key: "perf_gpu_spectral_ops"
11781178
command:
11791179
- "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
1180-
- "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl --device CUDA"
1180+
- "julia --color=yes --project=.buildkite test/Operators/spectralelement/benchmark_ops.jl"
11811181
env:
11821182
CLIMACOMMS_DEVICE: "CUDA"
11831183
agents:

test/Fields/field_multi_broadcast_fusion.jl

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -74,13 +74,9 @@ end
7474

7575
function benchmark_kernel!(f!, X, Y, device)
7676
println("\n--------------------------- $(nameof(typeof(f!))) ")
77-
trial = benchmark_kernel!(f!, X, Y, device)
77+
trial = BenchmarkTools.@benchmark ClimaComms.@cuda_sync $device $f!($X, $Y)
7878
show(stdout, MIME("text/plain"), trial)
7979
end
80-
benchmark_kernel!(f!, X, Y, ::ClimaComms.CUDADevice) =
81-
BenchmarkTools.@benchmark CUDA.@sync $f!($X, $Y);
82-
benchmark_kernel!(f!, X, Y, ::ClimaComms.AbstractCPUDevice) =
83-
BenchmarkTools.@benchmark $f!($X, $Y);
8480

8581
function show_diff(A, B)
8682
for pn in propertynames(A)

test/InputOutput/spectralelement2d.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using Test
22
using ClimaComms
3+
ClimaComms.@import_required_backends
34
using LinearAlgebra
45
import ClimaCore
56
import ClimaCore:

test/Operators/hybrid/cuda.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using Test
22
using StaticArrays
33
using ClimaComms, ClimaCore
4+
ClimaComms.@import_required_backends
45
import ClimaCore:
56
Geometry,
67
Fields,
@@ -11,7 +12,6 @@ import ClimaCore:
1112
Operators,
1213
Quadratures
1314
using LinearAlgebra, IntervalSets
14-
using CUDA
1515
using OrdinaryDiffEq
1616

1717
function hvspace_3D_box(

test/Operators/spectralelement/benchmark_ops.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ using Revise; using ClimaCore
1313
include(joinpath(pkgdir(ClimaCore), "test", "Operators", "spectralelement", "benchmark_utils.jl"))
1414
include(joinpath(pkgdir(ClimaCore), "test", "Operators", "spectralelement", "benchmark_kernels.jl"))
1515
kernel_args = setup_kernel_args(["--float-type", "Float64"]);
16-
device = kernel_args.device
16+
device = ClimaComms.device()
1717
trial = benchmark_kernel!(kernel_args, kernel_spectral_div_grad!, device; silent=true);
1818
trial = benchmark_kernel_array!(kernel_args.arr_args, kernel_spectral_wdiv_array!, device; silent=true);
1919
show(stdout, MIME("text/plain"), trial);
@@ -60,7 +60,7 @@ include(
6060

6161
function benchmark_all(kernel_args = setup_kernel_args(ARGS))
6262

63-
device = kernel_args.device
63+
(; device) = kernel_args
6464
#=
6565
# Run benchmarks for a single kernel with:
6666
trial = benchmark_kernel!(kernel_args, kernel_spectral_div_grad!, device)

test/Operators/spectralelement/benchmark_utils.jl

Lines changed: 12 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ function benchmark_kernel_array!(
5959
@test all(Array(ϕ_arr) .== Array(ψ_arr)) # compile and confirm correctness
6060

6161
# Perform benchmark
62-
trial = BenchmarkTools.@benchmark CUDA.@sync $kernel(
62+
trial = BenchmarkTools.@benchmark ClimaComms.@cuda_sync $device $kernel(
6363
$args,
6464
threads = $threads,
6565
blocks = $blocks,
@@ -74,21 +74,14 @@ end
7474
function benchmark_kernel!(
7575
args,
7676
kernel_fun!,
77-
::ClimaComms.AbstractCPUDevice;
77+
device::ClimaComms.AbstractDevice;
7878
silent,
7979
)
8080
kernel_fun!(args) # compile first
81-
trial = BenchmarkTools.@benchmark $kernel_fun!($args)
82-
if !silent
83-
show(stdout, MIME("text/plain"), trial)
84-
println()
85-
end
86-
return trial
87-
end
88-
89-
function benchmark_kernel!(args, kernel_fun!, ::ClimaComms.CUDADevice; silent)
90-
kernel_fun!(args) # compile first
91-
trial = BenchmarkTools.@benchmark CUDA.@sync $kernel_fun!($args)
81+
trial =
82+
BenchmarkTools.@benchmark ClimaComms.@cuda_sync $device $kernel_fun!(
83+
$args,
84+
)
9285
if !silent
9386
show(stdout, MIME("text/plain"), trial)
9487
println()
@@ -101,13 +94,6 @@ function initial_velocity(space)
10194
return @. Geometry.Covariant12Vector(Geometry.UVVector(uλ, uϕ))
10295
end
10396

104-
function ismpi()
105-
# detect common environment variables used by MPI launchers
106-
# PMI_RANK appears to be used by MPICH and srun
107-
# OMPI_COMM_WORLD_RANK appears to be used by OpenMPI
108-
return haskey(ENV, "PMI_RANK") || haskey(ENV, "OMPI_COMM_WORLD_RANK")
109-
end
110-
11197
function create_space(
11298
context;
11399
float_type = Float64,
@@ -142,14 +128,6 @@ end
142128
function setup_kernel_args(ARGS::Vector{String} = ARGS)
143129
s = ArgParseSettings(prog = "spectralelement operator benchmarks")
144130
@add_arg_table! s begin
145-
"--device"
146-
help = "Computation device (CPU, CUDA)"
147-
arg_type = String
148-
default = CUDA.functional() ? "CUDA" : "CPU"
149-
"--comms"
150-
help = "Communication type (Singleton, MPI)"
151-
arg_type = String
152-
default = ismpi() ? "MPI" : "Singleton"
153131
"--float-type"
154132
help = "Floating point type (Float32, Float64)"
155133
eval_arg = true
@@ -173,16 +151,8 @@ function setup_kernel_args(ARGS::Vector{String} = ARGS)
173151
end
174152
args = parse_args(ARGS, s)
175153

176-
device =
177-
args["device"] == "CUDA" ? ClimaComms.CUDADevice() :
178-
args["device"] == "CPU" ? ClimaComms.CPUSingleThreaded() :
179-
error("Unknown device: $(args["device"])")
180-
181-
context =
182-
args["comms"] == "MPI" ? ClimaComms.MPICommsContext(device) :
183-
args["comms"] == "Singleton" ?
184-
ClimaComms.SingletonCommsContext(device) :
185-
error("Unknown comms: $(args["comms"])")
154+
device = ClimaComms.device()
155+
context = ClimaComms.context(device)
186156

187157
ClimaComms.init(context)
188158

@@ -260,12 +230,10 @@ function setup_kernel_args(ARGS::Vector{String} = ARGS)
260230
f = @. Geometry.Contravariant3Vector(Geometry.WVector(ϕ))
261231

262232
s = size(parent(ϕ))
263-
array_kernel_args = if device isa ClimaComms.AbstractCPUDevice
264-
(; ϕ_arr = fill(FT(1), s), ψ_arr = fill(FT(2), s))
265-
else
266-
device isa ClimaComms.CUDADevice
267-
(; ϕ_arr = CUDA.fill(FT(1), s), ψ_arr = CUDA.fill(FT(2), s))
268-
end
233+
ArrayType = ClimaComms.array_type(device)
234+
ϕ_arr = ArrayType(fill(FT(1), s))
235+
ψ_arr = ArrayType(fill(FT(2), s))
236+
array_kernel_args = (; ϕ_arr, ψ_arr)
269237

270238
kernel_args = (; ϕ, ψ, u, du, f, ϕψ, nt_ϕψ, nt_ϕψ_ft, f_comp, f_comp2)
271239
# buffers cannot reside in CuArray kernels

test/Operators/spectralelement/rectilinear.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using Test
22
using StaticArrays
33
using ClimaComms
4+
ClimaComms.@import_required_backends
45
import ClimaCore.DataLayouts: IJFH, VF
56
import ClimaCore:
67
Geometry,

test/Operators/spectralelement/rectilinear_cuda.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using Test
22
using StaticArrays
33
using ClimaComms, ClimaCore
4+
ClimaComms.@import_required_backends
45
import ClimaCore:
56
Geometry,
67
Fields,
@@ -11,7 +12,6 @@ import ClimaCore:
1112
Operators,
1213
Quadratures
1314
using LinearAlgebra, IntervalSets
14-
using CUDA
1515

1616
FT = Float64
1717
domain = Domains.RectangleDomain(

test/Spaces/distributed_cuda/ddss2.jl

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import ClimaCore:
1313

1414
using ClimaComms
1515
ClimaComms.@import_required_backends
16-
using CUDA
1716

1817
# initializing MPI
1918
const device = ClimaComms.device()

test/Spaces/distributed_cuda/ddss3.jl

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import ClimaCore:
1313

1414
using ClimaComms
1515
ClimaComms.@import_required_backends
16-
using CUDA
1716

1817
# initializing MPI
1918
const device = ClimaComms.device()

0 commit comments

Comments
 (0)