LLM-Guided Strategy Synthesis for Scalable Equality Saturation (EggMind) 기술 요약

이 논문은 EggMind이라는 새로운 프레임워크를 제안하며, 이는 대규모 언어 모델 (LLM) 을 활용하여 확장 가능한 동등성 포화 (Equality Saturation, EqSat) 전략을 자동으로 합성하는 데 중점을 둡니다. EqSat 은 e-graph 를 통해 동등한 프로그램들의 방대한 공간을 탐색하고 최적의 코드를 추출하는 강력한 최적화 패러다임이지만, e-graph 의 급격한 폭발 (explosion) 로 인해 실용적인 자원 제약 하에서 적용하기 어렵다는 한계가 있습니다. EggMind 는 수동으로 설계된 전략의 한계를 극복하고, 재사용 가능한 전략 아티팩트를 생성하여 EqSat 의 효율성과 품질을 동시에 개선합니다.

1. 문제 정의 (Problem)

• EqSat 의 확장성 문제: 동등성 포화는 파괴적이지 않은 (non-destructive) 방식으로 모든 동등한 변형을 탐색할 수 있지만, 모든 규칙을 한 번에 적용하면 메모리와 실행 시간이 기하급수적으로 증가하여 (e-graph explosion) 최적화 자체가 불가능해집니다.
• 수동 전략 설계의 비효율성: 현재 EqSat 의 확장성을 제어하기 위해 규칙을 분할하고, 순서를 정하며, e-graph 를 단순화하는 '전략 (Strategy)'이 필요합니다. 그러나 이러한 전략 설계는 대부분 수동으로 이루어지며, 특정 도메인에 종속적이고 대규모 규칙 세트에 대해 잘 일반화되지 않아 자동화 compiler 의 주요 병목 현상이 됩니다.
• LLM 기반 코드 진화의 실패: 최근 LLM 을 이용한 코드 생성이 발전했지만, EqSat 전략을 직접적으로 진화시키는 것은 실패했습니다. 그 이유는 다음과 같습니다:
  1. 재사용 가능한 추상화 부재: 저수준의 백엔드 코드 (예: egglog) 를 직접 생성하면 고수준의 전략 의도를 표현하거나 재사용하기 어렵습니다.
  2. 피드백의 부재: 최종 비용 (cost) 같은 거시적 지표는 전략이 왜 실패했는지 설명하지 못하며, 전체 e-graph 상태는 LLM 에게 너무 방대합니다.
  3. 처리 가능성 (Tractability) 통제 부재: LLM 이 제안한 불안정한 전략은 e-graph 를 통제 불가능하게 성장시키거나 (OOM),过早에 최적해를 잘라내어 품질을 떨어뜨립니다.

2. 방법론 (Methodology)

EggMind 는 오프라인 합성 (offline synthesis) 과 온라인 재사용 (online reuse) 으로 구성된 엔드 - 투 - 엔드 프레임워크입니다.

2.1 EqSatL: 전략을 위한 도메인 특화 언어 (DSL)

LLM 이 생성할 수 있는 재사용 가능한 아티팩트로서 EqSat 전략을 명시적으로 표현하기 위해 EqSatL이라는 DSL 을 도입했습니다.

• 규칙 세트 분할 (Ruleset Partitioning): 규칙을 의미론적 태그 (semantic tags) 로 그룹화하여 어떤 규칙들이 상호작용할지 정의합니다.
• 스케줄 구성 (Schedule Construction): 규칙 적용을 '단계 (phase)', '시퀀스 (sequence)', '반복 (repeat)'으로 조직화하여 탐색의 시간적 흐름을 제어합니다.
• 단순화 제어 (Simplification Control): e-graph 의 크기를 줄이기 위해 휴리스틱 또는 LLM 힌트에 기반한 단순화 노드를 명시적으로 삽입합니다.

2.2 에이전트 워크플로우 (Agentic Workflow)

LLM 을 중심으로 한 반복적인 전략 탐색 프로세스를 수행합니다.

• Strategist: 현재 컨텍스트를 바탕으로 새로운 전략 구조를 제안하거나 기존 전략을 조정합니다.
• Tool Pool:
  • Generator: 새로운 EqSatL 전략을 생성합니다.
  • Evaluator: 다양한 예산 모드 (최소/중간/전체) 로 전략을 평가합니다.
  • Partitioner & Simplifier: 탐색의 안정성을 위해 규칙 분할 조언과 단순화 힌트를 제공합니다.

2.3 핵심 기술

