用AI芯片设计出更好的AI芯片

芯片设计既是一项工程壮举，也堪称一门艺术。逻辑和内存块的一切可能布局、连通每个元件的导线，共同构成了几乎无穷无尽的规划组合。而且信不信由你，芯片设计领域最出色的工程师们，都是在凭着经验和直觉在工作。没错，他们自己也说不清楚为什么某种模式有效，而其他模式却无效。

因此，芯片设计就成了一种高风险工作。研究人员一直努力把人类的推断猜测从其中剔除出去，真正朝着高度优化的方向迈进。即使如今我们已经转向小芯片设计时代，这个目标也仍然存在，毕竟计算引擎上的所有小芯片也都需要互连才能构成虚拟意义上的单体芯片，同时认真考虑整个电路综合体的延迟和功耗水平。

话说到这里，用AI技术协助设计芯片就成了最合乎逻辑的方案。几年之前，我们就曾经跟谷歌工程师聊过这个问题，这家云巨头也从来没有放弃探索。今年3月，谷歌研究院的科学家们发布了PRIME。这是一种深度学习方法，能够利用现有数据（例如蓝图、功率/延迟等指标）建立起比传统工具成果更快、体积更小的芯片设计。

他们在报告中写道，“要利用先前积累的加速器数据库实现硬件设计，最简单的方法也许就是使用监督机器学习来训练预测模型，让该模型以给定的加速器性能目标为输入，之后不断优化学习模型的性能输出，最终找到符合设计要求的加速器布局方案。”

一年前，谷歌开始使用强化学习（RL）技术设计自己的TPU AI加速器布局。而且除谷歌之外，Synopsys和Cadence等芯片设计工具开发商，也开始在自家产品组合中引入AI技术。

如今，英伟达推出一种新方法。根据三位深度学习科学家的解释，“使用AI设计出更小、更快、更高效的电路，以在每一代芯片中提供更高性能。庞大的算术电路阵列为英伟达GPU提供动力，从而让AI、高性能计算和计算机图形获得了前所未有的加速支持。因此，对这些算术电路的设计改进，将在GPU性能与效率的提升之路上发挥至关重要的作用。”

于是，英伟达选择了自己的强化学习测试方式，名为PrefixRL，并表示该技术证明AI不仅有能力从零开始学习电路设计，而且最终成果要比最新EDA工具的设计方案更小、更快。英伟达的“Hopper” GPU架构已经于今年3月推出，让该公司进一步扩大对AI、机器学习和神经网络的广泛关注，期间共产生近13000种由AI技术设计而成的电路实例。

图一：PrefixRL AI设计的64b加法器电路（左），要比当前最先进EDA工具设计的电路（右）小25%，且速度与功能相同。

在一篇关于PrefixRL的六项研究论文中，研究人员表示他们专注于所谓“并行前缀电路”的算术电路，加法器、增量器和编码器等都属于这类前缀电路，而且全都可以在宏观上被定义为前缀图。英伟达想了解AI代理能否设计出高质量的前缀图，并补充称“所有前缀图都具有O(2^n^n)的极大状态空间，因此无法用暴力方式硬性尝试。”

论文写道，“使用电路生成器，即可将前缀图转换为带有导线和逻辑门的电路。这些生成的电路随后通过物理综合工具加以进一步优化，即使用门尺寸、复制和缓冲器插入等方法实现物理综合优化。”

算术电路是由NAND、NOR和XOR等逻辑门，再配合大量导线构建而成的，同时既要小到能够在指定的芯片上被大量封装，又要具备极快的速度以降低性能延迟，同时在可能的范围内将功率控制在最低。对于PrefixRL，研究人员的关注重点是电路的大小和速度（用于降低延迟），因为这两个指标往往会相互冲突。芯片设计的挑战，就在于如何高效找到尺寸与性能最均衡的设计点。“简而言之，我们希望在满足延迟要求的条件下，尽可能缩小芯片尺寸。”

图二：PrefixRL的一次4b电路迭代示例

研究人员们写道，“但这些物理综合优化的效果，即最终电路属性（包括延迟、面积和功率等指标）无法直接由原始前缀图属性（例如电平和节点数）转换得出，所以AI代理才需要学习设计前缀图，尝试突破这道人类难以逾越的从前缀图、到最终电路属性的优化之路。我们将算术电路设计视为一种强化学习任务，训练一个代理来优化算术电路的面积和延迟特性。对于前缀电路，我们还专门设计了一个环境，允许强化学习代理在其中添加或删除前缀图内的节点。”

接下来，设计流程要使前缀图合法化，确保其中始终保持正确的前缀和计算方法，而后通过合法化的前缀图创建电路。再引入物理综合工具优化电路，最终测量电路的面积和延迟特性。在整个过程中，强化学习代理会通过添加或删除节点的一系列步骤构建起前缀图。

英伟达研究人员在自己的工作中使用了完全卷积神经网络和Q学习算法（强化学习算法的一种）。该算法使用前缀图的网格表示对电路设计代理进行训练，网格中的每个元素都唯一地映射至前缀节点。这种网格表示作为Q网络的输出和输出，输入网格中的每个元素表示节点是否存在，输出网格中的每个元素表示添加或删除节点的Q值。

图三：4b前缀图表示（左）与完全卷积Q学习代理架构（右）

PrefixRL的计算需求很大。据研究人员介绍，整个物理模拟任务需要为每个GPU配备256个CPU，而且训练时间超过32000个GPU小时。为了满足如此严苛的资源需求，英伟达专门创建了一个名为“Raptor”的分布式强化学习平台，专门利用英伟达硬件实现这一级别的强化学习用例。

研究人员们写道，“Raptor提供多项功能，可以提高可扩展性和训练速度，例如作业调度、自定义网络和GPU感知数据结构。在PrefixRL用例当中，Raptor能够将工作负载混合分配给CPU、GPU和Spot实例。这个强化学习应用涉及多种网络类型，其中Raptor能够在英伟达集体通信库（NCCL）之间切换以进行点对点传输，从而将模型参数直接从学习GPU传输至推理GPU。”

由Redis存储异步及较小消息，例如奖励和统计信息。大容量和低延迟请求由JIT编译的RPC负责处理，例如上传经验数据。最后，Raptor还提供GPU感知数据结构，例如具有多线程服务的重放缓冲器，能够接收来自多个worker的经验、并行批处理数据并将结果预先载入至GPU上。

图四：研究人员使用Raptor对并行训练和奖励计算进行解耦，借此克服电路综合延迟

研究人员表示，强化学习代理能够单凭来自综合电路属性的反馈中学习设计能力，PrefixRL也成功设计出了64b加法器电路。在延迟相同的情况下，PrefixRL最佳加法器设计方案比EDA工具设计成果的面积小25%。

研究人员们最后总结道，“据我们所知，这是第一种使用深度强化学习代理进行算术电路设计的方法。我们希望这种方法能够带来启发，推动AI在现实世界电路设计问题中的应用，包括构建动作空间、状态表示、强化学习代理模型、针对多个竞争目标进行优化，以及克服物理综合过程等慢速奖励计算过程。”