在摩尔定律逐渐放缓的背景下，要在相同工艺节点上持续提升芯片性能，必须寻找新的突破口。英伟达采取了一项创新策略——让AI自主设计更小、更快、更高效的算术电路，为芯片带来切实的性能提升。简单来说，就是利用深度强化学习来“训练”电路设计，取代传统工程师手动调参的方式。

GPU之所以能够加速AI、高性能计算与图形渲染，关键在于其内部集成了大量的算术电路阵列。因此，优化这些电路的设计成为提升GPU性能与效率的核心。那么，如果让AI学会设计这些电路，会产生怎样的效果？在最新发表的论文《PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning》中，英伟达研究团队给出了答案：AI不仅能够从零开始设计电路，而且所设计的电路比最先进的电子设计自动化（EDA）工具更小、更快。

论文地址：https://arxiv.org/pdf/2205.07000.pdf

在最新的英伟达Hopper GPU架构中，已经集成了将近13000个由AI设计的电路实例。下图中，左侧是PrefixRL AI设计的64位加法器电路，右侧是EDA工具设计的电路——左侧面积减少了25%。

电路设计概览

计算机芯片中的算术电路，本质上是由逻辑门（如NAND、NOR、XOR）和导线构成的网络。一个理想的电路需要满足三个条件：面积小（以便集成更多电路）、速度快（降低延迟、提升性能）以及功耗低。在本研究中，英伟达团队主要关注面积和延迟，他们发现功耗与面积密切相关。面积和延迟往往难以兼得——需要寻找帕累托边界，即在每个延迟点实现面积最小化的设计。

为此，研究者聚焦于一类常见的算术电路——并行前缀电路。GPU中的加法器、增量器、编码器等关键部件均属于前缀电路，它们可以在更高层次上抽象为“前缀图”。那么，AI能否设计出高质量的前缀图？所有可能的前缀图状态空间规模为O(2^n^n)，穷举搜索完全不可行。下图展示了一个4位电路实例的PrefixRL迭代过程：

研究者首先使用电路生成器将前缀图转换为包含逻辑门和导线的实际电路，再通过物理综合工具进行优化（例如门尺寸调整、复制、缓冲器插入）。值得注意的是，经过物理综合优化后，最终电路的属性（延迟、面积、功耗）并非简单地从原始前缀图的属性（如节点数、层级数）直接换算而来。因此，AI智能体必须学会基于最终电路的属性来优化前缀图，而非仅仅关注图的统计指标。

他们将算术电路设计定义为一个强化学习任务：训练智能体以优化电路的面积和延迟。针对前缀电路，他们构建了一个环境，智能体可以添加或删除前缀图中的节点，随后依次执行：标准化前缀图（确保前缀和计算正确）、从图中生成电路、使用物理综合工具优化、测量面积和延迟特性。每一步，智能体获得的奖励是电路面积和延迟的改进值。下面的动图展示了智能体通过逐步添加或删除节点来构建前缀图：

原图为可交互版本

完全卷积Q学习智能体

研究者采用Q学习算法训练智能体。他们将前缀图分解为网格表示——网格中的每个元素唯一对应一个前缀节点，并用作Q网络的输入和输出。输入网格表示节点是否存在，输出网格的每个元素代表添加或删除该节点的Q值。由于输入和输出均为网格结构，因此采用了完全卷积神经网络。此外，智能体分别预测面积和延迟的Q值，因为这两类奖励在训练过程中可以分开观察。

图3：4位前缀图表示（左）以及完全卷积Q学习智能体架构（右）。

Raptor进行分布式训练

PrefixRL的计算量非常庞大——物理模拟环节中，每个GPU需要256个CPU协同工作；训练64位任务需要超过32000个GPU小时。为应对这种工业级强化学习需求，英伟达开发了内部分布式强化学习平台Raptor，充分利用了自身硬件优势。Raptor能够提升训练模型的可扩展性和速度，例如作业调度、自定义网络以及GPU感知的数据结构。在PrefixRL场景中，Raptor支持跨CPU、GPU和Spot实例的混合资源分配。

该强化学习应用涉及多种网络通信，Raptor在多个方面展现出优势：在NCCL之间切换进行点对点传输，将模型参数直接从学习器GPU传输到推理GPU；使用Redis处理异步小消息，例如奖励或统计信息；利用JIT编译的RPC处理高容量、低延迟的请求，如上传经验数据。此外，Raptor提供了GPU感知的数据结构，例如带多线程服务的重放缓冲器，可以接收多个worker的经验，并行批处理数据并提前加载到GPU上。

图4：研究者采用Raptor进行解耦并行训练与奖励计算，以克服电路合成延迟。

奖励计算

研究者使用权衡权重w（范围0到1）来组合面积和延迟目标，训练了多个不同权重的智能体，以获得帕累托边界，平衡两者之间的取舍。在RL环境中进行物理综合优化，能够产生多种方案来权衡面积和延迟。研究者采用与训练智能体相同的权衡权重来驱动物理综合工具。

在奖励计算循环中执行物理综合优化有两个好处：第一，RL智能体可以直接优化目标技术节点和库下的最终电路属性；第二，RL智能体在物理综合过程中包含了目标电路周边的逻辑，从而能够联合优化它们。然而，物理综合本身速度较慢——64位加法器大约需要35秒，这会严重拖慢RL的训练和探索进程。

于是，研究者将奖励计算和状态更新解耦：智能体只需知道当前前缀图状态即可采取动作，无需等待电路合成或之前的奖励。借助Raptor，他们将耗时的奖励计算转移到CPU worker池中并行执行物理综合，而actor智能体无需等待即可继续探索。当CPU worker返回奖励时，转换被嵌入重放缓冲器。同时，综合奖励会被缓存，避免重复遇到同一状态时进行冗余计算。

结果及展望

下图展示了使用PrefixRL设计的64位加法器电路，与最先进EDA工具设计的帕累托最优加法器在面积和延迟方面的对比。最好的PrefixRL加法器，在相同延迟下面积比EDA工具设计的低了25%。这些经过物理综合优化后映射到帕累托最优的前缀图，结构相当不规则，并非人类工程师通常设计的样式。

图5：PrefixRL设计的算术电路比最先进的EDA工具设计的电路更小、更快。(左)电路架构；(右)对应的64位加法器电路特性图

据目前所知，这是首个利用深度强化学习智能体设计算术电路的研究。英伟达还描绘了一个蓝图：希望这种方法能将AI应用到真实的电路设计中——包括构建动作空间、状态表示、RL智能体模型、针对多个竞争目标进行优化，以及克服缓慢的奖励计算过程。这无疑为芯片设计领域打开了一扇新的大门。