Axe: 머신러닝 컴파일러를 위한 단순하고 통합된 레이아웃 추상화

Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers

최신 딥러닝 워크로드의 성능 향상을 위해서는 데이터와 연산을 장치 네트워크, 메모리 계층 구조 및 이기종 가속기 전반에 걸쳐 조정하여 배치해야 합니다. 본 논문에서는 하드웨어 인식을 갖춘 추상화 기법인 Axe Layout을 제시합니다. Axe Layout은 논리적인 텐서 좌표를 명명된 축을 통해 다차원 물리 공간에 매핑합니다. Axe는 타일링, 샤딩, 복제 및 오프셋을 장치 간 분산 및 장치 내 레이아웃 전반에 걸쳐 통합하여, 장치 네트워크에서 스레드까지 일관된 방식으로 집단 연산을 표현할 수 있도록 합니다. Axe를 기반으로, 본 논문에서는 다중 수준의 분산 인식 DSL(Domain Specific Language)과 컴파일러를 설계하여, 단일 커널 내에서 스레드 로컬 제어를 집단 연산자와 결합합니다. 실험 결과, 본 논문의 통합 접근 방식은 최신 GPU 장치 및 멀티 장치 환경, 그리고 가속기 백엔드에서 수동으로 튜닝된 커널과 거의 동등한 성능을 제공함을 보여줍니다.

Scaling modern deep learning workloads demands coordinated placement of data and compute across device meshes, memory hierarchies, and heterogeneous accelerators. We present Axe Layout, a hardware-aware abstraction that maps logical tensor coordinates to a multi-axis physical space via named axes. Axe unifies tiling, sharding, replication, and offsets across inter-device distribution and on-device layouts, enabling collective primitives to be expressed consistently from device meshes to threads. Building on Axe, we design a multi-granularity, distribution-aware DSL and compiler that composes thread-local control with collective operators in a single kernel. Experiments show that our unified approach can bring performance close to hand-tuned kernels on across latest GPU devices and multi-device environments and accelerator backends.