wrap RHS to reduce latency with OrdinaryDiffEq.jl #1255
Conversation
| if integrator.f isa RHSWrapper | ||
| return integrator.f.semi | ||
| else | ||
| return integrator.p | ||
| end |
There was a problem hiding this comment.
OOC (out-of-curiosity): Is this type stable? Or would it make sense to dispatch on the type of integrator.f, which is probably a type parameter of integrator? Or am I completely off track since we always return a semi and thus the return type is stable?
There was a problem hiding this comment.
It should be type table since it should specialize on the type of integrator, which knows the type of integrator.f - but it may be worth checking to be sure - and running some benchmarks,of course
Edit: See below
There was a problem hiding this comment.
The wrapping with losing type information is done at the level of OrdinaryDiffEq.jl. This means that accessing the semidiscretization like this is indeed not type stable. Thus, I can think of the following options
- keep the code as it is in this PR, introducing type instabilities in callbacks accessing
semi - introduce a function barrier in all callbacks after accessing
semito mitigate this problem - store the semidiscretization
semiin the callbacks requiring it (if feasible, i.e., no full duplication, just a pointer) - keep the status quo without this PR
There was a problem hiding this comment.
Have you been able to assess the impact on latency/compilation times with this change yet? Maybe we should try to test this on Roci for a more stable test environment?
If it does improve things (in a relevant way) in terms of TTFX, I think keeping the status quo is not desirable. In that case, I would try to go for the least complex, then the least intrusive solution.
There was a problem hiding this comment.
Please see the post below: #1255 (comment)
|
Initial benchmarks julia -e '
@time using OrdinaryDiffEq;
@time using Trixi;
@time begin
equations = LinearScalarAdvectionEquation2D((0.2, -0.7));
solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs);
mesh = TreeMesh((-1.0, -1.0), (1.0, 1.0), initial_refinement_level=4, n_cells_max=30_000);
semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_convergence_test, solver);
ode = semidiscretize(semi, (0.0, 1.0))
end;
@time integrator = init(ode, SSPRK43());
'and julia -e '
@time using OrdinaryDiffEq;
@time using Trixi;
@time begin
equations = LinearScalarAdvectionEquation2D((0.2, -0.7));
solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs);
mesh = TreeMesh((-1.0, -1.0), (1.0, 1.0), initial_refinement_level=4, n_cells_max=30_000);
semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_convergence_test, solver);
ode = semidiscretize(semi, (0.0, 1.0))
end;
@time solve(ode, SSPRK43());
'yield the following preliminary results:
|
This sounds very promising! However, the loss of stable typing in the callbacks could potentially affect the runtime performance, couldn't it? Thus we probably should either run tests for these as well, or make sure we restore type stability through one of the methods you proposed. |
|
For now, I switched everything to function barriers after extacting the semidisretization in a type-unstable way. Compared to the latest release of Trixi.jl, I get the following results with the current state of this PR.
vs. in the latest release (with Julia 1.8.2).
vs. in the latest release. |
|
@sloede I think this is ready for a first review. It would be great to get some help running tests and benchmarks in the wild with this. |
|
There is a huge problem with plotting. If we use this, the ODE problem (and solution) will not carry any type information from Trixi.jl anymore. That means dispatching on the ODE solution type for our plotting will be type piracy and overwrite existing plotting methods. I guess this basically settles this - we cannot use this approach 😞 |
Just such that I understand it correctly: With the change proposed here, |
|
Basically yes - we could make this work by breaking |
Hm. I see three options
Maybe we should collect the changes in the runtime numbers and put this up for discussion at a meeting? I also tend to dislike breaking |
Sounds reasonable. Any help is welcome 🙂 |
|
Would it possibly help if we, instead of storing the full |
This could indeed work but would also require that we perform the precompilation of ODE solvers from OrdinaryDiffEq.jl within Trixi.jl. In particular, we can't benefit from their precompilation directly. |
Could this be fixed with upstream support? E.g., if they added a type parameter in an appropriate location such that it will not break precompilation benefits but allow us to to "tag" the return value somehow? I'm just brainstorming, probably what I'm trying to do is fundamentally impossible with the type system 😅 |
|
Yeah, I guess it's impossible. They would need to precompile for something that is created by Trixi.jl - which is impossible since they don't depend on us. |
|
I am not sure, but is this something where the relatively new https://github.com/SciML/SciMLStructures.jl could help in the long run? |
I'm not sure. |
Closes #1241
TODO
EulerAcousticsCouplingCallbackset_preferences!(UUID("1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"), "PrecompileLowStorage" => true))