기술 요약: Tensor Seeks Layout: ML 컴파일러를 위한 레이아웃 선택의 정형화

1. 문제 정의

현대적인 머신러닝(ML) 컴파일러는 실행 비용을 최소화하면서 하드웨어 제약 조건을 준수하도록 텐서의 메모리 레이아웃을 선택해야 합니다. 텐서 레이아웃은 다차원 배열의 물리적 표현(예: NHWC와 NCHW 같은 차원 순서, 또는 시스톨릭 어레이를 위한 파티션 차원)을 정의합니다. 이 선택 과정은 본질적으로 전역적(global)입니다. 즉, 특정 연산자는 하나의 레이아웃에서 가장 빠르게 실행될 수 있지만, 그 연산자의 소비자(consumer)는 다른 레이아웃을 선호할 수 있습니다. 이러한 선호도의 불일치를 조정하기 위해서는 명시적인 레이아웃 변환(transpose, copy 또는 reformatting)이 필요하며, 이는 상당한 오버헤드를 발생시킵니다.

현재의 컴파일러들은 로컬 그리디 패스(local greedy pass)나 규칙 기반 전파(rule-based propagation)와 같은 임시방편적인 휴리스틱에 의존하고 있습니다. 그러나 이러한 접근 방식은 이론적 근거가 부족하여 종종 최적이 아닌 성능을 초래합니다. 본 논문은 연산자 실행 비용과 레이아웃 변환 비용 사이의 트레이드오프를 고려하는 통합된 수학적 모델의 부재라는 결정적인 격차를 식별합니다.

2. 방법론 및 정형화

2.1 수학적 정식화

저자들은 레이아웃 선택을 데이터플로 그래프($G = (V, E)$) 상의 조합 최적화 문제로 정식화합니다.

• 입력: 연산자를 나타내는 정점(vertex)과 텐서를 나타내는 간선(edge)으로 구성된 유향 비순환 그래프(DAG). 각 텐서 $t$는 유한한 허용 레이아웃 집합 $L(t)$를 가집니다.
• 비용:
  • 연산자 비용 ($C^{op}_o$): 입력 및 출력 레이아웃 조합을 실행 비용(또는 불가능한 구성의 경우 $\infty$)으로 매핑하는 함수.
  • Transpose 비용 ($C^{tr}_t$): 텐서의 프로듀서(producer) 레이아웃에서 컨슈머(consumer)가 요구하는 레이아웃으로의 변환 비용을 나타내는 함수.
• 목표: 모든 연산자 $o$와 텐서 $t$에 대해 연산자 실행 비용과 필요한 모든 변환 비용의 합을 최소화하는 레이아웃 할당 $\ell(t, o)$를 찾는 것입니다.

이 정식화는 XLA의 레이아웃 할당(차원 순서), 시스톨릭 어레이 파티션 차원 선택(예: AWS Trainium), 모바일 GPU 최적화를 포함한 기존 전략들을 통합합니다.

2.2 이론적 분석

본 논문은 문제의 계산 복잡도를 확립합니다:

• NP-난해성(NP-Hardness): 레이아웃 선택은 2D 텐서에 대한 행렬 곱셈만을 포함하고 각 텐서당 두 개의 레이아웃을 가진 프로그램으로 제한되더라도 NP-난해임이 증명되었습니다.
• 근사 불가능성(Inapproximability): P=NP가 아니라면, 일반적인 경우에 대해 유한한 근사 비율을 갖는 다항 시간 근사 알고리즘은 존재하지 않습니다.
• 고정 매개변수 가동성(Fixed-Parameter Tractability): 이 문제는 **제한된 트리에드폭(bounded treewidth)**을 가진 데이터플로 그래프에 대해 최적의 다항 시간 알고리즘을 허용합니다. 저자들은 이 경우에 대한 명시적인 동적 계획법(DP) 알고리즘을 제공하며, 이는 고정된 트리에드폭에 대해 연산자 수에 선형적입니다.
• 일반 솔빙(General Solving): 일반적인 인스턴스의 경우, 문제를 가중 MaxSAT 인스턴스로 인코딩하여 기성 솔버(off-the-shelf solver)가 최적의 해를 찾을 수 있도록 합니다.

2.3 구현

저자들은 이를 AWS Trainium AI 가속기용 프로덕션 컴파일러에 구현했습니다. 그들은 네 가지 전략을 평가했습니다:

1. Local: 변환을 무시하고 연산자 비용만을 최소화합니다.
2. Greedy: 전체 목적 함수(연산자 + 변환 비용)를 최소화하기 위해 위상 정렬된 할당을 반복적으로 개선합니다.
3. Rule-Based: 도메인 특화 규칙과 전파를 사용하는 기존의 컴파일러 전략입니다.
4. Solver-Based: 최적의 해를 구하기 위해 정확한 Treewidth DP 및 MaxSAT 인코딩을 사용합니다.

3. 주요 결과

3.1 레이아웃 선택의 실질적 영향

