Flux agnostic @generated 3D conservative volume turbo kernel

src/Trixi.jl

@@ -255,7 +255,7 @@ export flux, flux_central, flux_lax_friedrichs, flux_hll, flux_hllc, flux_hlle,
        FluxRotated,
        flux_shima_etal_turbo, flux_ranocha_turbo,
        FluxUpwind,
-       FluxTracerEquationsCentral
+       FluxTracerEquationsCentral, FluxVolumeTurbo
 
 export splitting_steger_warming, splitting_vanleer_haenel,
        splitting_coirier_vanleer, splitting_lax_friedrichs,

src/auxiliary/math.jl

@@ -1,3 +1,4 @@
+using LoopVectorization: AbstractSIMD
 # By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
 # Since these FMAs can increase the performance of many numerical algorithms,
 # we need to opt-in explicitly.
@@ -135,7 +136,7 @@ end
     See also [`Trixi.set_log_type!`](@ref).
     """
     @inline log(x::Real) = x < zero(x) ? oftype(x, NaN) : Base.log(x)
-
+    @inline log(x::AbstractSIMD) = ifelse(x < zero(x), oftype(x, NaN), Base.log(x))
     @inline log(x::Float64) = ccall("llvm.log.f64", llvmcall, Float64, (Float64,), x)
     @inline log(x::Float32) = ccall("llvm.log.f32", llvmcall, Float32, (Float32,), x)
     @inline log(x::Float16) = ccall("llvm.log.f16", llvmcall, Float16, (Float16,), x)
@@ -205,6 +206,19 @@ Given ε = 1.0e-4, we use the following algorithm.
     end
 end
 
+@inline function ln_mean(x::AbstractSIMD, y::AbstractSIMD)
+    RealT = eltype(x)
+    epsilon_f2 = convert(RealT, 1.0e-4)
+    f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
+    return ifelse(f2 < epsilon_f2,
+                  (x + y) / @evalpoly(f2,
+                            convert(RealT, 2),
+                            convert(RealT, 2 / 3),
+                            convert(RealT, 2 / 5),
+                            convert(RealT, 2 / 7)),
+                  (y - x) / log(y / x))
+end
+
 """
     Trixi.inv_ln_mean(x::Real, y::Real)
 
@@ -231,6 +245,19 @@ multiplication.
     end
 end
 
+@inline function inv_ln_mean(x::AbstractSIMD, y::AbstractSIMD)
+    RealT = eltype(x)
+    epsilon_f2 = convert(RealT, 1.0e-4)
+    f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
+    return ifelse(f2 < epsilon_f2,
+                  @evalpoly(f2,
+                            convert(RealT, 2),
+                            convert(RealT, 2 / 3),
+                            convert(RealT, 2 / 5),
+                            convert(RealT, 2 / 7)) / (x + y),
+                  log(y / x) / (y - x))
+end
+
 # `Base.max` and `Base.min` perform additional checks for signed zeros and `NaN`s
 # which are not present in comparable functions in Fortran/C++. For example,
 # ```julia

src/equations/numerical_fluxes.jl

@@ -578,4 +578,59 @@ end
 end
 
 Base.show(io::IO, f::FluxUpwind) = print(io, "FluxUpwind(", f.splitting, ")")
+
+"""
+    FluxVolumeTurbo(volume_flux)
+
+Specialize the volume flux to use the SIMD instructions via LoopVectorization.jl
+"""
+struct FluxVolumeTurbo{VolumeFlux}
+    volume_flux::VolumeFlux
+    function FluxVolumeTurbo{VolumeFlux}(volume_flux) where {VolumeFlux}
+        return new{VolumeFlux}(volume_flux)
+    end
+end
+
+# Helper function for conservative systems.
+function FluxVolumeTurbo(volume_flux)
+    turbo_flux = combined_turbo_flux(volume_flux)
+    return FluxVolumeTurbo{typeof(turbo_flux)}(turbo_flux)
+end
+
+# By default the turbo flux has no specialization and re-uses
+# the numerical flux in terms of conservative variables.
+@inline combined_turbo_flux(volume_flux) = volume_flux
+
+# By default the turbo flux has the same number of precomputed variables
+# as the number of variables.
+@inline n_turbo_flux_aux_node_vars(volume_flux, equations) = Val(nvariables(equations))
+
+# Transform the conserved variables in precomputed auxiliary variables to speed up the computation
+# of the numerical flux. When no specialization is given, this gives cons2cons.
+@inline cons2fluxauxiliary(volume_flux, conserved_and_equations...) = Base.front(conserved_and_equations)
+
+# Numerical volume flux that recalls the plain volume flux when no specialization is given.
+@inline function volume_flux_turbo(volume_flux, aux_and_normals_and_equations...)
+    equations = last(aux_and_normals_and_equations)
+    volume_flux_turbo(volume_flux, have_nonconservative_terms(equations),
+                      aux_and_normals_and_equations...)
+end
+
+@inline function volume_flux_turbo(volume_flux, have_nonconservative_terms::False,
+                                   aux_and_normals_and_equations...)
+    equations = last(aux_and_normals_and_equations)
+    n = nvariables(equations)
+    u_ll = SVector(ntuple(v -> aux_and_normals_and_equations[v], Val(n)))
+    u_rr = SVector(ntuple(v -> aux_and_normals_and_equations[n + v], Val(n)))
+    normal_direction = SVector(aux_and_normals_and_equations[end - 3],
+                               aux_and_normals_and_equations[end - 2],
+                               aux_and_normals_and_equations[end - 1])
+    return volume_flux(u_ll, u_rr, normal_direction, equations)
+end
+
+# Allow LoopVectorization to use SIMD instructions on volume_flux_turbo and cons2fluxauxiliary
+LoopVectorization.can_turbo(::typeof(volume_flux_turbo), ::Val) = true
+LoopVectorization.can_turbo(::typeof(cons2fluxauxiliary), ::Val) = true
+
+Base.show(io::IO, f::FluxUpwind) = print(io, "FluxVolumeTurbo(", f.volume_flux, ")")
 end # @muladd

src/solvers/dgsem_structured/dg.jl

@@ -234,4 +234,5 @@ include("dg_2d_subcell_limiters.jl")
 # Specialized implementations used to improve performance
 include("dg_2d_compressible_euler.jl")
 include("dg_3d_compressible_euler.jl")
+include("dg_3d_turbo.jl")
 end # @muladd