🤖 AI Summary
This work addresses the lack of efficient, differentiable, and multi-GPU-scalable native sparse matrix operations in PyTorch for industrial-scale scientific computing. The authors propose an open-source PyTorch library that enables automatic differentiation via the adjoint method, achieving O(1) computational graph nodes and O(nnz) memory footprint. The library integrates Newton, Picard, and Anderson nonlinear solvers alongside eigenvalue computation routines. Leveraging domain decomposition and halo exchange, it supports multi-GPU scaling and successfully solves sparse linear systems with 400 million degrees of freedom across three GPUs. By seamlessly interfacing with PyTorch’s autograd system, the library drastically reduces memory overhead and provides efficient, native support for large-scale end-to-end differentiable scientific simulations.
📝 Abstract
Industrial scientific computing predominantly uses sparse matrices to represent unstructured data -- finite element meshes, graphs, point clouds. We present \torchsla{}, an open-source PyTorch library that enables GPU-accelerated, scalable, and differentiable sparse linear algebra. The library addresses three fundamental challenges: (1) GPU acceleration for sparse linear solves, nonlinear solves (Newton, Picard, Anderson), and eigenvalue computation; (2) Multi-GPU scaling via domain decomposition with halo exchange, reaching \textbf{400 million DOF linear solve on 3 GPUs}; and (3) Adjoint-based differentiation} achieving $\mathcal{O}(1)$ computational graph nodes (for autograd) and $\mathcal{O}(\text{nnz})$ memory -- independent of solver iterations. \torchsla{} supports multiple backends (SciPy, cuDSS, PyTorch-native) and seamlessly integrates with PyTorch autograd for end-to-end differentiable simulations. Code is available at https://github.com/walkerchi/torch-sla.