用AI芯片加速同态加密，让隐私计算更省电。
Leveraging ASIC AI Chips for Homomorphic Encryption
- 将高精度加密运算转为低精度矩阵计算，适配TPU架构
- 通过预处理重排数据，避免运行时耗能的内存操作
- 在谷歌TPU上实现比现有方案更高的能效，适合隐私计算场景
同态加密（HE）可保障云服务中的数据隐私，但计算开销巨大。尽管GPU已成加速HE的实用平台，与专用但昂贵的HE ASIC相比仍存在一个数量级的能量效率差距。本文探索新路径：利用现有AI加速器（如谷歌TPU）的粗粒度计算与内存结构，实现接近ASIC级别的能效。然而，当前主流HE算法针对GPU设计，与TPU架构存在根本冲突：(1) 高精度（32位）整数运算无法发挥TPU低延迟向量单元作用，导致其高吞吐的8位矩阵单元（MXU）闲置；(2) 细粒度数据重排在粗粒度内存子系统中效率低下。直接移植GPU优化的HE库会导致资源严重浪费和性能下降。为此，我们提出CROSS编译框架，系统性重构HE工作负载以匹配TPU架构。主要贡献包括：(1) 基础对齐变换（BAT），将高精度模运算转换为密集的8位整数矩阵乘法，充分激活并提升TPU MXU利用率；(2) 内存对齐变换（MAT），通过离线参数变换将重排嵌入计算核中，消除运行时昂贵的数据重排。CROSS（TPU v6e）在NTT和HE算子上的每瓦吞吐率优于WarpDrive、FIDESlib、FAB、HEAP和Cheddar，确立了AI ASIC作为当前最高效的HE平台。
原文摘要 · Abstract (English)
Homomorphic Encryption (HE) provides strong data privacy for cloud services but at the cost of prohibitive computational overhead. While GPUs have emerged as a practical platform for accelerating HE, there remains an order-of-magnitude energy-efficiency gap compared to specialized (but expensive) HE ASICs. This paper explores an alternate direction: leveraging existing AI accelerators, like Google's TPUs with coarse-grained compute and memory architectures, to offer a path toward ASIC-level energy efficiency for HE. However, this architectural paradigm creates a fundamental mismatch with SoTA HE algorithms designed for GPUs. These algorithms rely heavily on: (1) high-precision (32-bit) integer arithmetic to now run on a TPU's low-throughput vector unit, leaving its high-throughput low-precision (8-bit) matrix engine (MXU) idle, and (2) fine-grained data permutations that are inefficient on the TPU's coarse-grained memory subsystem. Consequently, porting GPU-optimized HE libraries to TPUs results in severe resource under-utilization and performance degradation. To tackle above challenges, we introduce CROSS, a compiler framework that systematically transforms HE workloads to align with the TPU's architecture. CROSS makes two key contributions: (1) Basis-Aligned Transformation (BAT), a novel technique that converts high-precision modular arithmetic into dense, low-precision (INT8) matrix multiplications, unlocking and improving the utilization of TPU's MXU for HE, and (2) Memory-Aligned Transformation (MAT), which eliminates costly runtime data reordering by embedding reordering into compute kernels through offline parameter transformation. CROSS (TPU v6e) achieves higher throughput per watt on NTT and HE operators than WarpDrive, FIDESlib, FAB, HEAP, and Cheddar, establishing AI ASIC as the SotA efficient platform for HE operators. Code: https://github.com/EfficientPPML/CROSS
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。