特斯拉Dojo架构大解析——道场微架构 - Hot Chips 34

Hot Chips 34 - 特斯拉的道场微架构

我是斜杠青年，一个热爱前沿科技的“杂食性”学者！

说特斯拉只对机器学习感兴趣是轻描淡写的。这家电动汽车制造商制造了一台名为Dojo的内部超级计算机，专为训练其机器学习模型进行了优化。与许多其他超级计算机不同，Dojo没有使用现成的CPU和GPU，例如来自AMD、Intel或Nvidia。相反，特斯拉根据他们的需求设计了自己的微架构，让他们做出更一般的架构无法做出的权衡。在本文中，我们将根据特斯拉在Hot Chips的演示文稿来了解该架构。该架构没有单独的名称，因此为了简单起见，每当我们进一步提到Dojo时，我们都会谈论架构。

方框图

在高水平上，Dojo是一个八宽核心，具有四向SMT，以保守2 GHz运行。它具有CPU风格的管道，使其比GPU更能容忍不同的算法和分支代码。Dojo的指令集在标量方面类似于RISC-V，但特斯拉的工程师添加了一套专注于加速机器学习的自定义矢量指令。

特斯拉将道场描述为“高吞吐量、通用CPU”。从绩效的角度来看，这当然是有道理的。但为了增加计算密度，特斯拉做出了牺牲，与我们在台式机、笔记本电脑和智能手机中熟悉的CPU相比，Dojo内核将非常难以使用。在某些方面，Dojo核心处理的更像IBM Cell中的SPE，而不是传统的通用CPU核心。

像细胞的SPE？

在21世纪初中期，IBM的Cell处理器具有八个“辛能处理元件”（SPE），由功能齐全的CPU核心（“电源处理元件”或PPE）控制。乍一看，道场与SPE有很多相似之处。Dojo和SPE都针对矢量处理进行了优化，并依赖于单独的主机处理器进行工作分配。在Dojo或SPE上运行的代码无法直接访问系统内存。相反，应用程序预计将主要在一小块本地SRAM中工作。此本地SRAM由软件管理，不能用作缓存。如果需要来自主内存的数据，则必须使用DMA操作将其带入。最后，Dojo和Cell的SPE都缺乏虚拟内存支持。我们稍后会更多地讨论这意味着什么，但简而言之，它使多任务处理变得非常困难。

Dojo在几个重要方面有所不同。由于Dojo的设计没有考虑到小规模部署，因此主机处理器驻留在单独的主机系统上。这些主机系统具有带有接口处理器的PCIe卡，然后通过高速网络链路连接到Dojo芯片。相比之下，Cell的主机处理器位于同一个模具上。这使得单独部署单个Cell芯片成为可能——这在Dojo中是不可能的。Dojo的1.25 MB本地块是SRAM比Cell SPE大得多，带宽也比Cell SPE高得多，Cell SPE只有256 KB的SRAM。Cell的256 KB SRAM只有一个端口，每个周期可以交付128B。Dojo的SRAM有五个64B端口。当然，设计目标非常不同。Dojo是宽而低时钟的，而Cell SPE具有为高时钟设计的狭窄、深的管道。但有趣的是，Cell的高级设计概念今天仍然具有现实意义。

前端：CPU舒适度等

让我们从前端开始，简单地浏览一下道场的管道。有一种分支预测器，因为特斯拉的图表显示了BTB（分支目标缓冲区）。它的预测能力可能无法接近我们在AMD、ARM和英特尔高性能内核上看到的，因为Dojo需要优先将模具区域用于矢量执行。但即使是基本的分支预测器，也是从根本没有预测器迈出的一大步，Dojo的分支预测能力在处理分支代码或更大的指令足迹时，它应该比GPU提供更好的性能。

