【论文分享精炼版】 sNPU: Trusted Execution Environments on Integrated NPUs

今天在COMPASS分享了之前写的一个博客，做了进一步的提炼总结，大家可以看看原文~

今天分享的论文《sNPU: Trusted Execution Environments on Integrated NPUs》来自2024年ISCA，共同一作为Erhu Feng以及Dahu Feng。并且， 这两位作者今年在加速器安全研究领域还有一篇来自ASPLOS的论文，即《sIOPMP: Scalable and Efficient I/O Protection for TEEs》。值得一提的是，这两篇论文均借助了基于RISC-V的Penglai Enclave，即2021年USENIX的《Scalable Memory Protection in the PENGLAI Enclave》，Erhu Feng为该工作的第一作者，大家感兴趣可以去了解原文。

Background & Motivation

由于移动设备上人工智能应用的高需求，当前的移动芯片在 SoC 中集成了 NPU 内核，以提高能效和性能。然而，NPU 在 SoC 中的集成引入了新的三个攻击面：

（1）利用受损的NPU攻击CPU侧资源。

（2）NPU的内部攻击。

（3）利用CPU攻击NPU。

为了解决以上安全问题，设计结合 CPU 和 NPU 的 TEE 变得至关重要。

目前针对 CPU-NPU 系统的TEE设计分两种：

（1）可信NPU。该方法暂时将整个 NPU 指定为安全设备，并将整个 NPU 驱动程序迁移到 TEE 中。然而，该解决方案导致 NPU 资源利用率严重不足，并且由于软件堆栈的复杂性导致 TCB 较大。

（2）加密保护。该方法旨在通过内存加密和完整性保护以抵御 DRAM 的物理攻击（例如冻结内存）。然而，这些方法缺乏对 NPU 内部结构的保护，因为 NPU 内部的数据仍然是明文。

此外，直接将 GPU 等其他加速器的 TEE 机制应用于集成 NPU 存在以下两个限制：

（1）GPU TEE 中采用的 IOMMU 等隔离机制对于集成 NPU 来说效率不高，因为 NPU 需要更大的内存带宽。

（2）NPU 具有 scratchpad 和 Networks-on-Chip(NoC) 等专门的硬件结构，这带来了新的攻击面。

因此，NPU TEE 应该防御上述三个攻击面，并满足两个基本要求：最小化运行时性能开销和实现更高的资源利用率。

sNPU Overview

为防御前面所提到的三个攻击面，以及兼顾性能开销的基本要求，sNPU 引入了以下新颖方法：

（1）NPU Guarder。为了防御利用 NPU 外部行为（例如内存访问）的攻击，sNPU 引入了基于 tile 的内存翻译和检查单元，专门为适应 NPU 内存访问模式的特征而设计。这种设计产生（几乎）零运行时开销，同时节省了检查能耗。

（2）NPU Isolator。为了解决利用 scratchpad 和片上网络等 NPU 内结构的新攻击面，sNPU 利用 scratchpad 与主存储器没有关联的观察结果，对其采用更细粒度和动态的隔离机制。同时，sNPU 还结合了 NoC 隔离机制和离线路由检查，保证了 NoC 网络的完整性。

（3）NPU Monitor。为了最大限度地减少 CPU 端恶意软件的潜在攻击，sNPU 减少了 NPU 堆栈的软件 TCB 。sNPU 在安全世界中引入了 NPU Monitor，仅用于必要的安全检查。与此同时，人工智能框架和 NPU 驱动程序等其他软件组件可能仍然不受信任。

Detailed Design

NPU Guarder: Memory Access Guarder for NPU

Challenges

sNPU 专注于在集成的 NPU 上构建 TEE，在设计访问控制机制时会面临以下挑战：

（1）集成的 NPU 具有不同的内存访问路径，这可能会使统一访问控制器的设计变得复杂。如图1所示，一些 NPU（Type-1 和 Type-2）利用单独的 IOMMU/MMU 来限制 NPU 访问，而其他 NPU（Type-3）则依赖 CPU 端的访问检查机制。

