Speaker
Description
There is a growing interest in differentiable simulations that have fast execution time and yield additional gradient information. Most differentiable accelerator codes use a single autodiff engine or truncated power series algebra library and then add standard optics and particle tracking equations on top. While benefitting from ease of use and fast development, such approaches give up low level control and can be limiting when computational and memory resources are heavily constrained. Intense space charge beams, especially in the context of ring tracking, are one such case. We present our prototyping work on a new framework for differentiable simulations, JACC (Jax for ACCelerators), that can do dynamic beamline code synthesis into optimized routines implemented in several libraries (Jax, PyTorch, NVIDIA Warp) or raw CUDA kernels. It can also stitch together custom pre-compiled or eager execution kernels. This helps with easy debugging while still enabling performance-critical kernel fusion, compilation, and distributed execution. Common Xsuite elements are implemented and a template for custom differentiable models provided. We present benchmarks with collective effects on CPU and GPU, with focus on scalable long-term tracking implementations that explore the compute/memory pareto front.