[Performance]: Try to 5-10x the performance on CPU

[QuentinWach]

This is more of a study / experiment.
Can we optimize BEAMZ to be extremely fast even on CPUs?

• 0.1 - 0.3 GCUPS (we already achieve this, making BEAMZ already faster than Meep).
• 1 - 2 GCUPS.
• 3 GCUPS for large simulations and a really good CPU.

[QuentinWach]

I dug into this from the JAX side and did some local measurements in the current compiled engine.

Summary of where we are:

• JAX does support multiple logical CPU devices. This can be enabled with JAX_NUM_CPU_DEVICES (or --xla_force_host_platform_device_count) before JAX initializes.
• I verified locally that JAX_NUM_CPU_DEVICES=8 gives 8 CpuDevices on this machine.
• That means CPU-side SPMD/domain decomposition is on the table. However, this is not the same thing as SIMD. pmap/shard_map gives us parallel execution across logical CPU devices; vectorization still has to happen inside each shard via XLA/LLVM.

What I measured so far in the current codebase:

• Naive "multiple steps at once" by using lax.scan(..., unroll=k) over the full domain did not help in our current layout.
• Example local results:
  • 48^3, 80 steps: fori_loop baseline ~0.290s, scan(unroll=2) ~0.350s
  • 64^3, 80 steps: fori_loop baseline ~0.337s, scan(unroll=2) ~0.515s
• This matches the HLO growth we saw: more slices / larger code, but not enough locality improvement to offset the extra live temporaries.
• The best CPU win I found so far was not temporal unrolling. It was eliminating the hot-loop single-slab dynamic-update-slice in source injection and using a dense single-slab path on CPU. That materially improved the compiled CPU path and removed the last dynamic-update-slice from HLO in the simple Gaussian-source case.

My current read on the optimization direction:

1. Multi-CPU-device decomposition is worth testing

A realistic CPU parallelism approach here is to shard the simulation domain across multiple logical CPU devices.

Most likely first version:

• shard along z into slabs
• keep one local FDTD kernel per shard
• exchange left/right halo regions between neighboring shards before H and E updates
• keep sources and monitors local where possible
• only reduce globally when needed

This could give a meaningful speedup on large enough simulations because:

• each shard works on a smaller working set
• we can use more host cores concurrently
• the memory traffic per kernel launch is more localized

Main caveat:

• all CPU devices still share the same memory subsystem, so scaling will stop once memory bandwidth saturates
• for small domains, halo overhead and scheduling overhead will dominate

So I do think JAX multi-CPU-device execution is worth exploring here, but only as domain decomposition for large problems, not as a universal speed button.

2. The likely 3-5x path is tiled temporal blocking, not whole-grid scan unrolling

The fact that naive scan unrolling regressed does not mean multi-step is wrong.
It means we are doing the wrong version of it.

The version that is plausible for big CPU gains is:

• spatial tiling
• then advancing each tile for k timesteps locally
• with halo depth k
• so the same tile data stays hot in cache across multiple steps

That is much closer to the classic high-performance CPU stencil / FDTD playbook.

For Yee-grid FDTD, this means:

• split each shard into 3D tiles
• for a k-step temporal block, each tile needs a spatial halo of depth k in the blocked directions
• compute k timesteps on tile+halo
• write back only the valid interior after the block

This is fundamentally different from just asking JAX to unroll the full-domain loop.
The goal is cache reuse, not just fewer loop iterations in HLO.

3. Proposed plan toward 1 GCUPS+

Phase A: establish the right baseline on large domains

Use the current compiled benchmark harness and collect numbers on at least:

• 128^3
• 160^3
• 192^3

with:

• current single-device compiled path
• CPU dense single-slab source path enabled
• profiler traces + HLO stats

Reason: 64^3 is too small to say much about CPU peak behavior.

Phase B: prototype multi-CPU-device sharding

Implement a first sharded version using JAX CPU devices:

• 1D slab decomposition along z
• 4 and 8 device configurations
• explicit halo exchange in the compiled step loop
• measure scaling vs single-device baseline

Success criterion:

• on large simulations, get at least 3x speedup and ideally 4-5x

If this does not happen, the halo/memory-system costs are probably too high and we should pivot harder to per-device temporal blocking.

Phase C: add tiled temporal blocking inside each shard

Implement a new blocked kernel for the hot field update path:

• choose a tile size (tz, ty, tx)
• choose k = 2 first for temporal depth
• operate on tile + halo
• write back only the interior after k steps
• keep sources/monitors/material updates out of the blocked kernel initially

This is the most likely path to the bigger CPU gains.

Phase D: reintroduce side features carefully

After the blocked stencil kernel is good:

• reintroduce sparse source injection
• reintroduce material model updates
• reintroduce monitor accumulation
• measure regressions after each step

4. My current expectation

I do not expect pmap alone to get us to 1 GCUPS.
I do think the combination of:

• multi-CPU-device sharding for large domains
• tile-local temporal blocking within each shard
• already-removing hot-path dynamic slices

is a plausible route to 1 GCUPS+.

Roughly:

• current optimized single-device path: baseline
• multi-CPU-device sharding: likely first large multiplier
• tile-local temporal blocking: second large multiplier
• feature-specific cleanup: final polishing to preserve gains in real simulations

If we want to pursue this issue seriously, I suggest we explicitly split it into:

1. CPU multi-device sharding prototype
2. tile-local temporal blocking prototype
3. target benchmark matrix / GCUPS scoreboard

That will make it easier to tell which part is actually buying us performance.