1. 증거 기반 재작성 모티프 캐싱 (Proof-Derived Rewrite Motif Caching):
   • 성공적인 EqSat 실행에서 얻은 방대한 증명 (proof) 대신, **재작성 모티프 (rewrite motifs)**를 추출하여 캐시합니다.
   • 이는 구체적인 규칙 대신 고수준의 의미론적 태그 체인 (예: vec-lifting → vec-mac-opt) 으로 표현되어 LLM 이 재사용 가능한 구조적 증거로 활용할 수 있게 합니다.
2. 처리 가능성 안내 (Tractability Guidance):
   • 의존성 기반 규칙 세트 위험 모델: 규칙 간의 의존성을 분석하여 e-graph 폭발을 유발할 수 있는 불안정한 단계 할당을 방지하고 안전한 분할을 조언합니다.
   • LLM 기반 단순화 힌트: 각 단계별로 어떤 패턴을 우선시하거나 제거할지에 대한 구조적 선호도를 LLM 이 제안하여, 단순화 과정이 국소적 최적화에 그치지 않고 전역 최적화를 방해하지 않도록 유도합니다.

3. 주요 기여 (Key Contributions)

1. EggMind 프레임워크: LLM 을 활용하여 재사용 가능한 EqSat 전략을 합성하는 오픈소스 프레임워크를 최초로 제안했습니다.
2. EqSatL DSL: 규칙 조직화, 스케줄 구성, e-graph 단순화를 위한 명시적이고 재사용 가능한 전략 아티팩트 경계를 정의했습니다.
3. 고급 에이전트 워크플로우: 증명 기반 모티프 캐싱과 처리 가능성 안내를 결합하여, e-graph 성장 하에서도 안정적이고 고품질의 전략을 탐색할 수 있는 메커니즘을 구축했습니다.

4. 평가 결과 (Results)

EggMind 는 벡터화 (Vectorization), XLA 기반 텐서 컴파일, 논리 합성 (Logic Synthesis) 등 세 가지 영역에서 평가되었습니다.

• 벡터화 벤치마크 (Vectorization):
  • 최종 비용 (Cost): 전체 EqSat 대비 45.1% 감소, 전문가가 튜닝한 Isaria 전략 대비 20.6% 감소.
  • 메모리 사용량: 전체 EqSat 대비 69.1% 감소 (최대 95% 이상 감소 사례 존재).
  • 실행 시간: Isaria 대비 2.21 배 속도 향상.
• XLA 기반 텐서 컴파일:
  • 17 개 사례 중 13 개에서 전체 EqSat 보다 낮은 비용 달성.
  • 실행 시간에서 11.89 배의 기하평균 속도 향상.
• 논리 합성 (EqMap Case Study):
  • 규칙 세트가 확장된 환경 (Augmented Rewrite Space) 에서도 안정적인 최적화 수행.
  • 기존 EqMap 기준 대비 비용 33.76% 감소, 메모리 51.94% 감소.
• Ablation Study:
  • 모티프 캐싱 (Motif Caching) 과 반복 구조 (Repeat) 제거 시 성능이 급격히 저하됨 (비용 78% 이상 증가 등) 을 확인하여 이 두 요소의 중요성을 입증했습니다.
  • Raw egglog 코드를 직접 생성하는 일반 LLM 에이전트와 비교했을 때, EqSatL 과 특수 워크플로우를 사용한 EggMind 가 훨씬 효율적이고 안정적인 결과를 보였습니다.

5. 의의 및 결론 (Significance)

• 컴파일러 자동화의 새로운 지평: EqSat 전략 설계라는 난제를 LLM 기반의 오프라인 합성 문제로 재정의함으로써, 도메인별 전략을 자동으로 생성하고 재사용하는 길을 열었습니다.
• 확장성과 품질의 트레이드오프 해결: e-graph 폭발을 억제하면서도 EqSat 이 가진 풍부한 탐색 공간의 이점을 최대한 활용하여, 자원 제약 하에서도 최적의 최적화 결과를 도출할 수 있게 합니다.
• 실용적 적용 가능성: 벡터화, 텐서 연산, 논리 합성 등 다양한 컴파일러 최적화 영역에서 성공적으로 적용 가능함을 입증하여, 현대 e-graph 기반 컴파일러의 전략 계층을 자동화하는 실질적인 솔루션을 제시했습니다.

결론적으로, EggMind 는 LLM 의 추론 능력과 EqSat 의 강력한 탐색 능력을 결합하여, 수동 설계의 한계를 넘어 재사용 가능하고 안정적이며 고품질의 컴파일러 최적화 전략을 자동으로 생성하는 새로운 패러다임을 제시합니다.