PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks

이 논문은 MLIR 기반의 모듈형 컴파일러 인프라인 PolyBlocks 를 제안하여, 분석적 비용 모델과 휴리스틱을 활용한 자동 최적화를 통해 PyTorch 와 JAX 에서 Torch Inductor 나 XLA 와 경쟁력 있는 성능을 달성하고 새로운 AI 칩으로의 이식을 용이하게 한다는 내용을 담고 있습니다.

핵심

PolyBlocks: AI 칩을 위한 '만능 요리사'와 '레시피 자동화 시스템'

이 논문은 PolyBlocks라는 새로운 도구를 소개합니다. 이 도구는 인공지능 (AI) 을 만드는 개발자들이 복잡한 하드웨어 (칩) 에서 최고의 성능을 내도록 도와주는 **'컴파일러 인프라'**입니다.

쉽게 말해, PolyBlocks 는 **"고급 요리사 (AI 모델) 가 어떤 주방 (하드웨어) 에 있든, 그 주방에 맞춰 최고의 요리를 자동으로 만들어내는 자동화 시스템"**이라고 생각하시면 됩니다.

---

1. 왜 PolyBlocks 가 필요한가요? (현재의 문제점)

지금까지 AI 모델을 실행할 때 두 가지 방식이 주로 사용되었습니다.

1. 즉석 실행 (Eager Execution): 요리사가 재료를 하나씩 꺼내서 바로 요리하는 방식입니다. 빠르고 직관적이지만, 효율이 떨어질 수 있습니다.
2. 기존 라이브러리 의존: 요리사가 "이 요리는 이미 유명한 식당에서 만든 레시피 (CuDNN, CuBLAS 같은 상용 라이브러리) 를 가져와서 쓰자"라고 생각하는 방식입니다.
   • 문제점: 이 레시피들은 특정 주방 (예: 엔비디아 GPU) 에만 최적화되어 있습니다. 새로운 주방 (새로운 AI 칩) 이 나오면 다시 레시피를 새로 만들어야 하거나, 아예 쓸 수 없게 됩니다. 또한, 여러 요리를 합쳐서 한 번에 하는 '융합 (Fusion)'이 어렵습니다.

PolyBlocks 의 해결책:
PolyBlocks 는 **"아예 처음부터 재료를 다 직접 손질하고, 주방 기구도 맞춰서 요리하는 방식"**입니다. 미리 만들어진 레시피를 쓰지 않고, **모든 코드를 처음부터 새로 생성 (Code Generation)**합니다. 그래서 어떤 새로운 주방이 나오든, 그 주방에 딱 맞는 최적의 요리법을 자동으로 만들어냅니다.

---

2. PolyBlocks 가 어떻게 작동하나요? (핵심 원리)

PolyBlocks 는 MLIR이라는 '만능 번역기'를 기반으로 합니다. 이 시스템은 5 단계의 공정을 거쳐 코드를 최적화합니다.

🏗️ 1 단계: 재료 준비 (Frontend)

PyTorch 나 JAX 같은 고급 프로그래밍 언어로 쓴 AI 모델을 받아서, 기계가 이해할 수 있는 '중간 언어 (Affine Loop Nests)'로 번역합니다. 마치 고급 요리 레시피를 "밀가루 100g, 계란 2 개" 같은 구체적인 재료 목록으로 바꾸는 것과 같습니다.

🔍 2 단계: 분석과 계획 (Mid-level Optimization)

이 단계에서 PolyBlocks 는 가장 중요한 일을 합니다.

• 조각 내기 (Tiling): 큰 요리를 한 번에 다 하지 않고, 작은 접시 (Tile) 에 나누어 담습니다. 이렇게 하면 주방의 작은 냉장고 (온칩 메모리) 에 재료를 미리 넣어두고 빠르게 조리할 수 있습니다.
• 합치기 (Fusion): 예를 들어, "양파 다지기"와 "양파 볶기"를 따로 하는 대신, 한 번에 "양파 다져서 볶기"로 합칩니다. 이렇게 하면 불필요한 이동 (메모리 접근) 을 줄여 속도가 빨라집니다.
• 새로운 기술: 기존에는 불가능했던 복잡한 요리 (Attention 레이어 등) 도 자동으로 합쳐서 최적화합니다.

🚀 3 단계: 특수 장비 활용 (Mapping to Matrix Units)

