用大模型自动优化张量加速器代码,性能超商用库5.6倍。
Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
- 通过分阶段提示词设计,结合领域知识与硬件反馈进行搜索优化。
- 在三类硬件上实现最高5.6倍于厂商库的加速,优于专家调优代码1.9倍。
- 生成的优化方案可复用,相同采样预算下提速提升24%。
硬件加速器,尤其是为张量处理设计的设备,已成为现代计算的核心。然而,即便编译器发展成熟,编程这些张量加速器仍具挑战性,导致其潜力未被充分释放。近期,基于大量代码训练的大语言模型(LLMs)在代码生成与优化任务中展现出巨大潜力,但生成低资源语言如专用张量加速器代码仍面临困难。本文提出Autocomp,一种通过自动化LLM驱动搜索,融合领域知识与硬件反馈来优化代码的方法。具体包括:1)将每轮优化封装为规划与代码生成两阶段结构化提示;2)在规划阶段引入简洁可适配的优化菜单以注入领域知识;3)在每次搜索迭代中集成代码正确性与性能指标作为硬件反馈。在三个不同硬件平台上的实验表明,Autocomp优化的代码比厂商提供的库(Gemmini)快5.6倍,比专家级手工调优代码(AWS Trainium)快1.9倍,并在GPU上达到机器学习成本模型(NVIDIA L40S)3.8倍的性能提升。此外,我们证明了Autocomp生成的优化调度可跨相似张量操作复用,在固定采样预算下进一步提升速度提升达24%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。