如何设计一颗AI芯片？AI芯片设计入门方案

1 介绍 

机器学习（ML）已经在在线活动中变得无处不在。近年来，这些模型的规模和复杂性大幅增长，这有助于提高预测的准确性和有效性。然而，与此同时，这种增长给用于大规模训练和推理这些模型的硬件平台带来了巨大挑战。总拥有成本（TCO）是在数据中心将模型投入生产的主要制约因素之一，而功率是这些平台TCO的重要组成部分。因此，单位TCO的性能（以及每瓦的性能）已成为针对机器学习的所有硬件平台的重要衡量标准。

深度学习推荐模型（DLRM）已成为Meta数据中心最主要的工作负载之一。这些模型将计算密集型的传统多层感知器（MLP）操作（有时称为全连接或FC）与将稀疏特征转换为密集表示的嵌入表相结合。这些表包含随机索引的宽向量，并将其简化为单个向量，然后将其与来自其他层的数据组合以产生最终结果。虽然嵌入表操作的计算要求相对较低，但由于数据访问模式的性质和表的大小，它们的内存占用和带宽要求依然相对较高。

图1显示了Meta生产数据中心中与推荐模型相关的推理工作负载的复杂性和内存占用的历史增长和预计未来增长。虚线显示了模型计算需求的估计增长，而实线显示了内存占用的增加。灰色实线捕捉了用于存储嵌入表的设备内存的占用空间，这是这些模型的重要组成部分。计算和内存需求的增长水平是一个需要解决的问题，尤其需要考虑到这些工作负载通常是如何在数据中心运行。

图1：推理模型的增长趋势

2 动机 

传统上，CPU被用作Meta生产数据中心中服务推理工作负载的主要工具，但它们在满足最新工作负载的需求方面并不具有成本效益。在某种程度上，硬件加速被认为是一种有吸引力的解决方案，它可以解决功率和性能问题，并提供一种更有效的方式来服务推理请求，同时为运行未来的模型提供足够的计算性能余量。

图2显示了过去几年中为数据中心内的推理工作负载部署的服务器的估计数量。浅色实线显示基于CPU的服务器数量，虚线显示配备第一代推理加速器Intel NNPI的服务器数量；深色实线显示基于GPU的服务器数量。虽然使用NNPI加速器暂时满足了对增加容量的最初需求，但对推理模型的要求很快超过了NNPI能力，并为使用GPU提供了动力。这带来了利用已经用于训练的现有生态系统的额外优势。因此，可以观察到，GPU作为加速器越来越多地满足了对模型复杂性的日益增长的需求。

虽然最近几代GPU提供了大量的内存带宽和计算能力，但它们在设计时没有考虑到推理，因此处理实际推理工作负载的效率很低。开发人员使用了无数的软件技术，如算子融合、图转换和内核优化，以提高GPU的效率。但是，尽管做出了这些努力，仍然存在效率差距，这使得在实践中部署模型具有挑战性且成本高昂。

图2：服务器对推理工作负载的需求增长

基于部署NNPI和GPU作为加速器的经验，很明显对于重要的推理工作负载，还有更优化的解决方案的空间。这种最佳解决方案基于内部加速器，该加速器从头开始构建，以满足推理工作负载的苛刻要求，尤其侧重于满足DLRM系统的性能要求。然而，在关注DLRM工作负载的同时（考虑到它们正在进行的演变，以及该架构是为这些工作负载的下一代有效构建的情况），很明显，除了性能之外，该架构还应该提供足够的通用性和可编程性，以支持这些工作负载的未来版本以及潜在的其他类型的神经网络模型。

虽然创建定制的硅解决方案为目标工作负载的充分创新和专业化打开了大门，但为数据中心的大规模部署创建加速器架构是一项艰巨的任务。因此，构建加速器时的重点和策略一直是采用和重用供应商和开源社区提供的技术以及工具和环境。这不仅缩短了上市时间，而且还利用了来自社区和供应商的支持和增强功能，减少了构建、启用和部署此类平台所需的资源量。

本文的其余部分解释了Meta的第一个针对推理工作负载的加速器芯片MTIA的架构，以及随之而来的学习。下一节详细介绍了加速器的架构及其提供的各种功能和组件。第4节介绍了将一个示例算子映射到此体系结构，展示了如何利用各种提供的功能来高效地运行。第5节概述了加速器的软件堆栈，第6节描述了我们的评估方法和结果。最后，第7节讨论了在这个开发周期中吸取的一些重要经验教训。

3加速器架构

