谷歌AI造芯新突破！6小时设计AI芯片，质量超最先进算法

原标题：谷歌AI造芯新突破！6小时设计AI芯片，质量超最先进算法

智东西（公众号：zhidxcom）

文 | 心缘

智东西4月24日，谷歌用人工智能（AI）设计芯片又有新进展！在昨日新发布的预印本论文中，谷歌团队描述了一种深度强化学习方法，能在6小时内完成芯片布局设计，而完成相同设计步骤，人类专家往往要花费数周时间。

人工智能的快速发展得益于计算机系统和硬件的显著进步，但随着摩尔定律趋缓，世界需要新的硬件架构来满足指数级增长的算力需求。

然而，芯片通常需要数年的时间来设计，芯片设计专家不得不预测未来2-5年的机器学习（ML）模型对芯片的需求，大幅缩短芯片设计周期将使硬件更好地适应快速发展的人工智能领域。

谷歌研究人员相信，AI将提供缩短芯片设计周期的手段，在硬件和AI之间建立一种共生关系，彼此推动彼此的进步。

与原有方法不同，谷歌的新算法能从过去的经验中学习，并随着时间的推移而改进。特别是当训练更多芯片块（block）时，其算法能快速地为以前从未见过的芯片块生成优化的设计。

研究人员表示，相较当前最高水平的基线，这一方法可以使真正的AI加速器（谷歌TPUs）达到更好的功率、性能和面积（PPA）。

这项研究适用于任何芯片布局优化，如果该技术可公开使用，那么它将使现金短缺的初创公司能够开发自己的芯片用于AI和其他专门用途。它还可以帮助缩短芯片设计周期，从而使硬件更好地适应快速发展的研究。

该研究的论文题目为《芯片布局与深度强化学习 （Chip Placement with Deep Reinforcement Learning）》，Azalia Mirhoseini和Anna Goldie为共同作者。

论文链接：https://arxiv.org/pdf/2004.10746.pdf

一、攻克芯片设计最耗时的环节

在芯片设计中，全局布线是最复杂和耗时的阶段之一，需要在日益复杂的电路中进行多目标优化。

在这项新研究中，研究人员提出了一种基于深度强化学习的芯片布局方法，目标是将宏（如SRAM）和标准单元（逻辑门，如NAND、NOR和XOR）的网表节点映射到一个芯片画布（一个有界的2D空间）上，从而优化功率、性能和面积（PPA），同时遵守对布局密度和布线拥塞的约束。

这些网表图的大小范围，从数百万到数十亿个节点不等，通常评估目标指标需要数小时到一天的时间。尽管对芯片布局问题已研究几十年，人类专家仍然需要用现有的布局工具迭代数周，才能产生满足多方面设计标准的解决方案。

为了破解这一挑战，谷歌研究人员提出将芯片布局视为一个增强学习（RL）问题。研究人员称，这是第一个具有泛化能力的芯片布局方法，它能利用从以前布局网表中学到的东西来为从未见过的新网表生成优化的芯片设计方案，从而为芯片设计人员带来了新的可能性。

研究人员表示：“我们的方法可以直接优化线长、密度和拥塞等目标指标，而不必像其他方法一样定义这些功能的近似值。我们的公式不仅使新的代价函数在可用时易于整合，而且还使我们能够根据给定芯片块的需求（如功耗受限）来权衡它们的相对重要性。”

据论文介绍，其新领域适应策略不仅产生了较好的效果，而且与从零开始的训练策略相比，训练时间缩短了8倍。

二、用10000个芯片布局训练模型

谷歌研究人员采用深度强化学习的方法，包含状态、动作、状态转移、奖励四个关键要素，通过训练一个智能体通过累计奖励最大化来优化芯片布局。

初始状态有一个空的芯片画布和一个未布局的网表，最后的状态对应于一个完全布局的网表。真实的奖励是一个商业EDA工具的输出，包括线路长度、布线拥塞、复杂度、功率、时间和面积。

智能体在每个步骤中布局一个宏，一旦布局了所有宏，标准单元格就会使用强制定向的方法来放置。奖励是近似的线路长度和拥塞的线性组合，被计算并传递给代理来优化它的参数，以便下一次迭代。

由于强化学习策略需要10万个示例才能有效地学习，因此奖励函数的评估速度必须非常快，最好在几毫秒内运行，且这些近似的奖励功能也必须与真实的奖励有正相关。因此研究人员定义了线路长度和布线拥塞的近似代价函数，通过加权和将多个目标合并成一个单一的奖励功能。

根据论文描述，谷歌研究人员提出一种新的神经网络架构，能训练领域自适应策略的芯片布局。

训练这样一个策略网络是一个具有挑战性的任务，因为包含所有可能放置的所有可能的芯片的状态空间是巨大的。此外，不同的网表和网格大小可能具有非常不同的属性，包括不同数量的节点、宏大小、图形拓扑以及画布的宽度和高度。

为了应对这些挑战，研究人员首先集中学习状态空间的丰富表示。其直觉是，一个能够跨芯片迁移布局优化的策略网络架构，也应该能够在推理时，将与一个未被发现芯片相关的新状态，编码为一个有意义的信号。

因此，谷歌研究人员建议训练一个能够预测新网表奖励的神经架构，生成输入网表的丰富特征嵌入，然后使用这个架构作为其策略和价值网络的编码器，以支持迁移学习。

▲策略和价值网络架构

嵌入层对当前给定网表要放置的邻接、节点特性、宏信息进行编码，然后策略网络输出可用放置位置上的概率分布，而价值网络则输出当前放置位置的预期奖励估值。

