Make FDTDX extremely fast

Our goal is to increase the simulation and gradient computation speed of fdtdx even further. Currently, the simulation is memory bound, because for a single fdtd update, the E and H field of 6 surrounding cells need to be loaded. This leads to the GPU spendig most of its time loading data instead of performing computations. To address this, I propose the following changes:

• Instead of performing a single time step at a time for the whole simulation volume, the volume should be divided into "DiamondCandy" structures. This allows performing more computations with a single data load and less communication. Therefore, the multi-GPU speed would also increase drastically from this change.
• However, this would require massive changes to the fdtdx codebase. Firstly, sources, detectors and boundaries need to be able to work with a single cell rather than the whole volume of the object. This also requires possibly splitting the ArrayContainer structure into multiple subcomponents (?).
• Additionally, different DiamondCandy "kernels" would need to be exectued depending on if there is a source, detector or boundary within the current kernel volume. This would mean that we need to first run a separate program to determine the computational graph of all kernels, then schedule the different kernels manually. Implementing this within JAX could be problematic because long computational graphs will also lead to long compile times. However, I believe that a fast simulation usually is more important than the compile time, which only needs to be executed once.
• Moreover, we should make sure that gradient computation still works under this scheme. This means, that we still should be able to perform time reversible simulations by reversing the computational graph of the kernel.

If someone has suggestions for a good way to implement this, feel free to post it here. Any help is appreciated. Due to other projects, I might only be able to work on this feature start of next year and its probably going to take a few month of implementation effort.

References:

• OpenEMS Notes on Temoral tiling: https://github.com/thliebig/openEMS-Project/discussions/154#discussioncomment-8183199
• DiamondCandy: https://dl.acm.org/doi/10.1145/3407947.3407951
• DiamondTiling: https://www.academia.edu/10021612/Diamond_Tiling_A_Tiling_Framework_for_Time_iterated_Scientic_Applications
• LRnLA with Dispersion: https://library.keldysh.ru/prep_vw.asp?lg=e&pid=3580
• 1D: https://www.mdpi.com/2297-8747/26/3/52