图3显示了加速器的高级架构，它被组织为连接在网格上的处理元素（PE）阵列。网格通过每侧的交叉开关连接到一组片上存储器块和片外存储控制器。有一个单独的控制子系统，带有专用处理器和外设，用于运行系统的控制软件。包含PCIe接口、相关DMA引擎和安全引导处理器的主机接口单元也位于该控制子系统旁边。

图3：加速器架构

图4显示了PE的内部组织。PE由两个RISC-V处理器核心和相关外设（左）、以及几个专门执行特定计算或数据移动的固定功能单元（右）组成。此外，每个PE都有128KB的本地存储空间。本地互连建立了处理器、外设和自定义硬件块之间的连接。

图4：PE的内部组织

3.1 固定功能单元

每个PE总共有五个固定功能块和一个命令处理器，该命令处理器协调这些固定功能块上的操作执行。功能单元在PE中形成了一个粗粒度的管道，数据可以从一个单元传递到下一个单元，以执行连续的操作。每个功能单元还可以直接访问PE本地存储器内的数据，执行必要的操作，并将结果写回，而无需将数据传递给其他功能单元。

3.1.1 内存布局单元（MLU）

该功能块执行、复制和更改本地存储器中的数据布局有关的操作。它可以对4/8/16/32位数据类型的张量进行操作。诸如转置、串联或数据整形之类的操作都是使用此功能块执行的。输出数据可以直接发送到下一个功能块以立即进行操作，也可以存储在PE的存储器中。例如，MLU可以转置矩阵，并将输出直接提供给DPE块用于矩阵乘法运算，或者它可以将数据正确格式化为深度卷积运算的一部分，并将其发送给DPE以执行实际计算。

3.1.2 点积引擎（DPE）

该功能块对两个输入张量执行一组点积运算。首先读取第一张量并将其存储在DPE中，然后将第二张量流式传输，并对第一张量的所有行执行点积运算。DPE每个周期可以执行1024次INT8乘法（32×32）或512次FP16/BF16乘法（32×6）。操作完全流水线化；执行两个最大矩阵的乘法需要32个时钟周期。在INT8乘法的情况下，结果输出以INT32格式存储，而在BF16或FP16乘法的情况中，结果以FP32格式存储。结果总是被发送到管道中的下一个功能单元进行存储和累积。

3.1.3 缩减引擎 (RE)

RE托管存储元件，这些存储元件跟踪矩阵乘法运算的结果并在多个运算中累积它们。有四个独立的存储库，可以独立地用于存储和累积来自DPE的结果。RE可以将初始偏置加载到这些累加器中，并且还可以通过专用缩减网络将它们的内容发送到相邻PE（在本节稍后讨论）。在通过缩减网络接收结果时，RE将接收到的值累积在本地存储体之一中的值之上。然后，它可以将结果发送给相邻的功能块或SE，或者将其直接存储在PE的本地存储器中。

3.1.4 SIMD 引擎 (SE)

该功能块执行诸如量化/去量化和非线性函数之类的操作。在内部，该块包含一组查找表和浮点运算单元，用于计算非线性函数的线性或三次近似，如指数、S形、tanh等。近似接受INT8或FP16数据类型作为输入，在输出处产生INT8或FP32结果。该单元可以直接从RE块接收其输入，或者从本地存储器读取它们。此外，该块还能够使用其浮点ALU执行一组预定义的元素运算，如加法、乘法、累加等。

3.1.5 结构接口（FI）

这个功能块充当进出PE的网关。它连接到加速器的片上网络并通过该网络进行通信。它制定并向片上和片外存储器以及系统寄存器发送存储器访问请求，并接收回数据或写入完成。它实现了一组类似DMA的操作，在PE的本地内存中传输数据。它还接收和传输来自处理器内核的缓存未命中和未缓存访问，并允许其他实体（其他PE或控制子系统）访问PE的内部资源。

3.1.6 命令处理器（CP）

除了托管PE的本地存储器和寄存器外，CP块还充当中央处理单元，同时协调固定功能块上各种操作的执行。它从PE中的两个处理器核心接收指令，对这些指令执行依赖性检查、调度和跟踪，并将它们分派到固定功能单元执行。它包含两个独立的调度器（每个处理器内核一个）、一组命令队列以及用于访问本地内存和寄存器资源的仲裁逻辑。

硬件提供一组原子基元，以允许内核之间（在PE内或在多个PE之间）的同步。这些基元由处理器制定，允许对预定义寄存器进行原子更新，并且可以暂停处理器，直到外部满足某些条件（例如，计数器达到某个值）。在更高的级别上，这些机制用于有效地实现软件结构，如锁、互斥锁和屏障等。执行原子操作的逻辑以及相关寄存器位于命令处理器内，并通过自定义接口与处理器内核紧密集成。

