Add extensions feature from Julia 1.9 with backward compatibility

.github/workflows/ci.yml

@@ -1,17 +1,17 @@
 name: CI
 on:
-  - push
-  - pull_request
-
+  pull_request:
+  push:
 jobs:
   test:
+    if: github.event.pull_request.draft == false
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         version:
-          - '1.8'
+          - '1.9'
         os:
           - ubuntu-latest
           - macOS-latest
@@ -45,6 +45,7 @@ jobs:
         with:
           file: lcov.info
   docs:
+    if: github.event.pull_request.draft == false
     name: Documentation
     runs-on: ubuntu-latest
     steps:

Project.toml

@@ -4,40 +4,49 @@ authors = ["Gabriel Weymouth <gabriel.weymouth@gmail.com>"]
 version = "1.0.3"
 
 [deps]
-AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+
+[weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
+[extensions]
+WaterLilyAMDGPUExt = "AMDGPU"
+WaterLilyCUDAExt = "CUDA"
+WaterLilyWriteVTKExt = "WriteVTK"
+
 [compat]
 AMDGPU = "^0.4.13,0.6"
 CUDA = "4.1.4,5"
 ForwardDiff = "^0.10.18"
 KernelAbstractions = "^0.9.1"
 Reexport = "^1.2.2"
+Requires = "1.3"
 StaticArrays = "^1.1.0"
 julia = "1.6"
 
 [extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
-JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 Makie = "537997a7-5e4e-5d89-9595-2241ea00577e"
-OutMacro = "0ae4d431-9932-4135-a8f1-51ee5e017775"
 PerformanceTestTools = "dc46b164-d16f-48ec-a853-60448fc869fe"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
+WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [targets]
-test = ["Test", "PerformanceTestTools", "JLD2", "OutMacro"]
+test = ["Test", "CUDA", "AMDGPU", "GPUArrays"]

ext/WaterLilyAMDGPUExt.jl

@@ -0,0 +1,28 @@
+module WaterLilyAMDGPUExt
+
+if isdefined(Base, :get_extension)
+    using AMDGPU
+else
+    using ..AMDGPU
+end
+
+using WaterLily
+import WaterLily: L₂
+
+"""
+    __init__()
+
+Asserts AMDGPU is functional when loading this extension.
+"""
+__init__() = @assert AMDGPU.functional()
+
+AMDGPU.allowscalar(false) # disallow scalar operations on GPU
+
+"""
+    L₂(a)
+
+L₂ norm of ROCArray `a` excluding ghosts.
+"""
+L₂(a::ROCArray,R::CartesianIndices=inside(a)) = mapreduce(abs2,+,@inbounds(a[R]))
+
+end # module

ext/WaterLilyCUDAExt.jl

@@ -0,0 +1,28 @@
+module WaterLilyCUDAExt
+
+if isdefined(Base, :get_extension)
+    using CUDA
+else
+    using ..CUDA
+end
+
+using WaterLily
+import WaterLily: L₂
+
+"""
+    __init__()
+
+Asserts CUDA is functional when loading this extension.
+"""
+__init__() = @assert CUDA.functional()
+
+CUDA.allowscalar(false) # disallow scalar operations on GPU
+
+"""
+    L₂(a)
+
+L₂ norm of CUDA array `a` excluding ghosts.
+"""
+L₂(a::CuArray,R::CartesianIndices=inside(a)) = mapreduce(abs2,+,@inbounds(a[R]))
+
+end # module

src/vtkWriter.jl → ext/WaterLilyWriteVTKExt.jl

@@ -1,11 +1,21 @@
-using WriteVTK
+module WaterLilyWriteVTKExt
+
+if isdefined(Base, :get_extension)
+    using WriteVTK
+else
+    using ..WriteVTK
+end
+
+using WaterLily
 using Printf: @sprintf
+import Base: close
+
 """
     vtk_grid(name;attrib,T)
 
 Generates a `vtkWriter` that hold the collection name to which the `vtk` files are written.
 The default attributes that are saved are the `Velocity` and the `Pressure` fields.
-Custom attributes can be passed as `Dict{String,Function}` to the `attrib` keyword.    
+Custom attributes can be passed as `Dict{String,Function}` to the `attrib` keyword.
 """
 struct vtkWriter
     fname::String
@@ -21,7 +31,7 @@
 """
     default_attrib()
 
-return a `Dict` containing the name and bound funtion for the default attributes. 
+return a `Dict` containing the name and bound funtion for the default attributes.
 The name is used as the key in the `vtk` file and the function generates the data
 to put in the file. With this approach, any variable can be save to the vtk file.
 """
@@ -49,7 +59,7 @@
 
 closes the `vtkWriter`, this is required to write the collection file.
 """
-Base.close(w::vtkWriter)=(vtk_save(w.collection);nothing)
+close(w::vtkWriter)=(vtk_save(w.collection);nothing)
 """
     components_first(a::Array)
 
@@ -59,4 +69,6 @@
 function components_first(a::Array)
     N=length(size(a)); p=[N,1:N-1...]
     return permutedims(a,p)
-end
+end
+
+end # module

src/WaterLily.jl

@@ -43,7 +43,7 @@ Constructor for a WaterLily.jl simulation:
   - `uλ`: Function to generate the initial velocity field.
   - `body`: Immersed geometry.
   - `T`: Array element type.
-  - `mem`: memory location. `Array` and `CuArray` run on CPU and CUDA backends, respectively.
+  - `mem`: memory location. `Array`, `CuArray`, `ROCm` to run on CPU, NVIDIA, or AMD devices, respectively.
 
 See files in `examples` folder for examples.
 """
@@ -105,7 +105,20 @@ end
 
 export Simulation,sim_step!,sim_time,measure!
 
-include("vtkWriter.jl")
-export vtkWriter,write!,close
+# WriteVTKExt functions to be extended
+function vtkWriter end
+function write! end
+
+# Backward compatibility for extensions
+if !isdefined(Base, :get_extension)
+using Requires
+end
+function __init__()
+    @static if !isdefined(Base, :get_extension)
+        @require AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" include("../ext/WaterLilyAMDGPUExt.jl")
+        @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("../ext/WaterLilyCUDAExt.jl")
+        @require WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" include("../ext/WaterLilyWriteVTKExt.jl")
+    end
+end
 
 end # module

src/util.jl

@@ -1,7 +1,4 @@
 using KernelAbstractions: get_backend, @index, @kernel
-using CUDA: CuArray
-using AMDGPU: ROCArray
-GPUArray = Union{CuArray,ROCArray}
 
 @inline CI(a...) = CartesianIndex(a...)
 """
@@ -44,8 +41,7 @@ size_u(u) = splitn(size(u))
 
 L₂ norm of array `a` excluding ghosts.
 """
-L₂(a) = sum(abs2,@inbounds(a[I]) for I ∈ inside(a))
-L₂(a::GPUArray,R::CartesianIndices=inside(a)) = mapreduce(abs2,+,@inbounds(a[R]))
+L₂(a::Array) = sum(abs2,@inbounds(a[I]) for I ∈ inside(a))
 
 """
     @inside <expr>
@@ -131,7 +127,7 @@ Base.front(I::CartesianIndex) = CI(Base.front(I.I))
 """
     apply!(f, c)
 
-Apply a vector function `f(i,x)` to the faces of a uniform staggered array `c` or 
+Apply a vector function `f(i,x)` to the faces of a uniform staggered array `c` or
 a function `f(x)` to the center of a uniform array `c`.
 """
 apply!(f,c) = hasmethod(f,Tuple{Int,CartesianIndex}) ? applyV!(f,c) : applyS!(f,c)