一旦分支预测器生成了下一个指令获取指针，Dojo可以将每个周期32字节从“小”指令缓存拉入每个线程获取缓冲区。此指令缓存可能有助于减少本地SRAM的指令带宽压力，确保数据端可以以尽可能少的争用访问SRAM。此外，指令缓存不连贯。如果新代码加载到本地SRAM中，则在分支到该新代码之前必须刷新指令缓存。

从获取缓冲区，Dojo的解码器每个周期可以处理来自两个线程的八条指令。我对“每个周期两个线程”的含义有点困惑，因为具有SMT的CPU通常每个周期获取和解码处理一个线程，并在循环边界上的线程之间切换。也许Dojo将解码器分成两个集群，并选择两个线程在每个周期喂养它们。这可能会减少从被提取的分支中解码吞吐量损失。

在解码时，某些指令，如分支、谓词操作和即时加载（“列表解析”）可以在前端执行并从管道中删除。这有点像较新的x86 CPU，消除了重命名器中的寄存器到注册副本。道场没有通过管道跟踪“消除”指令，以保持有序的退休。其他处理器将所有内容跟踪到退休，以便他们可以在任何指令边界停止，并保持恢复执行所需的所有状态。这种功能被称为“精确异常”，被现代操作系统用于交付各种好东西，例如分页到磁盘，或确切地告诉你代码在哪里搞砸了。

特斯拉不在乎提供精确的例外情况。Dojo确实具有调试模式，通过管道中有更多的指令来提供“更精确”的异常，但没有像正常的无序CPU那样的有序退出逻辑。

Dojo的执行引擎

在看到宽大的前端后，看到一个只有两个ALU和两个AGU的4宽整数执行引擎可能会有点奇怪。但这个漏斗形状的管道是有道理的，因为一些指令被执行并丢弃在前端。Dojo也没有进入客户端系统，在客户端系统中，标量整数性能很重要。因此，整数边提供了足够的吞吐量来压缩控制流和地址生成，以保持矢量和矩阵单元的进给。

Dojo的矢量和矩阵执行侧位于核心管道中的标量执行引擎之后，并有两个执行管道。两根管子听起来可能不太像，但道场在这些管道后面有非常宽的执行单元。一个管道可以执行512位矢量执行，而另一个管道可以执行8x8x4矩阵乘法。因此，只要指令暴露足够的显式并行性，Dojo就能够实现非常高的吞吐量——特别是在使用矩阵单元的情况下。特斯拉声称，一个拥有354个道场核心的模具可以在2 GHz时达到362 BF16 TFLOPS，这表明每个核心每个周期可以达到512 BF16 FLOPS。

我们不确定Dojo是否可以进行完全的无序执行。但特斯拉确实表示，整数侧可以远远领先于矢量端，这表明它可以执行过去停滞不前的指令，直到其中一个调度器填满。缺乏有序的退休也表明执行能力不合序。

通常，实现无序执行会带来很多复杂性。这是因为CPU必须看起来按顺序执行指令。AMD、ARM和英特尔的高性能CPU使用大型重新排序缓冲区（和其他结构）跟踪指令，以便其结果可以按程序顺序提交。这意味着，如果一个程序做了一些愚蠢的事情，比如除以零，这些核心可以准确显示哪个指令是错误的。而且，他们可以显示CPU状态，该状态反映了故障前所有指令中的操作，但之后没有。这意味着你可以修复任何使指令出现故障的东西，并恢复执行。Dojo放弃了这项能力。作为交换，Dojo可以避免与通过管道跟踪每条指令相关的电源和区域开销，以确保结果按程序顺序提交。

SRAM访问

通常，我们会在这里讨论缓存。但Dojo无法直接访问系统内存，因此我们将讨论1.25 MB的SRAM块。它每个周期可以处理两个512位负载，使其每个周期的带宽与支持AVX-512的英特尔CPU相匹配。特斯拉表示，SRAM有五个512位端口（2个负载，1个存储端口，两个到网格停止）。但标量端只有两个AGU，这可能意味着核心无法在每个周期都承受两个512位负载和一个512位存储。

