如果我们说特斯拉只对机器学习感兴趣是轻描淡写的。事实上这家电动汽车制造商建造了一台名为 Dojo 的内部超级计算机,针对训练其机器学习模型进行了优化。

与许多其他超级计算机不同,Dojo 没有使用现成的 CPU 和 GPU,例如 AMD、Intel 或 Nvidia 的。特斯拉根据他们的需求设计了自己的微架构,让他们做出更通用架构无法做出的权衡。

在本文中,我们将根据特斯拉在 Hot Chips 上的演示来了解该架构。该架构没有单独的名称,因此为简单起见,每当我们在后面提到 Dojo 时,我们都是在谈论架构。

框图

从高层次上看,Dojo 是一个 8 宽内核,具有四路 SMT,运行在保守的 2 GHz,具有 CPU 风格的pipeline,使其比 GPU 更能容忍不同的算法和分支代码。Dojo 的指令集在标量方面类似于 RISC-V,但 Tesla 的工程师添加了一组自定义向量指令,专注于加速机器学习。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(1)

特斯拉 Dojo 核心的简化框图

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

像Cell的 SPE?

在 2000 年代中期推出的IBM Cell 处理器具有八个“协同处理元件”(Synergistic Processing Elements)或 SPE,由一个功能齐全的 CPU 内核(“电源处理元件”或 PPE:Power Processing Elemen)控制。乍一看,Dojo 与 SPE 有很多相似之处。

Dojo 和 SPE 都针对矢量处理进行了优化,并且依赖于单独的主机处理器进行工作分配。在 Dojo 或 SPE 上运行的代码不能直接访问系统内存。相反,应用程序预计主要在一小部分本地 SRAM 中工作。此本地 SRAM 由软件管理,不能用作缓存。如果需要来自主存储器的数据,则必须使用 DMA 操作将其引入。

最后,Dojo 和 Cell 的 SPE 都缺乏对虚拟内存的支持。稍后我们将详细介绍这意味着什么,但简而言之,它使多任务处理变得非常困难。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(2)

Tesla Dojo 和 IBM Cell SPE 之间的非常简化的比较

Dojo 在几个重要方面有所不同。由于 Dojo 在设计时并未考虑小规模部署,因此主机处理器驻留在单独的主机系统上。这些主机系统具有带接口处理器的 PCIe 卡,然后通过高速网络链路连接到 Dojo 芯片。相比之下,Cell 的主处理器驻留在同一个芯片上。这使得单独部署单个 Cell 芯片成为可能——这在 Dojo 中是不可能的。Dojo 的 1.25 MB 本地块是 SRAM,它比只有 256 KB SRAM 的 Cell SPE 大得多,带宽也更高。Cell 的 256 KB SRAM 只有一个端口,每个周期能够提供 128B。Dojo 的 SRAM 有五个 64B 端口。当然,架构目标是非常不同的。Dojo 是宽时钟和低时钟,而 Cell SPE 具有专为高时钟设计的窄而深的pipelines。

前端:CPU Comforts等

让我们从前端开始简要介绍一下 Dojo 的pipelines。有某种分支预测器,因为特斯拉的图表显示了 BTB(branch target buffer:分支目标缓冲区)。它的预测能力可能无法达到我们在 AMD、ARM 和 Intel 的高性能内核上看到的水平,因为 Dojo 需要优先考虑在矢量执行上花费裸片区域。但即使是一个基本的分支预测器也比没有预测器有了很大的进步,Dojo 的分支预测能力在处理分支代码或更大的指令占用空间时应该比 GPU 提供更好的性能。

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

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(3)

来自特斯拉的 Hot Chips 演示文稿

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

在解码时,某些指令(如branches, predicated operations和immediate loads (“list parsing”) )可以在前端执行并从pipeline中删除。这有点像更新的 x86 CPU 在重命名器(renamer)中消除了寄存器到寄存器的副本。但是您没听错——Dojo 不会通过pipeline跟踪“已消除”的指令以维持in-order retirement。其他处理器跟踪所有内容以退出,以便它们可以在任何指令边界处停止,并保持恢复执行所需的所有状态。这种能力被称为“precise exceptions”,现代操作系统使用它来提供各种好东西,比如paging to disk,或者告诉你你的代码到底在哪里搞砸了。

特斯拉并不关心precise exceptions。Dojo 确实具有调试模式,其中更多指令通过pipeline以提供“更精确”的exceptions,但没有像正常无序 CPU 那样的有序退出逻辑。

Dojo 的执行引擎