3.2 处理器内核

每个PE包含两个RISC-V内核，它们运行应用程序的代码并向CP发出命令，用于将各种计算加载到固定功能单元。内核按序执行指令，且单周期只执行一个指令，具有五级流水线（AX25-V100，来自Andes Technology），并经过大量定制以适应所需的功能。该组自定义包括自定义接口、自定义寄存器、自定义指令和自定义异常。自定义接口将内核连接到CP，以向固定功能单元发出命令，并在内核和本地存储器之间来回移动数据。自定义寄存器存储在发出命令时发送给CP的命令信息，添加自定义指令以启动每个固定功能单元上的所需操作。最后，自定义异常确保了向CP发出的每个命令的正确性，并在命令中出现非法值时引发异常。

其中一个处理器内核配备了RISC-V矢量扩展，这为PE增加了额外的灵活性，并允许实现无法很好地映射到现有固定功能单元的操作。矢量处理单元包含32个矢量寄存器，每个寄存器的宽度为64B，所有矢量功能单元的宽度相同。它实现了RISC-V矢量扩展的0.8.1版本。

3.3 本地存储 (LS)

每个PE总共有128KB的本地内存供处理器和功能单元使用。CP实现存储器组的仲裁方案，并协调来自内核和固定功能单元的访问。本地存储器被映射到系统的地址空间，并且可以由内核通过常规加载/存储指令来访问。

在本地存储器之上引入了一个抽象层，以简化它们的操作之间的使用和依赖性检查。每个PE可以定义映射到现有本地存储器的循环缓冲区（CB）。每个CB都指定有一个ID，并有一对寄存器，用于指定其大小（深度）和本地存储器中的起始地址。此外，每个CB还实现一组读写指针，以实现硬件FIFO。

在CB中，读取操作总是从读取指针开始读取数据，而写入操作总是从写入指针开始写入数据。读写操作携带一个偏移量，这允许它们访问缓冲区当前头或尾以外的位置（图5）。固定功能单元使用CB ID作为其输入/输出操作数；例如，矩阵乘法运算使用两个CB作为其输入操作数。在允许操作开始之前，命令处理器检查输入CB中数据的可用性和输出CB中的空间。它只允许在必要的元素和空间检查通过时启动操作。因此，可以保证操作具有完成所需的资源，并且不会在执行过程中暂停功能单元。

图5：从环形缓冲区读取

命令处理器还使用CB ID来强制执行不同自定义指令之间的依赖性检查和互锁。它确保访问和修改特定CB的操作始终按程序顺序执行，而在不同CB或同一CB的不同区域上的操作可以并行执行。与使用绝对本地存储器地址来实施这种互锁相比，这显著简化了依赖性检查。

CB还简化了不同操作之间生产者-消费者执行模型的实现。这些操作可以由不同的内核或不同的固定功能单元发起。例如，程序可以向硬件发出一系列DMA操作（将数据从外部存储器移动到CB），然后进行一组使用该数据的自定义计算操作（例如MATMUL），而不需要两者之间的显式同步。MATMUL指令由命令处理器自动停止，直到之前的DMA操作将足够的数据带入CB，然后立即启动，从而使程序不再明确检查数据的可用性。

虽然一些指令（如DMA操作）会自动调整读和写指针（因为它们会将数据移入和移出CB，从而产生或消耗元素），但其他自定义指令不会移动指针。这允许CB内部的数据在被明确标记为已消耗之前被不同的操作多次重用。硬件提供了额外的自定义指令，可以调整每个CB中的读指针和写指针，允许在必要时将数据元素明确标记为生成或消耗的数据元素。

3.4 存储器子系统和互连

除了PE内的本地存储器外，加速器还具有128MB的片上SRAM，这些SRAM被组织为网格周围的切片。这种片上存储器可以用作可寻址的暂存存储器，也可以用作通用、共享的存储器侧缓存。网格的每一侧都有四个LPDDR5控制器，总共提供176 GB/s（理论上）的片外带宽。该加速器总共可以支持128GB的片外存储器容量。存储器地址分布在这些控制器之间，以及片上SRAM片之间。当片上SRAM被配置为高速缓存时，每四个高速缓存片都与单个存储器控制器相关联，并高速缓存其地址。

将所有PE和存储器连接在一起的片上网络基于具有特殊增强功能的AXI互连。互连由两个网络组成，用于分别承载存储器和寄存器访问。存储器访问网络配备了多播功能，允许将来自多个PE的请求合并为一个PE（如果它们被发送到同一组地址）。然后向存储器块发送单个请求以检索数据并将其返回给所有请求PE。然而，多播仅支持位于网格中同一行或列的PE，不能用于任意一组PE。

