[논문 리뷰] Stripe: Tensor Compilation via the Nested Polyhedral Model

Tim Zerrell, Jeremy Bruestle|arXiv (Cornell University)|2019. 03. 14.

Parallel Computing and Optimization Techniques참고 문헌 20인용 수 25

한 줄 요약

이 논문은 기계학습 컴파일러를 위한 도메인 특화 중간 표현(IR)인 Stripe를 소개한다. 이는 중첩 다각형 모델 기반으로, 다양한 하드웨어에서 자동으로 고성능 커널 생성을 가능하게 한다. 텐서 연산을 계층적인 다각형 반복 공간으로 모델링함으로써, 알고리즘 설계, 최적화 단계, 하드웨어 타겟 간의 분리를 가능하게 하여 모듈화되고 재사용 가능한 최적화와 현대적 기계학습 가속기에서의 확장 가능한 코드 생성을 가능하게 한다.

ABSTRACT

Hardware architectures and machine learning (ML) libraries evolve rapidly. Traditional compilers often fail to generate high-performance code across the spectrum of new hardware offerings. To mitigate, engineers develop hand-tuned kernels for each ML library update and hardware upgrade. Unfortunately, this approach requires excessive engineering effort to scale or maintain with any degree of state-of-the-art performance. Here we present a Nested Polyhedral Model for representing highly parallelizable computations with limited dependencies between iterations. This model provides an underlying framework for an intermediate representation (IR) called Stripe, amenable to standard compiler techniques while naturally modeling key aspects of modern ML computing. Stripe represents parallelism, efficient memory layout, and multiple compute units at a level of abstraction amenable to automatic optimization. We describe how Stripe enables a compiler for ML in the style of LLVM that allows independent development of algorithms, optimizations, and hardware accelerators. We also discuss the design exploration advantages of Stripe over kernel libraries and schedule-based or schedule-space-based code generation.

연구 동기 및 목표

• 기계학습 알고리즘과 하드웨어 가속기의 빠른 진화를 지원하기 위해 수작업으로 최적화된 커널의 조합 폭발 문제를 해결하기 위해.
• 알고리즘 개발, 최적화 단계, 하드웨어 기반 코드 생성 간의 분리를 통해 유지보수 오버헤드를 줄이고 확장성을 향상시키기 위해.
• 복잡한 병렬성과 메모리 접근 패턴을 체계적이고 확장 가능한 IR을 사용해 모델링함으로써, 기계학습 워크로드에 대한 자동 고성능 코드 생성을 가능하게 하기 위해.
• 최적화 단계가 서로 재사용 가능하고 조합 가능한, LLVM과 유사한 컴파일러 스택의 기반을 제공하기 위해.
• 일반적인 컴파일러와 기존의 기계학습 전용 프레임워크의 한계를 극복하기 위해, 텐서 연산의 수학적으로 탄탄한 계층적 표현을 도입하기 위해.

제안 방법

• 매우 병렬화 가능하고 상호 반복 간 의존성이 제한된 텐서 연산을 표현하기 위해 중첩 다각형 모델을 공식화한 형태로 제안한다.
• 다각형 제약 조건과 중첩 구조를 사용하여 반복 공간, 메모리 레이아웃, 계산 단위 추상화를 포괄하는 도메인 특화 IR인 Stripe를 설계한다.
• 계층적 타일링과 스케줄링을 지원하도록 IR을 구성함으로써, 여러 메모리 레벨과 계산 단위를 아우르는 최적화를 가능하게 한다.
• PlaidML 컴파일러 스택에 Stripe를 통합하여, Tile(고수준 IR)로부터 내림된 고수준 텐서 연산을 수신한 후, 하드웨어에 종속되지 않은 최적화 단계를 적용한다.
• 표준 컴파일러 기법—예를 들어 타일링, 벡터화, 메모리 레이아웃 변환—을 Stripe IR에 적용하며, 최적화는 타겟 하드웨어와 독립적으로 적용된다.
• 중첩 다각형 모델의 공식적 구조를 활용하여, 새로운 최적화 단계를 Stripe IR 기반으로 구축할 수 있도록 확장성을 제공한다.

실험 결과

연구 질문

• RQ1 복잡하고 계층적인 병렬성 및 메모리 접근 패턴을 가진 현대 기계학습 워크로드를 모델링하면서도 자동 최적화에 적합한 컴파일러 IR은 어떻게 설계할 수 있는가?
• RQ2 비직사각형 반복 공간과 컨volution과 같은 연산에서 흔히 발생하는 경계 조건을 효과적으로 표현할 수 있는가?
• RQ3 성능을 희생시키지 않고도 컴파일러 스택에서 알고리즘, 최적화, 하드웨어 간의 책임 분리를 어느 정도 달성할 수 있는가?
• RQ4 표현력, 모ularity, 성능 이식성 측면에서 Stripe IR은 Halide, TVM, 또는 MLIR와 같은 기존 접근 방식과 비교해 어떻게 다를 수 있는가?
• RQ5 각 운영자와 아키텍처 조합에 대해 수작업 최적화가 필요 없이도 다양한 하드웨어 플랫폼에서 자동으로 고성능 코드 생성이 가능한가?

주요 결과

• 중첩 다각형 모델은 히알로 영역과 경계 제약 조건을 다각형 프레임워크 내에서 모델링하여, 비점별 컨볼루션과 같은 복잡한 비직사각형 반복 공간을 성공적으로 포괄한다.
• Stripe의 중첩 블록 구조는 다수의 메모리 레벨과 계산 단위에 대한 효율적 매핑을 가능하게 하는 다수 수준의 타일링을 자연스럽게 지원한다.
• 다양한 기계학습 연산과 하드웨어 타겟 간에 최적화 단계를 독립적으로 개발하고 재사용할 수 있도록 IR이 허용되어, 수작업 최적화된 커널의 필요성을 줄인다.
• GPU 타겟에서의 초도 벤치마크 결과에 따르면, Stripe는 수작업 최적화된 구현과 경쟁 가능한 성능을 보이는 자동 코드 생성을 가능하게 한다.
• PlaidML의 컴파일러 파이프라인에 Stripe를 통합함으로써, 하드웨어에 종속되지 않은 컴파일링과 재사용 가능하고 조합 가능한 최적화 단계를 실현한다.
• 저자들은 향후 Stripe를 MLIR의 다이어렉트로 통합함으로써 상호 운용성을 향상시키고, 다양한 컴파일러 생태계 간 최적화 단계 공유를 보다 넓게 가능하게 할 것으로 기대한다.