Jeff Dean：机器学习在硬件设计中的潜力

作者｜Jeff Dean

翻译｜沈佳丽、胡燕君、贾川

为什么芯片设计需要很长时间？能不能加速芯片设计周期？能否在几天或几周之内完成芯片的设计？这是一个非常有野心的目标。过去十年，机器学习的发展离不开系统和硬件的进步，现在机器学习正在促使系统和硬件发生变革。

Google在这个领域已率先出发。在第58届DAC大会上，Google AI负责人Jeff Dean分享了《机器学习在硬件设计中的潜力》，他介绍了神经网络发展的黄金十年，机器学习如何影响计算机硬件设计以及如何通过机器学习解决硬件设计中的难题，并展望了硬件设计的发展方向。 

他的演讲重点在于Google如何使用机器学习优化芯片设计流程，这主要包括架构搜索和RTL综合、验证、布局与布线（Placement and routing）三大阶段。在架构搜索阶段，Google提出了FAST架构自动优化硬件加速器的设计，而在验证阶段，他们认为使用深度表示学习可提升验证效率，在布局与布线阶段，则主要采用了强化学习技术进行优化。

以下是他的演讲内容，由OneFlow社区编译。

1

神经网络的黄金十年

制造出像人一样智能的计算机一直是人工智能研究人员的梦想。而机器学习是人工智能研究的一个子集，它正在取得很多进步。现在大家普遍认为，通过编程让计算机变得“聪明”到能观察世界并理解其含义，比直接将大量知识手动编码到人工智能系统中更容易。

神经网络技术是一种非常重要的机器学习技术。神经网络一词出现于1980年代左右，是计算机科学术语中一个相当古老的概念。虽然它当时并没有真正产生巨大的影响，但有些人坚信这是正确的抽象。 

本科时，我写了一篇关于神经网络并行训练的论文，我认为如果可以使用64个处理器而不是一个处理器来训练神经网络，那就太棒了。然而事实证明，我们需要大约100万倍的算力才能让它真正做好工作。

2009年前后，神经网络技术逐渐火热起来，因为我们开始有了足够的算力让它变得有效，以解决现实世界的问题以及我们不知道如何解决的其他问题。2010年代至今是机器学习取得显著进步的十年。

是什么导致了神经网络技术的变革？我们现在正在做的很多工作与1980年代的通用算法差不多，但我们拥有越来越多的新模型、新优化方法等，因此可以更好地工作，并且我们有更多的算力，可以在更多数据上训练这些模型，支撑我们使用更大型的模型来更好地解决问题。 

在探讨设计自动化方面之前，我们先来看看一些真实世界的例子。首先是语音识别。在使用深度学习方法之前，语音识别很难得到实际应用。但随后，使用机器学习和神经网络技术，大幅降低了词语的识别错误率。 

几年后，我们将错误率降低到5%左右，让语音识别更加实用，而现在，在不联网的设备里，我们都可以做到仅仅4%左右的错误率。这样的模型被部署在人们的手机里面，随时随地帮助人们识别自己的语音。

计算机视觉方面也取得了巨大的进步。2012年左右，Alex Krizhevsky、Ilya Sutskever和Geoffrey Hinton在ImageNet比赛中首次使用了AlexNet，错误率得到显著降低，并在当年夺得桂冠。

后一年的ImageNet比赛中，几乎所有参赛者都使用深度学习方法，研究人员则进一步放弃了传统的方法。其中，2015年，由何恺明等微软研究人员提出ResNet更进一步降低了错误率。 

当时的斯坦福大学研究生Andrej Karpathy正在帮助运营ImageNet比赛，他想知道如果人工识别这项艰难的任务，错误率会是多少。在上千个类别中有40种狗，你必须能够看着一张照片说：“哦，那是一只罗威纳犬，不是一只大力金刚犬，或者其他品种的狗。” 经过一百个小时的训练，他将错误率降到了5%。

这是一项非常艰难的任务，将计算机识别错误率从2011年的26%降低到2017年的2%是一件很了不起的事，过去计算机无法识别的东西，现在已经可以识别。自然语言处理、机器翻译和语言理解中也经历了类似的故事。 

此外，开源框架确实使世界各地的许多人能够应用机器学习技术，TensorFlow就是其中之一。

大约在2015年11月，我们开源了TensorFlow以及供Google内部使用的工具。TensorFlow对世界产生了相当大的影响，它已经被下载了大约5000万次，当然也出现了很多其他框架，比如JAX、PyTorch等等。

世界各地的人们能够将机器学习用于各种了不起的用途，例如医疗保健、机器人技术、自动驾驶等等，这些领域都是通过机器学习方法来理解周围的世界，进而推动领域的发展。

