Automatic Differentiation in High Performance Hydrodynamic Simulation
Please login to view abstract download link
Recent advances in scientific computing for topology optimization and machine learning have intensified the need for scalable, robust gradient computation in large-scale simulations. Automatic differentiation (AD) provides a robust way to algorithmically compute derivatives and assemble gradients of complex compositions of functions. Frameworks such as JAX and PyTorch, originally developed for machine learning, have demonstrated that AD can scale to very large problems and enable rapid, flexible experimentation. This capability reduces development time and supports the use of more sophisticated mathematical models. In this work, we bring these ideas into the high-performance mechanics setting by integrating AD-based gradient computation into the production multiphysics code MARBL at Lawrence Livermore National Laboratory. We focus on time dependent, mechanics driven applications where gradients are needed for topology optimization, material parameter inference, and optimal control of initial value problems. By exploiting the sequential structure of time integration, we design AD algorithms that compute sensitivities efficiently with respect to both control and material parameters, while remaining compatible with MARBL’s existing data structures, parallelization strategy, and performance constraints. We will discuss the practical challenges associated with deploying this AD capability in a production environment. In particular, we will address issues related to supporting AD across heterogeneous compute environments and parallelization strategies, designing optimal checkpointing schemes for time dependent simulations, and using just in time (JIT) compilation for user defined inputs while preserving both performance and flexibility in MARBL. In addition, we will present preliminary results that highlight how this AD infrastructure can be leveraged for inertial confinement fusion (ICF) relevant problems, illustrating its potential impact on ICF design and analysis workflows. This work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344 and Project No. 21-SI-006, 24-ERD-005, and 26-SI-002.
