🤖 AI Summary
This work addresses the inefficiency of AI ASICs—such as TPUs—optimized for low-precision computation when executing high-precision number-theoretic transforms (NTTs) required by fully homomorphic encryption (FHE). Existing approaches suffer from precision mismatches that force NTT output reconstruction off the matrix engine’s dataflow, creating a performance bottleneck. To overcome this, the paper proposes a lightweight multi-precision systolic array architecture that, for the first time, seamlessly integrates low-precision matrix multiplication and high-precision NTT reconstruction within a unified dataflow, eliminating the need for vector processor intervention and avoiding dataflow disruption. Evaluated via matrix decomposition, shift-and-add reconstruction, 7nm OpenRoad synthesis, and SCALE-Sim simulation on a 128×128 matrix engine, the design achieves at least 1.33× speedup for NTT sizes ranging from 2¹² to 2¹⁶ with negligible hardware overhead, significantly enhancing FHE execution efficiency on standard AI ASICs.
📝 Abstract
Fully Homomorphic Encryption (FHE) ensures robust data privacy but suffers from prohibitive computational overhead. Accelerating FHE on AI hardware like Tensor Processing Units (TPUs) is promising, yet fundamentally limited by a precision mismatch: TPUs are optimized for 8-bit arithmetic, whereas FHE and its critical parts such as the Number Theoretic Transform (NTT), demand high precision. Current approaches bridge this gap using matrix decomposition to execute NTT computations on low-precision matrix engines. However, reconstructing the full-precision results requires shift-and-add accumulation that does not match the dataflow of matrix multiplication. This forces offloading full-precision reconstruction from matrix engines to vector processors that disrupts the matrix multiplication dataflow, creating significant performance bottleneck. To resolve this limitation, we propose a minimally modified multi-precision systolic array that performs full-precision output reconstruction natively within the array in sync with low-precision matrix multiplication under a uniform dataflow. Synthesized at 7nm with OpenRoad, our design incurs negligible hardware overhead. Cycle-accurate simulations using SCALE-Sim demonstrate that natively executing NTTs on the proposed architecture achieves at least 1.33x speedup, for transform sizes 2^12 to 2^16 on 128x128 matrix engines, successfully enabling standard AI hardware to support high-precision FHE acceleration.