[논문 리뷰] High Performance Code Generation in MLIR: An Early Case Study with GEMM

Uday Bondhugula|arXiv (Cornell University)|2020. 03. 01.

Algorithms and Data Compression참고 문헌 6인용 수 27

한 줄 요약

이 논문은 애फ인 및 폴리드랄 추상화를 갖춘 컴파일러 인프라인 MLIR를 사용하여, OpenBLAS 및 MKL과 같은 수작업 최적화된 라이브러리와 경쟁할 수 있는 고성능 GEMM 커널을 자동으로 생성할 수 있음을 보여주는 초기 사례 연구를 제시한다. MLIR의 memref, affine 다이얼ект, 그리고 폴리드랄 패assing(예: unroll-and-jam, 타일링, 벡터화)을 활용하여 인텔 스카라이크 CPU에서 정점 성능의 92%를 달성하였다. 이는 모듈러하고 자동화된 코드 생성이 고성능 선형 대수 커널에 대해 실현 가능하다는 것을 보여준다.

ABSTRACT

This article is primarily meant to present an early case study on using MLIR, a new compiler intermediate representation infrastructure, for high-performance code generation. Aspects of MLIR covered in particular include memrefs, the affine dialect, and polyhedral utilities and pass infrastructure surrounding those. This article is also aimed at showing the role compiler infrastructure could play in generating code that is competitive with highly tuned manually developed libraries, albeit in a more modular, reusable, and automatable way.

연구 동기 및 목표

• MLIR가 수작업 어셈블리 코드 없이도 고성능 GEMM 커널을 자동으로 생성할 수 있음을 보여주기 위해.
• MLIR의 폴리드랄 및 애फ인 인프라가 OpenBLAS 및 MKL과 같은 수작업 최적화 라이브러리와 경쟁 가능한 성능을 달성할 수 있는지 평가하기 위해.
• 밀도 있는 선형 대수 커널에 대해 수작업 최적화 대신 모듈러하고 재사용 가능하며 조합 가능한 컴파일러 기반 접근 방식을 도입할 수 있는지 타당성을 탐색하기 위해.
• 타일링, 벡터화, 스칼라 교체 등의 MLIR 패스 인프라가 정점 성능에 얼마나 기여하는지 평가하기 위해.

제안 방법

• 저자는 애फ인 다이얼ект를 사용하여 행렬 곱셈을 애फ인 연산으로 표현하며, 명시적인 타일링 매개변수(예: M_R=3, N_R=16, K_U=4)를 포함한다.
• 타일링, unroll-and-jam, 벡터화, 스칼라 교체 등의 MLIR 패스를 순차적으로 적용하며, 이 모든 작업은 애फ인 연산을 대상으로 수행된다.
• 코드 생성 파이프라인은 고수준의 MLIR 연산에서 시작하여 MLIR 인프라를 통해 점진적으로 타겟 특화 코드로 내림차순 변환된다.
• 메모리 레이아웃 분석을 정확히 수행할 수 있도록, 애फ인 레이아웃 맵을 사용한 memref 타입이 스트라이드, 전치, 메모리 접근 패턴을 표현하는 데 사용된다.
• 변환 시퀀스는 C++ 패스로 인코딩되어 최적화 선택의 프로그래밍 기반 탐색과 자동화를 가능하게 한다.
• 최종 코드는 MLIR의 인프라를 사용하여 생성되며, 수작업으로 작성된 C 코드나 인-line 어셈블리 코드를 피한다.

실험 결과

연구 질문

• RQ1 MLIR의 폴리드랄 및 애फ인 인프라는 OpenBLAS 및 MKL과 같은 수작업 최적화 라이브러리와 경쟁 가능한 성능을 달성하는 GEMM 커널을 자동으로 생성할 수 있는가?
• RQ2 unroll-and-jam, 벡터화, 타일링 등의 MLIR 내장 패스는 수동 튜닝 없이도 고성능을 달성하는 데 얼마나 효과적인가?
• RQ3 MLIR의 모듈러 인프라는 고성능 선형 대수에서 수작업 어셈블리 코드 및 도메인 전문 최적화 전문 지식의 필요성을 얼마나 줄일 수 있는가?
• RQ4 타일 크기가 행렬 크기를 나누지 못하는 비균일한 문제 크기(예: 타일 차원이 행렬 차원을 나누지 못할 경우)를 처리할 때 MLIR는 성능 손실 없이 어떻게 대처하는가?
• RQ5 타일 매개변수에 대한 분석 모델을 MLIR의 변환 파이프라인에 효과적으로 통합하여 자동 성능 최적화를 이끄는 데 유용한가?

주요 결과

• MLIR가 생성한 GEMM 커널은 인텔 i7-8700K CPU에서 67.49 GFLOPS를 기록하여 이론적 정점 성능(75.2 GFLOPS)의 92%를 달성하였다.
• MLIR가 생성한 커널의 성능은 OpenBLAS(67.49 GFLOPS)와 동일했으며, MKL(67.70 GFLOPS)와는 0.3% 이내로 근접하였다.
• Clang와 GCC가 -O3 및 -ffast-math 옵션을 사용했을 때는 각각 0.47 GFLOPS와 4.53 GFLOPS에 머물러 일반 컴파일러의 한계를 부각시켰다.
• MLIR의 폴리드랄 패스(특히 unroll-and-jam 및 벡터화)는 타일링과 결합되었을 때 고성능 달성에 핵심적인 역할을 하였다.
• MLIR의 애फ인 레이아웃 맵 지원 덕분에 스트라이드, 전치, 메모리 접근 패턴을 깔끔하게 모델링할 수 있었으며, 코드 생성을 단순화시켰다.
• 본 연구는 MLIR가 고수준 GEMM에서 고성능 코드로의 전체 최적화 파이프라인을 자동화할 수 있음을 보여주며, 수작업 저수준 튜닝에 대한 의존도를 감소시킬 수 있음을 시사한다.