2025-11-11T09:13:09.652713

Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators

Hong, Bhatia, Cheung et al.
Hardware accelerators, especially those designed for tensor processing, have become ubiquitous in today's computing landscape. However, even with significant efforts in building compilers, programming these tensor accelerators remains challenging, leaving much of their potential underutilized. Recently, large language models (LLMs), trained on large amounts of code, have shown significant promise in code generation and optimization tasks, but generating low-resource languages, such as specialized tensor accelerator code still poses a significant challenge. We tackle this challenge with Autocomp, an approach that empowers accelerator programmers to leverage domain knowledge and hardware feedback to optimize code via an automated LLM-driven search. We accomplish this by: 1) formulating each optimization pass as a structured two-phase prompt, divided into planning and code generation phases, 2) inserting domain knowledge during planning via a concise and adaptable optimization menu, and 3) integrating correctness and performance metrics from hardware as feedback at each search iteration. Across three distinct hardware platforms, we demonstrate that Autocomp-optimized code runs 5.6x faster than the vendor-provided library (Gemmini), outperforms expert-level hand-tuned code by 1.9x (AWS Trainium), and achieves 3.8x higher performance than a machine learning-based cost model for GPUs (NVIDIA L40S). Additionally, we demonstrate that optimization schedules generated from Autocomp can be reused across similar tensor operations, improving speedups by up to 24% under a fixed sample budget.
academic

Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators

基本信息

  • 论文ID: 2505.18574
  • 标题: Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
  • 作者: Charles Hong, Sahil Bhatia, Alvin Cheung, Yakun Sophia Shao (UC Berkeley)
  • 分类: cs.PL cs.AI cs.AR cs.LG
  • 发表状态: Preprint. Under review.
  • 论文链接: https://arxiv.org/abs/2505.18574

摘要

硬件加速器，特别是专为张量处理设计的加速器，已经在当今的计算环境中无处不在。然而，即使在编译器构建方面投入了大量努力，编程这些张量加速器仍然具有挑战性，使其潜力大部分未得到充分利用。本文提出Autocomp，一种通过自动化LLM驱动搜索来优化代码的方法，使加速器程序员能够利用领域知识和硬件反馈。该方法通过三个关键技术实现：1）将每个优化过程制定为结构化的两阶段提示，分为规划和代码生成阶段；2）通过简洁且可适应的优化菜单在规划期间插入领域知识；3）在每次搜索迭代中集成来自硬件的正确性和性能指标作为反馈。

研究背景与动机

核心问题

张量加速器编程面临的主要挑战包括：

  1. 编程复杂性：与通用CPU编程不同，张量加速器需要显式管理数据移动、配置状态和操作调度
  2. 编译器适配成本：为新硬件平台适配传统编译器需要大量工程工作，软件开发成本占新硬件开发成本的40-50%
  3. 优化调度问题：确定应用哪些优化以及应用顺序的组合爆炸问题
  4. 低资源语言挑战：专用加速器的指令集架构（ISA）和DSL在LLM训练语料中代表性不足

现有方法局限性

  1. 传统编译器：XLA、TVM、Triton等只支持少数硬件后端，主要是CPU和GPU
  2. DSL方法：Halide、Exo等提供了表达张量计算的原语，但优化负担仍在程序员身上
  3. 数据驱动方法：需要大量性能数据进行训练，这对于领域特定硬件加速器来说极其稀缺
  4. 直接LLM应用：对于低资源加速器语言，零样本代码生成高度不可靠

核心贡献

  1. 首个LLM驱动的低资源张量加速器代码优化方法：提出Autocomp框架，专门针对专用硬件加速器
  2. 高度可移植的优化框架：通过修改提示即可适配新硬件平台，大大降低工程成本
  3. 卓越的性能表现：在三个不同硬件平台上显著优于现有方法
  4. 调度复用机制：证明了优化调度可在相似张量操作间复用，提高样本效率

方法详解

任务定义

输入：未优化的张量加速器代码 输出：功能等价但性能优化的代码 约束：保持语义等价性，通过硬件验证确保正确性

