用统一坐标系简化多设备深度学习编译布局问题
Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers
- 通过命名轴将张量逻辑坐标映射到物理空间
- 在多GPU环境下性能接近手工优化内核
- 适合需要高效分布式训练的工程师和研究者
现代深度学习负载的扩展需要在设备网格、内存层级和异构加速器之间协调数据与计算的放置。我们提出Axe Layout，一种硬件感知的抽象，通过命名轴将逻辑张量坐标映射到多轴物理空间。Axe统一了跨设备分布和设备内布局中的分块、分片、复制和偏移，使集体操作可从设备网格一致地表达到线程级别。基于Axe，我们设计了一个多粒度、分布感知的DSL与编译器，将线程本地控制与集体操作在单个内核中组合。实验表明，该统一方法在最新GPU设备及多设备环境和加速器后端上，性能可接近手写优化内核。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。