人工智能代理能否设计出良好的前缀图

 

芯片设计既是一门艺术，也是一项工程壮举。

芯片设计的风险很高，研究人员一直在努力朝着更优化的设计迈进。随着我们转向小芯片设计，所有小芯片都需要互连才能成为虚拟单片芯片，并且必须考虑延迟和功耗对此类电路复合体的影响。

人工智能技术也必然会帮助我们设计芯片，这也是我们几年前与谷歌工程师讨论过的事情。谷歌这家云计算巨头也继续努力用人工智能技术来进行芯片设计：今年3月，谷歌研究院的科学家们推出了一种深度学习方法PRIME，它可以利用现有数据，造出比使用传统工具设计的芯片更快、更小的加速设计。

“也许使用先前设计的加速数据库进行硬件设计的最简单方法，是使用监督机器学习来训练预测模型，该模型可以预测给定加速的性能目标，”他们在一份报告中写道。“然后，通过优化这个学习模型与输入加速设计相关的性能输出，人们可能会设计新的加速。”

一年前，谷歌公司使用了一种名为强化学习（RL）的技术来设计TPU 人工智能技术加速。不仅仅是谷歌在做这些。Synopsys和Cadence等芯片设计工具制造商都在将人工智能技术技术应用到自己的产品中。

现在英伟达推出了一种方法，该公司的三位深度学习科学家最近写道，利用人工智能设计更小、更快、更高效的电路，能在每一代芯片中提供更高的性能。大量的算术电路阵列为英伟达 GPU提供了动力，使其在人工智能、高性能计算和计算机图形学方面实现了前所未有的加速。因此，改进这些算术电路的设计对于提高 GPU 的性能和效率至关重要。

英伟达将深度强化学习技术称为PrefixRL，并表示该技术证明人工智能技术不仅可以从头开始学习设计电路，而且人工智能设计的电路也比最先进的电子设计自动化（EDA）工具设计的电路更小更快。最新 NVIDIA Hopper GPU 结构 拥有近13000个人工智能设计电路实例。

 

由 PrefixRL AI（左）设计的64b加法器电路比由最先进的EDA工具（右）设计的电路小25% ，同时速度快，功能等效

在一篇关于PrefixRL的研究论文中，研究人员表示他们关注一类流行的算术电路，称为（并行）前缀电路。GPU中的各种重要电路，如加法器、增量器和编码器，都是前缀电路，可以在更高级别上定义为前缀图。

他们专注于一类称为并行前缀电路的算术电路，其中包括加法器、增量器和编码器等电路，所有这些都可以在更高级别定义为前缀图表。Nvidia特别提出了一个问题：人工智能代理能否设计出良好的前缀图？

计算机芯片中的算术电路是使用逻辑门（如NAND、NOR和XOR）和导线组成的网络构建的。理想电路应该很小，以便芯片上可以容纳更多，快速减少延迟，并尽可能减少功耗。对于PrefixRL，研究人员的重点是电路的大小和速度，他们说这往往是相互竞争的特性。挑战就在于找到最有效地利用两者之间的权衡的设计。

PrefixRL 的一次迭代与 4b 电路示例

由于这些物理合成优化，最终电路特性（延迟、面积和功率）不会直接从原始前缀图特性（如电平和节点数）转换。这就是为什么人工智能代理学习设计前缀图，但优化由前缀图生成的最终电路的属性。

研究人员将算术电路设计作为强化学习 （RL） 任务，我们训练一个代理来优化算术电路的面积和延迟特性。对于前缀电路，我们设计了一个环境，RL 代理可以在其中添加或删除前缀图中的节点。

设计过程将前缀图合法化，以确保它始终保持正确的前缀和计算，之后，从合法化的前缀图创建电路。然后，使用物理合成工具对电路进行物理合成优化，最后测量电路的面积和延迟特性。在整个过程中，RL代理通过添加或删除节点逐步建立前缀图。

英伟达研究人员在他们的工作中使用了完全卷积神经网络和Q学习算法（一种RL算法）。该算法对前缀图使用网格表示，其中网格中的每个元素唯一地映射到前缀节点。这种网格表示法用于Q网络的输入和输出。输入网格中的每个元素表示节点是否存在。输出网格中的每个元素表示用于添加或删除节点的Q值。

某些 4b 前缀图（左）和全卷积 Q 学习代理架构（右）的表示

运行PrefixRL的计算需求很大。据研究人员称，物理模拟每个GPU需要256个CPU，64b 案例需要 32000 GPU 小时。为了满足这些需求，英伟达创建了一个名为“Raptor”的分布式强化学习平台，该平台专门利用英伟达硬件来实现这一级别的强化学习。

Raptor 具有一些增强可扩展性和训练速度的功能，例如作业调度、自定义网络和 GPU 感知的数据结构。在PrefixRL的背景下，Raptor使得跨CPU、GPU和Spot实例的混合分配工作成为可能。

Raptor还包含GPU感知数据结构，用于并行批处理数据并将其预取到GPU。

使用 Raptor 进行解耦并行训练和奖励计算，以克服电路合成延迟

研究人员表示，RL代理能够仅基于从合成电路属性反馈中学习，来设计电路，在相同延迟下，最好的 PrefixRL加法器的面积比EDA工具加法器低 25% 。

“据我们所知，这是第一种使用深度强化学习代理来设计算术电路的方法，”研究人员写道。“我们希望这种方法可以成为将人工智能应用于现实世界电路设计问题的蓝图：构建动作空间、状态表示、RL代理模型、针对多个竞争目标进行优化，以及克服物理合成等缓慢的奖励计算过程。”　　

      审核编辑：彭静