Applying Recent Machine Learning Approaches to Accelerate the Algebraic Multigrid Method for Fluid Simulations

In this work, we describe our experiences trying to apply recent machine learning (ML) advances to the Algebraic Multigrid (AMG) method to predict better prolongation (interpolation) operators and accelerate solver convergence. Published work often reports results on small, unrepresentative problems, such as 1D equations or very small computational grids. To better understand the performance of these methods on more realistic data, we create a new, reusable dataset of large, sparse matrices by leveraging the recently published Thingi10K dataset of 3D geometries, along with the FTetWild mesher for creating computational meshes that are valid for use in finite element method (FEM) simulations. We run simple 3D Navier-Stokes simulations, and capture the sparse linear systems that arise. We consider the integration of ML approaches with established tools and solvers that support distributed computation, such as HYPRE, but achieve little success. The only approach suitable for use with unstructured grid data involves inference against a multi-layer message-passing graph neural network, which is too memory-hungry for practical use, and we find existing frameworks to be unsuitable for efficient distributed inference. Furthermore, the model prediction times far exceed the complete solver time of traditional approaches. While our focus is on inference against trained models, we also note that retraining the proposed neural networks using our dataset remains intractable. We conclude that these ML approaches are not yet ready for general use, and that much more research focus is required into how efficient distributed inference against such models can be incorporated into existing HPC workflows.