2x speedup for default_example()

[cwittens]

Yesterday evening I spent a few hours of compute on Fable 5 (when it was still available) to try to improve the speed of Trixi (just to test the new model) and got, among other things, an ~8x speedup for the weak volume kernel (overall ~2x speedup for default_example() on CPU) by, as far as I understand it, leveraging a better memory access pattern.

A similar thing, along with a few small improvements, was also done for the P4estMesh, also giving about a 2x speedup overall.
All tests passed locally (at least all those that were already passing before my changes) and the difference in the final solution is within machine precision.

My problem is that I don't understand the vibe-coded code well enough to want to make a PR myself andI dont want to have someone else look it over. I also don't know the design policies of the package well enough to judge whether, for example, a CPU-specialized function is acceptable, given code maintainability vs. speed trade-off.

But I thought I should still at least flag that a ~2x speedup is potentially possible, and even if the code Fable produced may not be the way to go, at least the idea behind it seems not completely wrong (and can maybe also be generalized to not just the two kernels).

Here the code: https://github.com/cwittens/Trixi.jl/tree/fabel5-improvements
And the diff: https://github.com/cwittens/Trixi.jl/pull/1/changes

And here some benchmarks:

# default example using tspan = (0 , 100) Fable 5 improved version:
────────────────────────────────────────────────────────────────────────────────────
             Trixi.jl                      Time                    Allocations      
                                  ───────────────────────   ────────────────────────
        Tot / % measured:              242ms /  83.1%           14.2MiB / 100.0%    

Section                   ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────────
rhs!                       9.00k    108ms   53.5%  12.0μs   11.7MiB   82.1%  1.33KiB
  interface flux           9.00k   31.0ms   15.4%  3.44μs   2.61MiB   18.3%     304B
  volume integral          9.00k   16.6ms    8.2%  1.84μs   2.47MiB   17.4%     288B
  reset ∂u/∂t              9.00k   14.0ms    6.9%  1.55μs     0.00B    0.0%    0.00B
  ~rhs!~                   9.00k   13.7ms    6.8%  1.53μs   4.78KiB    0.0%    0.54B
  surface integral         9.00k   12.8ms    6.4%  1.43μs   2.20MiB   15.4%     256B
  prolong2interfaces       9.00k   11.0ms    5.4%  1.22μs   2.47MiB   17.4%     288B
  Jacobian                 9.00k   6.69ms    3.3%   743ns   1.92MiB   13.5%     224B
  prolong2mortars          9.00k    533μs    0.3%  59.2ns     0.00B    0.0%    0.00B
  prolong2boundaries       9.00k    458μs    0.2%  50.9ns     0.00B    0.0%    0.00B
  mortar flux              9.00k    448μs    0.2%  49.8ns     0.00B    0.0%    0.00B
  boundary flux            9.00k    262μs    0.1%  29.2ns     0.00B    0.0%    0.00B
  source terms             9.00k    257μs    0.1%  28.5ns     0.00B    0.0%    0.00B
I/O                           20   51.1ms   25.4%  2.56ms   1.90MiB   13.3%  97.2KiB
  save solution               19   46.5ms   23.1%  2.45ms   1.87MiB   13.2%   101KiB
  ~I/O~                       20   4.63ms    2.3%   231μs   26.7KiB    0.2%  1.33KiB
  get node variables          19   10.5μs    0.0%   553ns     0.00B    0.0%    0.00B
  save mesh                   19   3.30μs    0.0%   174ns     0.00B    0.0%    0.00B
  get element variables       19   2.50μs    0.0%   132ns     0.00B    0.0%    0.00B
analyze solution              19   39.5ms   19.6%  2.08ms    500KiB    3.4%  26.3KiB
calculate dt               1.80k   3.04ms    1.5%  1.69μs    169KiB    1.2%    96.0B
────────────────────────────────────────────────────────────────────────────────────

# default example using tspan = (0 , 100) Trixi version:
────────────────────────────────────────────────────────────────────────────────────
             Trixi.jl                      Time                    Allocations      
                                  ───────────────────────   ────────────────────────
        Tot / % measured:              344ms /  88.3%           14.2MiB / 100.0%    

