【TEE论文】Confidential Computing within an AI Accelerator

论文
幻灯片

文章目录

摘要

我们推出了IPU可信扩展（ITX），这是一组硬件扩展，可在Graphcore的AI加速器中实现可信执行环境。ITX 能够以较低的性能开销执行具有强大机密性和完整性保证的 AI 工作负载。ITX 将工作负载与不受信任的主机隔离开来，并确保其数据和模型在加速器的芯片外始终保持加密状态。ITX 包括一个硬件信任根，可提供证明功能并编排可信执行，以及用于在 PCIe 带宽下对代码/数据进行身份验证加密的片上可编程加密引擎。

我们还以编译器和运行时扩展的形式提供 ITX 软件，这些软件支持多方训练，而无需基于 CPU 的 TEE。

我们在台积电7nm节点流片的Graphcore的GC200 IPU中加入了对ITX的实验性支持。它在使用标准 DNN 训练工作负载的开发板上的评估表明，与基于 AMD SEV-SNP 的基于 CPU 的机密计算系统相比，ITX 增加了 <5% 性能开销，性能提高了 17 倍。

1 介绍

支持AI加速器的可信计算的需求：

• 需要添加隔离、测试、侧信道弹性等安全特征
• 有最小的设计改变和性能下降
• 灵活性，能支持不同类型的主机（无TEE支持的CPU，基于进程的TEE支持的CPU，基于虚拟机的TEE支持的CPU）
• 制造和装配过程以及协议必须针对供应链攻击进行加强

本文通过硬件ITX支持和compiler、runtime支持，能在IPU上保证ML任务的可信性和一致性，同时性能开销小于5%。

ITX的主要部件包括：

• 一种新的执行模式，用于将所有安全敏感状态与主机隔离并安全地处理安全异常
• 能够以线速（32 GB/s双向吞吐量支持PCIe Gen4）加密/解密CPUIPU PCIe流量的可编程加密引擎
• 一种新颖的身份验证加密协议，用于确保代码/数据传输的机密性和完整性，而无需信任主机。

对ITX的信任植根于机密计算单元（CCU），这是IPU板上的一种新的硬件信任根（HRoT）。CCU 根据制造结束时在 CCU 中采样的硬件密钥为每个设备提供唯一标识。CCU固件负责管理IPU上TEE的整个生命周期，包括创建、发布任务特定属性的证明报告、TEE的密钥交换、启动和终止。我们的设计还具有协议，用于在潜在的恶意制造环境中安全地向IPU提供固件，用于颁发捕获所有可更新固件身份的证书，以及支持固件更新而无需重新认证的设备。

一些特殊属性导致IPU ITX比CPU TEE有更高的安全性：

• 独占资源，不可能在相同资源上并发运行而产生侧信道
• 大型片上SRAM，IPU的数据访问具有固定延迟，IPU内计算和内存通信不需要加密，也不需要缓存和推测来隐藏内存访问延迟
• 编译器静态管理IPU内资源（cores,memory,communication channels）分配，应用程序的二进制定义数据流和控制流。因此不需要主机软件栈来控制执行，也不需要保证主机的可信。

注意：由于实现的限制，论文使用离散的HRoT，不对IPU-IPU链路进行加密，容易受到物理攻击，可以通过片上集成HRoT解决，或者在IPU-IPU链路引入加密引擎。

论文贡献如下：

• 可信的多方ML，不需要信任CPU
• 基于离散的硬件信任根，支持远程验证和安全密钥交换
• 通过PCIE对代码数据进行身份验证加解密的应用级别协议
• 制造过程安全地为设备预配机密、固件、证书

可以作为其他专用设备或加速器添加TEE支持的蓝图。

2 背景

2.1 IPU硬件结构

Tiles：1472tiles；900MB SRAM；专门为AI而设计的指令集；每个内核最多可以执行六个静态调度的线程；由于片上存储器以固定的延迟访问，指令可以被编译器准确地调度。

