Adapting AlphaEvolve to Optimize Fully Homomorphic Encryption on TPUs

2026-05-14

AI Summary

This work addresses the challenge of scaling fully homomorphic encryption (FHE) due to its high computational overhead and the lack of efficient compiler optimizations tailored for specialized hardware such as TPUs. It introduces, for the first time, an LLM-driven automated evolutionary search framework—AlphaEvolve—for implementation-level FHE optimization. The approach leverages large language models to generate optimized code, forming a closed-loop system that incorporates execution feedback from TPUv5e hardware and correctness verification to co-optimize vector register scheduling and resource allocation across MXUs and VPUs. Within 24 hours, the method autonomously discovers implementations surpassing hand-tuned optimizations, achieving significant performance gains: a 2.5× speedup in TFHE bootstrapping latency, and 1.31× and 1.18× acceleration for CKKS rotation and multiplication operations, respectively, thereby enabling efficient co-design across cryptography, compilers, and hardware.

Abstract

The deployment of Fully Homomorphic Encryption (FHE) at scale is hindered due to its heavy computational overhead. While specialized hardware accelerators like Google Tensor Processing Units (TPUs) can help, mapping complex cryptographic kernels onto such architectures remains a challenge. Efficient execution requires co-optimization between the systolic array-based Matrix Multiplication Unit (MXU) and Vector Processing Units (VPUs), as well as the orchestration of data movement across the vector register files. Existing compiler stacks often abstract low-level hardware utilization, requiring developers to adopt a manual trial-and-error process that often results in fragmented execution and underutilized resources. To accelerate this development process, we use AlphaEvolve to automate the exploration of hardware-aware cryptographic-kernel optimizations. We frame optimization as an evolutionary search problem, utilizing the closed-loop system provided by AlphaEvolve, that leverages LLM-driven code generation. We use real-world feedback from hardware execution and rigorous correctness testing to guide the evolution process. We evaluate AlphaEvolve optimization on primitives for both the TFHE (Jaxite) and CKKS (CROSS) FHE schemes on Google Cloud TPUv5e, a contemporary TPU architecture. Within 24 hours of automated exploration, AlphaEvolve discovered implementation-level optimizations that improve TFHE bootstrap latency by 2.5x and CKKS rotation and multiplication latency by 1.31x and 1.18x, respectively, relative to human-engineered state of the art. These results demonstrate that AlphaEvolve can be used to enable researchers to navigate the optimization trade-offs between cryptography, compilers, and hardware accelerators.

Problem

• Fully Homomorphic Encryption
• Hardware Acceleration
• TPU
• Compiler Optimization
• Cryptographic Kernels

Innovation

• AlphaEvolve
• Fully Homomorphic Encryption
• TPU
• evolutionary search
• hardware-aware optimization