Adapting AlphaEvolve to Optimize Fully Homomorphic Encryption on TPUs

Technical Summary: Adapting AlphaEvolve to Optimize Fully Homomorphic Encryption on TPUs

Problem Statement
The scalable deployment of Fully Homomorphic Encryption (FHE) is currently hindered by severe computational and memory overheads. While specialized hardware accelerators like Google Tensor Processing Units (TPUs) offer a path forward, mapping complex cryptographic kernels to these architectures remains a significant engineering challenge. Efficient execution on TPUs requires precise co-optimization between the systolic array-based Matrix Multiplication Unit (MXU) and Vector Processing Units (VPUs), alongside intricate data movement orchestration across vector register files.

Current compiler stacks (e.g., XLA) abstract low-level hardware details, forcing developers into a slow, manual trial-and-error loop to tune high-level code (JAX/Pallas) for hardware impact. This process often results in fragmented execution and underutilized resources. Furthermore, the application of AI to cryptographic optimization has been largely unexplored due to concerns that Large Language Model (LLM) hallucinations could introduce security vulnerabilities or side-channel attacks.

Methodology
The authors adapt AlphaEvolve, an agentic framework combining LLMs with evolutionary search, to automate the discovery of hardware-aware optimizations for FHE kernels. The system operates as a closed-loop evolutionary search problem with the following components:

1. Evolutionary Framework: The system samples a database of parent programs and prompts to generate candidate code modifications via an LLM. These candidates are compiled and evaluated.
2. Hardware-in-the-Loop Evaluation: Candidates are executed on physical Google Cloud TPUv5e hardware. The system uses XProf to capture high-fidelity execution traces, measuring compute/memory bandwidth utilization and storage occupancy.
3. Rigorous Correctness and Security Checks: To mitigate LLM hallucinations and ensure cryptographic integrity, the system enforces multi-tiered validation:
   • Correctness: Randomized input vectors are tested across computational sub-layers, including end-to-end encryption, homomorphic computation, and decryption against expected plaintexts.
   • Security: Modifications are validated against industry-standard lattice estimators to ensure parameters maintain at least 128-bit security. Additional end-to-end tests with high-depth computations prevent decryption failures.
   • Reward Function: The primary reward is the microsecond-level runtime latency of the cryptographic primitive, excluding compilation overhead. Candidates failing correctness or security checks are rejected.
4. Target Schemes: The framework optimizes primitives for two major FHE schemes:
   • TFHE (Jaxite): Specifically the blind_rotate operation used in bootstrapping.
   • CKKS (CROSS): Specifically slot rotation (CKKS-Rot) and slot-wise multiplication (CKKS-Mult).

Key Contributions
The paper presents the first application of agentic AI to optimize FHE performance on hardware accelerators. Key technical contributions include:

• Formulation: Defining JAX/Pallas-level FHE kernel optimization as an evolutionary search problem guided by real-world hardware feedback.
• Validation Engine: Designing a hardware-in-the-loop evaluator that scores candidates based on TPU latency while strictly rejecting code that fails functional or security tests.
• Discovery of Optimization Patterns: Identifying specific, non-intuitive optimization patterns that outperform human-engineered baselines, including:
  • Loop unrolling to reuse common parameters within on-chip memory.
  • Fine-grained scheduling to hide off-chip memory latency behind computation.
  • Elimination of redundant datatype conversions (e.g., removing unnecessary bitcasts).
  • XLA-favorable tiling strategies that improve Vector Register (VReg) utilization.

Results
Within 24 hours of automated exploration on TPUv5e, AlphaEvolve achieved significant latency reductions compared to state-of-the-art (SoTA) human-engineered implementations:

• TFHE Bootstrap: Latency improved by 2.5× (from 10 ms to 4 ms). The blind_rotate primitive specifically saw a 2.85× speedup (9.4 ms to 3.5 ms), driven by loop unrolling, fine-grained scheduling, and the removal of redundant type casts.
• CKKS Rotation: Latency improved by 1.31× (from 4854 µs to 3714 µs) through tensor partitioning that triggered specialized XLA optimizations, reducing the AUTOMORPHISM stage's contribution to total latency from 12% to 4%.
• CKKS Multiplication: Latency improved by 1.18× (from 6340 µs to 5359 µs) via scheduling optimizations for parameter reuse.

The authors note that the most effective optimizations occurred when the search addressed mismatches between cryptographic data layouts and the TPU compiler/hardware execution grain, often exposing "sweet spots" in the compiler that are difficult for humans to discover manually.

Significance and Claims
The paper claims that AlphaEvolve enables researchers to navigate complex trade-offs between cryptography, compilers, and hardware accelerators. It demonstrates that AI-driven search can autonomously discover implementation-level optimizations that domain experts frequently miss.

The authors emphasize that this work is a "first application" of agentic AI to cryptographic performance. They maintain a modest stance regarding the scope, noting that the system currently optimizes existing kernels rather than discovering new cryptographic algorithms. The paper highlights that while the framework successfully improves performance, the bottleneck of manual code review remains, and future work is needed to integrate functional equivalence checkers to fully automate the validation pipeline. Additionally, the authors identify a key software bottleneck: existing programmable APIs (JAX/Pallas) often force tensor reshapes and layout reorganizations that introduce avoidable overhead, a limitation that automated search can help expose and mitigate.