Interconnects：tiles通过Internal Interconnect高带宽互联；Internal Interconnect通过一组exchange blocks和external exchange Interconnect连接，每个IPU有一对PCIe Links连接主机，有IPU-Links连接其他IPU。Tiles通过写请求包向主机写入数据，数据包通过tiles的标识符进行路由。

Exchange Address Spaces：三个地址空间用于主机和IPU以及IPU和IPU之间通信。Tileadress space是瓦片之间相互寻址的空间。Host PCI space被主机用来处理Host Exchange块中的tile存储器和片上页表。Tile PCI space用于tile通过PCI处理对主机内存的读请求。
Host-IPU Interface：IPU通过驻留在PCI Complex一组配置寄存器和主机通信。PCI Complex由主机和IPU之间外部同步的主机同步代理( Host Sync Proxy，HSP )、转换PCI格式和专有外部交换数据包格式的主机交换、用于保存未完成PCI读请求的元数据的片上查找表以及用于访问所有其他内部组件的配置寄存器的控制端口组成。主机交换子系统还包括一个称为自动装载器的组件，实现tile内存的高效刷新和初始化。

Host-IPU Synchronization：IPU执行模型基于带栅栏和超步的批量同步并行( Bulk Synchronous Parallel，BSP )范式。

**IPU Control Unit (ICU)**：ICU是集成在板上的微控制器，通过JTAG与IPU连接，并配有PCB外设，用于供电和环境监测。它负责IPU的初始化。

2.2 IPU软件栈

IPU软件栈编译并执行TensorFlow和PyTorch等ML框架编写的应用程序。它由编译器、主机运行时和一组由IPU设备驱动程序支持的库组成。
Compiler：编译器在瓦片之间对图的每一层进行划分；为图的每个节点分配资源(线程,内存)；将主机和IPU之间的所有数据传输映射到一个运行时支持的称为流的抽象，从主机到IPU (以及IPU对主机的影响)的数据传输被映射为输入(输出)流，并被编译为对Tile PCI地址空间的读(写)指令序列。

Host Runtime：运行时提供了加载IPU二进制文件的代码，以及在IPU内部和外部流式传输数据的代码。它通过在每个瓦片内存的预留部分中部署一个小的bootloader来加载IPU二进制文件。

3 威胁模型

4 整体概述

IPU中的可信执行使模型开发者能够安全地卸载一个ML作业(训练或推断)，同时保护其模型和数据免受托管平台的影响。卸载作业的工作流程包括创建TEE、生成证明报告、远程验证、代码/数据加密、加密密钥的安全交换、作业执行和结果的解密。

4.1 硬件扩展（ITX）

硬件组件：硬件信任根CCU，可编程的AES-GCM引擎（用于验加密验证主机和IPU之间的链路，存在于安全交换通道SXP），可信执行模式（通过写寄存器进入，进入后寄存器和tile内存只能通过ICU和CCU访问，只能通过芯片重置退出安全模式）

4.2 软件支持

离线模式可以在基于IPU的TEE中部署多方ML训练工作负载，而不需要基于CPU的TEE。

