기술 요약: 컴파일러 월드 모델을 향하여 (Toward Compiler World Models)

문제 정의

텐서 프로그램 최적화는 현대 머신러닝 시스템에 필수적이지만, 가장 효율적인 프로그램을 식별하는 과정은 거대한 조합 탐색 공간으로 인해 어려움을 겪습니다. 온디바이스(on-device) 측정 비용을 피하기 위해, 기존의 오토 스케줄러(예: Ansor, AutoTVM)는 후보 프로그램을 순위 매기기 위한 학습된 비용 모델에 의존합니다. 그러나 이러한 모델들은 일반적으로 후보를 정적 코드 스냅샷(예: 최종 TensorIR 텍스트)으로 취급합니다.

이러한 정적 관점은 두 가지 결정적인 한계를 초래합니다:

1. 변환 문맥의 상실 (Loss of Transformation Context): 이는 프로그램을 생성한 단계별 변환 과정을 무시합니다. 결과적으로, 평가기는 후속 단계의 스케줄링 결정이 이전 결정에 의해 구축된 문맥에 어떻게 의존하는지 이해할 수 없습니다.
2. 구문적 노이즈에 대한 민감도 (Sensitivity to Syntactic Noise): 정적 인코더는 표면적인 코드 변화에 쉽게 현혹됩니다. 예를 들어, 구조적으로 동일한 프로그램(예: 루프 크기가 $N=1024$인 행렬 곱셈과 $N=1023$인 경우)은 경계 처리(loop peeling)로 인해 매우 다른 TensorIR 텍스트를 생성할 수 있으며, 이로 인해 모델은 하드웨어 동작이 동일함에도 불구하고 성능 잠재력을 잘못 분류할 수 있습니다.

방법론

저자들은 스케줄 평가를 프로그램 상태에 대한 **액션 조건부 잠재 역학(action-conditioned latent dynamics)**으로 재정의하는 컴파일러 월드 모델(Compiler World Model) 프레임워크를 제안합니다. 이 프레임워크는 최종적인 정적 코드 산출물을 평가하는 대신, 연속적인 잠재 공간을 통해 프로그램 상태의 진화를 시뮬레이션합니다.

이 프레임워크는 TVM AutoScheduler 내에 구현된 세 가지 학습된 구성 요소로 이루어져 있습니다:

1. 텐서 프로그램 상태 표현 (Encoder)

• 입력: TensorIR 텍스트 (스케줄링 전 및 후 상태 모두).
• 아키텍처: 긴 입력을 위해 조정된 학습 가능한 CodeBERT 기반 아키텍처입니다. 중첩된 청크(chunk)를 토큰화하고 인코딩하기 위해 슬라이딩 윈도우를 사용하며, 이를 마스크드 평균 풀링(masked mean pooling)을 통해 밀집 벡터 표현 $z = E(s)$로 집계합니다.
• 학습: 유사한 TensorIR 상태(동일한 워크로드 소스에서 유래된)를 잠재 공간에서 가깝게 배치하고, 서로 다른 상태는 멀리 밀어내는 대조 학습(contrastive learning)을 사용합니다. 이를 통해 구문적 변형에도 불구하고 의미적 동등성이 보존되는 상태 공간을 생성합니다.

2. 액션 조건부 다단계 상태 전이 (Action-Conditioned Multi-Step State Transition)

• 목표: 중간 AST를 명시적으로 구현하지 않고, 일련의 스케줄링 액션에 의해 유도된 최종 상태 표현을 예측하는 것입니다.
• 메커니즘: 초기 상태 임베딩 $z_0$에서 시작하여, 스케줄링 액션 시퀀스 $a_{1:T}$의 효과를 단계별로 전개(rollout)합니다.
  • 액션은 직렬화되어 밀집 임베딩으로 인코딩됩니다.
  • 전이 모델 $F(\cdot)$은 원래 지식 그래프 표현 학습에서 유래된 TransH 기반 아키텍처를 사용하여 다음 잠재 상태를 예측합니다: $z_t = F(z_{t-1}, a_t)$.
  • 이는 스케줄링 액션을 관계 특화 하이퍼플레인(relation-specific hyperplane)에서의 기하학적 변환으로 모델링하여, 결정의 출처(provenance)를 추적하고 참조 모호성(예: 퓨전 후 루프 변수가 무엇을 나타내는지)을 해결합니다.
• 학습: 상태 매칭 항(그라운드 트루스 미래 상태와의 거리 최소화)과 마진 기반 랭킹 항(예측된 상태가 부정 샘플보다 그라운드 트루스에 더 가깝도록 보장)을 결합한 다단계 전개 손실을 사용합니다.

3. 랭킹 기반 후보 평가 (Ranking-Based Candidate Evaluation)

