Refactor macros, bump patch version

CliMA · May 30, 2024 · b69139c · b69139c
1 parent a22e4ac
commit b69139c
Show file tree

Hide file tree

Showing 8 changed files with 92 additions and 90 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -1,7 +1,7 @@
 agents:
   queue: new-central
   slurm_mem: 8G
-  modules: climacommon/2024_03_18
+  modules: climacommon/2024_05_27
 
 env:
   OPENBLAS_NUM_THREADS: 1

diff --git a/NEWS.md b/NEWS.md
@@ -1,6 +1,9 @@
 ClimaComms.jl Release Notes
 ========================
 
+v0.6.1
+- Macros have been refactored which fix some issues with code loading. We've also added some query functions to determine if users should load `MPI`/`CUDA` .
+
 v0.6.0
 -------
 

diff --git a/Project.toml b/Project.toml
@@ -8,7 +8,7 @@ authors = [
     "Jake Bolewski <[email protected]>",
     "Gabriele Bozzola <[email protected]>",
 ]
-version = "0.6.0"
+version = "0.6.1"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

diff --git a/docs/src/index.md b/docs/src/index.md
@@ -12,6 +12,8 @@ ClimaComms
 
 ```@docs
 ClimaComms.import_required_backends
+ClimaComms.cuda_is_required
+ClimaComms.mpi_is_required
 ```
 
 ## Devices

diff --git a/ext/ClimaCommsCUDAExt.jl b/ext/ClimaCommsCUDAExt.jl
@@ -15,4 +15,9 @@ end
 
 ClimaComms.array_type(::ClimaComms.CUDADevice) = CUDA.CuArray
 
+# Extending ClimaComms methods that operate on expressions (cannot use dispatch here)
+ClimaComms.cuda_sync(expr) = CUDA.@sync expr
+ClimaComms.cuda_time(expr) = CUDA.@time expr
+ClimaComms.cuda_elasped(expr) = CUDA.@elapsed expr
+
 end
diff --git a/src/context.jl b/src/context.jl
@@ -34,7 +34,7 @@ Behavior can be overridden by setting the `CLIMACOMMS_CONTEXT` environment varia
 to either `MPI` or `SINGLETON`.
 """
 function context(device = device(); target_context = context_type())
-    if target_context == :MPICommsContext && mpi_ext_is_not_loaded()
+    if target_context == :MPICommsContext && !mpi_ext_is_loaded()
         error(
             "Loading MPI.jl is required to use MPICommsContext. You might want to call ClimaComms.@import_required_backends",
         )

diff --git a/src/devices.jl b/src/devices.jl
@@ -76,7 +76,7 @@ The default is `CPU`.
 """
 function device()
     target_device = device_type()
-    if target_device == :CUDADevice && cuda_ext_is_not_loaded()
+    if target_device == :CUDADevice && !cuda_ext_is_loaded()
         error(
             "Loading CUDA.jl is required to use CUDADevice. You might want to call ClimaComms.@import_required_backends",
         )
@@ -129,6 +129,8 @@ macro threaded(device, loop)
     end
 end
 
+function cuda_time end
+
 """
     @time device expr
 
@@ -145,24 +147,19 @@ CUDA.@time expr
 for CUDA devices.
 """
 macro time(device, expr)
-    return esc(
-        quote
-            if $device isa $CUDADevice
-                @static if isnothing(
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt),
-                )
-                    error("CUDA not loaded")
-                else
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt).CUDA.@time $expr
-                end
-            else
-                @assert $device isa $AbstractDevice
-                $Base.@time $(expr)
-            end
-        end,
-    )
+    __CC__ = ClimaComms
+    return esc(quote
+        if $device isa $CUDADevice
+            $(__CC__).cuda_time($expr)
+        else
+            @assert $device isa $AbstractDevice
+            $Base.@time $(expr)
+        end
+    end)
 end
 
+function cuda_elasped end
+
 """
     @elapsed device expr
 
@@ -179,24 +176,19 @@ CUDA.@elapsed expr
 for CUDA devices.
 """
 macro elapsed(device, expr)
-    return esc(
-        quote
-            if $device isa $CUDADevice
-                @static if isnothing(
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt),
-                )
-                    error("CUDA not loaded")
-                else
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt).CUDA.@elapsed $expr
-                end
-            else
-                @assert $device isa $AbstractDevice
-                $Base.@elapsed $(expr)
-            end
-        end,
-    )
+    __CC__ = ClimaComms
+    return esc(quote
+        if $device isa $CUDADevice
+            $(__CC__).cuda_elasped($expr)
+        else
+            @assert $device isa $AbstractDevice
+            $Base.@elapsed $(expr)
+        end
+    end)
 end
 
