谷歌第三代ai芯片（6小时完成芯片布局）

机器之心报道

机器之心编辑部

在芯片设计过程中，芯片布局（chip placement）可以说是其中最复杂和耗时的步骤了。芯片设计周期的缩短有助于硬件设备适应机器学习领域的快速发展，那么，机器学习能否助力芯片设计呢？最近，谷歌提出了一种基于强化学习的芯片布局方法。

与已有方法不同，该方法能够基于之前的经验进行学习，并随着时间不断改进。具体而言，在大量芯片 block 上执行训练时，该方法能够更迅速地为新的芯片 block 进行优化布局。

那么，这一结果是如何实现的呢？

研究者将芯片布局看作一个强化学习问题，然后训练智能体将芯片网表（netlist）的节点放置在芯片画布（canvas）上。为了使强化学习策略泛化至新的芯片 block，研究者将表征学习置于预测芯片布局质量的监督任务中。通过设计能够在大量网表及其布局上准确预测奖励的神经架构，该研究生成输入网表的丰富特征嵌入。然后利用该架构作为策略和价值网络的编码器，实现迁移学习。

该研究旨在最小化芯片设计的 PPA（功耗、性能和面积）。研究者称，该方法能够在 6 小时内完成芯片布局设计，布局质量超过或匹配人类设计，而现有的基线方法需要人类专家参与，且往往需要数周时间才能完成。此外，谷歌还表示，该方法可以为谷歌加速器芯片（TPU）生成更优化的芯片放置方案，还适用于任意类型的芯片（ASIC）。

博客链接：http://ai.googleblog.com/2020/04/chip-design-with-deep-reinforcement.html

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

芯片布局规划问题

计算机芯片通常分为数十个模块，每个模块都是一个单独的模组，例如内存的子系统、计算单元以及控制逻辑的系统。这些模块可以通过网表以及宏（内存组件）和标准单元（逻辑门，例如 NAND、NOR 和 XOR）等电路组件图来描述，而所有这些组件均通过网格连接。

确定芯片如何布局（通常称为芯片的布局规划过程）是芯片设计过程中最复杂、最耗时的阶段之一，它涉及到将网表放置在芯片的画布（2D 网格）上，尽可能使得功率、性能和面积（PPA）降至最低，同时还要注意密度和布线拥塞方面的限制。

尽管对此方向进行了数十年的研究，但是行业内专家仍然需要迭代数周的时间才能完成一个满足多方面设计标准的解决方案。简单来说，其复杂性来自于几个主要层面：网表图的大小（数百万至数十亿个节点）、网表图放置的网格粒度，以及计算真实目标所产生的过高成本，如果使用行业标准的电子设计自动化工具这个计算过程可能要花费数小时（有时甚至超过一天）。

谷歌提出用强化学习来解决芯片布局问题

谷歌模型的输入包括芯片网表（节点类型和图邻接信息）、即将放置的当前节点 ID，以及导线总数、宏和标准单元集群等网表元数据。研究者使用谷歌去年提出的基于边的图神经网络传输网表图和当前节点，以编码输入状态，从而为部分放置的图和候选节点生成嵌入。

研究者提出利用一种新型神经架构来为芯片布局训练域适应策略。首先重点学习每个状态空间的丰富表征，并进而提出一个能够预测新网表上奖励的神经网络架构，最终目标则是使用这一架构作为策略网络的编码层。

训练这种监督模型需要一个包含芯片布局和相应奖励标签的大型数据集。因此，研究者创建了一个包含 10000 个芯片布局的数据集，其中输入为与给定布局相关的状态，标签是对该布局位置的奖励（导线长度和拥塞）。在创建数据集时，研究者首先选择了 5 个不同的加速器网表，然后为每个网表生成 2000 个不同的布局。

此外，为了训练能够准确预测导线长度和拥塞标签并泛化至未见过数据的监督模型，研究者创建了一个嵌入网表信息的新型图神经网络架构。为了将该架构集成至策略网络，研究者去除了预测层，然后使用它作为策略网络的编码器组件，具体如下图所示：