Section                   ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────────
rhs!                       9.00k    210ms   69.3%  23.4μs   11.7MiB   82.1%  1.33KiB
  volume integral          9.00k    109ms   35.8%  12.1μs   2.47MiB   17.4%     288B
  interface flux           9.00k   30.0ms    9.9%  3.33μs   2.61MiB   18.3%     304B
  surface integral         9.00k   20.8ms    6.9%  2.31μs   2.20MiB   15.4%     256B
  prolong2interfaces       9.00k   19.3ms    6.4%  2.14μs   2.47MiB   17.4%     288B
  ~rhs!~                   9.00k   14.4ms    4.8%  1.60μs   4.78KiB    0.0%    0.54B
  reset ∂u/∂t              9.00k   9.32ms    3.1%  1.04μs     0.00B    0.0%    0.00B
  Jacobian                 9.00k   6.23ms    2.1%   692ns   1.92MiB   13.5%     224B
  prolong2mortars          9.00k    446μs    0.1%  49.5ns     0.00B    0.0%    0.00B
  mortar flux              9.00k    427μs    0.1%  47.5ns     0.00B    0.0%    0.00B
  prolong2boundaries       9.00k    416μs    0.1%  46.2ns     0.00B    0.0%    0.00B
  boundary flux            9.00k    246μs    0.1%  27.4ns     0.00B    0.0%    0.00B
  source terms             9.00k    236μs    0.1%  26.2ns     0.00B    0.0%    0.00B
I/O                           20   55.4ms   18.3%  2.77ms   1.90MiB   13.3%  97.2KiB
  save solution               19   49.6ms   16.3%  2.61ms   1.87MiB   13.2%   101KiB
  ~I/O~                       20   5.77ms    1.9%   289μs   26.7KiB    0.2%  1.33KiB
  get node variables          19   16.1μs    0.0%   847ns     0.00B    0.0%    0.00B
  save mesh                   19   4.80μs    0.0%   253ns     0.00B    0.0%    0.00B
  get element variables       19   1.80μs    0.0%  94.7ns     0.00B    0.0%    0.00B
analyze solution              19   34.9ms   11.5%  1.84ms    499KiB    3.4%  26.2KiB
calculate dt               1.80k   2.90ms    1.0%  1.61μs    169KiB    1.2%    96.0B
────────────────────────────────────────────────────────────────────────────────────


# /p4est_2d_dgsem/elixir_advection_basic.jl example using tspan = (0 , 100) Fable 5 improved version:
────────────────────────────────────────────────────────────────────────────────────
             Trixi.jl                      Time                    Allocations      
                                  ───────────────────────   ────────────────────────
        Tot / % measured:              366ms /  87.8%           13.2MiB / 100.0%    

Section                   ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────────
rhs!                       9.00k    209ms   65.1%  23.2μs   9.62MiB   72.9%  1.09KiB
  interface flux           9.00k   66.8ms   20.8%  7.42μs   1.92MiB   14.6%     224B
  volume integral          9.00k   45.7ms   14.2%  5.07μs   2.47MiB   18.7%     288B
  prolong2interfaces       9.00k   36.0ms   11.2%  3.99μs    844KiB    6.2%    96.0B
  surface integral         9.00k   17.2ms    5.4%  1.91μs   2.33MiB   17.7%     272B
  ~rhs!~                   9.00k   16.4ms    5.1%  1.82μs   4.78KiB    0.0%    0.54B
  reset ∂u/∂t              9.00k   12.5ms    3.9%  1.39μs     0.00B    0.0%    0.00B
  Jacobian                 9.00k   12.3ms    3.8%  1.36μs   2.06MiB   15.6%     240B
  mortar flux              9.00k    588μs    0.2%  65.3ns     0.00B    0.0%    0.00B
  prolong2mortars          9.00k    552μs    0.2%  61.3ns     0.00B    0.0%    0.00B
  prolong2boundaries       9.00k    533μs    0.2%  59.3ns     0.00B    0.0%    0.00B
  source terms             9.00k    288μs    0.1%  31.9ns     0.00B    0.0%    0.00B
  boundary flux            9.00k    287μs    0.1%  31.9ns     0.00B    0.0%    0.00B
I/O                           20   56.2ms   17.5%  2.81ms   1.90MiB   14.4%  97.0KiB
  save solution               19   49.3ms   15.4%  2.60ms   1.87MiB   14.2%   101KiB
  ~I/O~                       20   6.88ms    2.1%   344μs   21.3KiB    0.2%  1.07KiB
  get node variables          19   14.3μs    0.0%   753ns     0.00B    0.0%    0.00B
  save mesh                   19   4.60μs    0.0%   242ns     0.00B    0.0%    0.00B
  get element variables       19   3.20μs    0.0%   168ns     0.00B    0.0%    0.00B
analyze solution              19   38.6ms   12.0%  2.03ms   1.19MiB    9.0%  64.0KiB
calculate dt               1.80k   17.3ms    5.4%  9.62μs    507KiB    3.7%     288B
────────────────────────────────────────────────────────────────────────────────────



# /p4est_2d_dgsem/elixir_advection_basic.jl example using tspan = (0 , 100) Trixi version:
────────────────────────────────────────────────────────────────────────────────────
             Trixi.jl                      Time                    Allocations      
                                  ───────────────────────   ────────────────────────
        Tot / % measured:              500ms /  91.5%           13.2MiB / 100.0%    

