  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.