核心架构：两阶段优化框架

阶段1：规划生成（Plan Generation）

提示结构包含：

  1. 加速器ISA描述：指令语义、内存地址规范、硬件结构描述
  2. 当前代码：待优化的代码
  3. 性能反馈：延迟（周期数）、内存利用率等指标
  4. 优化菜单：预定义的高级优化选项（如循环平铺、重排序、融合等）
  5. 搜索迭代信息：当前迭代数，用于指导优化选择

阶段2：代码实现（Implementation）

提示结构包含：

  1. 加速器ISA描述：同阶段1
  2. 当前代码：同阶段1
  3. 生成的计划：阶段1输出的具体优化计划
  4. 上下文学习示例：针对复杂优化（如平铺）的代码示例
  5. 实现指令：应用计划并输出优化代码的自然语言指令

束搜索优化策略

  • 宽度B=6的束搜索，并行探索多个优化轨迹
  • 正确性过滤：通过功能测试套件验证候选代码
  • 性能筛选：只保留性能优于父节点的候选
  • 迭代优化：固定预算T次迭代的搜索过程

关键技术创新

1. 多样性增强技术

  • 优化菜单Dropout：每次规划时随机移除部分菜单选项（70%概率）
  • LLM集成：在多个LLM间分配请求以增加响应多样性

2. 硬件反馈集成

  • 实时性能指标（延迟、内存利用率）指导下一步优化选择
  • 周期精确模拟或芯片级性能测量

3. 调度复用机制

  • 记录高质量调度序列
  • 对相似张量操作（相同纵横比或共享维度）复用已知调度
  • 轻量级搜索后进行进一步优化

实验设置

硬件平台

  1. Gemmini：开源加速器生成器，支持脉动阵列和向量风格的张量加速器
  2. AWS Trainium：商用高性能张量加速器，使用Neuron Kernel Interface (NKI)
  3. NVIDIA L40S GPU：现代数据中心GPU，包含专用Tensor Core

基准测试

  • Gemmini：ResNet-50的GEMM和卷积，TinyMPC模型预测控制
  • Trainium：教程级和高级深度学习算子（RMSNorm、LayerNorm、GEMM、Mamba等）
  • GPU：KernelBench Level 1基准测试

对比方法

  1. 高级软件库：Gemmini软件库、PyTorch NeuronX、PyTorch
  2. 未优化低级代码：Exo未优化、nki-samples教程代码
  3. 手工优化代码：专家级手动调优实现
  4. ML成本模型：TVM MetaSchedule（GPU）
  5. 硬件FSM：Gemmini硬件有限状态机（参考上限）

实验结果

主要性能表现

Gemmini平台

  • GEMM基准：相比Gemmini软件库提升5.6×，超越专家手工优化代码1.4×
  • 卷积基准：相比软件库提升2.6×，超越手工优化1.1×
  • 细粒度线性代数：超越未优化代码2.7×，甚至超越专家优化的硬件FSM实现1.6×（前向传播）

AWS Trainium平台

  • 教程工作负载：超越手工优化代码1.36×（几何平均），超越PyTorch NeuronX编译代码13.52×
  • 高级工作负载：超越专家级优化代码1.9×（几何平均），1D深度卷积提升高达17.37×

NVIDIA L40S GPU

  • KernelBench基准：超越PyTorch 2.05×（几何平均），超越TVM MetaSchedule 3.8×
  • 在所有基准测试中均优于PyTorch，而TVM仅在2个基准中优于PyTorch

消融实验分析

通过详细的消融实验验证了各组件的重要性：

  • 加速器ISA：移除后性能显著下降，但仍能实现改进
  • 优化菜单：完全必要，移除后优化性能完全恶化
  • 菜单Dropout：显著影响性能，防止模型偏向有限的菜单选项
  • LLM集成：提供重要的多样性，单一模型性能较低
  • 硬件性能反馈：有帮助但效果有限，因为优化菜单已包含相关指标

