Use preferences to switch between SIMD and KernelAbstractions

Project.toml

@@ -1,14 +1,15 @@
 name = "WaterLily"
 uuid = "ed894a53-35f9-47f1-b17f-85db9237eebd"
 authors = ["Gabriel Weymouth <[email protected]>"]
-version = "1.1"
+version = "1.1.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"

src/WaterLily.jl

@@ -134,15 +134,22 @@ export restart_sim!
 
 # 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).
+A warning is shown when running in serial (JULIA_NUM_THREADS=1) with KernelAbstractions enabled.
 """
-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
+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
 
 # Backward compatibility for extensions
 if !isdefined(Base, :get_extension)
@@ -155,7 +162,7 @@ function __init__()
         @require WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" include("../ext/WaterLilyWriteVTKExt.jl")
         @require ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3" include("../ext/WaterLilyReadVTKExt.jl")
     end
-    check_nthreads(Val{Threads.nthreads()}())
+    check_nthreads()
 end
 
 end # module

src/util.jl

@@ -65,6 +65,20 @@ macro inside(ex)
     end |> esc
 end
 
+# Could also use ScopedValues in Julia 1.11+
+using Preferences
+const backend = Symbol(@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>
 
@@ -99,7 +113,7 @@ macro loop(args...)
     setdiff!(sym,[I]) # don't want to pass I as an argument
     @gensym(kern, kern_) # generate unique kernel function names for serial and KA execution
     return quote
-        function $kern($(rep.(sym)...),::Val{1})
+        function $kern($(rep.(sym)...),::Val{:SIMD})
             @simd for $I ∈ $R
                 @fastmath @inbounds $ex
             end
@@ -109,10 +123,10 @@ macro loop(args...)
             $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))
+        function $kern($(rep.(sym)...),::Val{:KernelAbstractions})
+            $kern_(get_backend($(sym[1])))($(sym...),$R[1]-oneunit($R[1]),ndrange=size($R))
         end
-        $kern($(sym...),Val{Threads.nthreads()}()) # dispatch to SIMD for -t 1, or KA otherwise
+        $kern($(sym...),Val{$(QuoteNode(backend))}()) # dispatch to SIMD or KA otherwise
     end |> esc
 end
 function grab!(sym,ex::Expr)