torch-sla: PyTorch를 위한 수반 솔버 및 희소 텐서 병렬화 기반의 미분 가능한 희소 선형 대수

torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch

산업용 과학 계산 분야에서는 유한 요소 메쉬, 그래프, 포인트 클라우드와 같은 비정형 데이터를 표현하기 위해 주로 희소 행렬을 사용합니다. 본 논문에서는 GPU 가속과 확장성, 그리고 미분 가능한 희소 선형 대수를 가능하게 하는 오픈 소스 PyTorch 라이브러리인 torch-sla를 제안합니다. 이 라이브러리는 세 가지 핵심 과제를 해결합니다: (1) 희소 선형 풀이, 비선형 풀이(Newton, Picard, Anderson), 고유값 계산을 위한 GPU 가속; (2) 헤일로 교환(halo exchange)을 동반한 영역 분할(domain decomposition)을 통한 다중 GPU 확장으로, 3개의 GPU에서 4억 자유도(DOF)의 선형 풀이 달성; (3) 솔버의 반복 횟수와 무관하게 O(1)의 계산 그래프 노드(autograd용)와 O(nnz) 메모리를 달성하는 수반(Adjoint) 기반 미분입니다. torch-sla는 다양한 백엔드(SciPy, cuDSS, PyTorch-native)를 지원하며, 엔드투엔드 미분 시뮬레이션을 위해 PyTorch autograd와 원활하게 통합됩니다. 코드는 https://github.com/walkerchi/torch-sla 에서 확인할 수 있습니다.

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.