除了基于AXI的主要互连外，PE还通过一个称为缩减网络的专用网络相互连接。这是一个单向网络，只从北向南和从西向东传播。它将一个PE的RE块中的累加器的部分和传送到另一个PE。使用该网络，PE可以方便地累积其计算结果，而不必将其保存和恢复在存储器中。在行或列中的最后一个PE可以在累积所有部分值之后将最终结果存储在存储器中。

3.5 并行性与数据复用

在任何深度学习加速器中，并行性、局部性和数据复用在有效利用有限的硬件资源方面都发挥着重要作用。MTIA体系结构提供了一组功能，以允许在神经网络模型和算子中实现多个并行度和最大限度地利用时间和空间数据复用，如下所述。

并行性：该体系结构支持多种级别的并行性和各种操作的重叠。数据级并行（DLP）是通过在固定功能单元中使用宽矢量以及矢量处理器来实现的。多个PE也可以以数据并行的方式对同一任务进行操作。命令处理器利用指令级并行性，允许不同的固定功能块同时处理多个未完成的操作。内存级并行（MLP）是通过允许每个PE对片上和片外存储器发出许多未完成的请求来实现的。最后，线程级并行（TLP）可以通过利用多个PE（或PE组）运行并行线程以及在每个PE内有两个独立的线程来实现。PE内的线程可以协作执行给定的任务，通过一个线程编排数据移动，而另一个线程则编排计算。

缓存：在硬件的各个功能块中有多个级别的缓存，以提高局部性并减少内存带宽消耗。这包括处理器核心中的指令和数据高速缓存、大型片上末级高速缓存以及DPE块中输入操作数的高速缓存。DPE级别的缓存允许引擎保存操作数A和操作数B的数据，并在命中时保存对本地内存的访问。

循环缓冲区/本地存储器：当PE执行计算时，循环缓冲区为保存输入操作数提供存储。调整指针以及偏移到循环缓冲区内任何位置的灵活性来允许程序在决定将每行数据标记为已消耗之前多次访问每行数据。

专用缩减：拥有专用缩减网络不仅可以从系统的主要片上网络中避免大部分数据传输，而且还提供了一种将PE分组在一起并以聚合形式使用其本地存储器的方法。这又允许在PE中存储更大部分的输入操作数，并减少从片外存储器加载它们的带宽要求。此外，DPE块利用归约树（空间和）来计算乘法运算的输出，这是已知的更节能的方法。

多播：如前所述，当多个PE访问内存中的同一组地址时，系统的NoC允许合并来自多个PE的请求。这减少了内存带宽，并通过允许共享数据来提高数据移动的能效，同时只从内存中读取一次数据并将其传递给所有请求者.

图6显示了由片上SRAM和片外DDR控制器包围的PE网格的芯片示意图，而表I列出了芯片特征和参数的摘要。

图6:MTIA die示意图

表I MTIA特征和参数概述。

4映射FC层

为了演示上述所有特征是如何协同工作的，让我们考虑一个FC算子，它以CT=a×BT的形式执行矩阵乘法运算，并看看它是如何映射到PE的子网格的。以转置方式执行操作的原因是保持k作为两个张量的内部维度，以提高存储器访问的效率。矩阵A假设为m×k，矩阵B假设为k×n（因此BT将为n×k），产生输出C，输出C将为m×n矩阵（或CT为n×m矩阵）。假设输入具有行主存储器布局。当内部维度（k）不是32B的倍数时，外部维度（m或n）步幅与32B边界对齐，以实现有效的数据移动。为了简单起见，我们将假设所有元素都是INT8数据类型。

如前所述，DPE对32（m）×32（k）×32个（n）输入的块进行操作，生成RE中累积的32（n）×32（m）个部分结果。此操作需要32个时钟周期。为了馈送DPE的流水线，必须在32个周期内将32（m）×32（k）个块矩阵A和32（n）×32（k）个块矩阵BT从外部存储器引入PE的本地存储器，需要64B/周期的带宽。为了缓解这种带宽压力，RE块中的四个累加器用于累加2×2个部分结果块，总共保持输出矩阵的64（n）×64（m）个元素。通过以这种方式使用累加器，每个32×32输入块我们使用两次，从而将外部带宽需求减少到32B/周期。

张量维度m、n和k分别以64、64和32的倍数分布在PE网格上。因此，每个PE以数据并行方式在较大结果矩阵的不同子块上工作。缩减维度（k）沿着行（或列）分布在多个PE上。这便于在乘法完成之后使用归约网络来累积部分结果。PE将计算出的部分结果相互传递，以累积并传递给下一个PE。当沿给定行或列的两个或多个PE使用来自任一输入矩阵的相同输入数据块时，芯片网络的多播功能用于合并来自多个PE的请求，并向内存发送单个请求，从而进一步降低内存带宽需求。

