Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add preference to disable LoopVectorization #2295

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
19 changes: 13 additions & 6 deletions src/Trixi.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ See also: [trixi-framework/Trixi.jl](https://github.com/trixi-framework/Trixi.jl
"""
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_LOOPVECTORIZATION = @load_preference("loop_vectorization", true)

# Include other packages that are used in Trixi.jl
# (standard library packages first, other packages next, all of them sorted alphabetically)

Expand Down Expand Up @@ -53,7 +59,13 @@ using FillArrays: Ones, Zeros
using ForwardDiff: ForwardDiff
using HDF5: HDF5, h5open, attributes, create_dataset, datatype, dataspace
using LinearMaps: LinearMap
using LoopVectorization: LoopVectorization, @turbo, indices
if _PREFERENCE_LOOPVECTORIZATION
using LoopVectorization: LoopVectorization, @turbo, indices
else
using LoopVectorization: LoopVectorization, indices
include("auxiliary/mock_turbo.jl")
end

using StaticArrayInterface: static_length # used by LoopVectorization
using MuladdMacro: @muladd
using Octavian: Octavian, matmul!
Expand Down Expand Up @@ -81,11 +93,6 @@ using SimpleUnPack: @pack!
using DataStructures: BinaryHeap, FasterForward, extract_all!

using UUIDs: UUID
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)

# finite difference SBP operators
using SummationByPartsOperators: AbstractDerivativeOperator,
Expand Down
15 changes: 15 additions & 0 deletions src/auxiliary/math.jl
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,21 @@ function set_polyester!(toggle::Bool; force = true)
@info "Please restart Julia and reload Trixi.jl for the `polyester` change to take effect"
end

"""
Trixi.set_loop_vectorization!(toggle::Bool; force = true)

Toggle the usage of [LoopVectorization.jl](https://github.com/JuliaSIMD/LoopVectorization.jl).
By default, LoopVectorization.jl is enabled, but it can
be useful for performance comparisons to switch to the Julia core backend.

This does not fully disable LoopVectorization.jl,
but only its internal use as part of Trixi.jl.
"""
function set_loop_vectorization!(toggle::Bool; force = true)
set_preferences!(TRIXI_UUID, "loop_vectorization" => toggle, force = force)
@info "Please restart Julia and reload Trixi.jl for the `loop_vectorization` change to take effect"
end

"""
Trixi.set_sqrt_type!(type; force = true)

Expand Down
28 changes: 28 additions & 0 deletions src/auxiliary/mock_turbo.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
macro turbo(exprs...)
body = nothing
for expr in exprs
if Meta.isexpr(expr, :for)
body = expr
end
end
@assert body !== nothing

function insert_loopinfo!(expr)
recurse = Meta.isexpr(expr, :for) || Meta.isexpr(expr, :block) ||
Meta.isexpr(expr, :let)
if recurse
foreach(insert_loopinfo!, expr.args)
end
if Meta.isexpr(expr, :for)
# TODO: Should we insert LLVM loopinfo or `julia.ivdep`?
push!(expr.args, Expr(:loopinfo, Symbol("julia.simdloop")))
end
Comment on lines +16 to +19
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a few comments here, please? Is this a public API, or are you doing something that may break at any time?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's an internal API (https://github.com/JuliaLang/julia/blob/9f7cdfdc28290063ee021a5c6e180e82feee32c8/doc/src/devdocs/ast.md?plain=1#L464).

Most notably, it is used for the @simd macro https://github.com/JuliaLang/julia/blob/9f7cdfdc28290063ee021a5c6e180e82feee32c8/base/simdloop.jl#L79

There is also an old PR of mine that exposes some more of the knobs available through this API, that is used by a couple of consumers JuliaLang/julia#31376

It is an internal API that directly interacts with the LLVM optimizer. It has not seen recent change, and I am not predicting any, worst case these annotations may be silently dropped.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. It's fine with me in this case. Could you please add a few comments/links to the code, check the TODO note, and let me know when this PR is finished from your point of view?

end
insert_loopinfo!(body)

body = Expr(:block,
Expr(:inbounds, true),
body,
Expr(:inbounds, :pop))
return esc(body)
end
3 changes: 3 additions & 0 deletions src/callbacks_step/summary.jl
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,9 @@ function initialize_summary_callback(cb::DiscreteCallback, u, t, integrator;
if !_PREFERENCE_POLYESTER
push!(setup, "Polyester" => "disabled")
end
if !_PREFERENCE_LOOPVECTORIZATION
push!(setup, "LoopVectorization" => "disabled")
end
if mpi_isparallel()
push!(setup,
"#MPI ranks" => mpi_nranks())
Expand Down
2 changes: 1 addition & 1 deletion src/solvers/dg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -672,7 +672,7 @@ end
nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache)
end
# See comments on the DGSEM version above
if LoopVectorization.check_args(u_ode)
if _PREFERENCE_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
Expand Down
Loading