@@ -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,
7474function 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ϕ))
10295end
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-
11197function create_space (
11298 context;
11399 float_type = Float64,
142128function 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
0 commit comments