사용자 스케줄 언어로 매트릭스 곱셈의 벽을 넘는 텐서 가속기 활용법

초록

이 논문은 텐서 가속기의 활용 범위를 기존의 행렬 곱셈(MatMul) 중심 작업을 넘어서 확장하는 방법을 제시합니다. 사용자 스케줄 가능 언어인 Halide와 컴파일러 기반 기술(equality saturation)을 결합한 HARD BOILED 시스템을 개발하여, 개발자가 간단한 알고리즘과 스케줄만으로도 다양한 응용 프로그램(예: 이미지 처리 필터링, 리샘플링)을 텐서 가속기에 효율적으로 매핑할 수 있게 합니다. 이를 통해 Nvidia Tensor Core 등을 활용한 최대 6.1배의 성능 향상을 달성했습니다.

상세 분석

본 논문의 핵심 기술적 기여는 기존의 딱딱한 패턴 매칭 방식이 아닌, 유연한 프로그램 최적화 기법인 ’equality saturation(EqSat)‘을 기반으로 한 텐서 명령어 선택기 ‘HARD BOILED’를 Halide 컴파일러에 통합한 점에 있습니다.

기존 Halide의 텐서 가속기 지원(예: Intel AMX)은 제한된 문법 패턴에 하드코딩된 매칭 로직에 의존했으며, 이를 확장하거나 수정하는 데 많은 개발자 리소스가 필요했습니다. 이는 새로운 가속기나 연산 패턴을 지원하는 데 걸림돌이 되었습니다. HARD BOILED는 EqSat을 활용하여 이러한 문제를 해결합니다. EqSat은 E-graph라는 데이터 구조를 사용해 다양한 재작성 규칙을 적용한 프로그램 변형들을 효율적으로 탐색하고, 최적의 형태를 찾아냅니다. 논문에서는 Halide의 벡터화 중간 표현(IR: Ramp, Broadcast, vector_reduce_add 노드)을 텐서 연산을 표현하는 데 재활용하고, 이를 EqSat 규칙 집합에 연결합니다. 이를 통해 컴파일러는 사용자가 작성한 알고리즘의 다양한 벡터화된 형태로부터, 하드웨어 텐서 명령어(예: Intel AMX의 TDPBF16PS, Nvidia Tensor Core 연산)와 정확히 매칭되는 패턴을 유연하게 발견할 수 있습니다.

또한, HARD BOILED는 Halide의 기존 스케줄링 원칙(예: producer-consumer fusion, tiling)과 완전히 호환됩니다. 사용자는 store_in(MemoryType::AMXTile)과 같은 스케줄링 프리미티브로 계산이 텐서 레지스터에서 이루어져야 함을 표시하기만 하면, 나머지 복잡한 매핑 작업은 HARD BOILED가 자동으로 처리합니다. 이는 개발자에게 높은 수준의 추상화를 제공하면서도, 하드웨어 특화된 최적화를 가능하게 하는 강력한 조합입니다.

이 기술을 통해 논문은 텐서 가속기의 적용 영역을 전통적인 ML/AI를 넘어 선형 변환으로 표현 가능한 다양한 신호 처리 작업(1D/2D 컨볼루션, 정수/비정수 배율 리샘플링, 재귀적 필터링, DCT 기반 디노이징 등)으로 성공적으로 확장했음을 실험을 통해 입증했습니다. 이는 단순히 라이브러리 API를 호출하는 방식으로는 불가능했던, 계산 그래프 수준의 통합 최적화(퓨전 등)가 가능하기 때문입니다. 결과적으로, 텐서 가속기라는 특화된 하드웨어 자원의 활용도를 극대화하고, 더 넓은 범주의 고성능 컴퓨팅 문제에 그 이점을 가져올 수 있는 길을 열었다는 점에서 중요한 의미를 가집니다.