AI助力芯片设计效率革命!Jeff Dean领衔推出最新方案,6小时内完成布局设计,新一代TPU已用上

设计一款高端芯片有多难?想象一下要将上百亿个晶体管集成到指甲盖大小的区域上,还要进行各类功能区域规划,这是一项多么微观且浩瀚的工程。

如今,最基本的芯片晶体管部件都是微米、纳米级尺度,如何铺设和利用这几百亿个晶体管,如何通过模块设计组合迸发出一款芯片的极致性能,如何满足日益增长的终端需求,这是如今芯片设计工程师们面临的严苛挑战,而且如此复杂的设计,不能有纰漏缺陷,一旦投入量产无法补救,必须复盘再来,而如果一款芯片设计方案足够好,甚至能弥补一定程度的制程差距。

现在,随着 AI 技术深入到芯片设计领域,一场效率革命可能要再次席卷而来。

来自由 Jeff Dean 领衔的谷歌大脑团队以及斯坦福大学计算机科学系的科学家们,在一项联合研究中证明,一种基于深度强化学习的芯片布局规划方法,能够生成可行的芯片设计方案。

在不到 6 小时的时间内,由该方法自动生成的芯片平面图在所有关键指标上(包括功耗、性能和芯片面积等参数)都优于或与人类专家生成的设计图效果相当,而人类工程师往往需要数月的紧张努力才能达到如此效果。