在看到一个宽前端之后,看到一个只有两个 ALU 和两个 AGU 的 4 宽整数执行引擎可能会有点奇怪。但是这个漏斗形的流水线是有意义的,因为一些指令是在前端执行和丢弃的。

Dojo 也不会进入客户端系统,其中标量整数性能很重要。因此,整数端提供了足够的吞吐量来处理控制流和地址生成,以保持向量和矩阵单元的馈送。

Dojo 的向量和矩阵执行端放置在核心pipeline中的标量执行引擎之后,并且有两个执行pipeline。两个pipeline可能听起来不多,但 Dojo 在这些pipeline后面有非常宽的执行单元。一个pipeline可以执行 512 位向量执行,而另一个pipeline执行 8x8x4 矩阵乘法。因此,只要指令公开足够的显式并行性,Dojo 就能够实现非常高的吞吐量——尤其是在使用矩阵单元的情况下。Tesla 声称具有 354 个 Dojo 内核的芯片在 2 GHz 时可以达到 362 BF16 TFLOPS,这表明每个内核每个周期可以执行 512 BF16 FLOPS。

我们不确定 Dojo 是否可以完全乱序执行。但特斯拉确实表示,整数端可以在向量端之前运行很远,这表明它可以执行过去停滞的指令,直到其中一个调度程序被填满。缺乏按序引退也指向无序执行能力。

通常,实现无序执行会带来很多复杂性。这是因为 CPU 必须按顺序执行指令。AMD、ARM 和 Intel 的高性能 CPU 使用大型重新排序缓冲区(和其他结构)跟踪指令,以便它们的结果可以按程序顺序提交。这意味着如果一个程序做了一些愚蠢的事情,比如除以零,这些内核可以准确地显示出哪条指令出错了。而且,它们可以显示 CPU 状态,该状态反映了故障之前所有指令的动作,但之后没有。这意味着您可以修复导致指令错误的任何原因,并恢复执行。Dojo 放弃了这种能力。作为交换,Dojo 可以避免与通过其pipeline跟踪每条指令相关的功率和面积开销,以确保按程序顺序提交结果。

SRAM 访问

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

因为 Dojo 的本地 SRAM 块不是缓存,所以它不需要与数据一起存储的标记和状态位。SRAM 前面也没有 L1D 缓存,因此它必须足够快以处理所有加载和存储指令而不会造成瓶颈,即使它的大小更接近 L2 缓存。不将 SRAM 实现为缓存可能是特斯拉保持低延迟的方式。如果我们回顾很久以前的 AMD 的 Hammer 架构,我们可以看到在检测到 L1D 未命中后访问 1 MB L2 需要 8 个周期(总延迟为 12 个周期)。如果直接寻址 1 MB 的 SRAM 而不是用作缓存,则可能会删除三个阶段,从而将延迟降低到 5 个周期:

来自较早的 Hot Chips 演示文稿的幻灯片,显示了 L2 缓存访问所涉及的pipeline阶段。如果 L2 不是缓存,则可以跳过的阶段标记为红色。

考虑到几十年的工艺节点改进,甚至比 Athlon 更低的时钟速度目标,很容易看出特斯拉如何以类似 L1 的延迟访问 L2 大小的 SRAM 块。跳过一级缓存当然可以节省面积和功耗。

为了进一步减少延迟、面积和核心复杂性,Dojo 没有虚拟内存支持。因此,它没有 TLB 或页面遍历机制(page walk mechanisms)。现代操作系统利用虚拟内存为每个进程提供自己的内存视图。程序使用的内存地址不是直接访问物理内存地址,而是由 CPU 使用操作系统设置的分页结构转换为物理地址。这就是现代操作系统如何使程序相互隔离,并防止一个行为不端的应用程序导致整个系统瘫痪。

虚拟内存也是您可以运行比物理内存更多的程序的方式。当您用完实际内存时,操作系统会取消映射页面,将其写入磁盘,并为您的程序提供所需的内存。当其他糟糕的程序尝试访问该内存时,CPU 会尝试将虚拟地址转换为物理地址,但发现转换不存在。CPU 抛出一个页面错误异常,操作系统通过将被驱逐的页面读回物理内存并填写页表条目来处理该异常。

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

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

内存存取

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

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(4)

具有 25 个独立芯片的 Dojo tile 可以访问 160 GB 的 HBM 内存。

Tesla 表示,他们可以跨 tile 边界从每个芯片边缘传输 900 GB/s,这意味着可以以 4.5 TB/s 的链路带宽访问接口处理器及其 HBM。因为访问 HBM 需要通过单独的芯片,所以访问延迟可能非常高。