调度复用效果

  • 在100个样本预算下：复用调度达到4.6×加速，无复用仅3.7×
  • 在200个样本预算下：复用调度达到5.0×加速，无复用仅4.2×
  • 证明了调度的可泛化性，能有效减少相似基准的搜索成本

相关工作

张量加速器代码优化

  • 性能模型：Timeloop、MAESTRO使用高级硬件架构模型
  • 自动化方法：机器学习、线性规划、黑盒优化、强化学习
  • 局限性：现有抽象忽略了实现特定和指令级优化

LLM代码优化

  • 应用范围：进化搜索、检索增强生成、迭代优化、模型后训练
  • 系统级优化：CUDA、SIMD内在函数
  • 研究空白：缺乏针对专用硬件（非CPU/GPU）的LLM代码优化工作

结论与讨论

主要结论

  1. LLM驱动优化的有效性：Autocomp在多个硬件平台上显著超越传统方法
  2. 极高的可移植性：仅通过修改提示即可适配新硬件，工程成本极低
  3. 调度复用的价值：优化调度具有良好的泛化性，能显著提高样本效率

技术洞察

  1. 两阶段设计的必要性：分离规划和实现阶段提高了复杂优化任务的成功率
  2. 领域知识的重要性：优化菜单提供的领域专业知识对性能至关重要
  3. 硬件反馈的价值：实时性能指标有效指导优化方向选择

局限性

  1. 依赖LLM能力：方法性能受限于底层LLM的代码生成和推理能力
  2. 搜索成本：需要多次LLM调用和硬件仿真，计算成本较高
  3. 领域特定性：优化菜单需要针对不同硬件平台手工设计
  4. 评估范围：主要集中在张量计算工作负载，其他类型计算的适用性未知

未来方向

  1. 自动菜单生成：研究自动构建优化菜单的方法
  2. 跨平台调度迁移：探索不同硬件平台间的调度知识迁移
  3. 成本效率优化：减少搜索过程中的LLM调用和硬件仿真次数
  4. 更广泛的应用：扩展到非张量计算的其他专用加速器

深度评价

优点

  1. 创新性强：首次将LLM应用于低资源张量加速器代码优化，技术路线新颖
  2. 实用价值高：解决了实际的工程痛点，大幅降低了新硬件的软件开发成本
  3. 实验充分：在三个不同硬件平台上进行了全面评估，结果具有说服力
  4. 方法通用：框架设计具有良好的可扩展性和可移植性
  5. 性能卓越：在多个基准测试中显著超越现有最佳方法

不足

  1. 计算成本：需要大量LLM调用和硬件仿真，可能限制实际应用
  2. 人工设计依赖：优化菜单仍需专家知识手工设计，自动化程度有限
  3. 评估局限：主要集中在特定类型的张量计算，泛化性有待验证
  4. 理论分析不足：缺乏对方法收敛性和最优性的理论保证

影响力评估

  1. 学术价值：开创了LLM在专用硬件编译优化领域的应用，具有重要学术意义
  2. 产业影响：有望显著降低新硬件的软件栈开发成本，具有重要产业价值
  3. 可复现性：作者承诺开源实现和提示，有利于后续研究
  4. 启发性：为其他专用硬件的编译优化提供了新的技术路径

适用场景

  1. 新硬件原型开发：快速为新设计的张量加速器生成优化代码
  2. DSL编译器构建：作为传统编译器的补充或替代方案
  3. 性能调优工具：帮助开发者优化现有的加速器代码
  4. 研究和教学：为加速器编程和优化提供自动化工具

参考文献

论文引用了大量相关工作，主要包括：

  • 硬件加速器设计（Gemmini、TPU、Trainium等）
  • 编译器和DSL（XLA、TVM、Halide、Exo等）
  • LLM代码生成（CodeGen、Codex等）
  • 自动化优化方法（强化学习、进化算法等）

总体评价：这是一篇高质量的研究论文，在LLM应用于专用硬件编译优化这一新兴交叉领域做出了重要贡献。方法创新性强，实验评估全面，实用价值显著。虽然在计算成本和自动化程度方面仍有改进空间，但为该领域的发展开辟了新的方向，具有重要的学术和产业价值。