图7显示了在4×4 PE子网格上分布尺寸为512（m）、1024（k）和256（n）的FC算子的示例。缩减尺寸（k）分布在沿着同一行的两个PE上，尺寸m分布在四行上。列0和2中的PE以及列1和3中的PE参与矩阵A的行多播读取。类似地，每列中的所有PE都参与矩阵BT的列多播读取。

图7：将FC操作映射到子网格

在PE中，操作以生产者-消费者的方式在两个内核之间进行划分。图8显示了对应于PE中每个内核的伪代码。Core0发出一组DMA操作，将数据从主存储器移动到CB_A和CB_B，用于本地存储矩阵A和B。在并行线程中，Core1发出一组矩阵乘法（MML）指令，分别从CB_A和CB_B读取数据，并将结果存储在累加器寄存器中。可以观察到，每个数据块被使用两次，以在每个累加器寄存器中产生部分结果。如果操作是最后一次迭代，则通过发出POP指令将数据标记为在CB中消耗，否则相应的CB偏移量将增加，以在下一次迭代中移动到下一个数据块。最后，调用归约运算（REDUCE）来累积PE之间的所有部分和。缩减链中的最后一个PE使用DMA操作将数据发送回主存储器。

图8:PE中运行的FC运算的伪代码

PE中的两个内核必须在操作开始时同步，因为它们中只有一个执行必要的初始化任务（例如，设置要使用的CB）。但之后，没有显式的每次迭代同步；生产者-消费者同步由硬件负责：如果消费者（MML操作）试图使用没有足够数据的CB，硬件会暂停操作，直到生产者（DMA操作）在CB中放置足够的数据，此时它允许进行矩阵乘法。这种异步性使生产者和消费者线程解耦，并允许生产者继续前进，为以后的迭代引入更多数据。

5软件栈

MTIA的软件堆栈在设计时考虑到了两个主要目标：高效生产，这意味着实现比其他同类最佳解决方案更高的性能/TOC；同时，使用起来简单明了，甚至比现有的替代方案更简单。MTIA的软件堆栈是围绕PyTorch设计和构建的，以受益于其功能，并实现与生产环境中可用的ML的其他基础组件的无缝集成。本节的其余部分提供了软件堆栈的每个组件的概述，如图9所示。

图9：MTIA的软件栈

ML服务平台：在软件堆栈的顶部，我们有专用于ML模型服务平台（如图9所示的应用层）。这些服务平台在PyTorch之上运行，大多与硬件无关，支持在异构硬件系统上执行，包括CPU、GPU和MTIA等加速器。

PyTorch Runtime：为MTIA开发了PyTorch Run集成，提供了必要的功能和特性，包括MTIA张量、主机端内存分配器和类似CUDA的流式API，用于在设备上调度所需的操作。运行时支持不同的模型执行模式，包括热切模式，以及全图编译和执行，以最大限度地提高性能。它还支持将模型拆分为跨多个卡的分区，从而在它们之间提供必要的同步和通信通道。

编译器：软件堆栈中的下一个重要组件是一组编译器，它由多个部分组成：

一个基于PyTorch FX的ML模型编译器，它对表示为FX IR的PyTorch图进行了多次转换和模型级优化，并逐渐将其转换为LLVM IR。它负责利用PE网格和MTIA的内存子系统进行图优化。它实现了一个张量放置方案，该方案采用尽最大努力的方法将生产者-消费者数据保存在片上存储器中。它还可以将模型拆分为子图，这些子图旨在跨多张卡运行，甚至跨同一芯片内的子网格运行。

用于ML内核开发的基于DSL的编译器（代号KNYFE），它对ML内核进行了简短的高级描述，并生成了低级优化的C++代码。它使用底层硬件特定的API来实现ML算子，并被广泛用于开发MTIA中使用的许多ML内核。

基于LLVM的编译工具链将LLVM IR转换为外设的可执行文件。LLVM的使用主要是因为它提供了RISC-V支持，并负责最低级别的优化，如寄存器分配、排队和代码生成。大多数主要的优化，如PE之间的工作和数据的平铺或调度，都是由前面提到的高级编译器执行的。

ML内核库：另一个重要组件是内核库和ML算子库，用于构建在设备上执行的ML模型。这些内核中的许多都是使用前面提到的DSL编译器开发的，但一些对性能要求最高的内核，例如全连接（FC）层和嵌入包（EB）层，是由专家直接在低级C++中使用公开的内部函数开发的，以确保它们可以在硬件上实现最高级别的性能。