图神经网络生成嵌入，该嵌入与网表元数据嵌入一道成为策略和价值网络的输入。策略和价值网络整体架构如上所示，其中嵌入层对网表邻接信息、节点特征和即将放置的当前宏的信息进行编码。然后，策略和价值网络分别输出可用布局位置的概率分布以及对当前布局预期奖励的估计。

接下来，边、宏和网表元数据嵌入结合在一起形成单个状态嵌入（single state embedding），并传递至前馈神经网络。该前馈网络输出一个学得表征（learned representation）。该表征捕捉到有用的特征，并成为策略和价值网络的输入。然后，策略网络对当前节点可以放置其上的所有可能网格单元生成概率分布。

研究者采用强化学习方法来解决芯片布局问题，其中 RL 智能体（即策略网络）依次放置宏。宏放置完成后，利用密歇根大学研究者在论文《VLSI Cell Placement Techniques》提出的力引导（force-directed）方法来完成标准单元的大致布局（tough placement），如下图所示：

谷歌利用 RL 智能体放置宏、标准单元的流程图。放置完成后，计算近似导线长度和拥塞的线性组合并将其传递至智能体，为下次训练迭代进行参数优化。

谷歌新方法的效果如何？

谷歌在今日发布的博客中表示，该方法是首个具有泛化能力的芯片布局方法，这表示它可以利用布局之前的网表时所学的知识，为新的未见过的网表生成更好的布局。研究者表示，随着执行预训练的芯片网表的数量增加（即该方法在布局优化方面更加有经验），其策略能够更好地泛化至新的网表。

比如，预训练策略有机地确定了一种布局：将宏放置在芯片的边缘附近，将标准单元放置在芯片中心的凸起空间。这将导致宏和标准单元之间的线长变短，不会带来过度的布线拥塞。相比之下，从头开始训练的策略从随机布局开始，需要更长的时间才能收敛得到高质量的解决方案，而且后知后觉地发现需要在芯片画布的中心留出一个位置。参见以下动图：

开源 RISC-V 处理器 Ariane 的宏布局随着训练进程的变化情况。左图中策略是从头开始训练的，而右图中使用预训练策略对芯片进行微调。每个矩形代表一个单独的宏布局。注意，从头开始的策略发现的空洞在一开始就存在于预训练策略的布局中的。

研究者观察到预训练使得样本效率和布局质量均有所提升。该研究对比了使用预训练策略生成的芯片布局质量和从头开始训练策略得到的布局质量。

为了为之前未见过的芯片 block 生成芯片布局，研究者使用了 zero-shot 方法，即只使用预训练策略（不进行微调）来放置新的 block，并在不到一秒的时间内完成了布局。而经过微调后的结果还有待进一步提升。

从头训练的策略所需的收敛时间较长，甚至在 24 小时后其芯片布局质量仍然不如微调策略 12 小时后的结果。

两种策略在 Ariane 芯片 block 上的收敛轨迹，其中一个策略从头开始训练，另一个基于预训练策略进行微调。

谷歌提出的方法在较大的数据集上实现了更好的性能。据介绍，随着训练集从 2 个 block 上升到 5 个再到 20 个，RL 策略生成了更好的芯片布局，在相同的训练时钟时间下以上两种策略均实现了这样的效果。

训练数据规模 vs 微调后的性能。

谷歌还将自己的方法分别与模拟退火（Simulated Annealing，SA）方法、当前 SOTA 方法 RePlAce 和人工基线方法进行了对比，结果如下所示：

与 SA 方法相比，谷歌的方法不超过 6 小时即完成了收敛，而 SA 方法需要 18 个小时。并且，SA 方法生成高质量布局时需要的导线长度更大，布线拥塞也更高。

谷歌方法与当前 SOTA 方法 RePlAce 和人工基线方法的对比结果。

在此项研究工作中，研究团队始终是面向芯片布局这个核心点，此次提出基于强化学习的方法，而该方法也同时支持迁移学习，这表明强化学习的智能体会在越来越多的芯片网表中获得学习经验，从而在处理芯片布局方面变得更快更好。而新方法已经证明了优于 SOTA 标准，同时此方法是端到端的，并且可以在 6 个小时内生成布局位置。