Nvidia 开发 PrefixRL 设计更小更快的并行前缀电路

Nvidia 开发了基于强化学习（Reinforcement Learning, RL）的 PrefixRL 方法，用于设计比现有电子设计自动化（Electronic Design Automation, EDA）工具更小、更快的并行前缀电路（parallel-prefix circuits）。这些电路在 GPU 中非常重要，广泛应用于加法器、增量器和编码器等高性能数字设计中。

并行前缀电路的重要性与挑战

并行前缀电路是高性能数字设计的基础，可以定义为前缀图。PrefixRL 专注于这类算术电路，目标是探索 AI 代理是否能够设计出优秀的前缀图。由于问题的状态空间为 O(2^n^n)，无法通过暴力方法解决，因此需要新的设计方法。

PrefixRL 的设计目标

理想的电路应具备面积小、速度快和功耗低的特点。Nvidia 发现功耗与电路面积密切相关，但面积和延迟往往是相互竞争的特性。PrefixRL 的目标是找到设计中的帕累托前沿（Pareto frontier），在面积和延迟之间取得最佳平衡，从而在更小的面积内容纳更多电路，减少芯片延迟，提升性能并降低功耗。

Hopper GPU 中的 AI 设计电路

Nvidia 最新的 Hopper GPU 架构中，有近 13,000 个电路是由 AI 设计的，展示了 AI 在电路设计中的广泛应用。

PrefixRL 的技术实现

PrefixRL 使用全卷积神经网络（Q-learning 代理）进行训练。输入和输出采用网格表示前缀图，每个网格元素唯一映射到一个前缀节点。输入网格中的元素表示节点是否存在，而输出网格中的元素表示添加或删除节点的 Q 值。PrefixRL 代理分别预测面积和延迟的值，因为这些属性在训练过程中是独立观察的。

强化学习过程

RL 代理可以在前缀图中添加或删除节点，在强化学习任务的每个回合中，代理会根据电路面积和延迟的改进获得奖励。设计过程包括：合法化前缀图以确保正确的前缀和计算，从合法化的前缀图生成电路，最后通过物理综合工具优化电路并测量其面积和延迟特性。

帕累托前沿的优化

通过训练多个具有不同权重的代理（权重从 0 到 1），PrefixRL 在 RL 环境中生成多种解决方案，从而找到面积和延迟之间的最佳平衡。该综合过程耗时较长（64 位加法器约 35 秒）且计算密集，物理仿真需要为每个 GPU 配备 256 个 CPU，训练 64 位案例需要超过 32,000 个 GPU 小时。

Raptor 分布式强化学习平台

为了加速此类 RL 任务，Nvidia 开发了 Raptor，这是一个内部分布式强化学习平台，充分利用了 Nvidia 硬件。其核心特性包括任务调度、GPU 感知数据结构和自定义网络，以提高可扩展性和训练速度。Raptor 通过并行化 CPU 工作者的物理综合过程，避免冗余计算，并缓存奖励数据，从而显著提升训练速度。

RL 设计电路的成果

RL 代理设计的加法器在相同延迟下比 EDA 工具设计的电路面积减少了 25%，并且具有不规则结构。这一成就得益于 RL 代理从零开始设计电路，并通过综合电路属性的反馈进行学习。

总结

PrefixRL 通过强化学习技术，成功设计了更小、更快的并行前缀电路，展示了 AI 在高性能数字设计中的巨大潜力。Nvidia 的 Raptor 平台进一步提升了 RL 任务的训练效率，为未来 AI 在电路设计中的应用提供了强大支持。