[논문 리뷰] Automatic Kernel Generation for Volta Tensor Cores

Somashekaracharya G. Bhaskaracharya, Julien Demouth외 1명|arXiv (Cornell University)|2020. 06. 22.

Parallel Computing and Optimization Techniques참고 문헌 19인용 수 13

한 줄 요약

이 논문은 NVIDIA Volta 텐서 코어를 대상으로 하여 행렬 곱셈과 ReLU, 편향 덧셈과 같은 포인트와이즈 연산을 융합한 고도로 최적화된 CUDA 커널을 자동으로 생성하기 위한 다각형 편집 기반 접근법을 제시한다. mma.sync.m8n8k4 PTX 명령어와 매크로-MMA 조합을 활용하여, 수작업으로 최적화된 라이브러리 구현보다 최대 2.55배 빠른 성능을 달성하며, 텐서 코어 워크로드를 위한 자동 커널 생성의 효과성을 입증한다.

ABSTRACT

A commonly occurring computation idiom in neural networks is to perform some pointwise operations on the result of a matrix multiplication. Such a sequence of operations is typically represented as a computation graph in deep learning compilers. When compiling to a GPU target, these computations can be individually mapped to manually tuned implementations provided by libraries such as cuBLAS and cuDNN. These libraries also provide off-the-shelf support for targeting tensor cores in NVIDIA GPUs, which can lead to huge performance boosts through their specialized support for mixed-precision matrix math. Alternatively, tensor cores can be programmed directly using CUDA APIs or inline assembly instructions, which opens up the possibility of generating efficient CUDA kernels automatically for such computations. Automatic kernel generation is particularly crucial when it is beneficial to generate efficient code for an entire computation graph by fusing several operations into a single device function instead of invoking a separate kernel for each of them. Polyhedral compilation techniques provide a systematic approach for the analysis and transformation of a sequence of affine loop-nests. In this paper, we describe a polyhedral approach to generate efficient CUDA kernels for matrix multiplication using inline assembly instructions for programming tensor cores on NVIDIA Volta GPUs. Furthermore, we build on this approach to generate fused kernels for computation sequences involving matrix multiplication and pointwise operations such as bias addition, ReLU activation etc. Experimental evaluation of these techniques show that automatically generated kernels can provide significantly better performance than manually tuned library implementations, with speedups ranging up to 2.55X.

연구 동기 및 목표

• 각 연산에 대해 별도의 커널을 사용할 경우 반복적인 글로벌 메모리 접근으로 인한 성능 저하 문제를 해결하기 위해.
• 혼합 정밀도 행렬 곱셈과 융합된 포인트와이즈 연산을 위한 Volta 텐서 코어를 활용하는 효율적인 CUDA 커널을 자동 생성할 수 있도록 하기 위해.
• 세부 제어를 통해 텐서 코어 사용을 최적화하기 위해 다각형 편집 기법을 mma.sync.m8n8k4 PTX 명령어에 직접 적용하기 위해.
• 자동으로 생성된 커널이 수작업으로 최적화된 라이브러리 구현(예: cuBLAS, cuDNN)을 초월할 수 있음을 입증하기 위해.
• 텐서 코어 가속을 위한 딥러닝 워크로드용 고성능 커널을 생성하기 위한 체계적이고 재사용 가능한 프레임워크를 제공하기 위해.

제안 방법

• 다각형 프레임워크를 사용하여 행렬 곱셈과 포인트와이즈 연산을 나타내는 애프린 루프 네스트를 분석하고 변환한다.
• 8×8×8 행렬 片각을 스레드 퀠트를 통해 처리하는 mma.sync.m8n8k4 PTX 명령어를 직접 사용하여 Volta 텐서 코어를 프로그래밍하는 커널을 생성한다.
• 낮은 수준의 mma.sync.m8n8k4 연산을 16×16×8 및 32×32×8 크기의 고차원 매크로-MMA 연산으로 조합하여 메모리 대역폭 활용도를 향상시킨다.
• ReLU, 시그모이드, 편향 덧셈과 같은 포인트와이즈 연산을 행렬 곱셈의 상류(프로로그) 및 하류(에필로그)에 융합할 수 있다.
• Diesel DSL 컴파일러와 통합되어 텐서 코어 커널의 타일 크기와 레지스터 사용을 최적화하기 위한 커스텀 스케줄링을 사용한다.
• 최소한의 수작업 조정으로 고성능 CUDA 커널을 생성하기 위한 체계적인 코드 생성 파이프라인을 구현한다.

실험 결과

연구 질문

• RQ1 다각형 편집 기법이 Volta 텐서 코어용 고성능 CUDA 커널을 효과적으로 생성하는 데 적용될 수 있는가?
• RQ2 자동 커널 생성 기법이 행렬 곱셈과 포인트와이즈 연산을 융합한 경우, cuBLAS 및 cuDNN과 같은 수작업 최적화 라이브러리 대비 얼마나 뛰어난 성능을 낼 수 있는가?
• RQ3 mma.sync.m8n8k4 PTX 명령어의 사용이 융합된 커널에서 고도의 산술 집중도와 메모리 대역폭 효율성을 달성하는 데 얼마나 효과적인가?
• RQ4 다양한 워크로드에서 타일 크기와 융합 전략의 성능 상충 관계는 어떠한가?
• RQ5 16×16×8 및 32×32×8 크기의 매크로-MMA 조합이 다각형 분석을 통해 자동으로 생성되고 최적화될 수 있는가, 원자적 mma 연산 대비 성능 향상이 이루어지는가?

주요 결과

• 제안된 방법은 행렬 곱셈과 포인트와이즈 연산을 융합한 경우, cuBLAS 및 cuDNN 대비 최대 2.55배 빠른 성능을 달성했으며, 특히 더 큰 문제 크기에서 뚜렷한 성능 향상을 보였다.
• 단일 행렬 곱셈 커널의 경우, cuBLAS 대비 1.46배, cuDNN 대비 1.35배의 성능 향상을 기록하여 융합 없이도 효과적인 성능 향상을 입증했다.
• 여러 개의 행렬 곱셈을 포함한 융합 커널은 더 높은 레지스터 압박을 유발하여 타일 크기를 작게 조정해야 했지만, 성능 향상은 여전히 유지되었다.
• Halide 및 TVM과 같은 기존 프레임워크에 비해 텐서 코어 커널 생성에서 뛰어난 성능을 보였으며, 이는 직접 mma.sync.m8n8k4 명령어 사용과 매크로-MMA 조합 덕분이었다.
• 다양한 문제 크기에서 뚜렷한 성능 향상을 달성했으며, 특히 행렬 곱셈과 ReLU 또는 편향 덧셈을 포함한 융합 계산 그래프에서 가장 큰 성능 향상이 관찰되었다.
• 다각형 편집 기법이 저수준 텐서 코어 명령어를 직접 타겟으로 삼을 수 있음을 입증하며, 고도로 최적화된 커널을 자동 생성할 수 있음을 보여주었다.