[논문 리뷰] High Performance GPU Code Generation for Matrix-Matrix Multiplication using MLIR: Some Early Results

Navdeep Katel, Vivek Khandelwal외 1명|arXiv (Cornell University)|2021. 08. 23.

Parallel Computing and Optimization Techniques참고 문헌 1인용 수 6

한 줄 요약

이 논문은 NVIDIA A100 및 RTX 3090 GPU의 텐서 코어를 대상으로 하이퍼퍼포먼스 행렬-행렬 곱셈을 위한 MLIR 기반 자동 코드 생성 파이프라인을 제시한다. 최적화된 다이얼렉트 변환을 통해 고수준 IR를 저수준으로 내림내림으로써, 이 접근법은 FP32에 대해 95–119%, FP16 혼합 정밀도 워크로드에 대해 80–160%의 CuBLAS 성능을 달성하며, 수작업 최적화된 라이브러리에 필적할 수 있음을 입증한다. 이는 수작업 코딩 없이도 IR 기반 컴파일러가 수작업 최적화된 라이브러리 수준의 성능을 낼 수 있음을 보여준다.

ABSTRACT

This report presents some early results on code generation targeting tensor cores on NVIDIA GPUs using the MLIR compiler infrastructure. The state-of-the-art in high-performance deep learning today is primarily driven by manually optimized highly tuned libraries. The approach to develop such libraries is often not modular or reusable to the same extent that compiler infrastructure like LLVM is. Manual optimization typically does not use a standard intermediate representation (IR), although the optimizations performed can be encoded as a sequence of transformation steps and customized passes on an IR. Hand tuning may also miss exploration of design points only reachable easily by automatic code generation. We believe that until the recent introduction of MLIR (Multi-level intermediate representation), IR infrastructure was not geared to tackle the problem of automatic generation of domain-specific libraries in an effective manner. In particular, it was hard to represent and transform compute abstractions at high, middle, and low levels using a single IR. With suitable abstractions in MLIR, we build an experimental lowering pipeline that is able to automatically generate code for matrix-matrix multiplication on NVIDIA GPUs targeting its tensor cores. On a set of problem sizes we evaluated, initial performance results show that we are able to attain performance that is 95-119% and 80-160% of CuBLAS for FP32 and FP16 accumulate respectively on NVIDIA's Ampere microarchitecture-based Geforce 3090 RTX. We believe that these results could be used as motivation for further research and development on automatic code and library generation using IR infrastructure for similar specialized accelerators.

연구 동기 및 목표

• MLIR를 사용하여 NVIDIA GPU에서 행렬-행렬 곱셈을 위한 모듈러하고 자동화된 코드 생성 파이프라인을 개발하기.
• 수작업 커널 튜닝 없이도 텐서 코어를 효율적으로 타겟팅하여 전문가 최적화 라이브러리 의존도를 줄이기.
• IR 기반 컴파일이 수작업 최적화된 라이브러리(예: CuBLAS)와 유사한 성능을 달성할 수 있는지 평가하기.
• 모듈러한 IR 파이프라인에서 체계적인 A/B 테스트를 통해 개별 최적화의 영향을 탐색하기.
• 통합 중간 표현을 사용하여 도메인 특화 고성능 라이브러리의 자동 생성 기반 마련하기.

제안 방법

• 이 방법은 MLIR의 다수준 IR을 사용하여 행렬 곱셈을 고수준, 중간수준, 저수준 추상화 수준에서 표현함으로써 체계적인 변환 가능성을 확보한다.
• Warp Matrix Multiply Accumulate (WMMA) 연산을 위한 새로운 MLIR 다이얼렉트를 도입하였으며, 이는 LLVM/NVPTX IR로 내림내림된다.
• 파이프라인은 다이얼렉트 내림내림 및 최적화 단계의 시퀀스를 적용하며, 이에는 타일링, 공유 메모리로의 타일링, 레지스터 수준 최적화가 포함된다.
• 피이프라인화된 글로벌 메모리 로딩 및 타일 크기 선택과 같은 최적화는 성능 영향을 연구하기 위해 점진적으로 적용된다.
• 코드 생성 파이프라인은 고수준 행렬 곱셈 연산에서 시작하여 텐서 코어를 타겟으로 하는 효율적인 CUDA 코드로 컴파일되는 엔드 투 엔드 구조를 취한다.
• 성능 평가는 표준 문제 크기와 혼합 정밀도/반정밀도 설정을 사용하여 NVIDIA의 Ampere 아키텍처 기반 RTX 3090 GPU에서 평가된다.

실험 결과

연구 질문

• RQ1 IR 기반 코드 생성 파이프라인이 텐서 코어에서 행렬-행렬 곱셈에 대해 수작업 최적화된 라이브러리(예: CuBLAS)와 경쟁 가능한 성능을 달성할 수 있는가?
• RQ2 타일 크기 및 메모리 액세스 피이프라인화와 같은 개별 최적화가 모듈러하고 IR 기반 파이프라인에서 성능에 어떤 영향을 미치는가?
• RQ3 자동 코드 생성이 다양한 정밀도 형식에서 메이저 벤더 최적화 라이브러리의 성능을 얼마나 잘 따라잡거나 초월할 수 있는가?
• RQ4 MLIR과 같은 통합 IR 인fra가 텐서 코어와 같은 특수 가속기의 고수준 계산 추상화를 효과적으로 표현하고 변환할 수 있는가?
• RQ5 다양한 문제 크기와 데이터 유형에서 자동으로 생성된 커널의 성능은 CuBLAS와 비교해 어떻게 되는가?

주요 결과

• 자동으로 생성된 커널은 RTX 3090에서 FP32 혼합 정밀도 행렬 곱셈에 대해 CuBLAS 성능의 95–119%를 달성하며, 최대 효율은 장치 이론적 최대 성능의 95.4%에 이르렀다.
• FP16 혼합 정밀도 행렬 곱셈에서는 성능이 항상 CuBLAS의 95–119% 수준을 유지했으며, 더 작은 문제 크기에서는 더 작은 스레드 블록 타일 덕분에 더 높은 점유율로 인해 CuBLAS를 초월하는 결과를 보였다.
• FP16 전용 행렬 곱셈에서는 생성된 코드가 CuBLAS 성능의 80–160%를 달성했으며, CuBLAS가 성능이 일관되지 않게 나타나고 타일 선택이 비최적화된 더 큰 문제 크기에서는 CuBLAS를 뛰어넘는 성능을 보였다.
• 연구 결과에 따르면 CuBLAS는 더 작은 타일 크기(예: 128×128×32)와 5단계 피이프라인을 사용하지만, 더 높은 글로벌 메모리 스탠드로 인해 지연 숨기 기술이 덜 효과적이라는 점을 드러냈다.
• 더 작은 타일 크기(예: 64×64×64)는 작은 문제 크기에서 점유율 증가로 인해 성능 향상을 이룬 반면, 더 큰 행렬에서는 더 나은 데이터 재사용 덕분에 더 큰 타일이 더 효과적이었다.
• 점진적인 최적화 연구 결과, 피이프라인화와 타일 크기 선택이 성능에 상당한 영향을 미치며, 최종 최적화된 파이프라인이 최고 성능에 가까운 성능을 달성함을 확인했다.