由于Dojo的本地SRAM块不是缓存，因此它不需要与数据一起存储的标签和状态位。SRAM前面也没有L1D缓存，因此它必须足够快，才能处理所有加载和存储指令，而不会造成瓶颈，即使其大小更接近L2缓存。不将SRAM实现为缓存可能是特斯拉保持低延迟的方式。如果我们回顾很久以前的AMD锤子架构，我们可以看到，在检测到L1D丢失后，访问1 MB L2需要8个周期（对于12个总延迟周期）。如果直接解决1 MB的SRAM而不是用作缓存，则可能会删除三个阶段，使延迟减少到5个周期：

考虑到几十年来的过程节点改进，甚至比Athlon更低的时钟速度目标，很容易看到特斯拉如何访问具有类似L1延迟的L2大小的SRAM块。当然，跳过一个级别的缓存可以节省面积和电力。

为了进一步降低延迟、面积和核心复杂性，Dojo没有虚拟内存支持。因此，它没有TLB或页面行走机制。现代操作系统利用虚拟内存为每个进程提供自己的内存视图。程序使用的内存地址不是直接访问物理内存地址，而是由CPU使用操作系统设置的分页结构将它转换为物理地址。这就是现代操作系统如何将程序彼此隔离，并防止一个行为不端的应用程序导致整个系统崩溃的方法。

虚拟内存也是你运行比物理内存更多的程序的方式。当你用完真实内存时，操作系统会解对映页面，将其写入磁盘，并向你的程序提供它所需的内存。当另一个糟糕的程序试图访问该内存时，CPU试图将虚拟地址转换为物理地址，但发现翻译不存在。CPU抛出页面故障异常，操作系统通过将被驱逐的页面读回物理内存并填写页表条目来处理该异常。

这些在道场上都不可能做到。核心的4向SMT功能更适合单个应用程序公开显式并行性，而不是提高多任务处理性能。例如，一个线程可以执行矢量计算，而另一个线程将数据从系统内存异步加载到SRAM（通过DMA）。

为了进一步简化设计，Dojo只需21个地址位即可寻址SRAM，这可以简化Dojo的AGU和寻址总线。这些权衡可能允许特斯拉以足够低的延迟访问该SRAM，以避免在其面前实现单独的L1数据缓存。

DRAM访问

说到系统内存，Dojo芯片不会直接连接到内存。相反，它们连接到配备HBM的接口处理器。这些接口处理器还负责与主机系统通信。

带有25个独立芯片的Dojo瓷砖可以访问160 GB的HBM内存。特斯拉表示，他们可以跨越瓷砖边界从每个模具边缘传输900 GB/s，这意味着可以使用4.5 TB/s的链路带宽访问接口处理器及其HBM。由于访问HBM需要通过单独的芯片，因此访问延迟可能非常高。

Dojo它很小。你觉得如何？

Dojo是一个8宽的核心，具有至少一些OoO执行能力、体面的矢量吞吐量和矩阵乘法单元。但即使有1.25 MB的本地SRAM，它最终也是一个非常小的核心。相比之下，富士通的A64FX占据了同一工艺节点的两倍多区域。

处理器设计就是做出正确的权衡。特斯拉希望通过将大量核心包装到模具上来最大限度地提高机器学习的吞吐量，因此单个内核必须很小。为了实现其区域效率，道场使用了一些熟悉的技术。它以保守的2 GHz运行。较低的时钟电路往往占用更少的面积。它可能有一个基本的分支预测器和一个小的指令缓存。如果程序具有较大的代码足迹或许多分支，则会牺牲一些性能。

但特斯拉还通过削减运行内部工作负载不需要的功能，在尽量减少电力和面积使用方面也走得更远。他们不进行数据端缓存，不支持虚拟内存，也不支持精确的异常。

其结果是一个处理器核心，提供了现代CPU核心的性能灵活性，同时在许多方面不如英特尔8086对用户和程序员友好。除了核心之外，特斯拉还通过专门为大规模部署设计Dojo芯片来节省模具区域。