图|人类设计的微芯片平面图(a)与机器学习系统设计(b)的不同(来源:Nature

据论文描述,如今这种方法已经被用于下一代 Google 张量处理单元(TPU)加速器的产品中,有望为今后的每一代计算机芯片迭代节省数千小时的人力,相关论文于 6 月 10 日发表在顶级科学期刊 Nature 上,谷歌 AI 掌门人 Jeff Dean 是核心作者之一。

(来源:Nature

芯片设计工程师的AI助理

事实上,科学家们一直在探索让芯片设计过程更高效的路径。

重大的进步出现在 1980 年代,电子设计自动化(EDA)技术诞生,开发人员能将整个设计过程自动化,利用计算机辅助设计(CAD)软件,完成超大规模集成电路(VLSI)的功能设计、综合、验证、物理设计等复杂流程,EDA 的不断完善架起了芯片设计与制造之间的桥梁,在半导体产业链中有着举足轻重的作用。

尽管已经经过了 50 年的相关研究,但芯片布局规划仍无法实现自动化,依然需要物理设计工程师数月精心布置才能完成可制造的芯片布局。

具体而言,一个芯片往往会被分成几十个模块,每个模块都有单独的功能属性,例如内存子系统、计算单元或控制逻辑系统,这些模块可以用网表(也称连线表)、电路组件的超图和标准单元(逻辑门,例如 NAND、NOR 和 XOR)来描述,所有这些模块再通过导线连接起来。

芯片布局规划则涉及将网表放置在芯片画布(二维网格)上,以便优化性能指标(例如功耗、定时、面积和线长等),同时遵守密度和路由拥塞的硬性约束条件。

其实,自 20 世纪 60 年代以来已有许多芯片布局规划方法被提出,大概分为三类:

  • 基于分区的方法(partitioning-based methods);
  • 随机/爬山方法(stochastic/hill-climbing approaches);
  • 解析解算器(analytic solvers)。

但是,这些方法目前都无法达到人类专家级水平的性能,因为各有各的局限。

  • 基于分区的方法为了扩展到更大的网表而牺牲了全局解决方案的质量,而一个糟糕的早期分区可能会导致一个无法解决的最终结果;
  • 随机/爬山方法具有较低的设计收敛率,不能扩展到具有数百万乃至数十亿节点的现代芯片网表;
  • 解析解算器是一种主流方法,但它只能优化可微损失函数,这意味着它们不能有效地优化关键指标,如路由拥塞或时序冲突等。

由于这些先前方法的局限性,人类物理设计师必须使用商业 EDA 工具迭代数月,将芯片网表的寄存器传输级(RTL)描述作为输入,在芯片画布上进行手动放置,并等待长达 72 小时的时间,让 EDA 工具来评估该位置是否最佳。

在这个反馈的基础上,人类设计师要么得出设计标准达到目标的结论,生成一个更新的层级平面图进行评估,要么得向上游 RTL 设计师提供反馈,然后上游 RTL 设计师修改低级代码以使放置任务更容易。

为了解决这类复杂的芯片布局规划问题,研究人员开发了一种强化学习方法,该方法能够跨芯片进行推广,这意味着它可以从经验中不断学习,在放置新芯片模块时变得更好更快,从而使芯片设计师能够得到 AI 的协助,且这个 AI 助理的“经验”比任何人都丰富。

但培训跨芯片的 AI 布局策略非常具有挑战性,因为它需要学习优化所有可能的芯片网表到所有可能的芯片画布上的布局。

这类似于一个棋盘游戏,元件是“棋子”(例如网表拓扑、宏计数、宏大小和纵横比)、放置元件的画布是“棋盘”(不同的芯片画布大小和纵横比)和获胜条件(不同的评估指标或不同的密度和路由拥塞约束的相对重要性)。即使是这个游戏的一个实例(将一个特定的网表放到一个特定的画布上)也有一个巨大的状态-动作空间,对全局形成影响。

为了实现泛化,研究人员重点研究了芯片的可转移表征学习,并将表征学习应用于预测芯片放置质量的有监督任务中。通过设计一种能够准确预测各种网表及其位置的神经结构,研究人员能够生成输入网表的丰富特征嵌入,然后,使用这个结构作为编码器,以实现转移学习。

除了对芯片布局规划的直接影响外,基于 AI 的芯片布局设计方案概括和快速生成高质量解决方案的能力也具有重大意义,为与芯片设计过程早期阶段的协同优化创造了机会。

在更抽象的层次上,研究人员表示,这种形式的位置优化策略还可出现在广泛的科学和工程应用中,例如硬件设计、城市规划、疫苗测试和分发以及大脑皮层布局研究等。

因此,研究人员相信,这种 AI 布局优化方法策略可以应用于芯片设计以外更广泛的有效布局问题。

设计域的自适应策略

为芯片布局规划开发领域自适应策略非常具有挑战性,因为每个芯片可以包含数千万个逻辑门(称为标准单元)以及数千个内存块(称为宏块或宏),单元和宏块通过数十公里的布线相互连接,以实现设计的功能,其中蕴藏着惊人的复杂性,工程师们在专业软件工具的帮助下,才能控制一部分复杂性。

其中,标准单元和宏块在芯片中的位置对设计结果至关重要,它们的位置决定了导线必须跨越的距离,从而影响元件之间的布线是否成功,以及逻辑门之间的信号传输速度。因此,随着逻辑电路设计的发展,人类工程师必须反复调整宏块的位置,每次迭代都需要手动执行,需要几天或几周的时间。

众所周知,计算机行业是由摩尔定律驱动的,每一块芯片的元件数量大约每两年翻一番,相当于一个芯片上的元件数量每周增加约 1%。因此,无法自动完成布局规划是一个迫切的问题,这不仅是因为相关的时间成本会增加,也限制了芯片开发计划中可以探索的解决方案的数量。

如何让 AI 代理成功地将宏块一个接一个合理地放入芯片布局中?

具体而言,研究人员首先集中学习状态空间的丰富表征,他们的直觉是,一个能够完成芯片放置一般任务的 AI 策略,也应该能够在推理时将与一个新的不可见芯片相关的状态编码成一个有意义的信号。

为此,研究人员首先训练了一个神经网络架构,能够预测新网表放置时的回报,最终目标是将此架构用作整个策略的编码层,训练这个架构需要一个大的芯片位置数据集和相应的奖励标签,研究人员创建了一个包含 10000 个芯片放置的数据集。

另外,为了准确预测奖赏标签并推广到未知数据,研究人员又开发了一种神经网络结构:Edge-GNN (基于边缘的图形神经网络),该网络的作用是嵌入网表,将节点类型和连通性的信息提取到低维向量表示中,以用于下游任务。

图|策略和价值网络架构(来源:该论文)

嵌入层会对有关网表邻接、节点特征和要放置的当前宏的信息进行编码,然后,策略和值网络分别输出可用网格单元上的概率分布和当前布局的预期回报估计。

图|整体方法和训练方案概述(来源:该论文)

在每个训练迭代中,强化学习代理一次放置一个宏(一种批量处理的称谓),放置完所有宏后,将使用强制方法放置标准单元,在每次迭代结束时,奖励被计算为近似线长、拥塞和密度的线性组合,并作为反馈提供给强化学习代理,以便为下一次迭代优化其参数。

域自适应结果

在后续的实验中,研究人员比较了使用预先训练策略生成的放置质量,以及通过从头开始训练策略生成的放置质量,训练数据集由 TPU 块和开源 Ariane RISC-V CPU 块组成,在每一个实验中,都对除目标块外的所有块预先训练策略。

研究人员展示了零样本模式的结果,以及在特定设计上对预训练策略进行 2 小时和 12 小时微调后的结果。

图|从无到有的训练与不同时间量的微调,从表中可以看出,预先训练的策略网络始终优于从无到有训练的策略网络,证明了从离线训练数据中学习的有效性(来源:该论文)

图|Ariane RISC-V CPU30 从零开始训练与从预先训练的策略网络开始训练的收敛图(来源:该论文)

结果发现,从零开始训练的策略需要更长的时间才能收敛,即使在 24 小时之后,结果(由奖励函数评估)也比微调策略在 12 小时内达到的效果差。这表明,在预训练期间接触许多不同的设计可以更快地为新的看不见的块生成更高质量的放置方案。

预先训练的策略不仅具有较低的放置成本,而且能比从头开始训练的策略收敛速度快 30 小时以上。

接下来,研究人员探讨了培训数据对策略学习能力的影响,例如从更大的数据集中学习。TPU 芯片块非常多样化,他们仔细选择了一系列具有代表性的功能块(例如片上和片间网络块、计算核心、内存控制器、数据传输缓冲区和逻辑以及各种接口控制器)、饱和(宏的总面积与画布的总面积之比,<30%,30%-60% 和 >60%)和宏计数(最多几百个),小训练集包含 2 个块,中训练集包含 5 个块,大训练集包含 20 个块。

图|训练前数据集大小的影响(来源:该论文)

当研究人员对更多的芯片块进行预训练时,能够更快地为新的看不见的芯片块生成更高质量的放置。当他们将训练集从 2 个块增加到 5 个块,最后增加到 20 个块时,策略网络在微调后都会生成更好的放置位置。

这表明,将策略暴露于更多不同的芯片设计时,它变得不太容易过度拟合,并且更善于推广到新的未知设计中。

研究人员也探讨了 AI 芯片布局设计与基线方法的数据比较。在评估芯片布局图的质量时,有几个重要指标,它们相互权衡。没有一个度量可以用来捕获放置的总体质量,因此报告中的关键度量,例如包括总的线长、计时、路由拥塞(水平和垂直)、面积和功率等。

为了进行公平比较,研究人员确保所有方法都有相同的实验设置,包括相同的输入和相同的EDA工具设置,使用一个在最大的数据集(20 个 TPU 块)上预先训练的 AI 策略,然后在 5 个目标不可见块上微调它,时间不超过 6 小时。

图|与基线的比较,AI 布局芯片的方法与最先进的方法,以及使用行业标准 EDA 工具手动放置进行比较。对于本表中的所有指标,越低越好(来源:该论文)

结果证明,利用 AI 进行芯片设计布局是有效的,不仅能生成高质量的模块放置方案,也能满足实际设计标准要求,在面积、功率和线长方面优于或不亚于人类专家手动放置的效果。

加利福尼亚大学圣迭戈分校电气与计算机工程系专家 Andrew B. Kahng 对这项成果评论称,对于芯片设计和设计自动化领域的长期实践者来说,这项研究成果确实很神奇。

最重要的启示可能是,该技术方案已经被纳入谷歌下一代 AI 处理器的芯片设计中,这意味着这种解决方案足够好、足够成熟,可以在昂贵、尖端的硅片上批量复制,因此,未来可以期望半导体行业从业者加倍关注此类相关工作,并在整个芯片设计过程中寻求大量类似的应用落地。

参考链接:

https://www.nature.com/articles/s41586-021-03544-w

https://www.nature.com/articles/d41

发表评论

Fill in your details below or click an icon to log in:

WordPress.com 徽标

您正在使用您的 WordPress.com 账号评论。 注销 /  更改 )

Google photo

您正在使用您的 Google 账号评论。 注销 /  更改 )

Twitter picture

您正在使用您的 Twitter 账号评论。 注销 /  更改 )

Facebook photo

您正在使用您的 Facebook 账号评论。 注销 /  更改 )

Connecting to %s