（2）NPU 需要更高的内存带宽，这就需要更高效的检查逻辑。

（3）IOMMU/MMU 机制为 NPU 引入了不可忽略的开销。

Solutions

图 2 展示了 NPU Guarder 设计，这是一种利用 NPU 中特定内存访问模式的轻量级内存访问控制器。与 MMU 或 IOMMU 相比，它有两个主要优点：

（1）它具有轻量级设计，无需检查开销。

（2）它可以集成在 NPU 内部。

首先，为了消除传统基于分页的内存访问控制的运行时开销，NPU Guarder 采用粗粒度内存检查和细粒度翻译机制。对于存储器检查，它利用记录连续存储器区域的访问权限的检查寄存器，因为移动系统中的敏感数据通常存储在预先分配的安全存储器区域（例如，TrustZone 安全存储器区域）中。至于地址转换，NPU Guarder 在 tile 级别（例如输入 tile 和输出 tile）提供细粒度的转换寄存器。每个转换寄存器将特定区域从虚拟地址映射到相应的物理地址。与很少修改的检查寄存器不同，转换寄存器可以在 NPU 计算之前更新（如果需要）。

其次，NPU Guarder 将这些检查和转换寄存器集成在 NPU 内核内部，位于 DMA 引擎之前。与独立模块（例如IOMMU）相比，集成设计降低了SoC的复杂性。此外，由于内存检查和转换是在 DMA 请求级别而不是内存数据包级别执行的，因此与 IOMMU 相比，NPU Guarder 可以节省额外的消耗。当接收到 DMA 请求时，DMA 引擎将其分为多个固定大小的内存数据包（例如 64 字节）。因此，NPU Guarder 仅检查一次（并且节省能源），而 IOMMU 需要在内存数据包级别检查 O(N) 次。

NPU Isolator: Inner Resource Isolation for NPU

Challenges

Scratchpad 和 NoC 是 NPU 中用于加速 AI 工作负载的专用结构，但它们也引入了新的威胁：

（1）Scratchpad威胁。例如，NPU驱动程序可以分配一个已被另一个任务使用的scratchpad条目，并且NPU编译器可以强制读取scratchpad中的内容而不需要之前写入。

（2）NoC威胁。例如，如果 NPU 调度程序受到损害，它可以将恶意任务调度到错误的 NPU 核心。因此，攻击者可以拦截从源核心传输的秘密中间结果，或者向受害核心发送恶意 NoC 数据包。通过篡改 NoC 网络的路由完整性，攻击者可以操纵整个 ML 任务（见图3）。

Solutions

为解决Scratchpad 和 NoC所面临的威胁，NPU Isolator分别提出了基于ID的隔离机制以及NoC身份验证机制。

（1）基于ID的隔离机制。如图4（b）所示。关键的见解是scratchpad条目和系统内存之间没有地址关联，允许sNPU在任何scratchpad条目中存储数据。因此，scratchpad可以采用比缓存和内存更细粒度、更动态的隔离机制。此外，由于每个scratchpad条目具有较大的有效负载（例如，≥128b），因此一位 ID 状态的资源开销可以忽略不计（<1%）。

（2）NoC身份验证机制，即peephole。如图5所示。peephole机制为NoC数据包生成身份（head flit）位于源核心，在NoC网络中传播。当目标核收到此 NoC 数据包时，目标核中的peephole根据其身份验证此 NoC 请求。NPU核心的ID状态可以作为peephole机制中的有效身份。为了确保全面的NoC保护，sNPU还需要考虑在多个NPU核上运行的ML任务的路由完整性。除了计算 ML 任务代码的哈希值并将其与预期测量值进行比较的代码完整性检查之外，路由完整性检查还确保实际的 NoC 路由与用户的期望一致。

