[논문 리뷰] DISC: A Dynamic Shape Compiler for Machine Learning Workloads

Kai Zhu, Wenyi Zhao외 9명|arXiv (Cornell University)|2021. 03. 09.

Parallel Computing and Optimization Techniques참고 문헌 9인용 수 4

한 줄 요약

DISC는 기계 학습 워크로드를 위한 동적 형상 컴파일러로, MLIR에 완전히 동적 형상 IR(DHLO)를 추가하여 런타임 해석 오버헤드를 방지하기 위해 컴파일 타임에 효율적인 런타임 제어 흐름을 생성하고, 형상 전파 및 제약 조건 수집을 통해 호스트-디바이스 공동 최적화를 가능하게 한다. 이는 트랜스포머와 같은 동적 형상 워크로드에서 텐서플로우/파이토치 대비 최대 3.3배 빠르고, 니블 기반 대비 1.8배 빠르다.

ABSTRACT

Many recent machine learning models show dynamic shape characteristics. However, existing AI compiler optimization systems suffer a lot from problems brought by dynamic shape models, including compilation overhead, memory usage, optimization pipeline and deployment complexity. This paper provides a compiler system to natively support optimization for dynamic shape workloads, named DISC. DISC enriches a set of IR to form a fully dynamic shape representation. It generates the runtime flow at compile time to support processing dynamic shape based logic, which avoids the interpretation overhead at runtime and enlarges the opportunity of host-device co-optimization. It addresses the kernel fusion problem of dynamic shapes with shape propagation and constraints collecting methods. This is the first work to demonstrate how to build an end-to-end dynamic shape compiler based on MLIR infrastructure. Experiments show that DISC achieves up to 3.3x speedup than TensorFlow/PyTorch, and 1.8x than Nimble.

연구 동기 및 목표

• 기존 컴파일러가 고비용 컴파일 오버헤드, 메모리 팽창, 융합 기회 제한 등의 문제로 인해 발생하는 동적 형상 기계 학습 워크로드의 성능 저하 문제를 해결한다.
• XLA 및 유사 시스템에서의 정적 형상 최적화의 한계를 극복하여, 동적 형상을 효율적으로 처리하지 못하고 이러한 워크로드에 대해 융합을 비활성화하는 문제를 해결한다.
• 런타임 해석이나 팯딩 기반의 대체 방법에 의존하지 않고도 동적 형상 계산을 원천적으로 지원하는 솔루션을 제공한다.
• 컴파일 타임에 제어 흐름 코드를 생성하여 동적 형상 환경에서 효율적인 커널 융합 및 호스트-디바이스 공동 최적화를 가능하게 한다.
• MLIR 인프라를 기반으로 한 종단 간 동적 형상 컴파일링을 구현하며, 다양한 프레임워크를 지원하고 가능할 경우 정적 최적화로 자동 후퇴한다.

제안 방법

• MLIR의 HLO 다이어ект를 확장하여 DHLO(Dynamic HLO)를 생성함으로써, 새로운 IR를 새로 설계하지 않고도 동적 형상 계산을 원천적으로 지원하는 완전히 동적 형상 표현을 구현한다.
• 런타임 제어 흐름(형상 추론, 버퍼 관리, 호스트 측 로직 포함)을 컴파일 타임에 전체적으로 생성하여 VM 해석 오버헤드를 방지하고 호스트-디바이스 공동 최적화를 가능하게 한다.
• 생산자와 소비자 간의 형상 전파를 통해 형상이 호환되는 인접 연산 간의 융합 후보를 식별한다.
• DHLO로 내림내림 단계에서 형상 제약 조건을 수집함으로써, 컴파일 타임에 전체 형상 정보를 확보하지 못하더라도 형상 호환성에 대해 추론할 수 있도록 하여 더 큰 규모의 융합을 가능하게 한다.
• 하나의 서브그래프가 알려진 정적 형상을 가질 경우 자동으로 정적 컴파일링으로 후퇴함으로써 정적 및 동적 최적화를 혼합으로 지원한다.
• DHLO 허브를 통해 텐서플로우 및 파이토치와 통합하여 동적 형상 모델의 종단 간 컴파일링을 가능하게 한다.

실험 결과

연구 질문

• RQ1 런타임 해석에 의존하지 않고도 기계 학습 컴파일러가 동적 형상 계산을 효율적으로 표현하고 최적화할 수 있는 방법은 무엇인가?
• RQ2 컴파일 타임에 전체 형상 정보를 확보할 수 없는 동적 형상 워크로드에서 커널 융합을 가능하게 하는 기법은 무엇인가?
• RQ3 동적 형상 실행에서 런타임 오버헤드를 줄이기 위해 호스트-디바이스 제어 흐름을 어떻게 공동 최적화할 수 있는가?
• RQ4 동적 형상 컴파일러가 유연성을 유지하면서도 정적 컴파일러 수준의 성능을 달성할 수 있는 정도는 어느 정도인가?
• RQ5 MLIR 인프라를 기반으로 한 컴파일러 시스템이 최소한의 런타임 비용으로 종단 간 동적 형상 컴파일링을 원천적으로 지원할 수 있는가?

주요 결과

• DISC는 트랜스포머와 같은 동적 형상 워크로드에서 텐서플로우 및 파이토치 대비 평균 2.27배 빠르며, 일부 모델에서는 최대 3.3배 빠르다.
• 트랜스포머의 메모리 집약적인 연산에서는 DISC가 니블 대비 2.61배 빠르며, 주로 형상 힌트 및 제약 조건 덕분에 더 효율적인 융합 덕분이다.
• DISC는 커널 수가 다소 감소했음에도 불구하고, 효율적인 코드 생성 및 공동 최적화 덕분에 런타임 제어 흐름에 소요되는 CPU 시간을 니블의 36.6%로 줄였다.
• DISC는 평균적으로 니블 대비 1.8배 빠르며, VM 기반 해석 대비 컴파일 타임 흐름 생성의 이점을 입증한다.
• 정적 최적화로의 후퇴를 비활성화한 경우 DISC는 평균적으로 정적 컴파일링 성능의 85%를 기록했으며, 범위는 74.5%에서 91.4%까지 다양했다. 이는 전체 형상 지식이 없을 경우 융합 기회가 제한되어 성능 격차가 발생함을 시사한다.
• 트랜스포머 모델에서 커널 수는 니블의 13,924개에서 DISC의 10,734개로 감소하여, 명시적인 커널 수 감소가 주요 성능 요인는 아니지만 융합 효율성이 향상됨을 보여준다.