Technical Area: Solution of Linear and Nonlinear Systems

Solution of systems of algebraic equations is undoubtedly one of the most common computational kernels in scientific applications of interest to the Department of Energy. Efficient, scalable, and reliable algorithms are crucial for the success of large-scale simulations. This FASTMath area includes both iterative and direct linear solvers as well as nonlinear solution methods. Our work is focused on two primary themes. First, we are developing new algorithms, including improved preconditioners, FFTs, advanced multigrid techniques, and hybrid nonlinear systems, to better solve physics problems. Second, we will improve and enhance our software performance by increased GPU support and better communication strategies. We will take advantage of machine-learning tools to achieve some of these goals. This work has application to fusion, nuclear structure calculation, quantum chemistry, accelerator modeling, climate, and astrophysics.

Direct Solvers and Fast Fourier Transforms:

We design factorization-based direct linear solvers and preconditioners for large-scale simulations. We will port spectral nested dissection to GPUs, which will benefit SuperLU, STRUMPACK and symPACK, and provide GPU-enabled matrix-vector multiplications in ButterflyPACK. We will harden symPACK’s capability to solve symmetric indefinite systems. We will investigate machine learning approaches for graph partitioning and leveraging butterfly structures for inventing new neural networks. We develop FFTs, including special kernels, provided in the new FFTX library and in ButterflyPACK to support SciDAC application codes.

GPU-enabled Kernels and Reduced Communication Strategies:

We are developing linear algebra and graph kernels that are needed by multigrid methods and domain decomposition methods to perform well on GPU architectures. They will be integrated in FASTMath solvers libraries and their performance on GPUs, including tensor cores, will be evaluated. We will also investigate various communication strategies in Trilinos’ linear algebra kernels and solvers.

Multigrid and multilevel methods:

Since GPUs favor regular compute patterns to achieve high performance, we are developing and implementing a variety of multigrid methods capable to take advantage of structure in problems, including semi-structured algebraic multigrid methods in hypre and region based structured multigrid methods in Trilinos/MueLu. We are also designing physics-based preconditioners motivated by Tokamak Disruption simulations. Since high-order finite element methods are becoming essential for scientific simulations on GPU-based computer architectures, we are also designing multilevel methods for high-order systems.

Hybrid nonlinear systems:

We are developing methods for neural ODE's and hybrid differential models. We have been developing an integrated framework that uses both PyTorch and PETSc. It leverages scalable GPU based differential equation solvers and the adjoint sensitivity analysis capabilities from PETSc and allows efficient hybridization of traditional differential simulations and ML tools at scale in SciDAC applications. We will combine simulation and ML to develop reduced order models with required conservation properties. The reduced order models are used to complete simulations or generate sensitivities for computationally expensive applications such as optimization.