2

机器学习改变计算机设计方式

ML研究社区中的许多成功源自使用更多算力和更大的模型，更多的算力促进了机器学习研究领域中重要成果的产生。深度学习的发展正在深刻改变计算机的结构。现在，我们想围绕机器学习计算类型构建专门的计算机。

近年来，我们已经在Google做了很多类似的工作，其中TPU（张量处理单元）是我们构建定制处理器的一种方法，这些处理器专为神经网络和机器学习模型而设计。

TPU v1是我们第一个针对推理的产品，当你拥有经过训练的模型，并且只想获得已投入生产使用的模型的预测结果，那它就很适合，它已经被用于神经机器翻译的搜索查询、AlphaGo比赛等应用中。

后来我们还构建了一系列处理器。TPU v2旨在连接在一起形成称为Pod的强大配置，因此其中的256个加速器芯片通过高速互联紧紧连接在一起。TPU v3则增加了水冷装置。

TPU v4 Pod不仅可以达到ExaFLOP级的算力，它还让我们能够在更大的模型训练中达到SOTA效果，并尝试做更多的事情。 

以ResNet-50模型为例，在8块P100 GPU上训练完ResNet-50需要29小时，而在2021年6月的MLPerf竞赛中，TPU v4 pod仅耗时14秒就完成了训练。但我们的目的不仅仅是在14秒内训练完ResNet，而是想把这种强大的算力用于训练其他更先进的模型。 

可以看到，从一开始的29小时到后来的14秒，模型的训练速度提高了7500倍。我认为实现快速迭代对于机器学习非常重要，这样才能方便研究者试验不同想法。

基于机器学习的计算方式越来越重要，计算机也正在往更适应机器学习计算方式的方向上演进。但深度学习有可能影响计算机的设计方式吗？我认为，答案是肯定的。

3

机器学习缩短芯片设计周期

目前，芯片的设计周期非常长，需要几十甚至几百人的专业团队花费数年的努力。从构思到完成设计，再到成功生产，中间的时间间隔十分漫长。但如果将来设计芯片只需要几个人花费几周时间呢？这是一个非常理想的愿景，也是研发人员当前的目标。

如上图所示，芯片设计包含四个阶段：架构探索→RTL综合→验证→布局和布线。完成设计之后，在制作生产环节需要进行布局和布线（Placement & Routing），有没有更快、更高质量的布局和布线方法？验证是非常耗时的一步，能不能用更少的测试次数涵盖更多的测试项目？有没有自动进行架构探索和RTL综合的方法？目前，我们的芯片架构探索只针对几种重要的应用，但我们终将要把目光扩大。

布局与布线

首先，关于布局和布线，Google在2020年4月发表过一篇论文Chip Placement with Deep Reinforcement Learning，2021年6月又在Nature上发表了A graph placement methodology for fast chip design。

我们知道强化学习的大致原理：机器执行某些决定，然后接收奖励（reward）信号，了解这些决定带来什么结果，再据此调整下一步决定。

因此，强化学习非常适合棋类游戏，比如国际象棋和围棋。棋类游戏有明确的输赢结果，机器下一盘棋，总共有50到100次走棋，机器可以根据最终的输赢结果评定自己和对手的整套走棋方法的有效性，从而不断调整自己的走棋，提高下棋水平。 

那么ASIC芯片布局这项任务能不能也由强化学习智能体来完成呢？

这个问题有三个难点。第一，芯片布局比围棋复杂得多，围棋有10^{360}种可能情况，芯片布局却有10^{9000}种。

第二，围棋只有“赢”这一个目标，但芯片布局有多个目标，需要权衡芯片面积、时序、拥塞、设计规则等问题，以找到最佳方案。

第三，使用真实奖励函数（true reward function）来评估效果的成本非常高。当智能体执行了某种芯片布局方案后，就需要判断这个方案好不好。如果使用EDA工具，每次迭代都要花上很多个小时，我们希望将每次迭代所需时间缩减为50微秒或50毫秒。

利用强化学习进行芯片布局的步骤如下：首先从空白底座开始，运用分布式PPO算法（强化学习的常用算法）进行设计，然后完成每个节点的布局放置，最后进行评估。

评估步骤使用的是代理奖励函数（proxy reward function），效果和真实奖励函数相近，但成本低得多。在一秒或半秒内就可以完成对本次布局方案的评估，然后指出可优化之处。

构建奖励函数需要结合多个不同的目标函数，例如线长、拥塞和密度，并分别为这些目标函数设定权重。