+function cuda_sync end
+
 """
     @sync device expr
 
@@ -233,26 +225,17 @@ to synchronize), then you may want to simply use [`@cuda_sync`](@ref).
 """
 macro sync(device, expr)
     # https://github.com/JuliaLang/julia/issues/28979#issuecomment-1756145207
-    return esc(
-        quote
-            if $device isa $CUDADevice
-                @static if isnothing(
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt),
-                )
-                    error("CUDA not loaded")
-                else
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt).CUDA.@sync begin
-                        $(expr)
-                    end
-                end
-            else
-                @assert $device isa $AbstractDevice
-                $Base.@sync begin
-                    $(expr)
-                end
+    __CC__ = ClimaComms
+    return esc(quote
+        if $device isa $CUDADevice
+            $(__CC__).cuda_sync($expr)
+        else
+            @assert $device isa $AbstractDevice
+            $Base.@sync begin
+                $(expr)
             end
-        end,
-    )
+        end
+    end)
 end
 
 """
@@ -272,22 +255,13 @@ for CUDA devices.
 """
 macro cuda_sync(device, expr)
     # https://github.com/JuliaLang/julia/issues/28979#issuecomment-1756145207
-    return esc(
-        quote
-            if $device isa $CUDADevice
-                @static if isnothing(
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt),
-                )
-                    error("CUDA not loaded")
-                else
-                    $Base.get_extension($ClimaComms, :ClimaCommsCUDAExt).CUDA.@sync begin
-                        $(expr)
-                    end
-                end
-            else
-                @assert $device isa $AbstractDevice
-                $(expr)
-            end
-        end,
-    )
+    __CC__ = ClimaComms
+    return esc(quote
+        if $device isa $CUDADevice
+            $(__CC__).cuda_sync($expr)
+        else
+            @assert $device isa $AbstractDevice
+            $(expr)
+        end
+    end)
 end
diff --git a/src/loading.jl b/src/loading.jl
@@ -1,22 +1,38 @@
 import ..ClimaComms
 
-export import_required_backends
+export @import_required_backends
 
-function mpi_is_required()
-    return context_type() == :MPICommsContext
-end
+"""
+    mpi_is_required()
 
-function mpi_ext_is_not_loaded()
-    return isnothing(Base.get_extension(ClimaComms, :ClimaCommsMPIExt))
-end
+Returns a Bool indicating if MPI should be loaded, based on the
+`ENV["CLIMACOMMS_CONTEXT"]`. See [`ClimaComms.context`](@ref) for
+more information.
 
-function cuda_is_required()
-    return device_type() == :CUDADevice
-end
+```julia
+mpi_is_required() && using MPI
+```
+"""
+mpi_is_required() = context_type() == :MPICommsContext
 
-function cuda_ext_is_not_loaded()
-    return isnothing(Base.get_extension(ClimaComms, :ClimaCommsCUDAExt))
-end
+"""
+    cuda_is_required()
+
+Returns a Bool indicating if CUDA should be loaded, based on the
+`ENV["CLIMACOMMS_DEVICE"]`. See [`ClimaComms.device`](@ref) for
+more information.
+
+```julia
+cuda_is_required() && using CUDA
+```
+"""
+cuda_is_required() = device_type() == :CUDADevice
+
+cuda_ext_is_loaded() =
+    !isnothing(Base.get_extension(ClimaComms, :ClimaCommsCUDAExt))
+
+mpi_ext_is_loaded() =
+    !isnothing(Base.get_extension(ClimaComms, :ClimaCommsMPIExt))
 
 """
     ClimaComms.@import_required_backends
@@ -26,13 +42,15 @@ If the desired device is CUDA (as determined by `ClimaComms.device()`), try load
 """
 macro import_required_backends()
     return quote
-        @static if $mpi_is_required() && $mpi_ext_is_not_loaded()
+        @static if $mpi_is_required() && $(!mpi_ext_is_loaded())
             import MPI
-            @info "Loaded MPI.jl"
+            @debug "Loaded MPI via `import MPI`"
+            @assert $(mpi_ext_is_loaded()) "Importing MPI did not succesfully load the ClimaCommsMPIExt"
         end
-        @static if $cuda_is_required() && $cuda_ext_is_not_loaded()
+        @static if $cuda_is_required() && $(!cuda_ext_is_loaded())
             import CUDA
-            @info "Loaded CUDA.jl"
+            @debug "Loaded CUDA via `import CUDA`"
+            @assert $(cuda_ext_is_loaded()) "Importing CUDA did not succesfully load the ClimaCommsCUDAExt"
         end
     end
 end