主机驱动程序和固件接口：MTIA平台软件使主机能够访问加速器。它管理设备的生命周期和资源，并帮助启动和跟踪设备上的运行时操作。堆栈的这一部分大致分为两部分：主机软件和固件。主机软件包括Linux驱动程序、用于提供统一设备接口的设备访问库、用于与PyTorch接口的流式API，以及用于管理和监控设备的软件工具和实用程序。

设备固件：设备固件包括基于ROM的预引导固件、在其自己的处理器上运行的安全引导固件、运行在执行运行时和管理操作的控制子系统上的控制核心处理器固件，以及最后运行在计算网格中的PE上的PE监视器，该监视器调度和监视在PE上运行的工作负载。主控制固件是基于Zephyr实时操作系统工作的。

6结论

我们通过将MTIA与基线加速器（NNPI）和最近部署的GPU进行比较来评估其性能。需要注意的是，我们报告了使用正在开发的软件堆栈收集的结果，因为我们认为这反映了端到端的性能，并且代表了生产环境。然而，这个堆栈目前并不像GPU的软件堆栈那样优化。因此，在某些情况下，GPU更高效，但我们希望随着时间的推移能够缩小这一差距，并让MTIA软件堆栈在所有DLRM工作负载空间中发挥架构的全部优势。我们评估了复杂性、规模和准确性方面不同的完整DLRM模型。由于这些加速器都基于不同的硬件平台，我们首先比较了它们的系统级硬件规范（表II）。这些平台如下：Yosemite V2服务器（带有六个NNPI加速器卡），Zion4S服务器（带有八个Nvidia A100 GPU），Yosemite V3服务器（带有十二个MTIA加速器卡）。

表II: 测试所用的推理硬件平台

虽然我们可以比较MTIA、NNPI和GPU的绝对性能，但每个设备在计算吞吐量、内存带宽和内存容量方面都有不同的能力，它们也在不同的电力预算下运行。因此，在我们的研究中，我们报告了性能/W（来代表性能/TCO，充分考虑到TCO的敏感性），因为功率是数据中心部署资源调配的一个重要因素。我们使用总平台功率除以加速器卡的数量来确定为每个加速器提供的功率，而不是使用卡的最大TDP。

6.1 基准测试表现

我们首先评估了几个重要的操作符和内核的性能，这些操作符和内核突破了体系结构的极限，是生产DLRM中主要组件的代表。表III显示了batch大小为64和256的代表性DLRM中请求的延迟细分。该模型有大约750层，其中近550层由EB算子组成。对于batch大小为64的情况，FC主导执行时间，其次是EB，而对于batch大小为256的情况，EB略微主导FC，两者合计占执行时间的62%。应该注意的是，对于较大的输入，内核能够更好地分摊设置成本，并更多地重用数据，从而实现硬件中固定功能单元的更高利用率。

表III-中等复杂度DLRM中各算子细分

基于细分，我们使用一组基准来评估MTIA硬件的效率。这些基准测试虽然不是全面的工作负载，但允许重要算子使用各种形状和大小，并揭示硬件中可能存在的潜在缺陷。GemmBench用于评估密集计算；它创建了一个由FC层链组成的模型。在我们的基准测试中，我们关注FP16和INT8（量化）数据，这需要额外的量化和去量化层。TBEBench用于评估稀疏计算，并允许我们配置TBE算子的batch大小、表数、每个表的行数、嵌入维度和池化因子。BatchGEMMBench、ConcatBench和TransposeBbench用于有效覆盖推荐模型中常见的其他重要算子。我们还评估了几个元素内核，包括量化、反量化和tanh。

密集计算：我们评估了INT8和FP16全连接（FC）层（图10和图11）。当精度足够时，INT8量化可以使FC吞吐量提高2倍。对于我们评估的一组形状，趋势线大致跟踪INT8和FP16上的MTIA和GPU，表明软件实现在一系列算术强度上得到了很好的优化。在许多情况下，MTIA实现了每瓦2倍或更高的性能，并且对于低batch特别有效，这有助于在严格的延迟要求下为请求提供服务。对于大batch，GPU能够随着工作量的增加而实现更高的利用率，因此MTIA的性能/W增益更低。请注意，当张量可以直接从SRAM流式传输时，MTIA是最有效的，这意味着图优化和管理数据局部性对于模型级别的良好性能非常重要。

图10:INT8 FC性能

图 11: FP16 FC 性能