如上图所示，布局优化采取的是混合方式。强化学习智能体每次放置宏（macro），然后通过力导向方法（force-directed method）放置标准单元。

上图来自前面提到的Nature论文，展示了更多芯片架构的细节。 

上图展示了一个TPU设计块的布局与布线结果。白色区域是宏，绿色区域是标准单元群（standard cell clusters）。

图中左边是人类专家完成的设计，从中可以看出一些规律。人类专家倾向于把宏沿边缘放置，把标准单元放在中间。一名人类专家需要6~8周完成这个布局，线长为57.07米。图中右边是由智能体（ML placer）完成的布局，耗时24小时，线长55.42米，违反设计规则的地方比人类专家略多，但几乎可以忽略。

可以看出，智能体并不像人类专家一样追求直线布局。为了优化布局，智能体更倾向于弧形布局。我们也希望能训练智能体高效地创造一些前所未有的布局方式。

为了实现这个目标，首先，我们运用强化学习算法优化某个芯片设计块的布局，期间需要经历上万次迭代；然后，重复前一步骤，在多个不同的设计块上预训练出一套布局规则，最终让算法在面对前所未见的新设计块时也能给出布局方案。

预训练好的策略有助于在推理时做更少的迭代，进行“零次（zero-shot）布局”。我们实际上还没有新算法来优化这个特定的设计，当然我们可以做数百次迭代以得到更好的结果。

上图展示了使用不同方法时的布局成本。蓝线表示不经过预训练的从头训练策略的布局成本，绿线表示用已预训练的算法优化新设计块的布局，X轴表示训练时间，可以看到，蓝线在经过20多小时的训练后方可大幅降低布局成本，此后仍需经过一段时间才能达到收敛。而绿线只用了极少的训练时间就达到了更低的布局成本并很快收敛。

最令我感到兴奋的是圈红部分。调优预训练策略在短时间内就可实现相当不错的布局，这种实现就是我所说的，在一两秒内完成芯片设计的布局。

 上图展示了更详细的不同设计的情况。Y轴表示线长成本（越低越好）。绿色表示使用商业工具的线长成本，可以看到，从浅蓝色（零次布局）→蓝色（2小时微调）→深蓝色（12小时微调），线长成本越来越低。深蓝色一直比黄色的线长成本要低，因为通过12小时的微调能从其他设计中学到最佳布局。

上图中，左边和右边分别展示了从头训练的策略和经过微调预训练策略的芯片布局过程。每个小方块表示一个宏的中心，空白部分表示为标准单元预留的位置。可以看到，右边从一开始就将宏放在边缘，将大片中间区域留空。而左边则要经过很多次迭代才能形成这样的格局。

我们利用强化学习工具针对TPU v5芯片的37个设计块进行了布局与布线。其中，26个设计块的布局与布线质量优于人类专家，7个设计块的质量与人类专家相近，4个设计块的质量不如人类专家。目前我们已经把这个强化学习工具投入到芯片设计流程中了。

总的来说，用机器学习进行芯片布局与布线的好处包括：可以快速生成多种布局方案；即使上游设计有重大改动也可以迅速重新布局；大幅减少开发新ASIC芯片所需的时间和精力。

验证

接下来是芯片设计的验证阶段。我们希望用较少的测试次数覆盖多个测试项目。验证是阻碍芯片设计提速的主要瓶颈。据估计，芯片设计过程中，80%的工作量在于验证，而设计本身仅占20%。因此，验证技术的任何一点进步都会产生重大作用。

Google在2021年NeurIPS（神经信息处理系统大会）上发表了论文《Learning Semantic Representations to Verify Hardware Designs》，我们能不能运用机器学习生成在更短时间内覆盖更广状态空间的测试用例？ 

验证阶段的基本问题是可达性（reachability）。目前的芯片设计能否让系统达成需要的状态？我们的想法是，根据当前的芯片设计生成一个连续的表示，从而预测对系统的不同状态的可达性。

我们可以通过RTL将芯片设计抽象为一张图，然后运用基于图的神经网络去了解该图的特性，从而了解其对应芯片设计的特性，继而决定测试覆盖率和测试用例，这给了我们一个很好的设计的抽象表示。 

当然，如何将这种方法应用到实际芯片设计中将是另外一个重要话题。用RTL生成图表示之后，我们在图神经网络中运用一种叫Design2Vec的技术进行深度表示学习，从而帮助我们作出预测。 

目前，芯片的验证环节需要大量人力，例如，找bug、查找测试覆盖率漏洞、分析和解决bug等，还需要经历多次如上图所示的流程循环。我们希望上述步骤可以实现自动化，自动生成新的测试用例以解决重要的问题。