Section                   ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────────
rhs!                       9.00k    342ms   74.8%  38.0μs   9.62MiB   72.9%  1.09KiB
  volume integral          9.00k    139ms   30.4%  15.4μs   2.47MiB   18.7%     288B
  interface flux           9.00k   82.9ms   18.1%  9.21μs   1.92MiB   14.6%     224B
  prolong2interfaces       9.00k   48.4ms   10.6%  5.37μs    844KiB    6.2%    96.0B
  Jacobian                 9.00k   21.9ms    4.8%  2.43μs   2.06MiB   15.6%     240B
  surface integral         9.00k   21.1ms    4.6%  2.34μs   2.33MiB   17.7%     272B
  ~rhs!~                   9.00k   15.5ms    3.4%  1.72μs   4.78KiB    0.0%    0.54B
  reset ∂u/∂t              9.00k   11.1ms    2.4%  1.23μs     0.00B    0.0%    0.00B
  mortar flux              9.00k    569μs    0.1%  63.2ns     0.00B    0.0%    0.00B
  prolong2mortars          9.00k    508μs    0.1%  56.4ns     0.00B    0.0%    0.00B
  prolong2boundaries       9.00k    501μs    0.1%  55.7ns     0.00B    0.0%    0.00B
  boundary flux            9.00k    303μs    0.1%  33.6ns     0.00B    0.0%    0.00B
  source terms             9.00k    285μs    0.1%  31.7ns     0.00B    0.0%    0.00B
I/O                           20   57.3ms   12.5%  2.86ms   1.90MiB   14.4%  97.0KiB
  save solution               19   51.8ms   11.3%  2.72ms   1.87MiB   14.2%   101KiB
  ~I/O~                       20   5.49ms    1.2%   274μs   21.3KiB    0.2%  1.07KiB
  get node variables          19   10.2μs    0.0%   537ns     0.00B    0.0%    0.00B
  save mesh                   19   3.60μs    0.0%   189ns     0.00B    0.0%    0.00B
  get element variables       19   3.10μs    0.0%   163ns     0.00B    0.0%    0.00B
analyze solution              19   42.0ms    9.2%  2.21ms   1.19MiB    9.0%  64.0KiB
calculate dt               1.80k   16.1ms    3.5%  8.92μs    507KiB    3.7%     288B
────────────────────────────────────────────────────────────────────────────────────

Here some AI slop explaining the change:

What was optimized and how

weak_form_kernel! — the volume integral (biggest win)

This kernel computes the DG volume contribution: for each element, it maps node values through the flux function and applies the derivative operator. The original code used a scatter-accumulate pattern: for each source node (i,j), it looped over all output nodes and added partial contributions to du. That means each du entry got 2n read-modify-write hits through memory.

The rewrite uses a flux buffer + register gather pattern:

• First pass: compute all fluxes for the element and store them in a stack-allocated MMatrix{n,n,SVector} buffer (two buffers for x and y directions).
• Convert the n×n derivative matrix to SMatrix{n,n} once per element — this lets the compiler fully unroll the inner loops.
• Second pass: for each output node (i,j), accumulate its full contribution du_node in a register using one fused loop over k, then write to du[v,i,j,element] exactly once.

Result: O(n²) memory accesses instead of O(n³). This is dispatched only for du::PtrArray, u::PtrArray (the CPU threading path), so GPU and AD paths fall back to the generic method untouched. Applied to TreeMesh 1D/2D/3D and curved meshes (StructuredMesh, P4estMesh, etc.).

[ranocha]

Thanks a lot! It would be great if you could measure the latency introduced by this change. We had some issues some time ago (#516) that were fixed by using standard Matrix types instaead of SMatrix types for serivative matrices.

@inbounds is another story @sloede can talk about more. We know that there are some parts that can be made faster by using @inbounds, but it also increases code complexity a bit while reading the code. Thus, we have to make some choice. Maybe we can reconsider this decision.

Changing the access pattern of the weak_form_kernel! to compute all fluxes first and then update them may be an interesting option to check. Could you please benchmark how this influences the kernel for more complicated fluxes, e.g., the compressible Euler equations?

[ranocha]

@MarcoArtiano suggested to also check the influence of each of these modifications alone.

[cwittens]

Yes I also already thought about that, see how much eg @inbounds actually does compare to the other stuff, I just haven't gotten around to trying it out. I will try to do it before my vacation but can't promise

[cwittens]

nvm the last commend lol

I will check the results more closely tomorrow but the TL;DR is
for Euler it only did ~10% speedup because it was compute and not memory
the @inbounds did basically nothing and can be dropped
and the new SMatrix stuff did not significantly increase the startup time (only ~2% max) and no scaling with polynomial size

[ranocha]

for Euler it only did ~10% speedup because it was compute and not memory

It would be great if you could always test three different resolutions (like "fits definitely in the (L2) cache", "becomes a bit too big to have everything in the cache", "definitely not completely in the cache"). This should move the barrier from compute- to memory-bound (if it is compute-bound for small resolutions).