稀疏计算：虽然一个典型的推荐模型可能包括数百个EmbeddingBag（EB）算子，但它们可以合并到一个或多个TableBatchedEmbedding（TBE）算子中，以分摊内核启动开销，并增加可以在设备上并行化的工作。图12显示了在一组具有代表性的算子的应用中，在MTIA和GPU上运行的TBE基准测试的性能（以GB/s/W为单位）。请注意，我们在这里以GB/s为单位报告性能，因为此基准测试主要是内存绑定的，而测量带宽（而不是查找/秒）可以更好地了解硬件利用率。在这里，我们利用片上SRAM的缓存配置来利用跨batch和batch内的局部性。在这些示例中，所有表条目都使用8位量化，图中显示的三元组描述了算子的池化因子、表中的行数和嵌入维度（每行元素）。MTIA实现了当前内核实现的GPU性能/W的0.6倍至1.5倍。

图 12: TBE 性能表现

考虑到软件堆栈的演变性质，我们观察到存在显著的改进空间：MTIA仅达到其内存带宽的10-20%，而GPU实现了其HBM带宽的60%左右。为了确保硬件没有缺陷，我们使用了为RTL验证而开发的手写内核，并且在SRAM中具有足够的位置时，可以观察到高达500 GB/s（超过最大限度的60%）或6 GB/s/W的性能水平。我们希望通过改进TBE内核的软件流水线和指令调度来弥补这一差距。

其他算子：虽然FC和TBE往往主导执行时间，但我们发现其他算子也同样重要，特别是考虑到优化前者所花费的精力。我们评估了BatchMatMul、Concat、Transpose和M=256、K=128、N=32的几个元素内核，张量数据放在SRAM和DRAM中（图13）。这些算子往往是内存绑定的，例如BatchMatMul和Tanh，它们分别达到SRAM带宽的90%和80%以上。当数据被放置在DRAM中时，由于更难隐藏额外的内存延迟，效率平均下降到40%左右。我们相信，实现数据布局优化、算子融合和最大限度地减少DRAM中的数据提取可能会缓解这一问题。

图13：其它算子的性能表现

6.2模型性能

我们测试了表IV中的五种不同的具有代表性DRLM的性能，它们的复杂性从低到高不等。MTIA可以运行与NNPI和GPU上运行的推荐模型相同的模型。以目前软件堆栈的成熟度水平，MTIA实现了与GPU接近的性能/W，并超过了NNPI的性能/W，而建模表明，随着软件堆栈的进一步成熟，还有很大的改进空间。

表IV: 用于评估的DLRM模型。

图14显示了上述一组DLRM的性能（单位为TFLOPS/s/W）。与NNPI相比，MTIA实现了1.6倍的效率，而与GPU相比，它达到了0.9倍的效率。在这些结果中有两个重要因素需要考虑：模型特性和实现中的软件优化水平。对于低复杂度模型，MTIA比GPU具有显著优势，因为这些模型由具有较小输入形状的FC层主导，并且MTIA非常有效地处理了这一点，例如，LC2显示出近3倍的改进。对于中等复杂度的模型，MTIA仍然比GPU更有效率增益，但它更低，因为FC不太占主导地位，并且GPU软件堆栈提供了其他算子的更高效实现（具有TBE和积极的算子融合）。对于高复杂度模型，我们看到GPU软件堆栈针对大型输入进行了更好的优化，MTIA需要类似的优化才能实现相同或更高级别的效率。这些初步结果使我们深入了解了软件堆栈中未来应该考虑关注的领域（例如，大型FC、TBE优化、算子融合等），并为我们接下来讨论的下一代架构提供了重要的参考。

图14：DLRM的性能

7讨论

构建芯片总是一个困难、漫长和耗时的过程，尤其是当初次尝试时。对于MTIA，所产生的芯片需要实现高性能，处理广泛的推荐模型，并提供一定程度的可编程性，以允许在生产中快速部署模型。本节重点介绍了我们对体系结构选择的重要观察和思考，以及它们如何影响软件堆栈、性能和开发人员效率。这些经验教训也为改进和加强未来几代架构提供了指导。

双核PE：选择在PE内有两个独立的处理器内核、并允许两者控制固定的功能单元，在线程级别提供了很大程度的并行性和灵活性，它能够允许计算与数据传输解耦。虽然这种解耦简化了编程，并减轻了特定算子受指令约束时的性能问题（通过提供两倍于总指令的吞吐量），但在软件中正确有效地使用这两个内核需要付出一些努力。在第一次运行之前，两个内核之间用于初始化和清理的同步等细节很难正确处理，但之后通过软件堆栈中的适当集成，在所有工作负载中都能够得到很好的利用。

