[논문 리뷰] Tenspiler: A Verified Lifting-Based Compiler for Tensor Operations (Extended Version)

Jie Qiu, C.R. Cai외 4명|arXiv (Cornell University)|2024. 01. 01.

Parallel Computing and Optimization Techniques인용 수 3

한 줄 요약

Tenspiler는 Tensir라는 새로운 통합 중간언어를 사용해 순차적 C++ 또는 Python 코드를 최적화된 텐서 연산으로 변환하는 검증된 리프팅 기반 컴파일러이다. 코드 번역을 문법 유도 프로그램 합성 문제로 정식화하고 정확성 보장을 위해 정리 증명 기법을 활용함으로써, 6개의 소프트웨어 및 하드웨어 백엔드에서 69개의 실제 벤치마크에 대해 평균 105×의 커널 성능 향상과 9.65×의 엔드투엔드 성능 향상을 달성한다.

ABSTRACT

Tensor processing infrastructures such as deep learning frameworks and specialized hardware accelerators have revolutionized how computationally intensive code from domains such as deep learning and image processing is executed and optimized. These infrastructures provide powerful and expressive abstractions while ensuring high performance. However, to utilize them, code must be written specifically using the APIs / ISAs of such software frameworks or hardware accelerators. Importantly, given the fast pace of innovation in these domains, code written today quickly becomes legacy as new frameworks and accelerators are developed, and migrating such legacy code manually is a considerable effort. To enable developers in leveraging such DSLs while preserving their current programming paradigm, we introduce Tenspiler, a verified lifting-based compiler that uses program synthesis to translate sequential programs written in general-purpose programming languages (e.g., C++ or Python code) into tensor operations. Central to Tenspiler is our carefully crafted yet simple intermediate language, named TensIR, that expresses tensor operations. TensIR enables efficient lifting, verification, and code generation. Currently, Tenspiler already supports $ extbf{six}$ DSLs, spanning a broad spectrum of software and hardware environments. Furthermore, we show that new backends can be easily supported by Tenspiler by adding simple pattern-matching rules for TensIR. Using 10 real-world code benchmark suites, our experimental evaluation shows that by translating code to be executed on $ extbf{6}$ different software frameworks and hardware devices, Tenspiler offers on average 105$ imes$ kernel and 9.65$ imes$ end-to-end execution time improvement over the fully-optimized sequential implementation of the same benchmarks.

연구 동기 및 목표

• 일반적인 코드에서 도메인 특화 언어(DSL) 타깃으로의 자동 번역을 가능하게 하여 텐서 컴퓨팅 환경에서의 레거시 코드 이식성 문제를 해결한다.
• 패턴 매칭 기반 컴파일러와 LLM 기반 코드 생성 기법의 취약성, 오류 발생 가능성, 정확성 보장 부족 문제를 해결한다.
• 효율적인 합성과 검증을 가능하게 하는 통합적이고 확장 가능한 중간 표현(Tensir)을 설계한다.
• Tensir에 단순한 패턴 매칭 규칙을 추가함으로써 새로운 백엔드에 대한 지원을 원활하게 구현함으로써, 새로운 프레임워크나 가속기로의 이식 작업을 최소화한다.

제안 방법

• 입력 프로그램을 등가인 Tensir 프로그램으로 리프팅함으로써 코드 번역 작업을 문법 유도 프로그램 합성(SGS) 문제로 정의한다.
• 정리 증명기를 사용해 입력 프로그램과 합성된 Tensir 프로그램 간의 의미적 동치성을 검증함으로써 정확성을 확보한다.
• 일반적인 벡터 및 행렬 연산을 포괄하는 최소한의 텐서 대수 기반 IR인 Tensir를 설계하여 효율적인 합성과 검증을 가능하게 한다.
• 검색 공간을 줄이고 확장성을 향상시키기 위해 연산자 제한 및 프로그램 상태 제한과 같은 최적화 기법을 적용한다.
• 사용자가 제공한 코드 생성기를 사용해 검증된 Tensir 프로그램에서 목표 코드를 생성하며, MLX, Gemmini, Gaudi를 포함한 여러 백엔드를 지원한다.
• 표현 트리와 벡터화 히우리스틱을 활용해 고차원 텐서에서 효율적이고 벡터화된 연산의 합성을 이끈다.

실험 결과

연구 질문

• RQ1 검증된 리프팅 기반 접근법이 다양한 텐서 DSL 및 하드웨어 백엔드로 일반 코드를 정확하고 효율적으로 컴파일하는 데 성공할 수 있는가?
• RQ2 통합된 중간 표현(Tensir)이 이질적인 텐서 처리 타깃 간의 확장 가능한 합성과 검증을 얼마나 효과적으로 지원하는가?
• RQ3 Tenspiler의 합성 및 검증 파이프라인은 1D, 2D, 3D 텐서 연산을 포함한 실제 복잡한 벤치마크에 대해 얼마나 잘 스케일링되는가?
• RQ4 커널 및 엔드투엔드 실행 시간 측면에서 Tenspiler의 성능은 완전 최적화된 순차적 구현과 비교해 어떻게 되는가?
• RQ5 Tenspiler는 최소한의 엔지니어링 노력으로 새로운 백엔드를 쉽게 확장할 수 있는가?

주요 결과

• Tenspiler는 69개의 실제 벤치마크에서 완전 최적화된 순차적 구현 대비 평균 105×의 커널 실행 시간 향상과 9.65×의 엔드투엔드 실행 시간 향상을 달성한다.
• 합성 과정은 효과적으로 스케일링되며, 깊이가 5인 인위적인 3D 텐서 벤치마크는 184초 내에 합성되어 1시간 타임아웃 이내에 처리된다.
• Tenspiler는 MLX, Gemmini, Gaudi를 포함한 6개의 서로 다른 백엔드로 코드를 성공적으로 트랜스파일하여 오픈소스 및 상용 프레임워크 모두에 넓은 호환성을 보였다.
• 최근 출시된 MLX 프레임워크를 위한 Tenspiler 구현은 200줄 이내의 코드로 완료되어 Tensir IR의 확장성과 유연성을 입증했다.
• LLM가 생성한 코드는 동일한 벤치마크에서 일관되게 잘못되었고 기능을 하지 못했으며, 이는 LLM이 낯선 DSL에 대해 검증된 코드를 생성하는 데 한계가 있음을 시사한다.
• Tensir의 사용은 효율적인 검증과 코드 생성을 가능하게 하며, 연산자 제한 및 프로그램 상태 제한과 같은 최적화 기법이 합성의 검색 공간을 크게 줄여 성능 향상에 기여한다.