[논문 리뷰] Equality Saturation for Tensor Graph Superoptimization

Yichen Yang, Phitchaya Mangpo Phothilimtha외 4명|arXiv (Cornell University)|2021. 01. 05.

Parallel Computing and Optimization Techniques참고 문헌 29인용 수 31

한 줄 요약

이 논문은 Tensat를 소개한다. 텐서 그래프의 등식 포화를 사용하여 e-그래프를 통해 모든 Rewrite를 동시에 적용하고, 기존의 순차적 방법 대비 최대 16%의 런타임 속도 향상과 48배 빠른 최적화를 달성한다.

ABSTRACT

One of the major optimizations employed in deep learning frameworks is graph rewriting. Production frameworks rely on heuristics to decide if rewrite rules should be applied and in which order. Prior research has shown that one can discover more optimal tensor computation graphs if we search for a better sequence of substitutions instead of relying on heuristics. However, we observe that existing approaches for tensor graph superoptimization both in production and research frameworks apply substitutions in a sequential manner. Such sequential search methods are sensitive to the order in which the substitutions are applied and often only explore a small fragment of the exponential space of equivalent graphs. This paper presents a novel technique for tensor graph superoptimization that employs equality saturation to apply all possible substitutions at once. We show that our approach can find optimized graphs with up to 16% speedup over state-of-the-art, while spending on average 48x less time optimizing.

연구 동기 및 목표

• 텐서 그래프 최적화에서 순차적 재작성 전략의 비효율성에 대한 동기 부여.
• e-그래프를 이용한 등식 포화를 도입해 동등한 그래프의 더 큰 공간을 탐색한다.
• 등식 포화를 비지역적이고 다중 패턴 재작성으로 텐서 그래프에 확장한다.
• 표준 ML 모델에서의 실용적 속도 향상을 시연하고 최적화 시간을 분석한다.
• 확장 가능한 추출 기법(ilp)과 순환 필터링 전략을 제공해 탐색 공간을 관리한다.

제안 방법

• 연산자 기반 DAG를 등식 포화에 확장한 텐서 그래프 표현.
• e-그래프를 통해 파괴적 치환 없이 모든 동등한 그래프를 생성하는 재작성 규칙 세트를 적용한다.
• 다중 패턴 재작성 처리를 가진 예전으로 포화된 e-그래프를 탐색하는 방법을 canonical S-expression 매칭 접근으로 수행한다.
• 루트 e-class에서 최저 비용 그래프를 추출하는 방법으로 탐욕적 방법이나 위상 제약이 있는 ILP 기반 형식을 사용한다.
• 하드웨어(GPU)에서의 연산자 런타임에 기초한 비용 모델을 추출 가이드로 사용한다.
• 추출 시 최적의 노드를 선택하기 위해 ILP 해결 옵션을 제공하고 위상 정렬을 통한 순환 회피 메커니즘을 포함한다.

실험 결과

연구 질문

• RQ1 등식 포화가 텐서 그래프 재작성에서 위상 정렬 문제를 극복해 전역적으로 최적의 그래프를 찾을 수 있는가?
• RQ2 e-그래프를 비지역적이고 다중 패턴의 텐서 재작성에 효과적으로 확장하는 방법은 무엇인가?
• RQ3 런타임 비용을 최소화하기 위한 텐서 그래프의 실용적 추출 방법(탐욕적 vs ILP)은 무엇인가?
• RQ4 다중 패턴 재작성 적용 시 성능과 최적화 시간 간의 트레이드오프는 어떤가?
• RQ5 Tensat은 다양한 모델에서 기존의 순차 검색(TASO)과 어떻게 비교되는가?

주요 결과

모델	검색 시간(초) TASO	실행 시간 속도 향상률(%) TASO	검색 시간(초) Tensat	실행 시간 속도 향상률(%) Tensat
NasRNN	177.3	0.5	45.4	68.9
BERT	13.6	1.4	8.5	9.2
ResNeXt-50	25.3	0.7	5.5	8.8
NasNet-A	1226	10.6	1.9	7.3
SqueezeNet	16.4	0.3	6.7	24.5
VGG-19	8.9	0.4	8.9	8.9
Inception-v3	68.6	5.1	6.3	10.0

• Tensat은 평가된 모델에서 최첨단 TASO에 비해 최대 16%의 런타임 속도 향상을 달성한다.
• 같은 그래프의 수많은 동등한 조합을 압축된 e-그래프 표현으로 표현해 최적화 시간이 대폭 감소(최대 300배 빠름).
• 등식 포화는 순차적 백트래킹보다 훨씬 큰 탐색 공간을 커버하여 대부분의 벤치마크에서 더 나은 속도 향상을 제공한다.
• 일부 경우 다중 패턴 반복을 늘리면 속도 향상이 더 커지지만 e-그래프 폭발과 ILP 타임아웃의 위험이 있어 순환 필터링으로 이를 완화한다.
• Tensat은 속도 향상과 최적화 시간 간에 실용적인 트레이드오프를 제공하며 컴파일 워크플로우에 실용적으로 통합될 수 있다.
• NasRNN, SqueezeNet, Inception-v3와 같은 모델에 대한 평가에서 광범위한 적용 가능성을 보여준다.