通用计算：以RISC-V矢量支持的形式添加通用计算是一个合适的选择：有些算子是在架构定义阶段之后开发或被发现重要性的，因此架构不包括对它们的任何卸载支持。像LayerNorm和BatchedReduceAdd这样的算子很容易用向量实现，而且这些实现被证明优于使用标量核和固定函数单元的版本。

自动代码生成：关于如何在PE中集成和操作固定功能单元的一些架构选择使编译器的自动代码生成变得困难。处理器必须安装并发出明确的命令来操作任何固定的功能块。虽然这是通过向处理器添加自定义指令和寄存器来完成的，但它仍然需要组装许多参数并将其传递给每个目标引擎，以指定操作的细节。从程序中控制一组异构的固定函数单元并平衡它们之间的数据流对编译器来说是一个具有挑战性的工作。在各种输入形状和大小的固定功能块上实现期望的利用水平也是很困难的。虽然我们基于DSL的KNYFE编译器使编写内核变得更容易并能够自动处理其中许多问题，但它仍然需要学习一种新的DSL。

冲区：添加循环缓冲区抽象极大地简化了在同一内存区域上工作的自定义操作之间的依赖性检查，因为循环缓冲区ID被用作依赖性检查的单元（类似于处理器核中的寄存器ID）。它们还简化了固定功能块和处理器之间生产者-消费者关系的实现，因为硬件会推迟操作，直到循环缓冲区中有足够的数据（或空间）可用，而不需要在软件级别进行任何显式同步。灵活的寻址机制还允许对循环缓冲区内的任何位置进行任意访问，这简化了数据重用，因为不同的操作可以多次访问循环缓冲区中的不同段。然而，这需要软件明确管理缓冲区内的空间，并决定何时将数据标记为已消耗，如果执行不当，这可能会产生难以调试的问题。

内存延迟：PE和片上SRAM内存的访问延迟都比典型的要长。让许多客户端访问PE存储器使仲裁方案复杂化并增加了延迟周期。对于固定功能块，这种延迟会被计入操作的延迟中，但当处理器试图使用本地内存时，软件必须采用软件流水线等技术来隐藏延迟，这会增加寄存器压力。当在矢量处理器上开发内核时，这种情况会加剧，因为它限制了可以执行的寄存器分组的数量，从而增加了动态指令数。

在外围放置片上SRAM，同时在多个片上均匀地分配请求，造成了内存访问延迟的很大程度的不均匀性。由于请求总是在最后一段数据到达后完成，因此当发出大于最小请求时，总体延迟会受到影响。虽然在大多数情况下，这种延迟可以通过将数据预取到PE的内存中来隐藏，但在具有小池化组的EmbeddingBag算子等情况下，延迟会凸显出来，因为内存访问模式事先未知，并且没有足够的未处理请求来隐藏延迟。

缓存一致性：虽然系统实现了共享内存范例，但没有硬件支持缓存一致性。在这个共享存储器系统中，不需要PE间的一致性，因为不同的PE以数据并行的方式对数据集的专用部分进行操作。然而，PE内部的一致性有时会导致准确性问题：如果两个处理器内核、固定的功能单元和一个内核接触到同一组内存地址，或者同一个内核在不同的算子之间重用地址，则必须从缓存中显式清除这些地址的缓存副本，否则可能会使用过时的数据。

架构层次：加速器的推荐模型在其使用的层和算子的大小和复杂性方面有很大差异。虽然映射到PE网格上的大层可以从可用硬件资源中提取所需的利用率水平，并分摊创建和调度的开销，但较小的层或较低的batch处理大小必须采用诸如利用子图并行性和融合等技术来获得相同的利用率。尽管在软件级别有足够的空间来执行此类优化或减少部署作业的开销，但我们相信，在架构级别进行一些资源调配会使解决此问题变得更容易。这是因为对于较小的任务，网格必须划分为较小的子网格，以便每个子网格都能处理较小的任务。而设置和拆除这些子网格的任务是系统固件的一部分。我们认为，在体系结构本身中拥有另一个层次，例如PE集群，可能会使这个问题更容易解决，因为与PE的单片网格相比，它提供了自然的隔离和管理单元。

在第一代MTIA中，我们构建了一种架构，与GPU和其他加速器相比，该架构可以为DLRM的工作显著提升效率。随着软件堆栈的成熟，我们希望随着时间的推移继续提高这种效率。为这种体系结构编写内核、构建编译器和优化模型的经验让我们深入了解了哪些功能有更大的影响，我们希望这些项目硬件和软件方面的经验教训能够在未来几代架构中得到整合和利用。

 　　审核编辑：汤梓红