A Sub-linear Low-Rank Solver for Poisson's Equation using Machine Learning Frameworks for GPU Acceleration

In this paper we explore a fast Poisson solver for problems with a solution that is known to be low-rank.

We use an adaptive and warm started cross approximation called Cross-DEIM that iterates between index selection and and cross approximation to generate a low-rank solution.

This paper focuses on leveraging a modern machine learning framework, PyTorch, as a general purpose array language to implement low-rank solvers based on Cross-DEIM.

PyTorch enables native access to GPUs and accelerators but with a user-friendly high-level interface.

We investigate statistical leverage scores for the index selection for the cross approximation due to the cost associated with the pivoted algorithms used with the discrete empirical interpolation methods (DEIM and QDEIM) which are historically preferred.

The cross approximation is naturally paired with a Discrete Sine Transform (DST) Poisson solver.

This allows the Fast Fourier Transform (FFT) to be evaluated in batches along dimensions independently without any global transpose even in higher dimensions.

We present performance results running on a A100 GPU and AMD EPYC CPU demonstrating the usefulness of the approach that enables problems sizes that previously were not feasible.