是小,如何做到的?

Dojo 是一个 8 宽的内核,具有至少一些 OoO 执行能力、不错的向量吞吐量和一个矩阵乘法单元。但即使有 1.25 MB 的本地 SRAM,它最终还是一个非常小的内核。相比之下,富士通的A64FX在同一工艺节点上占据的面积是其两倍以上。

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

但特斯拉还通过削减运行内部工作负载所不需要的功能来进一步减少功耗和面积使用。它们不进行数据端缓存,不支持虚拟内存,也不支持precise exceptions。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(5)

结果是一个处理器内核提供了现代 CPU 内核的性能灵活性,同时在许多方面不如 Intel 8086 对用户和程序员友好。除了内核之外,Tesla 还通过设计专门用于部署的 Dojo 芯片来节省芯片面积大规模。

物理实现

缩小,Dojo 核心是在一个非常大的 645 平方毫米的裸片上实现的,称为 D1。与我们熟悉的其他芯片不同,单个 Dojo 芯片不能自给自足。它没有 DDR 或 PCIe 控制器。die 边缘周围有 IO 接口,可以让 die 与相邻的 die 进行通信,延迟约为 100 ns。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(6)

要访问系统内存,Dojo D1 芯片必须与具有板载 HBM 的接口处理器通信。该接口处理器又通过 PCIe 连接到主机系统(接口处理器安装在 PCIe 卡上)。理论上,最小的功能性 Dojo 部署将涉及一个 Dojo 芯片、一个接口处理器卡和一个主机系统。但特斯拉将 Dojo die部署在每个包含 25 个die的模块中,以提供规模感。Dojo D1 die专门设计用于超级计算机的构建块,仅此而已。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(7)

这种专业化可以节省更多的裸片面积。Dojo D1 不会在 DDR 和 PCIe 控制器上花费空间。大部分裸片都被大量 Dojo 核心占据,除了外部设计用于与相邻裸片接口的定制 IO 连接器。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(8)

相比之下,设计时考虑到更多部署灵活性的芯片在 IO 上花费了大量空间。AMD 的 Zen 1 “Zeppelin” 芯片很好地展示了这一点。Zeppelin 可以直接连接到 DDR4 内存、PCIe 设备、SATA 驱动器和 USB 设备——非常适合满足客户要求。在服务器中,IFOP 接口让它与相邻的芯片进行通信。位于 IFOP 附近的大块 SRAM 可能是窥探过滤器(snoop filters),有助于在高核心数设置中有效地保持缓存一致性。Dojo 不会尝试跨内核保持缓存一致性,并且不会将任何 SRAM 用于窥探过滤。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(9)

来自 AMD 在 ISSCC 2021 上的小芯片演示

AMD 通过将大约 44% 的“Zeppelin”裸片区域用于除内核和缓存之外的逻辑来为这种灵活性付出代价。Dojo 仅将 28.9% 的裸片面积用于 SRAM 和内核以外的东西。

最后的话

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

但这些权衡在特斯拉的超级计算机中非常有意义。Tesla 不需要 Dojo 内核来同时处理多个正在运行的应用程序,Dojo 只需运行内部的可信代码。所以,特斯拉并不关心虚拟内存支持。同样,在 Dojo 上运行的机器学习程序将在编写时考虑到该特定系统。您不会有一批任意程序可能要求比可用内存更多的内存。这意味着您不需要precise exceptions(和虚拟内存)来允许诸如过度使用内存、内存压缩或在磁盘之间交换内存页面等技术。precise exceptions对于调试也很有用,但 Tesla 通过单独的调试模式以更便宜的方式进行调试。

特斯拉芯片意义:特斯拉自研芯片架构带来的启示(10)

可以肯定的是,特斯拉为获得高计算密度而采用的权衡取舍在消费者或服务器 CPU 中是不可能的。但他们在行动中非常有趣,我们必须感谢特斯拉花时间在 Hot Chips 上展示。

在过去的二十年里,进程节点的改进一直在放缓,导致单线程性能的进步更慢。在过去的五年中,电源和冷却限制一直在降低多线程性能。但对更多计算能力的需求并没有放缓,因此公司正在转向更专业的硬件以跟上步伐。

Tesla 的 Dojo 超级计算机中的架构是一个很好的例子,说明了如何进行权衡以提高计算密度,以及当前趋势如何有利于为吞吐量受限的应用程序引入专用硬件。

,