[논문 리뷰] Tensor Processing Primitives: A Programming Abstraction for Efficiency and Portability in Deep Learning & HPC Workloads

Evangelos Georganas, Dhiraj Kalamkar외 17명|arXiv (Cornell University)|2021. 04. 12.

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

한 줄 요약

이 논문은 딥러닝 및 HPC 워크로드를 위한 고성능이고 이식 가능한 구현을 가능하게 하기 위해 '가상 텐서 ISA'로서 최소한의 2D 텐서 연산을 노출하는 이식 가능하고 효율적인 프로그래밍 추상화인 텐서 처리 원자성(TPP)을 소개한다. 알고리즘 명세와 플랫폼에 특화된 최적화를 분리함으로써, TPP는 개발자와 컴파일러가 복잡한 연산자를 이식 가능하게 표현하면서도 다양한 아키텍처에서 뛰어난 성능을 달성할 수 있도록 한다. 이는 종합적인 딥러닝 워크로드와 독립적인 커널에서 최첨단의 구현보다 뛰어난 성능을 발휘한다.

ABSTRACT

During the past decade, novel Deep Learning (DL) algorithms, workloads and hardware have been developed to tackle a wide range of problems. Despite the advances in workload and hardware ecosystems, the programming methodology of DL systems is stagnant. DL workloads leverage either highly-optimized, yet platform-specific and inflexible kernels from DL libraries, or in the case of novel operators, reference implementations are built via DL framework primitives with underwhelming performance. This work introduces the Tensor Processing Primitives (TPP), a programming abstraction striving for efficient, portable implementation of DL workloads with high-productivity. TPPs define a compact, yet versatile set of 2D-tensor operators (or a virtual Tensor ISA), which subsequently can be utilized as building-blocks to construct complex operators on high-dimensional tensors. The TPP specification is platform-agnostic, thus code expressed via TPPs is portable, whereas the TPP implementation is highly-optimized and platform-specific. We demonstrate the efficacy and viability of our approach using standalone kernels and end-to-end DL & HPC workloads expressed entirely via TPPs that outperform state-of-the-art implementations on multiple platforms.

연구 동기 및 목표

• 딥러닝 워크로드와 하드웨어의 급속한 발전에도 불구하고 딥러닝 프로그래밍 방법론의 정체를 해결하기 위해.
• 단일화되고 이식성이 없는 딥러닝 라이브러리 커널과 새로운 연산자에 대한 성능이 열악한 레퍼런스 구현의 한계를 극복하기 위해.
• 저수준 하드웨어 전문 지식이 필요 없이도 효율적이고 플랫폼 최적화된 실행을 가능하게 하는 이식 가능하고 고생산성의 추상화를 제공하기 위해.
• 개발자와 컴파일러 모두에게 알고리즘 명세와 저수준 코드 생성 간의 명확한 분리 구조를 제공하기 위해.
• 표준화되고 최소한의 연산자 세트를 통해 프레임워크, 컴파일러, 하드웨어 플랫폼 간의 재사용과 최적화 공유를 촉진하기 위해.

제안 방법

• 텐서를 위한 가상 명령어셋 아키텍처(ISA)로 사용할 수 있는 컴act하고 정밀도를 고려한 2D 텐서 원자성(TPP)의 집합을 정의한다.
• TPP를 플랫폼에 종속되지 않은 명세로 정의하면서도 실행 구현에 있어서는 고도로 최적화되어, 이식 가능한 코드이면서도 최고 성능을 발휘하도록 설계한다.
• TPP를 사용해 딥러닝 및 HPC 워크로드의 복잡한 고차원 연산자를 표현한다.
• 딥러닝 프레임워크, 벤더 라이브러리(예: oneDNN), 텐서 컴파일러(예: MLIR 다이얼렉트)의 소프트웨어 스택에 내려가는 레이어로 TPP를 통합한다.
• 목표 아키텍처를 대상으로 인트리닉 명령어(예: AVX512, SVE, NEON)를 사용해 TPP 백엔드를 구현하여 저수준 성능을 확보한다.
• 기존 컴파일러 인프라구조(예: MLIR)를 활용해 TPP 기반 커널의 코드 생성 및 최적화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 최소한의 이식 가능한 2D 텐서 연산 집합이 복잡하고 고성능의 딥러닝 및 HPC 워크로드를 위한 효과적인 추상화로 기능할 수 있는가?
• RQ2 TPP를 통한 알고리즘 명세와 플랫폼에 특화된 최적화의 분리가 기존의 딥러닝 라이브러리 커널이나 일반적인 프레임워크 원자성과 비교해 성능과 이식성 면에서 어떻게 향상되는가?
• RQ3 다양한 하드웨어 플랫폼에서 실제 딥러닝 및 HPC 워크로드에서 TPP가 최첨단의 구현보다 얼마나 뛰어난 성능을 발휘할 수 있는가?
• RQ4 MLIR, 딥러닝 프레임워크, 벤더 최적화 라이브러리와 같은 기존 소프트웨어 스택에 TPP를 효과적으로 통합할 수 있는가?
• RQ5 가상 텐서 ISA(TPP)의 사용이 컴파일러와 시스템 간 최적화 공유를 가능하게 하고 중복을 줄이는 데 어떻게 기여하는가?

주요 결과

• TPP를 전적으로 사용해 구현한 종합적인 딥러닝 워크로드, 예를 들어 CNN, 확장 컨볼루션, DLRM, BERT, GNN 등이 다양한 플랫폼에서 최첨단의 구현보다 뛰어난 성능을 발휘한다.
• TPP를 사용한 독립적인 커널 구현은 최적화된 벤더 라이브러리와 레퍼런스 구현보다 뛰어난 성능을 달성한다.
• TPP 추상화는 플랫폼에 특화된 백엔드 최적화를 통해 고생산성과 이식성을 유지하면서도 저수준 성능을 확보한다.
• TPP는 MLIR에 내장된 내림림 다이얼렉트로 자연스럽게 통합되어, 다양한 컴파일러 스택에서 효율적인 코드 생성 및 최적화를 가능하게 한다.
• 데이터/머신러닝 과학자들이 수작업으로 저수준 최적화를 수행할 필요를 줄여, 새로운 연산자의 개발 및 구현을 가속화한다.
• x86, AArch64, ARMv8 SVE, 향후 GPU/POWER 지원 등 다양한 ISA 지원은 TPP 추상화의 확장성과 이식성의 증거이다.