物理实现

缩小后，道场内核在一个非常大的645毫米2模具上实现，称为D1。与我们熟悉的其他芯片不同，单个道场模具是不自给自足的。它没有DDR或PCIe控制器。模具边缘周围有IO接口，允许模具与相邻的模具通信，延迟约为100ns。

要获得系统内存，Dojo D1芯片必须与具有板载HBM的接口处理器通信。该接口处理器反过来通过PCIe连接到主机系统（接口处理器安装在PCIe卡上）。从理论上讲，最小的功能Dojo部署将涉及Dojo模具、接口处理器卡和主机系统。但特斯拉在包含25个模具的模块中部署了道场模具，以给人一种规模感。Dojo D1模具是专门为超级计算机设计的积木，仅此而已。

这种专业化可以节省更多的模具面积。Dojo D1不会将区域用于DDR和PCIe控制器。除了外部设计用于与相邻模具接口的定制IO连接器外，大多数模具都被道场岩心占据。

相比之下，设计时考虑到了更灵活的部署灵活性的芯片在IO上花费了大量区域。AMD的Zen 1“Zeppelin”模具是一个很棒的展示。Zeppelin可以直接连接到DDR4内存、PCIe设备、SATA驱动器和USB设备——非常适合满足客户要求。在服务器中，IFOP接口允许它与相邻的模具通信。放置在IFOP附近的大块SRAM可能是窥探过滤器，有助于在高核心计数设置中高效保持缓存一致性。Dojo不会试图保持跨核心的缓存一致性，也不使用任何SRAM进行窥探过滤器。

AMD通过将大约44%的“Zeppelin”模具面积花在内核和缓存以外的逻辑上来实现这种灵活性。Dojo只将28.9%的模具面积用于SRAM和内核以外的其他东西。

最后

特斯拉Dojo超级计算机背后的微架构展示了如何实现非常高的计算密度，同时仍然保持CPU使用分支代码良好运行的能力。为了实现这一点，你放弃了定义我们现代计算体验的大部分舒适感。如果你可以假设围绕Dojo核心构建桌面，那么任何使用过MS-DOS的人可能会熟悉这种体验。你无法同时运行多个应用程序。一个行为不端的应用程序可能会迫使你重新启动系统。如果你没有足够的RAM来运行某个程序，你可以完全忘记运行它（无需分页到磁盘）。

但这些权衡在特斯拉的超级计算机中是完全有意义的。特斯拉不需要Dojo内核来同时处理几个正在运行的应用程序，Dojo只需要运行内部的可信代码。因此，特斯拉并不关心虚拟内存支持。同样，在Dojo上运行的机器学习程序在编写时将考虑到该特定系统。你不会有一批任意程序，可能要求比可用内存更多。这意味着你不需要精确的异常（和虚拟内存）来允许过度提交内存、内存压缩或将内存页面切换到磁盘和从磁盘中切换等技术。精确的异常对调试也很有用，但特斯拉通过单独的调试模式使调试变得更便宜。这样，在编写和调试代码后，Dojo可以专注于运行它，而无需进行按顺序提交指令结果所需的簿记。

当然，特斯拉为获得高计算密度而采用的权衡在消费者或服务器CPU中是不可能的。但它们在行动中非常有趣，我们必须感谢特斯拉抽出时间在Hot Chips上展示。在过去的二十年里，流程节点的改进一直在放缓，导致单线程性能的进步放缓。在过去的五年里，电源和冷却限制一直在减缓多线程性能。但对更多计算能力的胃口并没有放缓，因此公司正在转向更专业的硬件来跟上步伐。特斯拉Dojo超级计算机的架构是一个很好的例子，说明如何进行权衡来提高计算密度，以及当前的趋势如何有利于为吞吐量绑定应用程序引入专用硬件。

了解更多前沿科技，关注我就是你最好的选择！