论文
太阳能:人工智能驱动的光速性能分析
SOLAR: AI-Powered Speed-of-Light Performance Analysis
摘要
深度学习模型在目标硬件上运行的速度有多快,目前的实现距离该限制还有多远?这些问题是软件、硬件和算法优化的核心。光速 (SOL) 分析通过计算给定架构上工作负载的理论最短执行时间来回答这些问题。然而,推导 SOL 边界仍然是手动的,容易出错,并且与快速模型开发脱节。为了弥补这一差距,我们引入了 SOLAR,这是一个可以从 PyTorch 和 JAX 源代码自动导出经过验证的 SOL 范围的框架。 SOLAR 在其流程中利用生成性和确定性组件:LLM 前端将任何源程序转换为可执行的 Affine Loop IR,并通过输出比较进行验证;确定性流程将 IR 提升为 einsum 图;分析后端计算未融合、融合和缓存感知的 SOL 边界。 SOLAR 提供全面的运算符和语言覆盖范围,生成经过验证的边界,观察到的 SOL 违规为零,并提供多保真度分析,以收紧边界和表面优化见解。我们跨 KernelBench、JAX/Flax 模型和机器人工作负载评估 SOLAR。这些实验展示了四个用例:多个保真度级别的余量分析、识别优化机会、跨平台探索和逆屋顶线硬件配置。