后来我们发现，可以把这个问题转化为一个监督学习问题。如果之前进行了一系列测试，并知道这些测试覆盖哪些测试点，就可以将这些数据用作监督学习中的训练数据。

然后，当出现新的测试点时，假设进行一个新的测试，我们需要预测这个测试能否覆盖新的测试点。我们希望能结合之前的训练数据以及芯片设计本身，来实现这种预测。 

我们有两个Baseline，其中一个能够看到测试点（test points）和覆盖点（cover points）的数据，这是一个黑盒测试。

而Design2Vec除了能够处理上述数据外，还能处理实际设计、设计的图结构等等。如果你使用一半的测试点作为训练数据，并且设置多个大小不同的训练集，然后对其它测试点进行预测，那么将会得到非常出色的结果，即使是对于相对较少的覆盖点，也能泛化得非常好。相比之下，Baseline这种方法就不能对此进行很好地泛化。 

但使用图神经网络来学习设计、覆盖率和测试属性的方法，实际上比NeurIPS论文中的其他所有Baseline都要好。

例如，我们常会遇到很多难以生成测试的覆盖点。工程师们发现使用RISC-V Design和TPU Design这两种不同的设计也很难为这些特定的覆盖点生成测试，于是我们又转向使用贝叶斯优化器来尝试生成测试。

上图右边这一列是贝叶斯优化器覆盖的不同测试点、覆盖点所需的模拟器调用数（simulator calls），中间一列是使用Design2Vec所需的模拟器调用数。从中可以看到，为覆盖这些有挑战性的覆盖点，Design2Vec生成的测试要少于贝叶斯优化器。所以Design2Vec非常好，相比之下它更快，能聚焦覆盖范围，还能节省在运行计算模拟器（本身很昂贵）上的开销。 

验证是芯片设计在理论和实践上长期面临的一个挑战。我们认为，深度表示学习能够显著提高验证效率和质量，并且在设计中实现泛化。

即使设计发生了一些改变，这个新设计的版本也能运用之前在众多设计上训练出来的系统，提高验证效率。正如在布局与布线阶段，经过训练后的算法即使面对新设计也能够预测不同测试的覆盖点，以带来好的结果。

架构探索和RTL综合

在芯片设计中，另一个比较耗时的方面是要清楚你究竟想要构建何种设计。此时你需要做一些架构探索（architectural exploration），然后做RTL综合。目前计算机架构师和其他芯片设计师等具有不同专业知识的人花费大量时间来构建他们真正想要的设计，然后验证、布局和布线，那么我们可以学习自动做架构探索和综合吗？ 

现在我们正在研究的就是如何为已知的问题实行架构探索。如果我们有一个机器学习模型，并且想要设计一个定制芯片来运行这个模型，这个过程能否实现自动化，并提出真正擅长运行该特定模型的优秀设计。

关于这项工作，我们在arXiv发表了论文《A Full-stack Accelerator Search Technique for Vision Applications》，它着眼于很多不同的计算机视觉模型。另外一个进阶版本的论文被ASPLOS大会接收了《A Full-stack Search Technique for Domain Optimized Deep Learning Accelerators》。

这里要解决的问题是：当你设计一个机器学习加速器时，需要考虑你想在哪个加速器上运行什么样的机器学习模型，而且这个领域的变化非常之快。 

上图中的红线是指引入的不同计算机视觉模型，以及通过这些新模型实现的ImageNet识别准确率提升。

但问题是，如果你在2016年想要尝试设计一个机器学习加速器，那么你需要两年时间来设计芯片，而设计出来的芯片三年后就会被淘汰。你在2016年做的决定将会影响计算，要保证在2018年-2021年高效运行，这真的很难。比如在2016年推出了Inception-v3模型，但此后计算机视觉模型又有四方面的大改进。

因此，如果我们能使设计周期变得更短，那么也许单个工作负载加速器能变得可用。如果我们能在诸多流程中实现自动化，那么我们或许能够得到正反馈循环，即：缩短机器学习加速器的上市时间，使其能更适合我们当下想要运行的模型，而不用等到五年后。

4

用机器学习探索设计空间

实际上，我们可以使用机器学习来探索设计空间。有两个因素影响加速器性能，一是设计中内置的硬件数据通道，二是工作负载如何通过编译器而不是更高级别的软件映射到该数据通道。通常，设计空间探索实际上只考虑当前编译器优化的数据通道，而不是协同设计的编译器优化和优化数据通道时可能会做的事。 