• 성능 이득: 단순한 휴리스틱은 복잡한 워크로드(예: ResNet-50, Vision Perceiver)에서 최적화된 전략에 비해 실행 시간을 최대 5배까지 저하시킬 수 있습니다.
• 워크로드 민감도: 정규적인 밀집 트랜스포머 아키텍처(예: BERT, ALBERT)의 경우, 레이아웃 선호도가 일관되기 때문에 단순한 "Local" 베이스라인도 최적 전략과 거의 유사한 성능을 보입니다. 그러나 이질적인 연산자(컨볼루션, 크로스 어텐션, 스킵 커넥션 등)가 포함된 복잡한 아키텍처에서는 전역 최적화가 필수적입니다.
• 성능 저하 메커니즘: 잘못된 레이아웃 선택은 삽입되는 transpose 연산의 급격한 증가를 초래합니다. Trainium과 같은 시스톨릭 어레이에서 이러한 변환은 소비 연산에 퓨즈(fuse)될 수 없는 온칩 파티션 간의 물리적 데이터 이동을 수반하며, 이는 상당한 DMA 오버헤드를 유발합니다.

3.2 정확한 솔버의 생존 가능성

• Treewidth: 경험적 분석 결과, 실제 ML 데이터플로 그래프의 92%는 트리에드폭 $\le 4$를 가집니다. Treewidth 기반 DP 알고리즘은 이러한 인스턴스를 수 초 내에 해결합니다. 그러나 나머지 8%(treewidth > 4)의 경우 메모리 제약으로 인해 실패합니다.
• MaxSAT: MaxSAT 인코딩은 대부분의 모델에 대해 실행 가능하며, 수 초 내에 인스턴스를 해결합니다. 솔버가 타임아웃되는 대규모 모델의 경우, 처음 몇 초 동안 찾아낸 최적에 가까운 해가 매우 유용합니다. 왜냐하면 솔버는 더 나은 할당을 찾는 것보다 최적성을 증명하는 데 대부분의 시간을 소비하기 때문입니다.

3.3 비용 모델 충실도(Fidelity)의 역할

핵심적인 발견은 비용 모델 하에서 증명된 최적의 솔루션이 항상 최고의 하드웨어 성능을 보장하는 것은 아니라는 점입니다.

• 비용 모델 오류: 솔버는 공식적인 목적 함수를 정확하게 최소화합니다. 솔버가 휴리스틱보다 성능이 떨어지는 경우(예: ResNet-50), 그 격차는 전적으로 비용 모델의 부정확성, 특히 레이아웃 선택 이후에 이루어지는 타일링 결정에서 발생하는 DMA 전송 비용의 과소평가에 기인합니다.
• 휴리스틱의 이점: Rule-Based 전략과 같은 휴리스틱은 때때로 솔버보다 뛰어난 성능을 보이는데, 이는 분석적 비용 모델이 포착하지 못하는 도메인 지식(예: 선호되는 파티션 차원)을 포함하고 있기 때문입니다. 이들은 관습적인 레이아웃 근처에 머물러 모델 오류를 효과적으로 "평균화"하는 반면, 솔버는 모델의 과소평가된 부분을 공격적으로 활용합니다.

4. 의의 및 주장

본 논문은 ML 컴파일러의 레이아웃 선택에 관한 최초의 정식 연구라고 주장합니다. 주요 기여는 다음과 같습니다:

1. 통합된 정식화: 기존의 다양한 "임시방편적" 레이아웃 최적화를 포괄하는 단일 수학적 목적 함수를 제공하여 설계 선택을 엄격하게 비교할 수 있게 합니다.
2. 복잡도 경계: NP-난해성과 근사 불가능성을 증명하는 동시에, 제한된 트리에드폭이 구체적인 동적 계획법 알고리즘을 가진 다항 시간 내 해결 가능한 영역임을 식별합니다.
3. 경험적 검증: 프로덕션 컴파일러에의 통합을 통해, 많은 워크로드에서 정확한 솔버가 실행 가능하지만, 궁극적인 성능 상한선은 탐색 알고리즘이 아닌 비용 모델의 충실도에 의해 결정된다는 것을 입증합니다.
4. 실무적 가이드라인: 저자들은 컴파일러 개발자들이 휴리스틱을 사용할 때도 레이아웃 선택을 최적화 문제로 정식화해야 한다고 주장합니다. 이는 디버깅과 최적성 격차 측정을 위한 참조점을 제공합니다. 또한, 휴리틱의 최적성 격차가 비용 모델의 오류보다 작아지면, 추가적인 투자를 탐색 알고리즘이 아닌 비용 모델의 충실도(예: 캘리브레이션 또는 학습된 모델)를 높이는 쪽으로 전환해야 한다고 제안합니다.

논문은 정확한 솔버가 정규적인 워크로드에서는 규칙 기반 전략과 대등하거나 더 나은 성능을 보일 수 있지만, 복잡한 워크로드에서의 "격차"는 더 정교한 탐색 알고리즘보다 더 정확한 비용 모델의 필요성을 강조한다는 점을 밝히며 결론을 맺습니다.