최신 AI 칩에는 '행렬 계산 전용 엔진 (Tensor Cores)'이 있습니다. PolyBlocks 는 일반 요리 (일반 계산) 를 이 특수 엔진이 가장 잘 할 수 있는 형태로 변환해 줍니다. 마치 일반 칼질 대신 전기 칼을 사용하는 것과 같습니다.

🛠️ 4 단계: 최종 조리 (Backend)

최종적으로 해당 칩 (GPU 등) 이 직접 실행할 수 있는 기계어 코드를 만들어냅니다.

---

3. 실제 성과는 어떨까요? (실험 결과)

연구진은 PolyBlocks 를 PyTorch와 JAX에 적용하여 엔비디아 GPU 에서 테스트했습니다.

• 기존 방식과의 비교:
  • Torch Inductor (현재 PyTorch 의 표준): 이 방식은 미리 만들어진 레시피 (라이브러리) 를 많이 사용합니다.
  • PolyBlocks: 모든 것을 직접 만들어냅니다.
• 결과:
  • 속도: PolyBlocks 는 Inductor 나 XLA(기존 JAX 컴파일러) 보다 더 빠르거나 비슷하게 작동했습니다. 특히 작은 배치 크기 (Batch Size 1) 나 복잡한 모델에서는 PolyBlocks 가 훨씬 뛰어났습니다.
  • 단일 연산: 행렬 곱셈 (Matmul) 이나 합성곱 (Convolution) 같은 기본 작업에서도 PolyBlocks 가 만든 코드는 수천만 원짜리 전문 요리사 (벤더 라이브러리) 가 쓴 코드와 맞먹는 성능을 냈습니다.
  • 새로운 칩 지원: 새로운 AI 칩이 나왔을 때, PolyBlocks 는 기존 인프라를 재사용해서 새로운 주방에 맞는 요리법을 아주 빠르게 만들어낼 수 있습니다.

---

4. 비유로 정리하기

• 기존 방식 (라이브러리 의존):

  • "이 요리는 '레미제라블'이라는 유명한 식당에서 만든 레시피를 가져와서 해."
  • 장점: 잘 만들어진 레시피라 실패 확률이 낮음.
  • 단점: 새로운 주방 (새 칩) 이 나오면 레시피가 안 맞고, 여러 요리를 합쳐서 한 번에 할 수 없음.

• PolyBlocks 방식 (완전 코드 생성):

  • "이 주방의 크기, 가스불 세기, 칼 종류를 분석해서 지금 당장 이 주방에 딱 맞는 최고의 레시피를 직접 짜서 요리해."
  • 장점: 어떤 주방이든 최적의 성능을 냄. 불필요한 이동 (메모리 접근) 을 줄여서 매우 빠름.
  • 단점: 처음 레시피를 짜는 데 시간이 걸릴 수 있음 (하지만 PolyBlocks 는 이를 자동화함).

결론

PolyBlocks는 AI 개발자들이 복잡한 하드웨어의 성능을 100% 끌어올릴 수 있게 해주는 **자동화된 '최적화 엔진'**입니다. 더 이상 특정 칩에 맞춰 수동으로 코드를 짜거나, 제한된 라이브러리에 의존할 필요가 없게 만들어, 어떤 AI 칩이 나오든 최고의 성능을 자동으로 뽑아내는 미래를 제시합니다.

기술 요약

PolyBlocks: AI 칩 및 프로그래밍 프레임워크를 위한 컴파일러 인프라

이 논문은 AI 프로그래밍 프레임워크 (PyTorch, JAX 등) 와 특수 목적 AI 칩 (GPU 등) 을 연결하는 모듈화되고 재사용 가능한 컴파일러 인프라인 PolyBlocks의 설계와 구현을 소개합니다. PolyBlocks 는 MLIR (Multi-Level Intermediate Representation) 기반이며, 완전한 코드 생성 (fully code-generating) 방식을 채택하여 기존 라이브러리 의존성을 배제하고 고성능 코드를 자동으로 생성합니다.

1. 문제 정의 (Problem)

현재 AI 컴파일러 생태계는 다음과 같은 주요 한계를 가지고 있습니다:

• 라이브러리 의존성: Torch Inductor, XLA 와 같은 주요 컴파일러는 CuDNN, CuBLAS, FlashAttention 과 같은 벤더 최적화 라이브러리에 크게 의존합니다. 이는 새로운 하드웨어 지원 시 재구현이 어렵고, 라이브러리가 없는 새로운 칩에서는 성능을 보장하기 어렵습니다. 또한, 라이브러리 기반 접근은 연산자 간 융합 (Operator Fusion) 을 제한하여 메모리 대역폭 병목 현상을 해결하는 데 한계가 있습니다.
• 프로그래밍 계층의 격차:
  • 고수준 (PyTorch, JAX): 생산성은 높지만 최적화 수준이 낮거나 라이브러리에 의존합니다.
  • 중수준 (Triton, Pallas): 타일링 (Tiling) 과 융합을 프로그래머가 직접 제어해야 하므로 개발 비용이 높습니다.
  • 저수준 (CUDA, C++): 최적화 성능은 높지만 개발 및 유지보수 비용이 매우 높고, 새로운 하드웨어에 대한 이식성이 낮습니다.
• 자동화 부족: 기존 방식은 수동 스케줄링이나 자동 튜닝 (Auto-tuning) 에 의존하는 경우가 많아, 모델과 하드웨어의 조합이 무한히 증가하는 상황에서 확장성이 떨어집니다.

2. 방법론 (Methodology)

PolyBlocks 는 MLIR 인프라를 기반으로 한 5 단계 패스 파이프라인 (Pass Pipeline) 구조를 채택하여 문제를 해결합니다.

핵심 설계 원칙

1. 완전 자동화 (Full Automation): 프로그래머의 수동 스케줄링이나 자동 튜닝 없이, 분석적 비용 모델 (Analytical Cost Models) 과 휴리스틱을 통해 최적화를 자동화합니다.
2. 완전 코드 생성 (Fully Code-Generating): CuDNN 등 외부 라이브러리를 사용하지 않고, 고수준 IR 에서부터 하드웨어 특화 인트린식 (Intrinsics) 까지 모든 저수준 코드를 컴파일러가 직접 생성합니다. 이는 새로운 하드웨어 지원 시 재사용성을 극대화합니다.
3. 재사용 가능한 패스 파이프라인: MLIR 의 패스 파이프라인을 기반으로 하여, 고수준 프레임워크 (PyTorch, JAX, TF) 와 저수준 타겟 (NVIDIA GPU, AMD GPU, CPU) 간의 변환을 표준화된 방식으로 처리합니다.

5 단계 컴파일러 아키텍처

1. S1-S2 (Frontend): 고수준 텐서 연산자를 메모리 참조 (memref) 및 아핀 (Affine) 루프 중첩으로 낮춥니다 (Lowering). 타겟에 무관한 단계입니다.
2. S3 (Mid-level Optimizer): 가장 핵심적인 단계입니다. 타겟 정보를 활용하여 다음과 같은 최적화를 수행합니다.
   • 다중 수준 타일링 (Multi-level Tiling): 로컬리티와 병렬성을 극대화합니다.
   • 슬라이스 기반 아핀 융합 (Slicing-based Affine Fusion): 전통적인 루프 융합을 넘어, 불완전한 중첩 (imperfect nests) 에서도 생산자와 소비자 간의 데이터 의존성을 분석하여 불필요한 임시 버퍼를 제거하고 융합합니다.
   • 메트릭스 유닛 매핑: 어핀 루프 중첩을 GPU 텐서 코어 (WMMA) 나 행렬 유닛에 매핑합니다.
   • Attention 레이어 최적화: reduce-reduce-fusion 및 wmma-fusion 패스를 통해 Attention 연산의 메모리 접근을 최소화하고 레지스터 재사용을 극대화합니다.
   • Convolution 의 On-the-fly Packing: 행렬 곱셈 (GEMM) 으로 변환 시, 메모리 오버헤드를 줄이기 위해 타일 단위로 데이터를 동적으로 패킹합니다.
3. S4-S5 (Backend): 병렬화 (GPU, OpenMP) 및 동기화 추상화를 처리하고, 최종적으로 LLVM IR 및 타겟-specific 어셈블리 (NVPTX 등) 로 변환합니다.

3. 주요 기여 (Key Contributions)

