自动计算深度学习模型在硬件上的理论最快运行时间,助力软硬件协同优化。
SOLAR: AI-Powered Speed-of-Light Performance Analysis
- 基于大模型与确定性流程,从PyTorch/JAX代码自动生成可验证的执行中间表示。
- 在多个基准和真实场景中实现零可观测性能越界,提供未融合、融合与缓存感知的理论极限。
- 适合算法工程师、芯片设计者用于分析性能余量、发现优化空间或指导硬件规划。
深度学习模型在目标硬件上能跑多快?当前实现距离理论极限还有多远?这些问题对软硬件与算法优化至关重要。速度光限(SOL)分析通过计算特定架构下的理论最小执行时间来回答。然而,现有方法依赖人工,易出错且难以跟上快速迭代的模型开发。为此,我们提出SOLAR框架,可自动从PyTorch和JAX源码中推导并验证SOL边界。SOLAR采用生成式与确定性相结合的流程:先由大模型前端将源码转为可执行的仿射循环中间表示,并通过输出比对验证;再将其提升为einsum图;最后由分析后端计算未融合、融合及缓存感知的SOL边界。SOLAR覆盖全面的算子与语言,实现零观测到的SOL违规,支持多精度分析,持续收紧边界并揭示优化洞察。我们在KernelBench、JAX/Flax模型与机器人工作负载上评估,验证了四大应用场景:多精度性能余量分析、优化机会识别、跨平台探索以及反向屋顶线硬件配置。
原文摘要 · Abstract (English)
How fast could a deep-learning model run on target hardware, and how far is today's implementation from that limit? These questions are central to software, hardware, and algorithm optimizations. Speed-of-Light (SOL) analysis answers them by computing a workload's theoretical minimum execution time on a given architecture. Yet deriving SOL bounds remains manual, error-prone, and disconnected from rapid model development. To close this gap, we introduce SOLAR, a framework that automatically derives validated SOL bounds from PyTorch and JAX source code. SOLAR leverages both generative and deterministic components in its flow: an LLM frontend translates any source programs into an executable Affine Loop IR, validated by output comparison; a deterministic flow lifts the IR into an einsum graph; and an analytical backend computes unfused, fused, and cache-aware SOL bounds. SOLAR provides comprehensive operator and language coverage, produces validated bounds with zero observed SOL violations, and offers multi-fidelity analysis that tightens bounds and surfaces optimization insights. We evaluate SOLAR across KernelBench, JAX/Flax models, and robotics workloads. These experiments demonstrate four use cases: headroom analysis at multiple fidelity levels, identifying optimization opportunities, cross-platform exploration, and inverse-roofline hardware provisioning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。