Sparse Computation의 새로운 접근법: Sandia 연구소의 LAPIS 컴파일러 분석

현업에서 대규모 데이터를 다루다 보면 필연적으로 ‘희소성(Sparsity)‘이라는 골칫거리와 마주하게 됩니다. Dense matrix 연산은 GPU의 병렬 처리 능력을 극한으로 끌어올리기 쉽지만, 데이터가 비어있는 Sparse matrix로 넘어가는 순간 메모리 접근 패턴은 엉망이 되고 캐시 미스가 폭발하며 연산 효율은 바닥을 칩니다. 보통은 cuSPARSE 같은 벤더 종속적인 최적화 라이브러리에 기대거나, 하드웨어 스펙으로 찍어 누르는 방식을 택하곤 하죠.

그런데 최근 Sandia National Lab에서 꽤 흥미로운 논문을 발표했습니다. 선형 대수학(Linear Algebra)을 이해하는 컴파일러 프레임워크인 LAPIS를 통해 이 문제를 해결하겠다는 접근입니다. 기존의 라이브러리 호출 방식이 아니라, MLIR을 활용해 코드 레벨에서 Sparse 연산을 컴파일 타임에 최적화하겠다는 건데, 아키텍처 설계가 꽤나 실용적입니다.

LAPIS의 핵심: MLIR과 Kokkos의 만남

이 프로젝트의 가장 큰 기술적 특징은 MLIR(Multilevel Intermediate Representation) 위에 구축되었다는 점입니다. 최근 컴파일러 생태계, 특히 AI/ML 인프라 쪽에서 MLIR은 사실상 표준 인프라로 자리 잡고 있죠.

LAPIS의 진짜 혁신은 Kokkos dialect 를 새로 정의했다는 데 있습니다. 웹 백엔드나 일반적인 ML 엔지니어들에게는 생소할 수 있지만, Kokkos는 HPC(High-Performance Computing) 진영에서 널리 쓰이는 C++ 기반의 성능 이식성(Performance portability) 프로그래밍 모델입니다. CPU, GPU 등 다양한 아키텍처에서 동일한 코드로 높은 성능을 내기 위해 사용됩니다.

  • High-level Abstraction: 고수준의 SciML(Scientific Machine Learning) 코드를 MLIR로 파싱합니다.
  • Dialect Lowering: 이를 Kokkos dialect로 변환하여 아키텍처에 독립적인 최적화를 수행합니다.
  • Codegen: 최종적으로 사람이 읽을 수 있는(readable) Kokkos C++ 커널 코드를 생성해 냅니다.

이 구조가 매력적인 이유는 기존 HPC 환경의 레거시 C++ 코드베이스에 최신 머신러닝 모델을 통합할 때 발생하는 거대한 마찰을 줄여주기 때문입니다. 블랙박스 같은 바이너리가 아니라, 기존 빌드 파이프라인에 태울 수 있는 C++ 코드를 뱉어준다는 건 실무자 입장에서 엄청난 장점입니다.

분산 메모리와 Partition Dialect

단일 노드에서의 최적화도 어렵지만, Sparse 연산을 분산 메모리 클러스터로 가져가면 통신 오버헤드(Communication overhead)라는 새로운 지옥이 열립니다. LAPIS 팀은 이를 해결하기 위해 새로운 Partition dialect를 만들었습니다.

Sparse tensor를 여러 노드에 어떻게 분배할지, 그리고 연산 중 발생하는 통신 패턴을 어떻게 최소화할지를 컴파일러 레벨에서 결정합니다. GraphBLAS 기반의 관계형 데이터베이스인 TenSQL이나 서브그래프 동형성(Subgraph isomorphism) 커널 같은 복잡한 그래프 알고리즘에서 성능 이식성을 증명했다는 점은 이 아키텍처의 범용성을 잘 보여줍니다.

Hacker News의 반응과 개인적인 생각

이번 발표를 두고 해커뉴스에서도 꽤 날카로운 토론이 오갔습니다. 몇 가지 흥미로운 포인트에 대해 제 생각을 덧붙여 보겠습니다.

1. “이거 그냥 MLIR 얹은 SciPy 아니냐?” 한 유저가 이런 질문을 던졌는데, 완전히 빗나간 비유입니다. 다른 유저가 지적했듯, 이건 오히려 OpenXLA나 PyTorch의 torch.compile()에 훨씬 가깝습니다. 차이가 있다면 타겟이 범용 ML 가속기가 아니라, 과학 연산에 특화된 Kokkos C++라는 점이죠.

2. “왜 Torch Inductor나 OpenXLA에 기여하지 않고 바퀴를 재발명했나?” 학계의 고질적인 ‘바퀴 재발명’ 아니냐는 비판이 있었습니다. 하지만 저는 다른 유저의 반박에 100% 동의합니다. 현재 PyTorch나 JAX 생태계의 내부를 들여다본 분들은 아시겠지만, 의존성은 끔찍하게 얽혀 있고 빌드 시스템은 거대합니다. 특정 제약 조건(여기서는 HPC 환경과 C++ 통합)을 풀기 위해 거대한 프레임워크의 늪에 빠지느니, MLIR이라는 훌륭한 레고 블록을 가져다 가볍고 목적에 맞는 독립적인 컴파일러를 구축하는 것이 훨씬 영리한 엔지니어링 선택입니다.

3. Kahan Summation과 부동소수점 오차 문제 한 유저가 “이런 컴파일러가 Kahan Summation 같은 부동소수점 오차 보정 알고리즘을 어떻게 처리할지 궁금하다”고 언급했습니다. 아주 날카로운 지적입니다. 종종 공격적인 컴파일러 최적화(예: 연산 순서 재배치)는 수학적으로는 동일할지 몰라도 부동소수점 정밀도 측면에서는 재앙을 초래하곤 합니다. 저 역시 과거에 커스텀 CUDA 커널을 작성하다가 Associative reordering 때문에 모델 학습이 발산해 버리는 문제를 겪은 적이 있습니다. LAPIS가 선형 대수학을 ‘이해’하는 컴파일러를 표방한 만큼, 이런 수치적 안정성(Numerical stability)을 어떻게 보장하는지가 향후 프로덕션 도입의 관건이 될 것입니다.

결론: 실무에 도입할 만한가?

당장 내일 여러분의 웹 서비스 추천 시스템에 LAPIS를 도입하라고 권하지는 않겠습니다. 이 프로젝트는 여전히 연구 단계에 가깝고, 타겟 유즈케이스도 HPC와 과학 연산에 강하게 맞춰져 있습니다. (여담으로 누군가 “Mojo는 어디로 갔나?”라고 농담을 던졌는데, 이런 틈새시장을 보면 Mojo가 가야 할 길이 아직 멀었다는 생각도 듭니다.)

하지만 우리가 여기서 얻어야 할 인사이트는 ‘도메인 특화 컴파일러(DSL)의 강력함’ 입니다. 복잡한 비즈니스 로직이나 연산 패턴을 범용 언어로 하드코딩하며 최적화의 한계에 부딪혔을 때, MLIR 같은 인프라를 활용해 문제를 중간 표현(IR) 레벨로 끌어올리고 타겟 아키텍처에 맞게 코드를 생성(Codegen)하는 패턴은 앞으로 시니어 엔지니어들이 반드시 숙지해야 할 강력한 무기입니다.

References: