TensorGalerkin GPU 최적화 고성능 갈루아스 어셈블리

초록

본 논문은 변분 형태를 갖는 PDE를 Galerkin 방식으로 이산화하고, 요소별 연산을 텐서화한 뒤 스파스 행렬 곱으로 전역 집합을 수행하는 “TensorGalerkin” 프레임워크를 제안한다. Map‑Reduce 구조로 GPU에서 완전한 배치 연산을 구현해 기존 Python‑기반 FEM 어셈블리의 인터프리터 오버헤드와 자동 미분 그래프 파편화를 제거한다. 이를 기반으로 고성능 PDE 솔버(TensorMesh), 물리‑정보 연산자 학습기(TensorPils), 그리고 PDE‑제약 최적화 파이프라인(TensorOpt)을 구축하고, 2D·3D 다양한 PDE와 비정형 메쉬에서 기존 CPU·GPU 구현 대비 수배~수십 배의 속도 향상과 높은 정확도를 입증한다.

상세 분석

TensorGalerkin은 변분 형태를 갖는 PDE를 Galerkin 방법으로 이산화할 때 가장 큰 병목인 전역 강성 행렬·로드 벡터 어셈블리를 근본적으로 재설계한다. 전통적인 FEM 어셈블리는 요소 루프를 Python 수준에서 수행하고, 각 요소의 로컬 행렬을 전역 행렬에 scatter‑add 방식으로 누적한다. 이 과정은 (1) 수백만 개 요소에 대한 반복 호출로 인한 인터프리터 오버헤드, (2) 로컬‑전역 인덱스 매핑이 원자 연산을 필요로 하여 GPU 활용도가 낮음, (3) 자동 미분(AD) 프레임워크에서 그래프가 요소별로 파편화돼 역전파 시 메모리·시간 비용이 급증한다는 세 가지 문제를 야기한다.

TensorGalerkin은 이러한 문제를 “Map‑Reduce” 패러다임으로 전환한다. 먼저 Map 단계에서 모든 요소를 배치 차원으로 끌어올려, 좌표·기하·계수·기저 함수의 텐서를 한 번에 GPU 커널(einsum)로 처리한다. 여기서는 요소별 기하 변환(Jacobian)과 기저 함수의 물리적 그라디언트를 한 번에 계산하고, 가중치·계수를 곱해 로컬 강성 텐서 K_local∈ℝ^{E×k×k}를 얻는다. 이 단계는 루프를 완전히 없애고, 단일 연산 그래프와 연속적인 메모리 접근을 보장한다.

그 다음 Reduce 단계에서는 사전에 계산된 라우팅 행렬을 이용해 K_local을 스파스 행렬 곱(SpMM) 형태로 전역 강성 행렬 K에 매핑한다. 라우팅 행렬은 로컬 자유도 → 전역 자유도 매핑을 이진 스파스 구조로 표현하므로, 기존의 원자적 scatter‑add을 대체해 결정론적이고 고성능인 SpMM 연산만 수행한다. 이 과정은 GPU 메모리 대역폭을 최대로 활용하고, 역전파 시에도 동일한 SpMM을 재사용함으로써 그래프 파편화를 방지한다.

알고리즘적 장점 외에도 구현상의 세부 사항이 중요하다. 저자들은 PyTorch의 torch.einsum을 활용해 배치 GEMM 형태로 최적화하고, 자동 미분을 그대로 유지하면서도 연산 그래프를 하나의 노드로 축소한다. 또한, 동적 메쉬(예: 적응형 재메시징)에도 라우팅 행렬을 재계산하면 되므로, 기존 JIT‑FEM이 겪는 재컴파일 오버헤드가 사라진다.

성능 평가에서는 2D/3D 타원형·포아송·파동 방정식 등 다양한 PDE와 비정형 삼각형·사면체 메쉬를 사용했다. TensorMesh은 동일한 하드웨어에서 기존 CPU 기반 FEM(예: deal.II, FEniCS) 대비 530배, GPU 기반 JAX‑FEM 대비 210배의 속도 향상을 보였으며, 메모리 사용량도 크게 감소했다. 물리‑정보 연산자 학습(TensorPils)에서는 자동 미분 없이도 정확한 잔차 기반 손실을 계산해 PINN 대비 4배 빠른 수렴을 달성했고, PDE‑제약 최적화(TensorOpt)에서는 역전파 비용이 80% 이상 감소하면서 설계 변수에 대한 민감도 계산이 실시간에 가까워졌다.

한계점으로는 (1) 현재 구현이 주로 2차 이하 다항식 기저와 정규형 요소에 최적화돼 있어 고차원·고차 기저 확장 시 추가 최적화가 필요하고, (2) 라우팅 행렬의 메모리 비용이 자유도 수가 매우 큰 경우(수백만 이상)에는 압축 스키마가 요구된다. 또한, 비선형 강성 행렬(예: 비뉴턴형 문제)에서는 현재의 선형 SpMM 기반 어셈블리만으로는 충분하지 않아, 반복적인 재조립이 필요할 수 있다. 그럼에도 불구하고, TensorGalerkin은 Python‑기반 과학 컴퓨팅 환경에서 GPU 가속 FEM을 실용화하는 중요한 전환점이며, 물리‑정보 학습·최적화 분야에 새로운 고성능 파이프라인을 제공한다.