TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing

异构协作计算是指利用NPU和CPU的优势进行计算，因其性能优势而受到广泛关注。为确保计算过程中的数据机密性和完整性，可信执行环境（TEE）被认为是一种有前途的解决方案，因为其开销相对较低。然而，现有的异构TEE设计对于协作计算来说效率较低，原因在于CPU和NPU之间存在细微和不同的内存粒度。首先，CPU TEE的缓存行粒度增加了内存压力，因为需要额外的内存访问；其次，NPU的缓存行粒度MAC加剧了有限内存存储的压力；第三，跨异构enclave的数据传输依赖于非安全区域的传输，导致繁琐的重新加密和调度。为解决这些问题，我们提出了TensorTEE，一个统一的张量粒度异构TEE，用于高效的安全协作张量计算。首先，我们在CPU TEE中虚拟支持张量粒度，通过检测和维护芯片上的张量结构来消除芯片外的元数据访问。其次，我们提出了张量粒度MAC管理和预测执行，以避免计算停顿，同时消除芯片外MAC存储和访问。此外，基于统一粒度，我们实现了直接数据传输，避免了重新加密和调度的困境。我们的评估基于增强版Gem5和一个精确到周期的NPU模拟器。结果表明，与现有工作相比，TensorTEE将大型语言模型（LLM）训练工作负载的性能提高了4.0倍，并且与非安全训练相比仅产生2.1％的开销，为LLM训练提供了实际的安全保证。

TensorTEE: A Unified Tensor-Granularity Heterogeneous Trusted Execution Environment for Efficient Secure Collaborative Tensor Computing

The paper proposes a unified tensor-granularity heterogeneous Trusted Execution Environment (TEE) called TensorTEE, which supports efficient secure collaborative tensor computing by eliminating off-chip metadata access, avoiding computational stalls, and enabling direct data transfer without re-encryption and scheduling dilemmas.

TensorTEE virtually supports tensor granularity in CPU TEE to eliminate the off-chip metadata access by detecting and maintaining tensor structures on-chip. Tensor-granularity MAC management with predictive execution is proposed to avoid computational stalls while eliminating off-chip MAC storage and access. Direct data transfer is enabled based on the unified granularity without re-encryption and scheduling dilemmas. The evaluation is built on enhanced Gem5 and a cycle-accurate NPU simulator, and the results show that TensorTEE improves the performance of Large Language Model (LLM) training workloads by 4.0x compared to existing work and incurs only 2.1% overhead compared to non-secure training.

Related work includes 'SecureML: A System for Scalable Privacy-Preserving Machine Learning' and 'TEEchain: Enabling Trustworthy and Private Blockchain via Trusted Execution Environment'.