Use preferences to switch between SIMD and KernelAbstractions

.github/workflows/ci.yml

@@ -26,14 +26,13 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.10'
+          - '1.11'
         os:
           - ubuntu-latest
           - macOS-latest
           - windows-latest
         arch:
           - x64
-          - x86
         nthreads:
           - '1'
           - 'auto'
@@ -56,10 +55,18 @@ jobs:
             ${{ runner.os }}-test-${{ env.cache-name }}-
             ${{ runner.os }}-test-
             ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
+      - name: WaterLily tests
         env:
           JULIA_NUM_THREADS: ${{ matrix.nthreads }}
+        shell: bash
+        run: |
+          if [ "${{ matrix.nthreads }}" = "auto" ]
+          then
+            printf "[WaterLily]\nbackend = \"KernelAbstractions\"" > LocalPreferences.toml
+          else
+            printf "[WaterLily]\nbackend = \"SIMD\"" > LocalPreferences.toml
+          fi
+          julia --proj --color=yes -e "using Pkg; Pkg.instantiate(); Pkg.test(; coverage=true);"
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
         with:

.gitignore

@@ -2,6 +2,7 @@
 .vscode/
 .DS_Store
 Manifest.toml
+LocalPreferences.toml
 /docs/build/
 
 # gfx

Project.toml

@@ -10,7 +10,9 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -19,8 +21,8 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -43,6 +45,7 @@ EllipsisNotation = "1.8"
 ForwardDiff = "^0.10.18"
 KernelAbstractions = "0.9.1"
 LoggingExtras = "1.1"
+Random = "1.11.0"
 Reexport = "^1.2.2"
 Requires = "1.3"
 StaticArrays = "^1.1.0"

src/WaterLily.jl

@@ -6,7 +6,7 @@
 using DocStringExtensions
 
 include("util.jl")
-export L₂,BC!,@inside,inside,δ,apply!,loc,@log
+export L₂,BC!,@inside,inside,δ,apply!,loc,@log,set_backend,backend
 
 using Reexport
 @reexport using KernelAbstractions: @kernel,@index,get_backend
@@ -166,15 +166,21 @@
 
 # Check number of threads when loading WaterLily
 """
-    check_nthreads(::Val{1})
+    check_nthreads()
 
 Check the number of threads available for the Julia session that loads WaterLily.
-A warning is shown when running in serial (`JULIA_NUM_THREADS=1`).
-"""
-check_nthreads(::Val{1}) = @warn("\nUsing WaterLily in serial (ie. JULIA_NUM_THREADS=1) is not recommended because \
-    it disables the GPU backend and defaults to serial CPU."*
-    "\nUse JULIA_NUM_THREADS=auto, or any number of threads greater than 1, to allow multi-threading in CPU or GPU backends.")
-check_nthreads(_) = nothing
+A warning is shown when running in serial (JULIA_NUM_THREADS=1) with KernelAbstractions enabled.
+"""
+function check_nthreads()
+    if backend == "KernelAbstractions" && Threads.nthreads() == 1
+        @warn """
+        Using WaterLily in serial (ie. JULIA_NUM_THREADS=1) is not recommended because it defaults to serial CPU execution.
+        Use JULIA_NUM_THREADS=auto, or any number of threads greater than 1, to allow multi-threading in CPU backends.
+        For a low-overhead single-threaded CPU only backend set: WaterLily.set_backend("SIMD")
+        """
+    end
+end
+check_nthreads()
 
 # Backward compatibility for extensions
 if !isdefined(Base, :get_extension)
@@ -191,7 +197,6 @@
         @require Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73" include("../ext/WaterLilyMeshingExt.jl")
         @require JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" include("../ext/WaterLilyJLD2Ext.jl")
     end
-    check_nthreads(Val{Threads.nthreads()}())
 end
 
 end # module

src/util.jl

@@ -89,6 +89,19 @@
     end |> esc
 end
 
+# Could also use ScopedValues in Julia 1.11+
+using Preferences
+const backend = @load_preference("backend", "KernelAbstractions")
+function set_backend(new_backend::String)
+    if !(new_backend in ("SIMD", "KernelAbstractions"))
+        throw(ArgumentError("Invalid backend: \"$(new_backend)\""))
+    end
+
+    # Set it in our runtime values, as well as saving it to disk
+    @set_preferences!("backend" => new_backend)
+    @info("New backend set; restart your Julia session for this change to take effect!")
+end
+
 """
     @loop <expr> over <I ∈ R>
 
