高科技产品都离不开作为「心脏」的芯片,芯片设计周期的缩短将有助于硬件设备适应机器学习领域的快速发展。最近,谷歌提出了一种基于强化学习的芯片布局方法。
人工智能的快速发展得益于计算机系统和硬件的显著进步,但随着摩尔定律趋缓,世界需要新的硬件架构来满足指数级增长的算力需求。然而,芯片通常需要数年的时间来设计,而在芯片设计过程中,芯片布局(chip placement)可以说是其中最复杂和耗时的步骤了。
“让天下没有难做的芯片”。为了解决这一困境,Google 做了一个大胆的决定,让 AI 来设计 AI 芯片,提出了一种基于强化学习的芯片布局方法。与已有方法不同,该方法能够基于之前的经验进行学习,并随着时间不断改进。
谷歌研究人员相信,AI将提供缩短芯片设计周期的手段,在硬件和AI之间建立一种共生关系,彼此推动彼此的进步。
该研究的论文题目为《芯片布局与深度强化学习 (Chip Placement with Deep Reinforcement Learning)》,Azalia Mirhoseini和Anna Goldie为共同作者。
| 论文链接: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 和人工基线方法的对比结果。
总体来看,谷歌这项新研究提出了一种基于深度强化学习的芯片布局方法,支持迁移学习,能产生比现有最先进方法性能更优的芯片布局,并且能在6小时内完成人类专家需要几周才能完成的设计。
随着 AI 技术的发展,不知道还会给未来带来多少便捷,给生活带来多少惊喜。科技改变生活,我们拭目以待。
思铺学院·人工智能原理与应用场景学术挑战营,讲透人工智能的基本概念和方法,实践企业产品应用场景和科学研究技能,以小组形式协作完成一个多阶段场景案例智能解决方案。从而进一步熟练掌握人工智能的相关知识与技能,提高分析和解决实际问题的综合能力,欢迎对人工智能感兴趣的同学们踊跃报名。