NPU Monitor: Trusted Software Module for NPU

Challenges

NPU 包含一个由各种组件组成的大型软件堆栈，例如 AI 框架（例如 TensorFlow、PyTorch）、编译器（例如 TVM、CANN）和 NPU 驱动程序。将整个软件堆栈包含在 TCB 中可能会带来潜在的漏洞和安全风险。

Solutions

如图6所示，NPU Monitor由几个模块组成：上下文设置器、可信分配器、代码验证器和安全加载器。值得注意的是，NPU Monitor仅适用于安全的 ML 任务。而对于非安全任务，仅依靠硬件机制来保证安全任务和非安全任务之间的隔离。

（1）上下文设置器。其负责设置NPU安全上下文，其中包括NPU的ID状态、安全任务的检查和翻译寄存器。NPU上下文决定了NPU可以访问的硬件资源，例如系统内存和scratchpad。

（2）可信分配器。其负责在保留的安全内存中分配内存缓冲区，例如输入/输出数据和安全任务模型。它还检查scratchpad是否有重叠。

（3）代码验证器。首先，它将安全任务的代码和敏感模型加载到安全任务队列中。然后，它根据用户的期望计算并验证任务代码的测量结果。

（4）安全加载器。首先，它保证 ML 任务的路由完整性。与传统的 CPU TEE 不同，ML 任务可以利用与 NoC 网络连接的多个 NPU 核心。安全加载程序验证调度的 NPU 内核是否与预期 NoC 网络的拓扑匹配。验证路由完整性后，安全加载器将 ML 任务上传到相应的 NPU 内核中。

除了以上模块之外，NPU Monitor 还有两个辅助组件：trampoline 和安全任务队列。Trampoline 充当非安全 NPU 驱动程序和 NPU Monitor 之间的中介（用于数据传输），而安全任务队列则存储用于调度的安全 NPU 任务。

Evaluation

Experimental Setup

硬件原型：使用Chipyard在FPGA上实现了sNPU的硬件原型。

NPU设计：参考了Gemmini和AuRORA，基于systolic-array的DNN加速器。

CPU侧TEE：基于Penglai，一个RISC-V TEE系统。

模拟工具：使用FireSim进行端到端DNN工作负载的性能评估。

Protected Memory Access for sNPU

（1）性能比较。比较了使用IOMMU（如TrustZone NPU中采用的）和NPU Guarder（sNPU采用）对DNN推理吞吐量的影响。结果显示，IOMMU引入了显著的性能开销，而NPU Guarder则几乎没有性能损失。

（2）检查次数。分析了IOMMU和NPU Guarder在内存事务中的检查请求次数。NPU Guarder由于能够处理连续地址块，因此请求次数大大减少，从而降低了能耗。

ID-based Scratchpad Isolation

（1）性能评估。评估了在不同刷新粒度下，sNPU的基于ID的动态隔离机制与当前TrustZone NPU中采用的粗粒度刷新和静态分区方法的性能影响。结果显示，sNPU的方法提供了更高的灵活性和适应性，允许更高的Scratchpad利用率。

NoC Isolation

（1）微测试。通过微测试比较了使用软件NoC（使用共享内存）、未经授权的NoC和带有peephole机制的NoC的数据传输成本。结果显示，peephole机制在保持安全性的同时，显著减少了延迟并提高了带宽。

（2）应用测试。在真实世界场景下测试了peephole机制的性能。结果表明，与使用共享内存的软件NoC相比，peephole机制能够减少约20%的执行时间。

Hardware Cost Analysis

（1）资源开销。在FPGA上实现sNPU，并与基线NPU和TrustZone NPU进行硬件资源消耗的比较。sNPU的额外资源开销很小，特别是与采用复杂页表遍历的IOMMU相比。

TCB Size Analysis

（1）软件TCB。sNPU设计只引入了一个小的软件TCB，主要由NPU Monitor组成，其代码行数远小于整个NPU软件栈。

​