@@ -118,26 +131,34 @@
 """
 macro loop(args...)
     ex,_,itr = args
-    _,I,R = itr.args; sym = []
+    _,I,R = itr.args
+    sym = []
     grab!(sym,ex)     # get arguments and replace composites in `ex`
     setdiff!(sym,[I]) # don't want to pass I as an argument
+    symT = symtypes(sym) # generate a list of types for each symbol
     @gensym(kern, kern_) # generate unique kernel function names for serial and KA execution
-    return quote
-        function $kern($(rep.(sym)...),::Val{1})
-            @simd for $I ∈ $R
+    @static if backend == "KernelAbstractions"
+        return quote
+            @kernel function $kern_($(rep.(sym)...),@Const(I0)) # replace composite arguments
+                $I = @index(Global,Cartesian)
+                $I += I0
                 @fastmath @inbounds $ex
             end
-        end
-        @kernel function $kern_($(rep.(sym)...),@Const(I0)) # replace composite arguments
-            $I = @index(Global,Cartesian)
-            $I += I0
-            @fastmath @inbounds $ex
-        end
-        function $kern($(rep.(sym)...),_)
-            $kern_(get_backend($(sym[1])),64)($(sym...),$R[1]-oneunit($R[1]),ndrange=size($R))
-        end
-        $kern($(sym...),Val{Threads.nthreads()}()) # dispatch to SIMD for -t 1, or KA otherwise
-    end |> esc
+            function $kern($(joinsymtype(rep.(sym),symT)...)) where {$(symT...)}
+                $kern_(get_backend($(sym[1])),64)($(sym...),$R[1]-oneunit($R[1]),ndrange=size($R))
+            end
+            $kern($(sym...))
+        end |> esc
+    else # backend == "SIMD"
+        return quote
+            function $kern($(joinsymtype(rep.(sym),symT)...)) where {$(symT...)}
+                @simd for $I ∈ $R
+                    @fastmath @inbounds $ex
+                end
+            end
+            $kern($(sym...))
+        end |> esc
+    end
 end
 function grab!(sym,ex::Expr)
     ex.head == :. && return union!(sym,[ex])      # grab composite name and return
@@ -149,6 +170,10 @@
 grab!(sym,ex) = nothing
 rep(ex) = ex
 rep(ex::Expr) = ex.head == :. ? Symbol(ex.args[2].value) : ex
+using Random
+symtypes(sym) = [Symbol.(Random.randstring('A':'Z',4)) for _ in 1:length(sym)]
+joinsymtype(sym::Symbol,symT::Symbol) = Expr(:(::), sym, symT)
+joinsymtype(sym,symT) = zip(sym,symT) .|> x->joinsymtype(x...)
 
 using StaticArrays
 """

test/alloctest.jl

@@ -1,5 +1,6 @@
 using BenchmarkTools, Printf
 
+backend != "SIMD" && throw(ArgumentError("KernelAbstractions backend not allowed to run allocations tests, use SIMD backend"))
 @testset "mom_step! allocations" begin
     function Sim(θ;L=32,U=1,Re=100,perdir=())
         function map(x,t)

test/maintests.jl

@@ -1,6 +1,7 @@
 using GPUArrays
-using ReadVTK, WriteVTK, JLD2
+using ReadVTK, WriteVTK, JLD2, Random
 
+backend != "KernelAbstractions" && throw(ArgumentError("SIMD backend not allowed to run main tests, use KernelAbstractions backend"))
 @info "Test backends: $(join(arrays,", "))"
 @testset "util.jl" begin
     I = CartesianIndex(1,2,3,4)
@@ -17,6 +18,10 @@ using ReadVTK, WriteVTK, JLD2
     WaterLily.grab!(sym,ex)
     @test ex == :(a[I, i] = Math.add(b[I], func(I, q)))
     @test sym == [:a, :I, :i, :(p.b), :q]
+    sym = [:a,:b,:c]
+    Random.seed!(99)
+    symT = WaterLily.symtypes(sym)
+    @test WaterLily.joinsymtype(sym,symT) == Expr[:(a::MDFZ), :(b::AXYU), :(c::RFIB)]
 
     for f ∈ arrays
         p = zeros(4,5) |> f