• MLIR 기반 모듈형 인프라: 새로운 AI 칩을 지원하기 위해 복잡한 최적화 로직을 처음부터 다시 작성할 필요 없이, S1-S2 와 표준 패스를 재사용하고 S3-S5 만 타겟에 맞게 커스터마이징하면 됩니다.
• 고급 최적화 기술의 자동화:
  • 슬라이스 기반 융합: 기존 컴파일러가 처리하기 어려웠던 불완전한 루프 중첩과 복잡한 데이터 의존성을 가진 연산자 (예: 컨볼루션, 스텐실) 를 자동으로 융합하여 글로벌 메모리 접근을 제거합니다.
  • Attention 레이어 자동 최적화: FlashAttention 과 유사한 성능을 내기 위해 Attention 연산의 모든 단계를 하나의 커널로 융합하고, WMMA 연산자와 매핑하여 레지스터 수준에서 데이터를 재사용합니다.
  • 임시 GEMM 변환: 컨볼루션 연산을 행렬 곱셈으로 변환할 때, 메모리 오버헤드를 줄이는 'On-the-fly packing'을 자동으로 수행합니다.
• 라이브러리 독립적 고성능: 외부 라이브러리 없이 생성된 코드가 벤더 튜닝 라이브러리 (CuDNN, CuBLAS) 와 경쟁 가능한 성능을 달성함을 입증했습니다.

4. 실험 결과 (Results)

NVIDIA A10 및 A100 GPU 를 사용하여 PyTorch 및 JAX 워크로드에 대해 평가했습니다.

• PyTorch 성능:
  • Batch Size 1 (실시간 추론): PolyBlocks 는 Torch Eager 대비 평균 2.15 배, Torch Inductor 대비 1.4 배, TensorRT 대비 2.4 배 빠른 성능을 보였습니다.
  • Batch Size 8: PolyBlocks 는 Eager 대비 1.8 배 빠르며, Inductor 와 **동등한 성능 (0.97 배)**을 보였습니다. Inductor 는 대용량 배치에 최적화된 라이브러리를 사용하므로 PolyBlocks 와 격차가 줄어듭니다.
• JAX 성능: JAX 모델에서 PolyBlocks 는 JAX Eager 대비 2.12 배, XLA 대비 1.15 배 빠른 성능을 기록했습니다.
• 단일 연산자 성능:
  • 컨볼루션: CuDNN 대비 수백 개의 다양한 컨볼루션 연산에서 경쟁력 있는 성능을 보였으며, 약 50 개 이상의 경우에서 CuDNN 보다 2 배 이상 빠른 성능을 달성했습니다.
  • Matmul: CuBLAS 및 Triton 과 비교했을 때 평균적으로 매우 근접한 성능 (Machine Peak 에 근접) 을 보였습니다.
  • Attention: 최적화 기법 (RR-fusion, WMMA-fusion) 을 적용한 PolyBlocks 는 Inductor 가 사용하는 FlashAttention 커널과 경쟁 가능한 성능을 보였습니다.
• Ablation Study:
  • 텐서 코어 사용은 평균 17 배의 속도 향상을 가져왔습니다.
  • 연산자 간 융합 (Cross-operator fusion) 은 최적화된 코드 대비 2.87 배의 추가 향상을 제공하여, 글로벌 메모리 왕복을 줄이는 중요성을 입증했습니다.

5. 의의 및 결론 (Significance)

PolyBlocks 는 AI 컴파일러 분야에서 다음과 같은 중요한 의의를 가집니다:

1. 새로운 하드웨어 지원의 용이성: 벤더 라이브러리에 의존하지 않는 완전한 코드 생성 방식은 NVIDIA GPU 외의 새로운 AI 가속기 (예: NPU, FPGA, 차세대 GPU) 를 지원할 때 필수적인 인프라를 제공합니다.
2. 최적화의 자동화 및 확장성: 수동 스케줄링 없이도 복잡한 연산자 융합과 메모리 계층 최적화를 자동화하여, 모델과 하드웨어의 조합이 복잡해져도 확장 가능한 솔루션을 제공합니다.
3. 성능 경쟁력: 기존에 "라이브러리 없이는 CuDNN 과 경쟁할 수 없다"는 편견을 깨고, 고수준 프레임워크에서 완전 자동화된 코드 생성으로도 벤더 최적화 라이브러리 수준의 성능을 달성할 수 있음을 증명했습니다.

결론적으로 PolyBlocks 는 AI 프로그래밍 프레임워크와 다양한 하드웨어 아키텍처 간의 격차를 해소하고, 차세대 AI 하드웨어에 대한 빠르고 효율적인 컴파일러 개발을 가능하게 하는 강력한 인프라입니다.