因此，我们能否同时优化数据通道、调度（schedule）和一些编译器决策，并创建一个搜索空间，探索出你希望做出的共同设计的决策。这是一种覆盖计算和内存瓶颈的自动搜索技术，探索不同操作之间的数据通道映射和融合。通常，你希望能够将事物融合在一起，避免内存传输的每次内存负载中执行更多操作。

根本上说，我们在机器学习加速器中可能做出的设计决策创建了一种更高级别的元搜索空间，因此，可以探索乘法的脉冲列阵（systolic array）在一维或二维情况下的大小，以及不同的缓存大小等等。 

如前所述，考虑编译器优化与硬件设计的协同设计也很重要，因为如果默认编译器不会更改，就无法真正利用处理器中底层设计单元的变化。实际上，不一定要考虑特定设计的所有效果和影响。

接下来看看这种方式产生的一系列结果，将这些结果与TPUv3芯片的baseline（上图蓝条）进行比较。实际上这是假定型TPUv3芯片，其中模拟器已停止了运行。我们已经将其缩小到了sub-10纳米工艺。我们还将研究TPUv3的软件效用，以及共同探索在设计空间中的编译器优化。

红条和蓝条表示的内容是一致的，但一些探索过的编译器优化不一定在蓝条中得以体现，而这里的绿条则表示的是为单一计算机视觉模型定制的假定型设计。EfficientNet-B0...B7表示相关但规模不同的计算机视觉模型。与蓝条baseline相比，（绿条的）Perf/TDP的改进大约在3到6倍之间。 

那么除EfficientNet-B0...B7外，其他模型的情况如何？在此前所述的ASPLOS论文中提出更广泛的模型集，尤其是那些计算机视觉以外的BERT-seq 128和BERT-seq 1024等NLP模型。 

实际上，定制化芯片不只是适用于单个机器学习模型，而是使其适用于一组机器学习模型。你可能不想使你的加速器芯片设计仅针对某一项任务，而是想涵盖你所关注的那一类任务。

上图的黄条代表为五种不同模型设计的定制化芯片，而我们想要一个能同时运行这五种模型（红色箭头所指）的芯片，然后就能看出其性能能达到何种程度。可喜的是，从中可以看到，黄条（单一负载）并不比绿条（多负载）的性能低多少。所以你实际上可以得到一个非常适合这五种模型的加速器设计，这就好比你对其中任何一个模型都进行了优化。它的效果可能不是最好的，但已经很不错了。

而且，如果你关注的点是性能而非Perf/TDP，得到的结果实际上会更好。所以结果如何取决于你关注的是什么，是绝对性能还是每瓦性能？在Perf//TDP指标中，性能结果甚至提升了2.6到8.8倍，而非Perf/TDP指标下的1.8到6.4倍。

因此，我们能够针对特定工作负载进行定制和优化，而不用构建更多通用设备。我认为这将会带来显著改进。如果能缩短设计周期，那么我们将能以一种更自动化的方式用定制化芯片解决更多问题。 

当前的一大挑战是，如果了解下为新问题构建新设计的固定成本，就会发现固定成本还很高，因此不能广泛用于解决更多问题。但如果我们能大幅降低这些固定成本，那么它的应用面将会越来越广。

5

总结

我认为，在计算机芯片的设计过程中，机器学习将大有作为。

如果机器学习在合适的地方得以正确应用，那么在学习方法（learning approaches）和机器学习计算的加持下，芯片设计周期能不能缩短，只需要几个人花费几周甚至几天完成呢？我们可以用强化学习使得与设计周期有关的流程实现自动化，我认为这是一个很好的发展方向。

目前人们正通过一组或多组实验来进行测验，并基于其结果来决定后续研发方向。如果这个实验过程能实现自动化，并且能获取满足该实验正常运行的各项指标，那么我们完全有能力实现设计周期自动化，这也是缩短芯片设计周期的一个重要方面。 

这是本次演讲的部分参考文献以及相关论文，主要涉及机器学习在芯片设计和计算机系统优化中的应用。 

机器学习正在很大程度上改变人们对计算的看法。我们想要的是一个可以从数据和现实世界中学习的系统，其计算方法与传统的手工编码系统完全不同，这意味着我们要采取新方式，才能创建出我们想要的那种计算设备和芯片。同时，机器学习也对芯片种类和芯片设计的方法论产生了影响。

我认为，加速定制化芯片设计过程中应该将机器学习视为一个非常重要的工具。那么，到底能否将芯片设计周期缩短到几天或者几周呢？这是可能的，我们都应该为之奋斗。

（原视频：https://www.youtube.com/watch?v=FraDFZ2t__A）

其他人都在看

点击“阅读原文”，欢迎体验OneFlow v0.8.0