| 赛题 | proj325 - 面向AI计算的可信GPU计算框架 |
|---|---|
| 小组成员 | 黄家俊 杨南 |
| 指导教师 | 郭力维 |
- 幻灯片：proj325.pptx
- 总结文档：docs文件夹
- 演示视频：https://pan.baidu.com/s/1yR0yPObunOJx0sxENbk2ng?pwd=ttqx 提取码: ttqx
RPlayer是一个可信GPU计算框架，通过记录GPU在完整软件栈中的执行过程，并将其参数化并迁移至可信执行环境（TEE）中重放，从而实现可信GPU计算。
RPlayer 主要参考了基于录制重放的设备驱动复用相关工作，gpureplay-linux-kernel 和 gpureplay-circle-tinystack。在此基础上，RPlayer在RPi5平台完整记录GPU计算过程，于可信执行环境中重放，实现TEE中端到端的AI推理，并验证得到可解释的计算结果（如图像分类标签）。此外，RPlayer还进一步考察了多种计算框架在该机制下的应用与适配，并将其首次延伸扩展至大模型场景。
目前AI计算应用广泛并且发展迅猛，由于需要处理大量的并发计算，目前的AI框架会调用GPU等加速器来提升AI算力。
然而AI框架运行的敏感数据以及模型本身都会受到攻击并泄露，如攻击者使用提权漏洞获取内核权限并直接读取传递给GPU的数据信息。
为了防止恶意的攻击者，需要构建一套可信 GPU 计算框架，保证AI计算安全。本项目旨在使用 CPU TEE 技术，利用软硬件协同手段为GPU构造隔离计算环境，并保证AI计算数据在内核中的安全流通。
- 扩展 CPU TEE 安全能力，并应用于GPU中，保证GPU中的数据传输与计算安全；
- 加固AI计算使用的内核模块，防止内核特权攻击；
- 可信GPU模块的安全高效切换，不干扰原有系统的运行；
TEE 旨在从硬件层面确保代码及数据的可信性，为CPU提供了运行可信应用的安全区域。因此，我们以 TEE 作为本项目的关键基础。
OP-TEE 是 TEE 的一种内核实现，为安全区域的可信应用提供基础内核服务，如内存分配与上下文切换。
为使复杂GPU软件栈得到来自 TEE 的强安全隔离防护，我们考虑了如下两类方法：
- 基于划分：HETEE、StrongBox、LDR。核心思想是分离负载、软件栈等层面的可信与不可信部分。
- 基于 Record & Replay (RnR)：GPUReplay、RT-TEE。核心思想是记录普通系统的行为，并在可信系统重放。
考虑到目标是AI计算，软件栈层级复杂，我们主要参考了基于 RnR 的方法，GPUReplay。
RnR的原理：利用设备内部确定性有限自动机，记录（Record）关键驱动行为并重放（Replay），以复现状态转移。基于两点核心观察：
- 设备通过软硬件接口（如中断、DMA和IO寄存器读写）与驱动交互，视作驱使GPU状态转移的关键事件
- 执行请求时，设备的状态转移与IO数据的内容无关
GPUReplay 分为两步：第一步称为 Record，预先记录GPU在完整的软件栈上执行计算的过程，得到 Recording；第二步称为 Replay，利用 recording 和新的输入数据，执行GPU计算。
GPUReplay 绕过了GPU软件栈，使GPU能够通过 recording 及应用提供的数据直接进行运算，大大减小了GPU暴露在外的攻击面，从而降低了GPU计算的安全威胁。
原作者通过向 OP-TEE 新增驱动的方式，在不大幅修改原有代码的情况下，实现了基于 Mali Bifrost GPU 的 replayer 在 OP-TEE 的移植。不使用 GPUReplay 的应用将在普通世界的GPU软件栈上运行，使用 GPUReplay 的应用将在安全世界借助 replayer 运行。这证明了 GPUReplay 方法的可移植性，为我们结合 TEE 与 GPUReplay 以实现可信GPU计算框架提供了思路。
我们认为，同 TEE 结合的 GPUReplay 方法能够有效地完成本项目的预期目标。
- GPUReplay 能够通过让CPU免于执行GPU软件栈的方式，扩展 CPU TEE 在GPU计算上的安全性。
- GPUReplay 能够通过重放GPU执行记录的方式，减少AI计算内容对外界的暴露程度。
虽然 GPUReplay 方法并不局限于某一类GPU，但为了记录GPU的行动并对GPU驱动进行插桩，必须知道GPU的接口信息并获取GPU驱动的代码。最终，我们选用了 Raspberry Pi 5（简称 RPi 5）作为我们的硬件目标环境。RPi 5 所使用的 VideoCore 7 GPU 在Linux系统有开源的 V3D 驱动，且 RPi 5 本身有较为完善的文档支持与庞大的开源社区，是比较理想的开发目标。详见 docs/environments.md。
由于 RPi 系列机器使用的内存均为DRAM，无法隔离安全世界与非安全世界，因此 RPi 系列机器并不能真正实现安全的可信执行环境。不过，由于 RPi 系列机器的易用性和广泛性，我们认为用 RPi 5 进行概念验证仍然是不错的选择。
目前，网络上还没有将 OP-TEE OS 移植到 RPi 5 的相关资料。我们参考了一篇将OP-TEE移植到RPi4的资料，将 OP-TEE OS 手动移植到了 RPi 5。详见 docs/optee-os.md
RPlayer 是本项目的主体，以 RPi 5 为平台，运用 GPUReplay 方法，分为 recorder 和 replayer 两部分。其中，recorder 是对原 V3D 驱动程序插桩后得到的新驱动程序；replayer 是 OP-TEE 应用，在 OP-TEE 的可信支撑下重放GPU计算，包含运行在 Linux 之上的 Client Application(CA) 及由 OP-TEE OS 提供防护的 Trusted Application(TA) 。
GPUReplay 方法中，实现 record（记录）的程序称为 Recorder。recorder 需要记录GPU的各项行动，包括存取寄存器、设置GPU页表、设置GPU内存映射、在CPU与GPU间传输数据、等待中断等。此外，recorder 还需要定位 recording 所涉及的输入与输出数据。由于 replay 减少了CPU运行的负荷，recorder 还需要同步CPU与GPU的行动。
本项目 recorder 是经过插桩的GPU驱动程序，运行在内核态。recorder 通过插桩，记录了GPU寄存器存取情况与GPU用到的内存数据，最终产生用于 replay 的 recording，即 GPU 活动记录。详见 docs/recorder.md。
此外，本项目的 recorder 集成了在 Linux 内核态下的 replay 功能，可用于测试重放的基本功能。
GPUReplay 方法中，实现 replay（重放）的程序称为 Replayer。replayer 提供了如下API：（1）获取与释放GPU；（2）加载 recording，验证 recording 的安全性，分配所需的GPU内存；（3）replay主逻辑，根据 recording 及应用提供的输入与输出缓存，执行GPU运算。
完整的 replayer 程序分为三个部分，分别是（1）静态验证器，用于验证 recording 的安全性；（2）解释器，用于整理和执行 recording；（3）模拟GPU硬件的微型驱动。replayer 既可以是用户态程序（1与2）与内核模块（3）的混合，也可以是一整个内核模块，整合了以上三个部分。
本项目分别包含内核模块形式的 replayer 及 OP-TEE 应用形式的 replayer。前者经测试已能完整复现GPU计算结果，后者包含 CA 及 以 OPTEE OS 接口形式实现的 Pseudo TA。详见 docs/replayer.md。
就计算框架而言，我们分别考察了 NCNN、PyTorch、ExecuTorch 等框架在 GPUReplay 方法下的适用性，还能提取 NCNN 的重放输出。详见docs/workloads.md。
目前，OP-TEE 仅对 RPi 3 有较为全面的支持，对 RPi 4 及 RPi 5 均无官方支持。RPi 3 并没有供 Mesa 等 GPU 运行时使用的 drm/v3d 驱动，所以无法适配当下的绝大部分 GPU 计算框架。我们最终选用了计算能力更强、支持 Mesa 的 RPi 5 作为主要开发平台，因此我们需要手动实现 OP-TEE 对 RPi 5 的适配。
在适配过程中，我们遇到了一些技术问题。首先，我们发现 OP-TEE OS 在 RPi 5 设备的初始化并不稳定，即使先前成功初始化，在没有修改任何配置的情况下，后续也可能无法完成初始化。经过 JTAG debug 后，我们发现 RPi 5 设备串口没有正常输出 OP-TEE 日志，因此陷入了无限循环。我们修改了冲刷函数，解决了该问题。
此外，我们还发现 OP-TEE OS 的内存对 GPU 计算而言是严重的性能瓶颈。为此，我们修改了 OP-TEE OS 的相关配置，使其拥有 128 MB 的内存。
在设备驱动之外，本项目主要考虑的 GPU 负载分为三部分：GPU 运行时、计算框架、AI 应用。如果设备不支持某一计算框架依赖的 GPU 运行时，就不能支持该计算框架及其上层的 AI 应用。本项目选用的 RPi 5 只能支持一小部分 GPU 运行时。在目前主流的运行时当中，RPi 5 仅支持 Vulkan 及其 Mesa 实现。这对我们考察计算框架带来了不小的挑战。
最终，我们在 RPi 5 上考察了 NCNN、PyTorch 和 ExecuTorch 的 GPU 计算能力。其中，NCNN 对 Vulkan 的适配最为完善；PyTorch 的 Vulkan 分支已被废弃；而 ExecuTorch 虽已能够支持包括 Llama 在内的大语言模型的 GPU 加速推理，但其支持程度仍不充分。
为了实现更具泛化性的 GPU 可信计算，本项目还需要支持动态的计算输入与输出。我们对 NCNN SqueezeNet 的计算过程进行了深入分析，发现 NCNN 提交给 GPU 的指令在不同图片输入下保持不变。基于这一特性，我们无需手动修改 GPU 计算记录即可在重放过程中提交不同的图片输入。
在输出方面，我们进一步分析了 NCNN 的内存结果，发现其存储图片标签的位置相对固定，因此能够较为直接地提取计算输出。
此外，我们还发现设备在重放过程中无法正确等待部分提交作业的中断。进一步分析后发现，这些中断对应的作业具有异常高的执行次数。将这些作业的执行次数调低后，重放端的输出结果并未发生变化，仍与记录端保持一致。
存放了 RPlayer 的源代码，包括：
- recorder：存放内核驱动形式的 recorder 及 parser。该驱动也包含了 replayer 的内核模块形式实现。
- replayer：存放 OP-TEE 应用形式的 replayer，包含 CA(host文件夹) 及 TA(ta文件夹) 两部分。
- optee：存放 OP-TEE 移植到 RPi 5 所需的一些文件。
存放 RPlayer 测试所用的部分 AI 应用，包括 SqueezeNet、MobileNet、NanoDet 及 Llama。其中，SqueezeNet 有较为完整的测试，包含以下文件夹：
- record：包含GPU计算记录的输入、日志输出、内存输出及示例程序。
- replay：包含重放后得到的内存，及辅助查看 record 结果的应用。
存放项目文档，包括：
- environment.md，对本项目软硬件环境的介绍。
- optee-os.md，对本项目所需 OP-TEE OS 的介绍。
- progress.md，记录了项目时间线及实现项目过程中遇到的困难。
- proj325.pptx，幻灯片。
- proj325-tiny.pptx，幻灯片，供演示视频用。
- project.md，原题目的描述。
- recorder.md，对本项目 recorder 的介绍。
- related-work.md，对其他相关方法的介绍。
- replayer.md，对本项目 replayer 的介绍。
- v3d.md，对原版 V3D 驱动的介绍。
- workloads.md，对本项目所涉及的GPU运行时及AI框架的介绍。
黄家俊负责 OP-TEE OS 及 recorder 相关；杨南负责 replayer 相关。
郭力维老师为我们确认了本次比赛的选题，并联系到了 GPUReplay 论文的原作者 Heejin Park 与 Felix Xiaozhu Lin。两位原作者提供了极大的帮助。
本次比赛对我们而言是绝佳的锻炼机会。项目可能很小，但已经能让我们对系统方向的研究有切身的体会了。就项目本身而言，我们也算是完整地考察了 GPU 计算的运作机制。
[1] J. Zhu et al., "Enabling Rack-scale Confidential Computing using Heterogeneous Trusted Execution Environment”
[2] Y. Deng et al., “StrongBox: A GPU TEE on Arm Endpoints”
[3] H. Yan et al., “LDR: Secure and Efficient Linux Driver Runtime for Embedded TEE Systems”
[4] H. Park et al., “GPUReplay: a 50-KB GPU stack for client ML”
[5] J. Wang et al., “RT-TEE: Real-time System Availability for Cyber-physical Systems using ARM TrustZone”
[6] Tencent, https://github.com/Tencent/ncnn
[7] Meta, https://github.com/pytorch/executorch
[8] F. N. Landola et al., “SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size”
[9] A. G. Howard et al., “MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications”
[10] Meta, https://github.com/meta-llama/llama