Accelerating Algebraic Multigrid Using Machine Learning
Citation (APA)
Nzoyem, R., Louw, T., McIntosh-Smith, S. & Deakin, T., (2022). Accelerating Algebraic Multigrid Using Machine Learning: Application to problems originating from fluid simulation.
Abstract
High-fidelity fluid simulations require solving extremely large and sparse linear systems, often in- volving millions of unknowns. With its fast convergence properties, Algebraic Multigrid (AMG) is the method of choice in several application areas. This said, classical AMG suffers from a number of issues, calling Machine Learning (ML) to the rescue. Recent years have seen a flourish of ML tactics to accelerate AMG. However, published work tends to focus on small and unrepresentative problems. Moreover, these methods tend to be developed dissociatively, not involving human end users in the process.
With insights from potential users, we aim to build a ML-augmented AMG framework that is: (i) robust i.e. reusable in a variety of problems arising from fluid simulation; (ii) computationally cheap i.e. converges faster than the classical AMG; (iii) works for structured as well as unstructured grids; (iv) works for small systems as well as large systems, while taking advantage of their sparsity; and (v) scalable i.e. data- and model-parallelisable with efficient memory management.
The main contributions of this project are as follows: (1) a Graph Neural Network methodology for learning prolongation operators, built around DGL and PyTorch; (2) a Cloud Formation stack on AWS to run experiments, containing all the dependencies our implementation requires; (3) ex- periments indicating that our current AI model is better than the classical approach in about 20% of cases; and (4) a survey that, among other remarks, endeared participants to using (AI-enabled) AMG for their various problems.
Our work entails much efficient usage of HPC resources. By noticing that ML can be leveraged not only when solving linear systems, we pave the way for ML to be efficiently used in other parts of the high-fidelity simulation pipeline. Moreover, the consideration of users and their needs is a positive step towards broader Design Space Exploration of prolongation operators, model hyperparameters, and physical parameters influencing mechanical engines in operation.