[논문 리뷰] Sketch-Guided Equality Saturation: Scaling Equality Saturation to Complex Optimizations of Functional Programs

Thomas Kœhler, Phil Trinder외 1명|arXiv (Cornell University)|2021. 11. 25.

Parallel Computing and Optimization Techniques인용 수 5

한 줄 요약

이 논문은 효율적인 람다 계산법 인코딩과 프로그래머가 제공한 프로그램 스케치를 통해 등가 포화를 안내하는 기법인 스케치 유도 등가 포화를 소개한다. 이 기법은 기능적 프로그램 최적화에서 복잡한 행렬 곱셈 최적화(예: 루프 블로킹, 벡터화, 다중 스레딩)를 1GB 미만의 메모리와 몇 초 내에 발견할 수 있게 하며, 반면에 안내되지 않은 등가 포화는 한 시간이 넘게 걸리고 60GB의 메모리가 필요해도 실패한다.

ABSTRACT

Generating high-performance code for diverse hardware and application domains is challenging. Functional array programming languages with patterns like map and reduce have been successfully combined with term rewriting to define and explore optimization spaces. However, deciding what sequence of rewrites to apply is hard and has a huge impact on the performance of the rewritten program. Equality saturation avoids the issue by exploring many possible ways to apply rewrites, efficiently representing many equivalent programs in an e-graph data structure. Equality saturation has some limitations when rewriting functional language terms, as currently naive encodings of the lambda calculus are used. We present new techniques for encoding polymorphically typed lambda calculi, and show that the efficient encoding reduces the runtime and memory consumption of equality saturation by orders of magnitude. Moreover, equality saturation does not yet scale to complex compiler optimizations. These emerge from long rewrite sequences of thousands of rewrite steps, and may use pathological combinations of rewrite rules that cause the e-graph to quickly grow too large. This paper introduces \emph{sketch-guided equality saturation}, a semi-automatic technique that allows programmers to provide program sketches to guide rewriting. Sketch-guided equality saturation is evaluated for seven complex matrix multiplication optimizations, including loop blocking, vectorization, and multi-threading. Even with efficient lambda calculus encoding, unguided equality saturation can locate only the two simplest of these optimizations, the remaining five are undiscovered even with an hour of compilation time and 60GB of RAM. By specifying three or fewer sketch guides all seven optimizations are found in seconds of compilation time, using under 1GB of RAM, and generating high performance code.

연구 동기 및 목표

• 등가 포화의 확장성 한계를 해결하여, 수천 개의 재작성 단계가 필요한 복잡한 기능적 프로그램 최적화에 대응한다.
• 결합 법칙과 교환 법칙과 같은 병리적인 재작성 규칙 조합으로 인해 발생하는 지수적 증가하는 e-그래프 성장을 극복한다.
• 고성능 계산 분야에서 사용되는 다형적 타입을 가진 기능적 언어인 Rise와 같은 언어에 등가 포화의 실용적 적용을 가능하게 한다.
• 재작성 전략에서 수동 단계 순서를 설정할 필요를 줄이기 위해 반자동적, 스케치 기반의 안내 메커니즘을 도입한다.
• 등가 포화를 실세계 기능적 프로그램 최적화에 적용하기 위해 람다 계산법의 효율적 인코딩이 필수적이라는 것을 입증한다.

제안 방법

• 프로그래머가 부분적인 프로그램 패턴(세부 사항이 생략된)을 제시하는 스케치 유도 등가 포화를 도입하여 최적화 과정을 안내한다.
• 복잡한 최적화를 작은, 다룰 수 있는 등가 포화 검색 단계들로 분해하고, 각 단계에서 검색 공간을 제약하기 위해 스케치를 사용한다.
• 데 브리쥐느 인덱스와 추출 기반 치환을 사용하여 다형적 타입을 가진 람다 계산법의 효율적인 인코딩을 개발하여 e-그래프의 과도한 증가를 최소화하고 성능을 향상시킨다.
• Risegg 등가 포화 프레임워크에 새로운 인코딩을 통합하여 이름 바인딩이 있는 기능적 언어의 표현식 재작성 기능을 지원한다.
• 포화 후 e-그래프에서 최고 성능을 내는 프로그램을 비용 모델을 사용해 추출하여 의미적 동치성과 최적화된 성능을 보장한다.
• 스케치 유도 등가 포화를 Rise 기능적 언어에서의 일곱 가지 행렬 곱셈 최적화에 적용하여, 무안내 등가 포화 및 TVM과의 성능과 자원 사용량을 평가한다.

실험 결과

연구 질문

• RQ1 스케치 유도 등가 포화는 등가 포화를 수천 개의 재작성 단계가 필요한 기능적 프로그램 최적화에 효과적으로 확장시킬 수 있는가?
• RQ2 효율적인 람다 계산법 인코딩은 기능적 프로그램 최적화에서 등가 포화의 메모리 소비와 런타임 오버헤드를 어떻게 줄이는가?
• RQ3 프로그래머가 제공한 스케치는 재작성 규칙의 수동 단계 순서 없이 복잡한 최적화를 얼마나 잘 안내할 수 있는가?
• RQ4 무안내 등가 포화 및 최신 기술의 컴파일러인 TVM과 비교해 스케치 유도 등가 포화는 어떤 성능 및 자원 사용량 향상을 달성하는가?
• RQ5 효율적인 인코딩과 스케치 안내의 조합이 e-그래프 폭발로 인해 도달할 수 없는 최적화를 발견하는 데 기여하는가?

주요 결과

• 스케치 유도 등가 포화는 1GB 미만의 메모리와 몇 초 이내에 일곱 가지 복잡한 행렬 곱셈 최적화(루프 블로킹, 벡터화, 다중 스레딩 포함)를 모두 발견했다.
• 무안내 등가 포화는 한 시간이 넘게 실행되고 60GB의 메모리가 투입된 후에도 일곱 가지 최적화 중 다섯 가지를 발견하지 못해, 난이도 높은 등가 포화의 확장성 한계를 드러냈다.
• 새로운 람다 계산법 인코딩은 난이도 높은 인코딩 대비 런타임과 메모리 소비를 수십만 배 이상 줄여 등가 포화를 기능적 프로그램 최적화에 실용적으로 만들었다.
• 스케치 유도 등가 포화에서 생성된 코드는 최신 기술의 TVM 컴파일러와 수준이 유사했으며, 재작성 전략을 통해 수작업 최적화된 코드와도 유사한 성능을 보였다.
• 세 개 이하의 스케치 안내만으로도 모든 복잡한 최적화를 성공적으로 탐색하여 스케치 기반 안내가 검색 공간을 효과적으로 제약할 수 있음을 입증했다.
• 이 방법은 기술적 스케치와 등가 포화를 조합한 하이브리드 최적화 전략을 가능하게 하여, 기능적 언어를 위한 상호작용적이고 고성능 코드 생성의 실용적인 길을 제시한다.