Support serial execution and introduce backend preference (#2476)

* Support serial backend

Co-authored-by: Valentin Churavy <[email protected]>

* Rebase post 2212 and KernelAbstraction support

* Update src/auxiliary/containers.jl

* Apply suggestions from code review

Co-authored-by: Hendrik Ranocha <[email protected]>

* add news and bump version

* Update NEWS.md

* Improve summary for threading backend

---------

Co-authored-by: afilogo <[email protected]>
Co-authored-by: Hendrik Ranocha <[email protected]>

Author: 3 people

NEWS.md

@@ -5,6 +5,17 @@ Trixi.jl follows the interpretation of
 used in the Julia ecosystem. Notable changes will be documented in this file
 for human readability.
 
+
+## Changes when updating to v0.13 from v0.12.x
+
+#### Changed
+
+- The `polyester` preference got merged with the `native_threading` preference and the `Trixi.set_polyester!` function got renamed to `Trixi.set_threading_backend!` ([#2476]).
+
+#### Deprecated
+
+#### Removed
+
 ## Changes in the v0.12 lifecycle
 
 #### Added

Project.toml

@@ -1,7 +1,7 @@
 name = "Trixi"
 uuid = "a7f1ee26-1774-49b1-8366-f1abc58fbfcb"
 authors = ["Michael Schlottke-Lakemper <[email protected]>", "Gregor Gassner <[email protected]>", "Hendrik Ranocha <[email protected]>", "Andrew R. Winters <[email protected]>", "Jesse Chan <[email protected]>", "Andrés Rueda-Ramírez <[email protected]>"]
-version = "0.12.8-DEV"
+version = "0.13.0-DEV"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

src/Trixi.jl

@@ -18,9 +18,8 @@ module Trixi
 using Preferences: @load_preference, set_preferences!
 const _PREFERENCE_SQRT = @load_preference("sqrt", "sqrt_Trixi_NaN")
 const _PREFERENCE_LOG = @load_preference("log", "log_Trixi_NaN")
-const _PREFERENCE_POLYESTER = @load_preference("polyester", true)
+const _PREFERENCE_THREADING = @load_preference("backend", :polyester)
 const _PREFERENCE_LOOPVECTORIZATION = @load_preference("loop_vectorization", true)
-const _PREFERENCE_USE_NATIVE_THREADING = @load_preference("native_threading", true)
 
 # Include other packages that are used in Trixi.jl
 # (standard library packages first, other packages next, all of them sorted alphabetically)

src/auxiliary/auxiliary.jl

@@ -206,16 +206,19 @@ and [https://discourse.julialang.org/t/threads-threads-with-one-thread-how-to-re
 macro threaded(expr)
     # !!! danger "Heisenbug"
     #     Look at the comments for `wrap_array` when considering to change this macro.
-    expr = if _PREFERENCE_POLYESTER
+    expr = @static if _PREFERENCE_THREADING === :polyester
         # Currently using `@batch` from Polyester.jl is more efficient,
         # bypasses the Julia task scheduler and provides parallelization with less overhead.
         quote
             $Trixi.@batch $(expr)
         end
-    else
+    elseif _PREFERENCE_THREADING === :static ||
+           _PREFERENCE_THREADING === :kernelabstractions
         # The following code is a simple version using only `Threads.@threads` from the
         # standard library with an additional check whether only a single thread is used
         # to reduce some overhead (and allocations) for serial execution.
+        # If we want to execute on KernelAbstractions, we use the static backend here to fallback on,
+        # for loops that do not yet support GPU execution.
         quote
             let
                 if $Threads.nthreads() == 1
@@ -225,6 +228,10 @@ macro threaded(expr)
                 end
             end
         end
+    elseif _PREFERENCE_THREADING === :serial
+        quote
+            $(expr)
+        end
     end
     # Use `esc(quote ... end)` for nested macro calls as suggested in
     # https://github.com/JuliaLang/julia/issues/23221

src/auxiliary/containers.jl

@@ -358,9 +358,9 @@ Return the computational backend for `x`, which is either a KernelAbstractions b
 If the backend is `nothing`, the default multi-threaded CPU backend is used.
 """
 function trixi_backend(x)
-    # TODO: https://github.com/trixi-framework/Trixi.jl/pull/2417
-    if (_PREFERENCE_POLYESTER && LoopVectorization.check_args(x)) ||
-       (_PREFERENCE_USE_NATIVE_THREADING && get_backend(x) isa KernelAbstractions.CPU)
+    if (_PREFERENCE_THREADING === :polyester && LoopVectorization.check_args(x)) ||
+       (_PREFERENCE_THREADING !== :kernelabstractions &&
+        get_backend(x) isa KernelAbstractions.CPU)
         return nothing
     end
     return get_backend(x)

src/auxiliary/math.jl

@@ -8,18 +8,27 @@
 const TRIXI_UUID = UUID("a7f1ee26-1774-49b1-8366-f1abc58fbfcb")
 
 """
-    Trixi.set_polyester!(toggle::Bool; force = true)
-
-Toggle the usage of [Polyester.jl](https://github.com/JuliaSIMD/Polyester.jl) for multithreading.
-By default, Polyester.jl is enabled, but it can
-be useful for performance comparisons to switch to the Julia core backend.
-
-This does not fully disable Polyester.jl,
-but only its use as part of Trixi.jl's [`@threaded`](@ref) macro.
+    Trixi.set_threading_backend!(backend::Symbol; force = true)
+
+Toggle and/or switch backend behavior used in multithreaded loops inside Trixi.jl.
+The selected backend affects the behavior of Trixi.jl's [`@threaded`](@ref) macro, which is used
+throughout the codebase for parallel loops. By default, Polyester.jl is enabled for
+optimal performance, but switching backends can be useful for comparisons or debugging.
+
+# Available backends
+- `:polyester`: Uses the default [Polyester.jl](https://github.com/JuliaSIMD/Polyester.jl)
+- `:static`: Uses Julia's built-in static thread scheduling via `Threads.@threads :static`
+- `:serial`: Disables threading, executing loops serially
+- `:kernelabstractions`: Preferentially use the [KernelAbstractions.jl](https://github.com/JuliaGPU/KernelAbstractions.jl)
+  kernels written in Trixi.jl, falling back to `:static` execution.
 """
-function set_polyester!(toggle::Bool; force = true)
-    set_preferences!(TRIXI_UUID, "polyester" => toggle, force = force)
-    @info "Please restart Julia and reload Trixi.jl for the `polyester` change to take effect"
+function set_threading_backend!(backend::Symbol = :polyester; force = true)
+    valid_backends = (:polyester, :static, :serial, :kernelabstractions)
+    if !(backend in valid_backends)
+        throw(ArgumentError("Invalid threading backend: $(backend). Current options are: $(join(valid_backends, ", "))"))
+    end
+    set_preferences!(TRIXI_UUID, "backend" => backend, force = force)
+    @info "Please restart Julia and reload Trixi.jl for the `backend` change to take effect"
 end
 
 """

src/callbacks_step/summary.jl

@@ -209,9 +209,7 @@ function initialize_summary_callback(cb::DiscreteCallback, u, t, integrator;
 
     # technical details
     setup = Pair{String, Any}["#threads" => Threads.nthreads()]
-    if !_PREFERENCE_POLYESTER
-        push!(setup, "Polyester" => "disabled")
-    end
+    push!(setup, "threading backend" => string(_PREFERENCE_THREADING))
     if !_PREFERENCE_LOOPVECTORIZATION
         push!(setup, "LoopVectorization" => "disabled")
     end

src/solvers/dg.jl

@@ -669,7 +669,7 @@ end
     # since LoopVectorization does not support `ForwardDiff.Dual`s. Hence, we use
     # optimized `PtrArray`s whenever possible and fall back to plain `Array`s
     # otherwise.
-    if _PREFERENCE_POLYESTER && LoopVectorization.check_args(u_ode)
+    if _PREFERENCE_THREADING === :polyester && LoopVectorization.check_args(u_ode)
         # This version using `PtrArray`s from StrideArrays.jl is very fast and
         # does not result in allocations.
         #
@@ -704,7 +704,7 @@ end
                 nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache)
     end
     # See comments on the DGSEM version above
-    if _PREFERENCE_POLYESTER && LoopVectorization.check_args(u_ode)
+    if _PREFERENCE_THREADING === :polyester && LoopVectorization.check_args(u_ode)
         # Here, we do not specialize on the number of nodes using `StaticInt` since
         # - it will not be type stable (SBP operators just store it as a runtime value)
         # - FD methods tend to use high node counts