在实践中，上述智能体从空芯片开始依次布局组件，直到完成网表为止，直到智能体波长（与功率和性能相关）的负加权总和和拥塞为负值时才收到奖励。为了指导智能体选择首先布局哪些组件，按大小递减的顺序对组件进行排序，首先布局较大的组件，可减少以后没有可行布局的机会。

训练该监督模型需创建一个包含10000个芯片布局的数据集，其中输入是与给定布局相关的状态，而标签是布局的奖励（即线长和拥塞）。研究人员通过首先选择五个不同的芯片网表来构建这个数据集，并用AI算法为每个网表创建2000个不同的布局位置。

在实验中，合著者报告说，当他们在更多芯片上训练框架时，他们能够加快训练过程并更快地产生高质量的结果。

三、6小时生成芯片布局，支持迁移学习

在实验评估阶段，研究人员着重评估几个问题：新方法是否使领域转移和从经验中学习成为可能？训练前的策略对结果的质量有什么影响?与最先进的基线相比，生成布局的质量如何?

1、预训练策略优于从零训练策略

下图对比了使用预训练策略和从零开始训练策略所产生的设计质量，包括zero-shot以及在2小时和12小时内对预训练策略网络的具体设计细节进行调整的结果。

▲领域适应结果

从图中可以看出，预训练策略网络始终优于从零开始训练的策略网络，证明了从离线训练数据中学习的有效性。

策略网络从零训练需要更长时间收敛，甚至24小时后的收敛效果还不如收敛12小时的微调策略网络，说明权重和专家设计经验有助于在更短时间内实现更高质量的新设计。

上图显示了从零开始训练与从Ariane RISC-V CPU预训练策略网络训练的收敛图。预训练策略网络在微调过程开始时布局代价较低，且能收敛到更低代价，比从零开始训练的策略网络快30多小时。

2、更大数据集能产生更好布局效果

训练数据集从内部TPU块创建，由各种块组成，包括内存子系统、计算单元和控制逻辑。当将训练集从2个块增加到5个块，最后增加到20个块时，zero-shot策略网络和微调策略网络在相同的时间内均能产生更好的结果。

▲布局代价与微调时间及数据集大小的关系

研究人员在三个不同的训练数据集上对策略网络进行预训练，其中小数据集是中等数据集的子集，中等数据集是大数据集的子集，然后在相同的测试块上对预训练策略网络进行调整，并在不同的训练期间报告代价。

随着数据集大小的增加，生成的布局质量和测试块的收敛时间都有所提高。下图显示了在每个数据集上训练的策略的评估曲线，对于小的训练数据集，策略网络在测试数据集上快速地超拟合训练数据和性能；对于最大的数据集，策略网络需要更长的时间来拟合，且在这个更大的数据集上，预训练的策略网络在测试数据上能更快产生更好的结果。

▲不同大小数据集的收敛情况（每个点显示了训练策略生成的布局代价）

这表明，如果将策略网络提供更多不同的芯片块时，可能需要更长的时间来进行预训练，策略网络变得不易过度拟合，并且更擅长为新的未知块找到优化的布局位置。

3、可视化芯片布局对比

下图显示了开源处理器Ariane RISC-V CPU的布局结果。 左边是zero-shot策略网络的结果，右边是微调策略网络的结果。

▲可视化布局

zero-shot的结果是在推理期间于一个以前未见过的芯片上产生的。zero-shot策略网络将标准单元置于画布的中心，周围环绕着宏，已经非常接近于最佳布局。微调后，宏的位置变得更有规律，中心的标准细胞区域变得更不受干扰。

▲随着训练进行，开源处理器Ariane的布局（左边是从头开始训练策略，右边是针对该芯片的预训练策略，每个矩形代表一个单独的宏位置。）

下图左边是人工布局的结果，右边是谷歌新方法的布局结果。白色区域为宏位置，绿色区域为标准单元位置，新方法在标准单元周围创建环状的宏布局，从而减少总线长。

4、不到6小时完成芯片布局设计

在本节中，研究人员将其新方法与模拟退火（Simulated Annealing）、RePlAce和人工布局方法进行比较。

谷歌的新方法，在最大的数据集（20个TPU块）上使用预训练策略，然后在块1到5表示的5个目标不可见块上对其进行微调。相比模拟退火算法需要18个小时来收敛，而谷歌端到端的学习方法耗时不超过6个小时。

研究人员进行多次模拟退火实验，结果显示，模拟退火算法很难产生高质量的芯片布局，较谷歌新方法平均线长增加14.4%，拥塞率增加24.1%。

▲谷歌新方法与模拟退火方法的实验结果对比

RePlAce的收敛速度更快，能在1-3.5小时内收敛，谷歌新方法则需要3-6小时来完成。

▲谷歌新方法与RePlAce方法、人工布局方法的实验结果对比（表中所有指标越低越好）

但相比最先进的RePlAce，谷歌新方法在生成满足设计要求的布局方面性能更优，其主要优点包括：能很容易地优化各种不可微分的成本函数，而不需要转换封闭形式或与那些成本函数等价的可微分形式；当有更多芯片块时，该新方法的性能会随着时间推移而改进；遵守对布局密度和布线拥塞的约束。

结语：新方法可拓展到芯片设计的其他环节

总体来看，谷歌这项新研究提出了一种基于深度强化学习的芯片布局方法，支持迁移学习，能产生比现有最先进方法性能更优的芯片布局，并且能在6小时内完成人类专家需要几周才能完成的设计。

研究人员提到，其新方法只是领域适应策略优化的一个例子，它可以被拓展到芯片设计的其他阶段，比如架构和逻辑设计、合成和设计验证。基于学习的方法还可以在组成芯片设计过程的任务级联中进一步探索设计空间和协同优化。

原文来自：Google AI