Adapting AlphaEvolve to Optimize Fully Homomorphic Encryption on TPUs Shruthi Gorantala, Jianming Tong† , Asra Ali, Baiyu Li, Jonathan Katz, Jeremy Kun Thomas Steinke‡ , Abhradeep Thakurta‡ , Julian Walker‡ , Amir Yazdanbakhsh‡
arXiv:2605.14718v1 [cs.CR] 14 May 2026
† Georgia Institute of Technology
‡ Google DeepMind
becomes a slow trial-and-error tuning loop where high-level code modification must traverse the compilation pipeline to observe its hardware impact. Current systems lack a systematic approach to map transformed cryptographic primitives onto commodity hardware designed for low-precision, GEMM-heavy workloads. AI coding agents such as AlphaEvolve [16] have demonstrated significant advances in multiple fields including algorithm design, improving code efficiency, architecture discovery, and enabling discovery across scientific domains. Further, the use of AI to improve cryptographic implementation remains largely unexplored. This is primarily because LLM hallucinations might introduce security vulnerabilities which are unacceptable in cryptographic implementations. In this paper, we use AlphaEvolve, an agentic framework, to optimize FHE kernels on hardware accelerators, such as Google’s TPU, with rigorous correctness checking. We formulate kernel optimization as an evolutionary search problem using AlphaEvolve, with real-world hardware feedback. By automating the discovery of micro-architectural improvements, AlphaEvolve enables cryptographic developers to explore architectural ideas and co-optimize cryptographic primitives on specialized hardware. The main contributions of our work are: • We formulate JAX/Pallas-level FHE-on-TPU kernel optimizations as an evolutionary search problem. • We design a hardware-in-the-loop evaluator that scores candidates using real TPU latency while rejecting candidates that fail robust functional correctness tests. • We identify concrete optimization patterns discovered by the search that outperform domain-expert baselines: loop unrolling to reuse common parameters, finer-grained scheduling to hide off-chip memory access latency, minimize datatype conversion, and XLA-favorable tiling that improves VReg utilization. • Within 24 hours of exploration, AlphaEvolve achieved a 2.5× reduction in TFHE bootstrapping latency and speedups of 1.31× and 1.18× for CKKS rotation and multiplication, respectively.
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 arraybased 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 cryptographickernel 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 implementationlevel optimizations that improve TFHE bootstrap latency by 2.5× and CKKS rotation and multiplication latency by 1.31× and 1.18×, 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.
I. I NTRODUCTION The rapid proliferation of personal data in AI ecosystems has created a demand for privacy-enhancing technologies (PETs). Fully Homomorphic Encryption (FHE) [9] stands as a foundational Privacy Enhancing Technology (PET) by enabling arbitrary computations on encrypted data without requiring decryption. Despite its privacy guarantees, Fully Homomorphic Encryption (FHE) incurs orders-of-magnitude compute and memory overheads, severely hindering its scalable deployment [10]. To mitigate these bottlenecks, the community has increasingly pivoted toward specialized accelerators like Tensor Processing Units (TPUs). Mapping FHE primitives to accelerators requires engineering optimizations that spans across tiling workloads to strictly align with the underlying vector register granularity, navigating frontend abstractions (e.g., JAX and Pallas APIs for TPUs) to bypass costly, implicit layout transformations and sizing tensor granularities to effectively reduce off-chip memory latency. Modern compiler stacks like XLA [17] abstract these microarchitectural details. However performance tuning then
II. BACKGROUND AND M OTIVATION A. AlphaEvolve AlphaEvolve [16] is an agentic framework designed for algorithmic and scientific discovery by combining LLMs with
1
Algorithm: ToeplitzMul(a, b) Input: a = (a0 , . . . , an−1 ), b = (b0 , . . . , bn−1 ) Construct a Toeplitz matrix T (a) ∈ Zn×n : a0 a1 · · · an−1 −an−1 a0 · · · an−2 T (a) = . .. .. .. .. . . . −a1 −a2 · · · a0
the evolutionary search. It could be used to automatically evolve code in any programming language through a flywheel of code generation and code evaluation. The workflow begins with a human programmer configuring the system based on expert-guided, task specific instructions, an initial code example to seed the evolutionary database and an evaluator that scores the generated code. Next, the system samples code database and a pool of prompts. These prompts are fed into LLMs to generate new code modification. The generated code blocks are then compiled and fed into the evaluators, where they are scored, ranked, and only top-rank correct programs are added back to the evolutionary database. The ranking and evolution employs structured evolutionary techniques, such as island-based models or MAP-Elites algorithms [16]. These techniques balance genetic diversity based on inherent metrics (e.g., program length, code complexity) and performance on the evaluator metrics (fitness score). This optimization loop improves the code to find better program based on the fitness score.
// Note: Negacyclic entries account for reduction modulo X n + 1. Compute the matrix-vector product c = T (a) · b ∈ Znq . i return c // Represents the polynomial c(X) = ∑n−1 i=0 ci X .
Fig. 1: Polynomial multiplication a(X) · b(X) ∈ RQ = ZQ [X]/(X n + 1) via Toeplitz matrix-vector product, where a and b are coefficient vectors of a(X) and b(X), respectively.
8 sublanes, and each sublane provides 32 32-bit registers. Registers are organized into (8, 128) blocks, called vector registers (VRegs), each formed by selecting one register from every sublane across all 128 lanes. The XLU performs data reorganization, while the VPU and MXU accelerate vector and matrix operations. Both computation and data movement execute in a single-instruction, multiple-data (SIMD) manner over one (8, 128) VReg at a time. Consequently, programs that fully utilize VRegs achieve higher efficiency, whereas operations requiring fine-grained or irregular data access incur significant overhead. For FHE, the implication is direct: mathematically equivalent implementations can differ sharply in performance depending on shapes, layout and scheduling choices of JAX/Pallas tensors.
B. Fully Homomorphic Encryption FHE [9] enables computation on encrypted data. A homomorphic operation transforms ciphertexts so that the decrypted results are the same as if the operation were performed on the underlying plaintext data. Existing FHE schemes can be broadly categorized into scalar schemes and vector schemes based on the type of data they encrypt. Scalar schemes such as TFHE [5] encrypt bits or short integers, and supports general-purpose computation on encrypted data. Vector schemes such as CKKS [4] pack a vector of real or complex numbers into a ciphertext, and are widely used for privacy-preserving machine-learning and numerical workloads as vectors of data are processed in a SIMD fashion. Both approaches lead to 2∼5 order-ofmagnitude memory and computation overhead. In this paper, we use AlphaEvolve to optimize implementations of both schemes on TPUs. 1) TFHE - TFHE-bootstrap: This is the key primitive in TFHE used in Google transpiler [11] and HEIR compiler [3]. TFHE bootstrap algorithms utilize the so-called RingGSW scheme defined over a polynomial residual ring RQ = ZQ [X]/(X n + 1) and a special, cryptographic accumulator to perform decryption homomorphically. The primary bottleneck in the bootstrap operation is blind_rotate, which consists of a sequence of multiplication-accumulation operations on polynomials. Note that we reformulate polynomial multiplications in RQ as Toeplitz matrix-vector products to take advantage of MXUs in TPUs. 2) CKKS - CKKS-Rot and CKKS-Mult: They perform homomorphic slot rotation and slot-wise multiplication, respectively, which are the backbone components in almost all CKKS use cases. We optimize their implementations from CROSS [18], combining automorphisms, key switching, rescaling, tensor multiplication, and modular reduction operations.
D. The TPU Compiler Stack We inherit the multi-layer TPU compilation stacks and profiling system to support FHE computations: 1) JAX: expresses high-level numerical and cryptographic computation. Both expert-crafted TPU libraries, Jaxite [2] and CROSS [18], use JAX / Pallas as the frontend. 2) Pallas: serves as a kernel language that allows for the creation of custom, low-level hardware kernels within JAX. Our system relies on Pallas kernels to establish fine-grained control over data movement and memory staging on the TPU [14]. 3) XLA: lowers JAX/Pallas programs to TPU executables and applies optimizations such as fusion and layout selection [12], [17]. XLA compilation passes are shape-sensitive. 4) Xprof: is an integrated profiling tool within the XLA ecosystem that captures high-fidelity hardware execution traces. It extracts metrics such as compute and memory bandwidth utilization as well as on-chip and off-chip storage occupancy. E. Challenge: Compiler-Obscured Performance Engineering TPU performance engineering demands rigorous optimizations: (1) tiling workloads to strictly align with VReg granularities, (2) navigating frontend abstractions to bypass implicit layout transformations, and (3) sizing tensors to effectively hide off-chip memory latency behind compute. Developers are strictly bound to high-level programming interfaces such as JAX and Pallas, limiting the easy access to manipulating the low level schedule. At this level of abstraction,
C. TPU Execution Model Fig. 2 summarizes the TPU features. On-chip vector memory (VMEM) is partitioned across 128 lanes. Each lane contains
2
HBM 10~200 GB/s 10~100 GB PCI-E ~1000 GB/s VMEM ~0.1 GB
Lane 0 0.125~1 MB
TPU
32
Lane 1
Lane 127
Initial Cryptography Program (Jaxite or CROSS)
CPU 1~10 TOPS
Island 1 XLU Layout Transform
Controller (AlphaEvolve) Scaffolding: LLM with Evolutionary Search Feedback: "Execution Trace" program 1 program 2 program 10
"SIMD" Coarse-Grained Control
32-bit regs 128
Compute
VPU SIMD 1~10 TOPS
MXU Systolic Array (MatMul) 100~1000 TOPS
Island 10
Score: "Latency"
~10 TB/s
VReg 4 KB 8
Island 2
Scalar Unit SReg, SMEM For control
SandBox Evaluator
SandBox Evaluator
SandBox Evaluator
TPU
TPU
TPU
Correctness Checking
Correctness Checking
Correctness Checking
Performance Profiling (XProf)
Performance Profiling (XProf)
Slow
Fast
Fig. 2: TPU microarchitecture overview – SIMD machine. every source modification must traverse hundreds of opaque compilation passes within stacks like XLA [17] before finally materializing as a low-level TPU execution trace. Because these intermediate layers obscure critical microarchitectural details, adapting an algorithm for TPU execution devolves into a sluggish trial-and-error process. Developers are forced to experimentally deduce the specific shapes, casts, loops, and tiling parameters at frontend JAX/Pallas that cause XLA to produce a good TPU program. This paper adapts AlphaEvolve to automate this compilersteering loop while preserving correctness and security guarantees. The use of AI to improve cryptographic implementation is challenging as LLM hallucinations may introduce security vulnerabilities prone to side channel attacks. We use optimizing FHE as an initial case study where the data is encrypted.
Parent Programs Prompt Pools
Loop until Manual Stop
Drop
program 2 Filtered fast programs Manual Stop Post-AlphaEvolve Optimized program
Fig. 3: Workflow for adapting AlphaEvolve to cryptographic kernel optimization. AlphaEvolve samples parent programs, prompts an LLM to generate candidate JAX/Pallas implementations, and evaluates each candidate in a sandbox on TPU hardware. Invalid candidates are rejected by correctness and security checks. Valid candidates are scored by latency and profiled through execution traces. Fast candidates are returned to the evolutionary population for further search.
III. A DAPTING A LPHA E VOLVE FOR C RYPTOGRAPHY Configuring AlphaEvolve requires users to specify initial prompts (which establish a persona, algorithm and hardware platform), initial programs (which will be optimized by AlphaEvolve), the AlphaEvolve framework itself and an evaluation strategy (which validates the correctness, and then feeds reward as feedback back to AlphaEvolve). Fig. 3 provides an overview of various parts of the system with details described below:
be used to evolve security parameters and kernels, which allows for tuning security parameters for accelerator-specific kernel super-optimization. C. Evaluation Engine The evaluation engine only rewards candidate kernels that satisfy these multi-tiered correctness and security checks, thereby pruning invalid or insecure code early in the evolutionary search process. 1) Correctness Checks: To mitigate the risks of reward hacking and code hallucination inherent in LLM-generated kernels, we enforce strict correctness tests using randomized input vectors across all computational sub-layers. The correctness checking requires correct compilation to mitigate the risks of unintended bugs and erroneous code. This is done by testing subset of sub-modules as well as incorporating end-to-end tests which initializes the cryptosystem with valid security parameters, generates randomized inputs, encrypts the inputs, optionally runs cryptographic computations on the ciphertexts such as homomorphic computations for FHE cryptosystem, decrypts the resultant ciphertexts and checks against the expected plaintext outputs. 2) Security Checks: Our system validates any modifications to the parameters of the cryptosystem maintain standardized security guarantees (example: at least 128-bit security for
A. Initial Prompts The initial prompts describe the task and role of AlphaEvolve as an expert in hardware-aware kernels engineering. The prompts also include a functional description of the cryptographic algorithms being optimized, architecture of the hardware accelerator, optimization strategies and a pool of potential optimization strategies. B. Initial Program We feed well-optimized reference codes into the evolutionary search framework as the initial programs. The choice of initial program needs a balance between providing enough code for AlphaEvolve to optimize while not allowing it to change the security guarantees provided by the algorithm. We use coevolution as a process to evolve multiple kernels in tandem to find kernel improvements within various aspects of the cryptosystem. In addition to co-evolution, co-optimization can
3
Fully Homomorphic Encryption utilizing acceptable parameters IV. A LPHA E VOLVE FOR OPTIMIZING FHE generated from industry standard lattice estimators [1]). Further- A. Setup more in case of fully homomorphic encryption, we implement 1) Initial programs: For TFHE, we use the bootstrap implesafeguards to prevent the selection of scheme parameters that mentation from jaxite [2] and allow AlphaEvolve to modify could lead to decryption failures by adding additional end-tovarious parts of blind_rotate. For CKKS, we adopt end test with high depth computation. CKKS-Rot and CKKS-Mult from the TPU-based CKKS 3) Target Hardware Latency as Reward: By deploying library, CROSS [18] that offers state-of-the-art throughput, as directly on the target platform, we use real-world latency as the initial programs and apply AlphaEvolve to search for faster the primary reward score. In this paper, candidate kernels in implementations. We expose the implementations of R ESCALE, Jaxite and CROSS are lowered through the XLA compiler T ENSOR M ULTIPLY, K EY S WITCH, A PPROX M OD D OWN, and and executed on physical TPU hardware. Our reward score AUTO M ORPHISM from both kernels to AlphaEvolve. uses microsecond-level execution latency of the resulting 2) Prompt Setup: We use the system prompt together with cryptographic primitive, excluding XLA compilation overhead a prompt sampled uniformly from §A. These prompts are to isolate runtime performance. designed to steer AlphaEvolve towards implementation-level A persistent bottleneck in FHE acceleration is the lack of improvements, especially scheduling and layout choices that parity between isolated primitive speedups and actual end-toincrease VReg utilization in TPU and minimize implicit layout end performance due to the bottlenecks introduced by data transformation, rather than toward new algorithms discovery. movement between host and device and also between devices. 3) Correctness, and Performance Tests: For TFHE, we To ensure primitive kernel speedup translates to end-to-end incorporate end-to-end tests where the cryptosystem is initiated speedup, AlphaEvolve evolves the one or more primitive kernels with parameters offering 128-bit security and random inputs are simultaneously, while using end-to-end latency across the entire encrypted and evaluated against the generated code. The outputs functional module as the score. Measuring end-to-end latency are then decrypted and compared against the expected plaintext proves to be a useful reward metric for co-evolution and covalues. All values of the 3-bit lookup tables are evaluated for optimization techniques. correctness. 4) Feedback: To facilitate the autonomous discovery of For CKKS, functional correctness is enforced using a hidden hardware-optimal execution patterns, our system captures high- test set with fixed ciphertext inputs and expected ciphertext fidelity execution traces during the evaluation cycle using xprof outputs. integrated with XLA. These metrics from profiling including For both schemes, performance is measured from XProf compute and memory bandwidth utilization as well as on-chip traces, and we return the negative kernel latency as the and off-chip storage occupancy are fed back to AlphaEvolve optimization score so that lower latency corresponds to a higher as diagnostic feedback. By analyzing these traces, the agent score. can algorithmically identify latency bottlenecks and instances We run the evolution process with 10 controllers and 10 where the XLA compiler needs additional signal through evaluators per controller, for a total of 100 TPUv5e chips. JAX (or Pallas), enabling targeted iterative refinements to the TABLE I: Latency Comparison to SoTA [2], [18] cryptographic kernels. Workload Blind Rotation (TFHE) + Unrolling + Fine-grained scheduling + Type Cast Removal Post-AlphaEvolve blind Rotation Post-AlphaEvolve bootstrap CKKS-Mult (CKKS) + XLA-favor tiling choice Post-AlphaEvolve CKKS-Mult CKKS-Rot (CKKS) + XLA-favor tiling choice Post-AlphaEvolve CKKS-Rot
D. Manual reviews / Functional equivalence The best program generated by the evolutionary loop can be integrated into the original library after human review. Currently, a human programmer reviews the best program and merges the improved code into the original library. To reduce human-in-the-loop reviews, in principle a module to check functional equivalence of the initial code and the AlphaEvolve generated code could be used, though doing this is currently an unsolved research problem.
Latency Improvement 9.4 ms (TPUv5e-1) 10 → 7.8 ms 7.8 → 6.3 ms 6.3 → 3.5 ms 9.4 ms → 3.5 ms 10 → 4 ms 6340 µs (TPUv5e-1) 5359 µs (TPUv5e-1) 6340 → 5359 µs 4854 µs (TPUv5e-1) 3714 µs (TPUv5e-1) 4854 → 3714 µs
Code/Analysis Code: jaxite [2] Code: [6] Reason: §IV-B1 Code: [7] Reason: §IV-B2 Code: [8] Reason: §IV-B3 2.85× speedup 2.5× speedup Code ↗ [18] Code: ↗ Reason: §IV-C 1.18× speedup Code: ↗ [18] Code: ↗ Reason: §IV-C 1.3× speedup
B. AlphaEvolve for TFHE AlphaEvolve explored several distinct classes of optimizations, ranging from standard loop transformations to complex kernel scheduling and data type optimizations. 1) Loop Unrolling for Parameter Reuse [6]: In the negacyclic vector-matrix polynomial multiplication (negacyclic_vector_matrix_polymul) routine [2], the system independently discovered that unrolling the primary for loop by a factor of 8 yielded significant latency reductions. In the baseline trace, the loop executes in serial, loading the common parameters repeatedly. By unrolling the loop, common parameters get reused across multiple loop
E. Putting it all together After configuring initial prompts, initial programs, evaluation engine, and setting up appropriate hardware environment for the evaluators, AlphaEvolve begins the discovery process of automatically optimizing the cryptographic kernels. After reaching convergence based on the fitness function, the best generated program can be further tested for correctness through human reviews. The final code can then be merged into the original library and deployed for use.
4
Data iterations. Because the target parameter size (3.51 MB) fit Initial Program Post-Evolved Program 2 Register Register within the available on-chip memory, this transformation 51 ... Wasted Wasted 8 8 Fully Used reduced latency from 10 ms to 7.8 ms. This highlights that 65536 128 128 AlphaEvolve could automatically identify and exploit loop unroll optimizations. Fig. 4: Effect of the discovered layout changes. The initial 2) Fine-Grained Memory and Compute Scheduling [7]: program leaves substantial VReg capacity unused because The system optimized data access patterns by adjusting the operations are expressed at a granularity that does not match tiling and scheduling choice of the kernel. Specifically, it tiled VReg. The evolved program exposes fuller VReg utilization. one big JAX tensor into two blocks with shape halved. This modification allowed the XLA compiler to schedule more TPU compiler/hardware execution grain. In TFHE, the search fine-grained data loading from off-chip memory, and hiding discovers both conventional locality improvements and a it behind computation. This optimization reduced the overall non-obvious cast elimination that depends on the ciphertext bootstrapping latency from 7.8 ms to 6.3 ms. decomposition range. In CKKS, the search finds tensor shapes 3) Elimination of Redundant Type Casts [8]: The most that steer XLA toward higher VReg utilization. Across all significant single optimization involved the removal of an ex- cases, the successful edits preserve the correctness and security plicit bitcast operation. The baseline implementation explicitly while improving how the algorithm is presented to the compiler. cast a right-hand side (rhs) operand to 8-bit integers (i8). By using AlphaEvolve to expose the performance headroom AlphaEvolve identified that the dynamic range of the input data in the current TPU deployment flow, we also identify a key (values from 0 to 255 due to RLWE decomposition in [8]) was software bottleneck: existing programmable APIs force tensor already sufficiently captured by the existing bf16 precision. reshapes and layout reorganizations that are not inherent to the By removing the redundant and expensive bitcast operator, the hardware, introducing avoidable overhead. system reduced the execution time from 6.29 ms to 3.35 ms. This optimization requires a deep understanding of both the V. R ELATED W ORK mathematical dynamics of TFHE and the specific hardware execution costs of bitcast operations, making it difficult for To our knowledge this is the first application of AI techhuman engineers to consistently identify. nologies to improve cryptographic performance. The closest work for improving kernel generation is done by autocomp C. AlphaEvolve for CKKS [13] and EvoX [15] where agentic infrastructure has been used • CKKS-Rot: On TPUv5e, AlphaEvolve reduces the lato improve machine learning kernels. tency of CKKS-Rot from 4874 µs to 3754 µs, yielding a 1.31× speedup. This performance gain is primarily VI. C ONCLUSION achieved by partitioning a tensor into two smaller subtensors. This strategic split triggers specialized XLA This paper presents the first application of agentic AI to optimizations that significantly improve Vector Register optimize FHE primitives on TPUs. We detail the design of (VReg) utilization during the AUTO M ORPHISM stage. using AlphaEvolve, a closed-loop system that combines LLMs Consequently, AUTO M ORPHISM’s contribution to the with real-world hardware execution feedback and rigorous total end-to-end latency drops from 12% to just 4%. correctness checks. We demonstrate the framework’s ability to These results demonstrate AlphaEvolve’s ability to expose autonomously discover performance improvements on TPUs. and exploit non-intuitive, compiler-specific performance Our results demonstrate that AI-driven search uncovers “sweet spots” that are exceedingly difficult for human optimizations frequently missed by domain experts. engineers to discover manually. • CKKS-Mult: On TPUv5e, AlphaEvolve reduces latency of CKKS-Mult from 6340 µs to 5359 µs, corresponding to A. Experience and Lessons Learned using AlphaEvolve a 1.18× speedup. Such performance improvement comes Some of our experiences and learnings in adapting AlphaEfrom the scheduling optimization to reuse common parameters. The modest improvement indicates that CKKS-Mult volve for cryptographic optimization are described below: 1) Granularity of Optimization: Targeting isolated operais less constrained by underlying compiler inefficiencies. tions, such as polynomial multiplication (polymul [2]), failed These results show that AlphaEvolve can improve CKKS kernel to produce end-to-end speedups. We found that the agent performance on TPU by discovering implementation changes requires visibility into broader operational scopes to make that increase effective VReg utilization and reduce inefficiencies a meaningful impact on overall latency. in the execution schedule. 2) Co-evolution and System-Level Scoring: Building on the D. Takeaways need for broader visibility, shifting to the co-evolution of the The three workloads show that AlphaEvolve is most effective external_product and scoring it against the complete for cryptographic kernels when the bottleneck is caused bootstrapping execution successfully reduced latency from by a mismatch between cryptographic data layout and the 10 ms to 5.6 ms.
5
3) Guided Domain Constraints: While evolving the blind-rotate operation [2], the system discovered that RLWE decomposition allowed input tensors to safely use int8 rather than int32, and that unrolling the loop by a factor of eight yielded significant benefits. This underscores the importance of domain experts selectively granting the AI permission to modify high-impact code blocks. 4) Comprehensive Correctness Tests: Initial attempts to evolve the AND-gate failed to cover all values of a 3-input lookup table (lut3 [2]). We had to update the evaluator to exhaustively verify all lookup table values, demonstrating that strict, comprehensive validation is essential to prevent the model from reward hacking. 5) Security Checks: Because performance cannot come at the expense of cryptographic integrity, we introduced a standardized parameter check using a lattice estimator to ensure at least 128-bit security when co-optimizing scheme parameters. 6) Feedback and Prompts: To effectively steer the evolutionary search, we integrated hardware profiling tools (xprof) directly into the feedback loop and provided hardware-specific optimization guidelines as automated compilation prompts.
R EFERENCES
[1] M. R. Albrecht, R. Player, and S. Scott, “On the concrete hardness of learning with errors,” Cryptology ePrint Archive, Paper 2015/046, 2015. [Online]. Available: https://eprint.iacr.org/2015/046 [2] A. Ali, E. Astor, B. Gipson, S. Gorantala, M. Guevara, J. Kun, W. Lam, R. Misoczki, R. Springer, J. Takeshita, J. Tong, C. Tew, and C. Yun, “Jaxite: A fully homomorphic encryption backend targeting TPUs and GPUs, written in JAX,” https://github.com/google/jaxite, 2024. [3] A. Ali, J. Choi, B. Gipson, S. Gorantala, J. Kun, W. Legiest, L. Lim, A. Viand, M. Z. Demissie, and H. Zheng, “Heir: A universal compiler for homomorphic encryption,” 2025. [Online]. Available: https://arxiv.org/abs/2508.11095 [4] J. H. Cheon, A. Kim, M. Kim, and Y. Song, “Homomorphic encryption for arithmetic of approximate numbers,” in Advances in Cryptology – ASIACRYPT 2017, ser. Lecture Notes in Computer Science, vol. 10624. Springer, 2017, pp. 409–437. [5] I. Chillotti, N. Gama, M. Georgieva, and M. Izabachène, “TFHE: Fast fully homomorphic encryption over the torus,” Journal of Cryptology, vol. 33, no. 1, pp. 34–91, Jan 2020. [6] code perspective, “Add an unroll factor of 8 (reduces bootstrap speedup on tpu by 2 ms),” https://github.com/google/jaxite/pull/88, 2026, gitHub Pull Request #88, Accessed: 2026-05-06. [7] code perspective, “Improvements to the vector_matrix_polymul provides a speedup of 1.7 ms,” https://github.com/google/jaxite/pull/89, 2026, gitHub Pull Request #89, Accessed: 2026-05-06. [8] code perspective, “Update i32 matmul unreduced cggi to use single bfloat16 (instead of 4) for int32 multiplications,” https://github.com/ google/jaxite/pull/90, 2026, gitHub Pull Request #90, Accessed: 202605-06. [9] C. Gentry, “Fully homomorphic encryption using ideal lattices,” in Proceedings of the Forty-First Annual ACM Symposium on Theory of Computing, ser. STOC ’09. New York, NY, USA: Association for Computing Machinery, 2009, p. 169–178. [Online]. Available: https://doi.org/10.1145/1536414.1536440 [10] S. Gorantala, R. Springer, and B. Gipson, “Unlocking the potential of fully homomorphic encryption,” Communications of the ACM, vol. 66, no. 6, pp. 68–75, jun 2023. [Online]. Available: https://doi.org/10.1145/3582500 [11] S. Gorantala, R. Springer, S. Purser-Haskell, W. Lam, R. J. Wilson, A. Ali, E. P. Astor, I. Zukerman, S. Ruth, C. Dibak, P. Schoppmann, S. Kulankhina, A. Forget, D. Marn, C. Tew, R. Misoczki, B. Guillen, X. Ye, D. Kraft, D. Desfontaines, A. Krishnamurthy, M. Guevara, I. M. Perera, Y. Sushko, and B. Gipson, “A general purpose transpiler for fully homomorphic encryption,” CoRR, vol. abs/2106.07893, 2021. [Online]. Available: https://arxiv.org/abs/2106.07893 [12] X. He, “Accelerated linear algebra compiler for computationally efficient numerical models: Success and potential area of improvement,” PLOS ONE, vol. 18, no. 2, p. e0282265, 2023. [13] C. Hong, S. Bhatia, A. Cheung, and Y. S. Shao, “Autocomp: A powerful and portable code optimizer for tensor accelerators,” 2025. [Online]. Available: https://arxiv.org/abs/2505.18574 [14] JAX Developers, “Pallas quickstart — JAX documentation,” https://docs. jax.dev/en/latest/pallas/quickstart.html, 2024. [15] S. Liu, S. Agarwal, M. Maheswaran, M. Cemri, Z. Li, Q. Mang, A. Naren, E. Boneh, A. Cheng, M. Z. Pan, A. Du, K. Keutzer, A. Cheung, A. G. Dimakis, K. Sen, M. Zaharia, and I. Stoica, “Evox: Meta-evolution for automated discovery,” 2026. [Online]. Available: https://arxiv.org/abs/2602.23413 [16] A. Novikov, N. Vũ, M. Eisenberger, E. Dupont, P.-S. Huang, A. Z. Wagner, S. Shirobokov, B. Kozlovskii, F. J. Ruiz, A. Mehrabian et al., “Alphaevolve: A coding agent for scientific and algorithmic discovery,” arXiv preprint arXiv:2506.13131, 2025. [17] OpenXLA Project, “OpenXLA: Accelerated linear algebra,” https://github. com/openxla/xla, 2026. [18] J. Tong, T. Huang, J. Dang, L. de Castro, A. Itagi, A. Golder, A. Ali, J. Jiang, J. Kun, Arvind, G. E. Suh, and T. Krishna, “Leveraging asic ai chips for homomorphic encryption,” ser. HPCA’26. Australia: 2026 IEEE International Symposium on High Performance Computer Architecture (HPCA), 2026.
B. Future Work Our initial results focus on TPUs but the framework is extensible and can potentially optimize FHE primitives as well as any cryptographic primitives on a variety of hardware accelerators. It can also be integrated with simulators to guide specialized hardware-software co-design. To ensure security guarantees are maintained during hardware optimization, future iterations can integrate parameter selection with noise tracking. Furthermore, the bottleneck of manual code reviews can be reduced by adding functional equivalence checkers directly into the validation pipeline. Finally, a fundamental open question remains in the design of reward functions to safely allow AI agents to discover novel cryptographic algorithms and protocols while preserving security guarantees. ACKNOWLEDGMENTS This work represents an equal contributions between the first two authors, with contributing co-authors listed alphabetically The authors express their gratitude to Charles Hong, Andrew Ferraiuolo, Ben Kreuter, and Cindee Madison for their valuable feedback during various phases of this work. We thank PoSen Huang, Ngân (NV) Vũ, and the AlphaEvolve team for their assistance in establishing the experimental framework. We thank Drew Angeloff, Ankita Malviya Bairaria, Naomi Black, Elie Burzstein, Bryant Gipson, Pankaj Rohatgi, Amanda Walker, Moti Yung and the Safeworks leadership team for their guidance throughout this project. Finally, we also thank our extended team at Google DeepMind for their support of this research direction. We also acknowledge the use of large language models (LLMs) in the preparation of this manuscript. Specifically, the models were used for proofreading, summarizing text, and generating LaTeX outlines. The authors have reviewed and verified all generated content and take full responsibility for the final manuscript.
6
A PPENDIX A. LLM Instruction-Prompt Pool • Prompt 1: Suggest a new way to minimize data reorganization in the program. • Prompt 2: Suggest a new way to keep the logical shapes of JAX arrays invariant throughout execution. • Prompt 3: Suggest a new way to hide off-chip memory latency behind computation. • Prompt 4: Suggest a new optimization idea based on your expert knowledge of TPU performance tuning. B. Role and Hardware Architecture Prompts • Role: You are an expert hardware-aware compiler optimizer specializing in TPU architectures. • Objective: Optimize the input computational graph for maximum performance on the TPU v5e architecture. You must make strategic decisions regarding operation mapping (Vector vs. Matrix units), memory hierarchy management, and parallelism. • Target Hardware Specifications: • Compute Capabilities (Peak): – Clock Frequency: 1.50 GHz – Matrix Units (MXU): 197 TFLOPS (bf16/fp8), 394 TOPS (int8), 788 TOPS (int4). – Vector Units (VPU): vmatmul (1.54 TB/s), vmatpush (6.16 TB/s). • Memory Constraints (CRITICAL): – CMEM: None (0 bytes). Do not tile data into CMEM. – Vmem (Vector Memory): 128 MB. Read: 18.5 TB/s, Write: 6.16 TB/s. – HBM (High Bandwidth Memory): ∼17.2 GB at 820 GB/s. • Interconnect: 2D Torus, 4 links/chip, 45 GB/s per link. • Optimization Strategies Required: • Compute Mapping: Map matrix multiplications to the Systolic Array (MXU). Map element-wise, reduction, and non-matmul ops to the Vector Unit (VPU). • Memory Tiling & Fusion: All functional tiling must target the 128 MB Vmem. Perform loop fusion to avoid spilling back to HBM. • Precision Tuning: Evaluate suitability for int8 quantization for 2× throughput. • Pipeline Parallelism: Overlap HBM-to-Vmem DMA transfers with VPU/MXU execution. •Output Requirements: Generate optimized code/IR and summarize: theoretical peak utilization, memory bandwidth utilization, and a bottleneck analysis.
7