Nvidia使用AI设计GPU运算电路，面积比最先进的EDA小25%，加速更快更高效。因此，改进这些运算电路的设计对于提高GPU性能和效率至关重要。如果人工智能学会设计这些电路会怎么样？在最近的Nvidia论文《PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning》中，研究人员证明了人工智能不仅可以从头开始设计这些电路，而且人工智能设计的电路比最先进的电子设计自动化（EDA）设计的电路更小、更快）工具。论文地址：https://arxiv.org/pdf/2205.07000.pdf最新的NVIDIAHopperGPU架构有近13000个AI设计的电路实例。下图1左侧PrefixRLAI设计的64b加法器电路比图1右侧最先进的EDA工具设计的电路小25%。电路设计概述计算机芯片中的算术电路由网络构建逻辑门（例如NAND、NOR和XOR）和电线。一个理想的电路应该具备以下属性：小：面积更小，芯片上可以封装更多的电路；快：延迟更低，提高芯片性能；更低的功耗。在这项Nvidia研究中，研究人员专注于电路面积和延迟。他们发现功耗与感兴趣的电路面积密切相关。电路面积和延迟通常是相互竞争的属性，因此需要为有效权衡这些属性的设计找到帕累托边界。简而言之，研究人员希望每个延迟的电路面积都最小化。因此，在PrefixRL中，研究人员专注于一类流行的算术电路——并行前缀电路。GPU中的加速器、增量器、编码器等各种重要电路都是前缀电路，在更高层次上可以定义为前缀映射。那么问题来了：AIagent能否设计出好的前缀映射？所有前缀图的状态空间都很大O(2^n^n)，无法使用暴力方法进行探索。下面的图2显示了PrefixRL与电路实例4b的迭代。研究人员使用电路发生器将前缀图转换为带有电线和逻辑门的电路。接下来，这些生成的电路通过物理综合工具进行优化，该工具使用门大小、复制和缓冲区插入等物理综合优化。由于这些物理综合优化，最终电路属性（延迟、面积和功率）不会直接从原始前缀映射属性（例如电平和节点数）转换而来。这就是为什么AI代理学习设计前缀图但优化从前缀图生成的最终电路的属性。我们将算术电路设计视为强化学习(RL)任务，其中训练代理以优化算术电路的面积和延迟特性。对于前缀电路，他们设计了一个RL代理可以在前缀图中添加或删除节点的环境，然后执行以下步骤：前缀图被归一化以始终保持正确的前缀和计算；电路是从归一化前缀图生成的；使用物理综合工具优化电路的物理综合；测量电路的面积和延迟特性。在下面的动画中，RL代理通过添加或删除节点逐步构建前缀图。在每一步中，代理都会在电路面积和延迟方面得到改善。原始图像是交互式版本。FullyConvolutionalQ-LearningAgentResearchers使用Q-learning算法来训练代理电路设计。如下图3所示，他们将前缀图分解为网格表示，其中网格中的每个元素都唯一映射到一个前缀节点。该网格表示Q网络的输入和输出。输入网格中的每个元素代表一个节点的存在或不存在。输出网格中的每个元素代表添加或删除节点的Q值。研究人员采用全卷积神经网络架构，因为Q学习代理的输入和输出都是网格表示。智能体分别预测面积和延迟属性的Q值，因为面积和延迟的奖励在训练期间是分别可观察到的。图3：4b前缀图表示（左）和全卷积Q学习代理架构（右）。Raptor对PrefixRL的分布式训练需要大量的计算。在物理模拟中，每个GPU需要256个CPU，训练64b任务需要32000多个GPU小时。此次，英伟达内部开发了分布式强化学习平台Raptor，充分利用了英伟达硬件的优势，从而可以进行这种工业级的强化学习（见下图4）。Raptor提高了训练模型的可扩展性和速度，例如作业调度、自定义网络和GPU感知数据结构。在PrefixRL的上下文中，Raptor支持跨CPU、GPU和Spot实例的混合分配。此强化学习应用程序中的网络多种多样，并受益于以下内容：Raptor在NCCL之间切换以进行点对点传输，将模型参数直接从学习GPU传输到推理GPU；Redis用于异步和较小的消息，例如奖励或统计；对于JIT编译的RPC，对于大容量和低延迟的请求，例如上传体验数据。最后，Raptor提供GPU感知数据结构，例如具有多线程服务的回放缓冲区，以接收来自多个工作人员的经验，并行批处理数据并将其预加载到GPU上。下图4显示PrefixRL框架支持并发训练和数据收集，并利用NCCL高效地将最新参数发送给参与者（下图中的参与者）。图4：研究人员使用Raptor解耦并行训练和奖励计算以克服电路综合延迟。奖励计算研究人员使用权衡权重w（在[0,1]范围内）来组合区域和延迟目标。他们训练具有不同权重的各种代理以获得帕累托前沿，从而平衡面积和延迟之间的权衡。RL环境中的物理综合优化可以生成各种权衡面积和延迟的解决方案。研究人员使用用于训练特定代理的相同权衡权重来驱动物理综合工具。在奖励计算的循环中进行物理综合优化具有以下优点：RL智能体学习直接优化目标技术节点和库的最终电路属性；RL代理在物理综合时包含目标算法电路的外围逻辑，从而共同优化目标算法电路及其周边逻辑的性能。然而，执行物理综合是一个缓慢的过程（64b加法器约35秒），这会显着减慢RL训练和探索。我们将奖励计算与状态更新分离，因为代理只需要当前前缀图状态即可采取行动，而无需电路合成或先前的奖励。多亏了Raptor，他们可以将冗长的奖励计算卸载到CPU工作池以并行执行物理合成，而演员代理无需等待即可在环境中执行。当CPU工作人员返回奖励时，转换可以嵌入重播缓冲区。组合的奖励被缓存以避免再次遇到状态时的冗余计算。结果和展望下面的图5显示了使用PrefixRL设计的64b加法器电路的面积和延迟，以及来自最先进的EDA工具的Pareto支配加法器电路。在相同的延迟下，最好的PrefixRL加法器的面积比EDA工具加法器低25%。这些经过物理综合优化后映射到帕累托最优加法器电路的前缀图具有不规则的结构。图5：PrefixRL设计的算术电路比最先进的EDA工具设计的电路更小、更快。（左）电路架构；（右）对应的64b加法器电路特性图据我们所知，这是第一种使用深度强化学习代理设计算术电路的方法。Nvidia设想了一个蓝图，将AI应用于现实世界的电路设计问题，构建动作空间、状态表示、RL代理模型，针对多个竞争目标进行优化，并克服缓慢的奖励计算。