Job Preparation：静态编译任务，在可信离线的clean room environment执行(1）；模型提供者加密二进制和参数，生成新鲜的公钥共享用于密钥交换；数据提供者预处理并加密输入数据和标签，并创建包括密钥共享和认证的数据包(2)；得到的包加载到IPU（3，4）
Job Initialization：运行时为作业创建TEE ( 5 )，获取它们的证明报告和额外的附属物，如设备特定证书( 6 )，以及从依赖方传递密钥交换消息到CCU ( 8 )等操作提供用户模式的API。这个运行时是不可信的。
Remote Attestation：在可信模式下，CCU可以发布远程可验证的证明，并将其传递给依赖方( 7 )，作为其工作负载的TEE配置的证明。该证明是一个从IPU制造商根CA到由CCU签署的终端证书的证书链。
Job Execution：在模型开发者和数据提供者向CCU发布密钥后，CCU将密钥部署到SXPs中，并通过使用自动装载器将bootloader安装到IPU tilea中启动作业( 9 )。在可信模式下，这些块在写入tile存储器之前由SXP解密并进行完整性校验

5 IPU里的可信执行

5.1可信计算单元CCU

CCU负责将每个设备与唯一的密码身份相关联，并管理其IPU中的可信执行。CCU是基于意法半导体的STM32H753微控制器的分立芯片[ 24 ]。CCU通过ICU与IPU相连。一个专用的引脚接收IPU在可信模式下产生的所有异常，赋予CCU固件对异常处理的完全控制权。

Firmware Architecture and Attestation：CCU 实现测量的引导协议，该协议是设备身份组合引擎 （DICE） 架构的变体。DICE确保为每个设备分配一个唯一的身份，同时最大限度地减少硬件机密的暴露。除了稳定的设备标识外，所有派生的密钥和密钥都会在固件（及其度量）更改时自动更改，从而确保低级固件攻击不会泄露其他固件中使用的密钥。CCU 固件由三层组成：一个不可变的主引导加载程序，在制造时配置在一次性可编程闪存中；一个可变的辅助引导加载程序，负责设备标识；一个管理 TEE 生命周期的机密计算引擎 （CCE）。

在制造过程中，CCU 配置了主引导加载程序固件。当设备首次退出重置时，此主引导加载程序接收来自 ROM 固件的控制，使用基于硬件的 TRNG 对唯一设备密钥 （UDS） 进行采样，将其存储在闪存区域中，并永久阻止其从任何其他固件层访问。UDS是IPU密钥层次结构的根，该协议确保它永远不会暴露在CCU之外，甚至不会暴露给制造商。
在每次后续启动时，CCU都会使用IPU制造商的固件签名密钥从闪存加载并验证辅助引导加载程序。接下来，它派生两个中间密钥：来自 UDS 的硬件设备标识符 （HDI） 和来自 UDS 的复合设备标识符 （CDI） 以及辅助引导加载程序的度量。

辅助引导加载程序进一步派生两个公钥对：来自 HDI 的卡身份密钥 （CIK） 和来自 CDI 的平台身份密钥 （PIK）。因此，CIK 为每个卡提供稳定的身份，而 PIK 对于每个卡和辅助引导加载程序都是唯一的。
辅助引导加载程序从 CDI 和 CCE 度量派生证明密钥 （AK）。因此，AK对于每个设备、辅助引导加载程序和CCE都是唯一的。

Firmware Update：根据DICE，辅助引导加载程序更新会使制造商颁发的PIK证书失效，并且由于UDS在每个设备中配置，因此IPU制造商无法独立派生和认证更新的PIK。相反，我们依靠 CIK 作为本地 CA 来签署更新的 PIK 证书。此外，制造商将颁发包含旧版本和新版本固件测量值的 TCB 更新证书。

5.2 TEE生命周期管理

TEE Initialization：将作业安全地卸载到IPU的第一步是为此作业创建一个新的TEE。TEE 初始化需要作业清单、公钥共享、每个信赖方的密钥共享和证书上的签名，以及指示作业是从检查点开始还是从检查点恢复的检查点计数器。在TEE初始化期间，CCU首先使IPU静默，确保主机和IPU之间没有正在进行的读写请求。然后，它将IPU切换到可信模式，使用自动加载器清理所有tile内存，并测量配置寄存器的状态。然后，它使用证书检查密钥共享上的签名，并生成自己的新 EC 共享，该共享用于在每个依赖方和 CCU 之间建立 ECDH 共享密钥
CCU生成由证明密钥签名的证明报告，其中包含各种特定于IPU的属性和特定于作业的属性。每个信赖方都可以与支持证书链一起查看证明，以验证设备以及CCU和IPU的初始状态，然后可以计算其ECDH共享密钥，并包装包含该方的数据加密密钥和随机数的密钥包来运行作业。

TEE Launch：从所有信赖方收集包装的加密密钥后，主机将启动作业的执行。首先，CCU 计算每一方的 ECDH 共享密钥，并使用它们来解包从每一方收到的密钥包。然后，它将随机数组合在一起，以派生此作业运行的检查点密钥和最终模型加密密钥（如果从另一次运行恢复，则从上一次运行中获取检查点密钥以还原其状态）。此密钥派生可确保此运行的检查点密钥是最新的（只要一个信赖方的随机数是最新的），并且一旦所有信赖方都同意从检查点恢复，就可以重新计算先前运行的检查点密钥。
接下来，CCU使用自动加载器在IPU tile上部署预定义的引导加载程序，并按照作业清单中指定的方式将第一组加密密钥部署到SXP（包括型号密钥）。
然后，它会在每个tile上激活引导加载程序（其度量包含在证明报告中），从而发出从主机内存中读取其加密应用程序二进制文件的请求。对这些读取请求的响应在复制到专用磁贴内存之前由 SXP 进行身份验证和解密。最后，CCU按照作业清单中的规定部署下一组加密密钥（包括数据密钥），并在IPUtile上触发主执行循环。

TEE Termination：在初始化 TEE 后的任何时候，主机运行时也可以请求终止 TEE。如果IPU引发安全异常（例如无法验证读取请求的响应），CCU也可以终止TEE。在TEE终止期间，CCU使IPU静默，使用自动加载器清理tile内存，并禁用所有SXP密钥。最后，CCU通过Newmanry复位将IPU切换到正常模式。TEE 也可以通过器件的硬复位来终止。在这种情况下，所有CCU状态都会被清除，IPU将恢复到正常模式。当它退出重置时，在重新启用主机链接之前，ICU 会清理tile内存，以确保在主机重新获得对设备的访问权限之前擦除上次执行遗留的任何机密。

6 加密的DMA

接下来，我们将介绍用于加密代码和与IPU tile之间的数据传输的ITX协议。该协议是应用程序级的，而不是传输级的。虽然它对ML框架是透明的，但它依赖于应用软件（例如IPU编译器）来分配IV进行身份验证加密，并对tile进行编程，以安全地加载代码、初始权重、训练数据，并保存/重新加载检查点和结果。

6.1 数据格式

应用软件将每个代码和数据流划分为大小相等的加密帧。每个帧由一个 128 位 IV 和一个 128 位身份验证标记组成，后跟一系列携带帧加密内容的密码块。应用软件可以自由地对不同的流使用不同的帧大小，只要总帧大小（包括 IV 和身份验证标签）是 128 字节的倍数，最大为 1KB，这是支持的最大 PCIe 读取。
至关重要的是，应用软件必须确保 IV 永远不会在使用相同密钥加密的帧之间重复使用，这对于 AES-GCM 来说将是灾难性的。编译器确保了这种不变性，编译器通过组合特定于流的标识符和帧索引来构造 IV，并且代码和数据流都是一次写入的抽象。这保证了除非关联的密钥已泄露，否则使用正确的 IV 进行经过身份验证的解密将产生正确的有效负载。

6.2 硬件支持

安全交换管道 （SXP）:SXP 是一个可编程硬件模块，支持 AES256-GCM 身份验证的帧加密和解密。每个 SXP 可实现 16 GBps 的单向吞吐量，对延迟的影响可以忽略不计。交换互连上放置了四个 SXP（每个方向两个），以支持 PCIe Gen4 线速（32GBps 双向）的加密/解密。在可信模式下，每个 SXP 都配置为拦截来自四个交换块的读/写请求。

AES-GCM Engine：SXP 的核心是一个完全流水线的 AESGCM 引擎，每个上下文都可以通过将 256 位密钥加载到通过内部控制总线暴露给 CCU 的控制寄存器中来编程。虽然帧可能是交错的，但为了功能的正确性，我们要求每个上下文一次处理一个帧。在每个循环中，核心对其上下文执行以下操作之一：( i )上下文空闲，核心接收帧的IV；( ii )上下文活动，核心接收一个数据块；( iii )上下文活动，核心接收一个MAC。核心通过比较连续周期之间的关键上下文标识符来检测上下文切换，从而可以在下一次操作之前获取下一个上下文。

Frame encryption/decryption：SXP接收3种类型的外部交换数据包：读请求( egress )；需要解密(入口)的读取完成；并写出需要加密的请求(出口)。它们的头部被扩展以携带额外的信息，以帮助SXP确定数据包应该如何处理：一个AES位表示读完成或写请求被加密；一个4比特的KEY _ INDEX字段标识了要使用的物理密钥上下文；而一个CC位表示该帧的最后一个数据包，并触发其认证标签的计算。
在写请求包(出站到主机PCIe域)中，AES和CC位由Tile设置，而KEY _ INDEX由SXP设置。在读完成包中，信息由PCIe复合体根据其维护的关于待处理读请求的可信状态进行设置。读取请求数据包和AES位未设置的数据包不需要加密/解密；它们是通过不变的。对于其他所有数据包，报头绕过AES核，然后AES核根据帧是否开始、帧是否继续或帧是否结束处理每个数据包(及其块)。

Key Selection：每个SXP支持多个物理密钥上下文，以实现并发I / O的加密/解密。SXP提供一组可编程(通过CCU )寄存器来定义数据包和用于加密/解密其有效载荷的物理密钥上下文之间的映射。编译器通过分配一组与交换块上下文相关的瓦片来访问单个流来产生这种映射。SXP在接收到数据包后，利用源瓦片标识符在包头中计算出的交换块上下文索引来查找物理密钥上下文。
一旦SXP推断出物理密钥上下文，它将更新请求包头中的KEY _ INDEX字段。对于写请求，SXP使用该字段将AES内核切换到推断的物理上下文以加密其有效负载。然而，对于读请求，情况则更多地涉及到，因为读请求绕过了AES核心，必须使用推断的物理密钥上下文来解密读请求处理后将由主机返回的读完成。当PCI复合体接收到读请求时，它将KEY _ INDEX和AES字段与其他元数据(如源瓦片标识符)一起缓存在片上查找表中。

7 软件扩展

我们现在描述了一组对IPU软件栈的扩展，以便在离线模式下使用ITX编译和执行机密的ML任务。该模式由Tensorflow配置选项触发。当启用时：( 1 ) XLA后端将计算图转换为一个新的抽象，称为机密数据流，用于所有的数据传输，包括初始权重、训练数据、检查点和训练好的模型。( 2 ) IPU编译器将计算图编译成一组IPU应用二进制码(每个IPU一个)，其中每个二进制码是瓦片专用二进制码的级联。编译器使用新采样的模型密钥将瓦片二进制加密成一组加密帧。一个帧被分配一个唯一的IV，包括编码类型，IPU / tile ID和帧索引。( 3 )通过扩展IPU运行时来安全地引导任务，然后在主机和IPU之间传输加密的二进制文件和数据。

7.1 可信数据流

机密数据流是一种编译器和运行时抽象，用于利用SXP向IPU传输数据，具有机密性和完整性保证。每条数据流都是经过相同对称加密密钥加密的数据实例序列。每个数据实例被划分为一系列帧，每一帧都使用由流类型(数据)、流标识符和流内帧的索引组成的唯一IV进行加密。
编译器和运行时实现对机密数据流的读写操作如下。正如第6.1节所讨论的那样，编译器首先为每个流分配一个瓦片PCI空间中的区域，其约束条件是不超过IPU环形缓冲区(例如, 256 MB)的总容量。
接下来，编译器为每个流分配一组可读或可写的瓦片，在每个瓦片上预留SRAM以容纳流的一部分，并生成SXP映射，约束条件为：( a )与这些瓦片相关的交换块上下文映射到分配给流的物理密钥上下文，( b )对于任何一个SXP，程序中任意一点正在使用的物理密钥上下文的数量不超过16。
为了在这些约束条件下最大化性能，编译器可以在现有密钥失效和加载新密钥的应用程序中引入同步点。编译器将这些同步点连同它们的关键标识符一起包含在作业清单中；而(不可信的) IPU运行时使用该部分显单，要求CCU在这些点将下一个解密密钥加载到SXP中。关键的变化只适用于输入流。在TEE发射时，输出流的密钥由CCU导出并加载，并且不改变吞吐量的寿命。
恶意运行时不遵循作业清单的密钥调度只能导致解密失败，最多导致拒绝服务。
接下来，编译器在每个瓦片上调度读写方法操作。该调度需要满足一个硬件约束，即在任意时刻，针对任意给定物理密钥上下文产生请求的瓦片都与单个交换块上下文相关联。这是因为，虽然交换块可以在每个交换块上下文(从而使其物理密钥上下文每次被一个瓦片使用)内动态地同步和调节请求，但是跨交换块上下文并不存在这种同步。最后，编译器在执行调度的每个瓦片上生成代码，对访问的帧发出读写方法请求。(由于空间限制,细节被省略。)

7.2 安全检查点

检查点保存到并从主机内存中恢复，并通过数据流实现。安全检查点通过一种特殊形式的机密数据流实现，其中IV捕获历元(统计这项工作的检查点恢复数量)和检查点标识符(统计一个历元内存储的检查点数量。)，CCU为每个历元使用单独的检查点密钥，并在应用程序启动时通过bootloader使历元计数器和检查点标识符可供IPU瓦片使用。(详情见附录A.7。)

7.3 安全辅助程序

Secure辅助程序是在作业开始时或从检查点恢复作业时将加密的应用程序二进制文件安全加载到IPU的过程。
辅助程序涉及以下步骤。首先，IPU运行时在主机内存中加载加密的IPU二进制文件，并使用CCU API创建一个TEE；这就将IPU切换到可信模式。接下来，CCU使用第2.1节中描述的自动装载器在每个IPU瓦片上安装一个bootloader (见附录A.5)，并使用模型解密密钥对SXP进行配置。每个瓦片上的bootloader通过发出读请求序列从主机内存中获取瓦片的二进制文件。从主机接收到的每一帧都被SXP截获，经过认证和解密，并复制到瓦片存储器中。然后，bootloader检查接收到的IV与bootloader逻辑内置的预期IV是否匹配；这种检查是在软件中进行的，因为SXP只保证了每一帧的真实性，而不是整个流的完整性。这种检查的失败表明主机试图篡改码流，例如通过重放/重新排序帧。在这种情况下，瓦片会产生一个安全异常，由CCU处理。如果所有的校验都通过，bootloader通过从所有帧中剥离IV和认证标签来重构原始二进制。
最后，bootloader计算瓦片二进制的哈希值；瓦片累加了整个应用程序的哈希值；并且CCU检查其与作业清单中的度量是否匹配，否则产生安全异常。该协议和bootloader完整性(它的度量包含在证明中)一起保证了应用程序的完整性。

8 评估

Implementation：我们实现了一个用于主机CPU服务器不可信的保密训练任务的软件原型。我们的原型包括ML框架中的实验支持、IPU编译器和运行时。在我们的原型中存在以下不足：( 1 )我们的实现目前只支持板上的一个IPU；( 2 )编译器只使用一个逻辑关键字区域，将代码、数据、标签、检查点和输出映射到该区域上；然而，每个加密帧被静态地分配一个唯一的IV，保留了每个IV只使用一次的不变量；( 3 )安全复工尚未落实；( 4 )部署在IPU瓦片上的bootloader在解密后没有对IPU二进制进行测量
Experimental Results：启用SEV - SNP的开销较小，从8 % (小模型)到14 % (大模型)不等，启用ITX的开销从3 % (大模型)到58 % (小模型)不等。ITX开销以一次性设置成本为主，在较大的训练次数下摊销；**该成本包括TEE初始化和认证( 40 % )、TEE启动包括SXP设置( 43 % )和TEE终止包括SXP刷新( 13 % )。运行时加密仅引入总开销的4 %**。