• 입력: 예측된 최종 상태 표현 $\hat{z}_T$, 구조화된 액션 특징(예: 액션 횟수, 시퀀스 길이), 그리고 선택적 하드웨어 특징.
• 출력: 동일 워크로드 내에서의 상대적 후보 우선순위에 대한 스칼라 점수.
• 아키텍처: LambdaRank 목적 함수로 학습된 XGBoost 랭킹 모델입니다.
• 학습 전략: 모델은 절대적인 실행 시간을 예측하지 않습니다. 대신, 측정된 실행 시간을 기반으로 워크로드 그룹 내에서 후보들의 순위를 매기는 법을 학습하며, 이를 학습 기반 랭킹(learning-to-rank) 작업으로 취급합니다.

주요 기여

1. 컴파일러를 위한 월드 모델 관점: 본 논문은 텐서 최적화를 정적 코드 평가가 아닌 액션 조건부 잠재 상태 진화로서의 동적 프로세스로 적응시켜, 후보 스케줄을 모델링하는 프레임워크를 구축했습니다.
2. 액션-상태 궤적 데이터셋: 다단계 예측을 위해 튜닝 로그를 정렬된 액션-상태 궤적(스케줄링 전 상태, 액션 시퀀스, 중간 상태, 스케줄링 후 상태)으로 구성하는 TVM/TenSet 기반 데이터셋을 구축했습니다.
3. 효율적인 탐색 통합: 이 프레임워크를 TVM AutoScheduler에 구현하여, 잠재 역학을 모델링하는 것이 높은 검색 품질을 유지하면서도 오토튜닝 비용을 크게 줄임을 입증했습니다.

실험 결과

본 방법은 Intel Xeon Gold 6430 CPU와 NVIDIA RTX 4090 GPU에서 7개의 신경망 모델(ResNet, BERT, GPT, OPT 포함) 및 103개의 추출된 서브그래프를 대상으로 평가되었습니다.

• Ansor 대비 성능: 64회의 시도(trial) 예산 하에서, 이 방법은 대표 서브그래프 레이턴시를 Ansor 대비 GPU에서 1.37배, CPU에서 1.54배 개선했습니다.
• 샘플 효율성: 이 방법은 단 1,000회의 시도만으로 Ansor-10K(10,000회 시도)와 대등한 검색 품질을 달성합니다(10배 적은 측정). 구체적으로, 대표 서브그래프에 대해 산술 평균에서 Ansor-10K를 약간 상회하며, 기하 평균에서는 2.2% 이내의 차이로 대등한 성능을 보입니다.
• 엔드 투 엔드 가속: PyTorch/PyTorch-opt (cuDNN)와 비교했을 때, 이 방법은 기하 평균 가속도 4.61배 (CPU) 및 **3.67배 (GPU)**를 전달하며, 피크 가속도는 최대 58.61배에 달합니다.
• 절제 연구 (Ablation Studies): 학습된 전이 모듈을 제거했을 때(정적 초기 또는 최종 상태 인코딩으로 회귀) 레이턴시가 크게 증가(중앙값 약 56~69%)하였으며, 이는 중간 상태의 진화를 모델링하는 것이 가장 핵심적인 구성 요소임을 확인시켜 줍니다.

의의 및 한계

의의:
본 논문은 컴파일러 최적화를 정적 코드 평가 문제가 아닌 상태 전이의 동적 프로세스로 취급함으로써, 평가기가 스케줄링 결정 사이의 의미적 의존성을 포착할 수 있게 한다고 주장합니다. 이 접근 방식은 구문적 노이즈를 효과적으로 억제하고 참조 모호성을 해결하여, 더 적은 하드웨어 측정으로도 더 효율적인 탐색을 가능하게 합니다.

한계 (저자들이 명시한 내용):

1. 탐색 공간에 대한 의존성: 이 방법은 후보를 평가하는 것이지 생성하는 것이 아닙니다. 따라서 그 효과는 강력한 스케줄을 제안하는 기본 스케줄러의 능력에 의해 제한됩니다. 즉, 제안되지 않은 스케줄에 대해서는 더 나은 평가기도 복구할 수 없습니다.
2. 오차 누적: 잠재 공간 내에서 길거나 복잡한 스케줄링 궤적을 따라 예측 오차가 누적되어, 까다로운 사례에서 신뢰도가 떨어질 수 있습니다.
3. 랭킹 vs 예측: 모델은 워크로드 내의 랭킹을 위해 학습되었으며, 절대적인 레이턴시 예측을 위한 것이 아닙니다. 이는 온라인 우선순위 지정에는 이상적이지만, 서로 다른 워크로드 간에 보정된 성능 추정치가 필요한 설정에는 적합하지 않을 수 있습니다.