arXiv:2505.01637cs.PLcs.LG2025-05

用动态规划+空间压缩,高效编译出高速神经网络推理代码。

Morello: Compiling Fast Neural Networks with Dynamic Programming and Spatial Compression

  • 通过动态规划分解与重写程序规格,系统探索优化空间。
  • 在Zen 1 CPU上实现1×2048×16384的半精度矩阵乘,速度领先。
  • 适合需要极致推理性能的编译器研发者与高性能计算工程师。

高吞吐神经网络推理需协调并行分块、微核选择和数据布局等众多优化决策,其组合形成的程序搜索空间通常难以穷尽。现有方法(如自动调度器)常通过启发式采样应对。本文提出一种基于动态规划的方法,将大型程序规格逐步分解为可通过重写规则到达的小规格,再从每条重写路径中构建最小化仿射代价模型的最终程序。为降低内存开销,采用新型记忆化表结构,按非负整数坐标索引规格,并压缩相邻相同解。该方法可覆盖比以往更广的程序集合。我们开发了Morello编译器,将类似单节点XLA计算图的规格降级至x86。实验表明,仿射代价模型已足够挖掘高性能程序。例如,Morello生成了针对Zen 1 CPU的矩阵乘基准测试,包括1×2048×16384、bfloat16转float32的向量-矩阵乘法,已集成至Google的gemma.cpp。

原文摘要 · Abstract (English)

High-throughput neural network inference requires coordinating many optimization decisions, including parallel tiling, microkernel selection, and data layout. The product of these decisions forms a search space of programs which is typically intractably large. Existing approaches (e.g., auto-schedulers) often address this problem by sampling this space heuristically. In contrast, we introduce a dynamic-programming-based approach to explore more of the search space by iteratively decomposing large program specifications into smaller specifications reachable from a set of rewrites, then composing a final program from each rewrite that minimizes an affine cost model. To reduce memory requirements, we employ a novel memoization table representation, which indexes specifications by coordinates in $Z_{\geq 0}$ and compresses identical, adjacent solutions. This approach can visit a much larger set of programs than prior work. To evaluate the approach, we developed Morello, a compiler which lowers specifications roughly equivalent to a few-node XLA computation graph to x86. Notably, we found that an affine cost model is sufficient to surface high-throughput programs. For example, Morello synthesized a collection of matrix multiplication benchmarks targeting a Zen 1 CPU, including a 1x2048x16384, bfloat16-to-float32 vector-matrix multiply, which was integrated into Google's gemma.cpp.

编译器动态规划矩阵乘高性能

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。