Abstract:Heterogeneous collaborative computing with NPU and CPU has received widespread attention due to its substantial performance benefits. To ensure data confidentiality and integrity during computing, Trusted Execution Environments (TEE) is considered a promising solution because of its comparatively lower overhead. However, existing heterogeneous TEE designs are inefficient for collaborative computing due to fine and different memory granularities between CPU and NPU. 1) The cacheline granularity of CPU TEE intensifies memory pressure due to its extra memory access, and 2) the cacheline granularity MAC of NPU escalates the pressure on the limited memory storage. 3) Data transfer across heterogeneous enclaves relies on the transit of non-secure regions, resulting in cumbersome re-encryption and scheduling. To address these issues, we propose TensorTEE, a unified tensor-granularity heterogeneous TEE for efficient secure collaborative tensor computing. First, we virtually support tensor granularity in CPU TEE to eliminate the off-chip metadata access by detecting and maintaining tensor structures on-chip. Second, we propose tensor-granularity MAC management with predictive execution to avoid computational stalls while eliminating off-chip MAC storage and access. Moreover, based on the unified granularity, we enable direct data transfer without re-encryption and scheduling dilemmas. Our evaluation is built on enhanced Gem5 and a cycle-accurate NPU simulator. 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, offering a practical security assurance for LLM training.

What problem does this paper attempt to address?

### Problems the Paper Aims to Solve The paper aims to address the inefficiency of heterogeneous Trusted Execution Environments (TEEs) in collaborative computing. Specifically, the existing heterogeneous TEE designs incur significant performance overhead due to different memory protection granularities between the CPU and NPU. These issues include: 1. **Cache Line Granularity of CPU TEE**: The cache line granularity of CPU TEE adds extra memory accesses, especially in memory-intensive workloads, which significantly increases memory pressure. 2. **Cache Line Granularity MAC Management of NPU TEE**: Although NPU TEE manages version numbers (VN) at tensor granularity, its cache line granularity MAC management introduces additional storage overhead, further exacerbating the pressure on limited memory storage. 3. **Data Transfer Across Heterogeneous Enclaves**: Due to different TEE granularities, data transfer between heterogeneous enclaves requires passing through non-secure areas, leading to cumbersome re-encryption and scheduling, which adds extra memory accesses and affects parallel execution. To address these issues, the paper proposes TensorTEE, a unified tensor-granularity heterogeneous TEE architecture to achieve efficient and secure collaborative tensor computation. Specific improvements include: 1. **Supporting Tensor Granularity in CPU TEE**: By virtually supporting tensor granularity, it eliminates off-chip metadata access and improves efficiency by detecting and maintaining on-chip tensor structures. 2. **Proposing Tensor Granularity MAC Management in NPU TEE**: By reducing storage overhead through delayed verification, avoiding computation stalls, and ensuring data integrity through tensor poisoning tracking and verification barrier mechanisms. 3. **Direct Data Transfer Protocol**: Based on unified granularity, it enables direct data transfer, avoiding transmission through non-secure areas, and reducing the overhead of re-encryption and scheduling. Through these improvements, TensorTEE significantly enhances performance in tasks such as large language model (LLM) training while maintaining low overhead.

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

Enabling Rack-scale Confidential Computing using Heterogeneous Trusted Execution Environment

Enabling Privacy-Preserving, Compute- and Data-Intensive Computing using Heterogeneous Trusted Execution Environment

Dtee: A Declarative Approach to Secure IoT Applications Using TrustZone

Confidential Computing on nVIDIA Hopper GPUs: A Performance Benchmark Study

Confidential Computing on nVIDIA H100 GPU: A Performance Benchmark Study

Empowering Data Centers for Next Generation Trusted Computing

TEEp: Supporting Secure Parallel Processing in ARM TrustZone.

EN-T: Optimizing Tensor Computing Engines Performance via Encoder-Based Methodology

CRONUS: Fault-isolated, Secure and High-performance Heterogeneous Computing for Trusted Execution Environment

Unified Enclave Abstraction and Secure Enclave Migration on Heterogeneous Security Architectures

Privacy-Preserving DNN Training with Prefetched Meta-Keys on Heterogeneous Neural Network Accelerators

Toward Scalable Fully Homomorphic Encryption Through Light Trusted Computing Assistance

Fastrack: Fast IO for Secure ML using GPU TEEs

Composite Enclaves: Towards Disaggregated Trusted Execution

Snpu: Trusted Execution Environments on Integrated NPUs

Efficient Processing of Sparse Tensor Decomposition via Unified Abstraction and PE-Interactive Architecture

Ascend-CC: Confidential Computing on Heterogeneous NPU for Emerging Generative AI Workloads

TSTC: Two-Level Sparsity Tensor Core Enabling Both Algorithm Flexibility and Hardware Efficiency

Hybrid Trust Multi-party Computation with Trusted Execution Environment.