💬 NLP

Leveraging ASIC AI Chips for Homomorphic Encryption

This paper introduces CROSS, a compiler framework that bridges the architectural gap between AI accelerators like TPUs and Homomorphic Encryption workloads by transforming high-precision arithmetic and data permutations into efficient low-precision matrix operations and offline parameter adjustments, thereby achieving state-of-the-art energy efficiency for HE on AI ASICs.

Original authors: Jianming Tong, Tianhao Huang, Jingtian Dang, Leo de Castro, Anirudh Itagi, Anupam Golder, Asra Ali, Jeremy Kun, Jevin Jiang, Arvind, G. Edward Suh, Tushar Krishna

Published 2026-04-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jianming Tong, Tianhao Huang, Jingtian Dang, Leo de Castro, Anirudh Itagi, Anupam Golder, Asra Ali, Jeremy Kun, Jevin Jiang, Arvind, G. Edward Suh, Tushar Krishna

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a secret recipe for a delicious cake. You want to send this recipe to a baker in a different city so they can bake the cake for you, but you don't want them to see the ingredients or the instructions.

Homomorphic Encryption (HE) is like a magical oven that allows the baker to bake the cake while it's still inside a locked, opaque box. They can mix, heat, and frost the cake without ever opening the box or seeing what's inside. When they send the box back, you open it, and there's a perfectly baked cake, but the baker never knew what the recipe was.

The problem? This "magic" is incredibly slow and energy-hungry. It's like trying to bake a cake by manually turning every single gear in the oven with your bare hands. It takes forever and tires you out.

The Current Solution: The "Generalist" Chef (GPUs)

To speed this up, scientists started using GPUs (the powerful chips in gaming computers). Think of a GPU as a team of 10,000 tiny, general-purpose chefs. They are fast at many things, but they are still trying to do this "locked-box baking" with a lot of manual, fine-tuned movements. It's much faster than doing it by hand, but it's still inefficient and uses a lot of electricity.

The New Idea: The "Specialist" Factory (AI Chips)

Then, there are AI Chips (like Google's TPUs). These are like massive, hyper-specialized factories designed to build millions of identical toy cars at once. They are incredibly efficient at doing one specific thing: crunching huge blocks of numbers (matrices) very quickly.

The Problem: The "locked-box baking" instructions (the HE algorithms) were written for the "Generalist Chefs" (GPUs). If you try to run those instructions on the "Specialist Factory" (AI Chips), it's a disaster.

The Mismatch: The instructions ask the factory to do tiny, precise tasks one by one (like threading a needle), but the factory is built to stamp out thousands of metal sheets at once.
The Result: The factory sits 90% idle, waiting for instructions, while the few workers who can do the tiny tasks get overwhelmed. It's like using a giant steamroller to push a single grain of sand.

The Solution: CROSS (The Translator)

The paper introduces CROSS, a smart compiler (a translator program) that rewrites the "locked-box baking" instructions so the "Specialist Factory" can actually understand and use them efficiently.

CROSS does two magical things:

1. BAT (Basis-Aligned Transformation): "The Puzzle Solver"

The Old Way: The instructions told the factory to multiply numbers in a way that created a giant grid full of empty spaces (zeros). It was like asking a factory to stamp a sheet of metal, but 50% of the sheet was empty space. The machine wasted time and energy stamping nothing.
The CROSS Way: BAT looks at the instructions before the factory starts working. It realizes, "Hey, we know these empty spaces are always empty!" It cuts them out and rearranges the remaining pieces into a tight, dense block.
The Analogy: Instead of asking the factory to stamp a sheet with holes in it, BAT gives them a solid block of metal that fits perfectly into the machine. Now, the factory is stamping at 100% capacity, using its super-fast "Matrix Engine" (MXU) instead of its slow, weak "Vector Engine."

2. MAT (Memory-Aligned Transformation): "The Pre-Organized Library"

The Old Way: The instructions required the factory to constantly stop, pick up a book, move it to a different shelf, and then put it back. In computing terms, this is "shuffling" data around. On these massive factories, moving data is slow and expensive.
The CROSS Way: MAT realizes that the factory always needs the books in a specific order. Instead of moving the books during the work, MAT rearranges the library before the workers even arrive.
The Analogy: It's like a chef who, instead of running back and forth to the fridge to grab ingredients, has the entire kitchen pre-stocked with every ingredient exactly where it needs to be. The workers never stop to move things; they just cook.

The Result: A Super-Efficient Bakery

By using CROSS, the researchers took a Google TPU (an AI chip) and made it run Homomorphic Encryption faster and more energy-efficiently than almost anything else on the market, including:

GPUs: It's significantly faster and uses less power.
FPGAs: Custom chips designed for this, but still slower.
Specialized HE ASICs: Even chips built specifically for this encryption job (which cost millions to design) are only slightly faster, but they cost a fortune to build. CROSS gets you 90% of that performance on a chip that already exists and is widely available.

Why This Matters

This is a big deal because it means we don't need to wait for expensive, custom-built "magic ovens" to protect our privacy. We can use the powerful AI chips that are already in our data centers today.

In short: CROSS is the translator that teaches a super-fast AI factory how to do the slow, complex work of secret math, turning a clumsy, energy-wasting process into a sleek, high-speed operation. It proves that we can have strong privacy protection without sacrificing speed or breaking the bank.

1. Problem Statement

Homomorphic Encryption (HE) enables computation on encrypted data, offering strong privacy for cloud services. However, its practical adoption is hindered by prohibitive computational overheads (often 1000× slower than plaintext on CPUs).

Current State: GPUs are the de-facto standard for accelerating HE, but they still lag significantly behind specialized HE ASICs in energy efficiency (approx. 33× gap).
The Gap: Dedicated HE ASICs are energy-efficient but expensive to design and fabricate, limiting widespread deployment.
The Challenge: Existing AI accelerators (like Google TPUs) offer massive compute and memory resources with superior energy efficiency compared to GPUs. However, SoTA HE algorithms designed for GPUs do not map efficiently to TPUs due to fundamental architectural mismatches:
1. Arithmetic Mismatch: HE relies on high-precision (32-bit) modular arithmetic. GPUs handle this well, but TPUs have a low-throughput Vector Processing Unit (VPU) for 32-bit ops and a high-throughput Matrix Multiplication Unit (MXU) for low-precision (8-bit) ops. Current HE libraries leave the powerful MXU idle.
2. Memory/Granularity Mismatch: HE kernels (like Number Theoretic Transform - NTT) require frequent, fine-grained data permutations (shuffling/transposing). TPUs use coarse-grained, SIMD-controlled memory (large registers), making these fine-grained operations extremely inefficient and costly.

2. Methodology: The CROSS Framework

The authors propose CROSS, a compiler framework that systematically transforms HE workloads to align with the coarse-grained, high-throughput architecture of AI ASICs (specifically Google TPUs). CROSS introduces two core transformations:

A. Basis-Aligned Transformation (BAT)

Goal: Convert high-precision (32-bit) modular arithmetic into dense, low-precision (8-bit) matrix multiplications to utilize the TPU's MXU.
Mechanism:
- Standard HE multiplication decomposes high-precision numbers into chunks, resulting in sparse matrix-vector multiplications with many redundant zeros (wasting compute).
- BAT recognizes that parameters like twiddle factors (in NTT) and evaluation keys are static and known offline.
- It applies modular reduction ( $mod\ q$ ) to these static parameters offline, "folding" high-basis terms into low-basis terms.
- This transforms the sparse matrix into a dense, smaller matrix, eliminating redundant zeros and allowing the operation to be executed as a dense 8-bit MatMul on the MXU.
- Result: ~2× theoretical reduction in compute and memory waste; enables the use of the MXU (which is ~58× faster than the VPU on TPUs).

B. Memory-Aligned Transformation (MAT)

Goal: Eliminate costly runtime data reordering (transpose/shuffling) required by HE algorithms.
Mechanism:
- Instead of performing data permutations at runtime (which causes low register utilization and high latency on TPUs), MAT treats permutations as multiplication by a permutation matrix.
- These permutation matrices are applied offline to the pre-known parameters (e.g., twiddle factors) during compilation.
- This creates a "layout-invariant" kernel where the computation naturally produces output in the desired order without explicit runtime memory movement.
- Application: Applied to the 4-step NTT algorithm, MAT eliminates the explicit matrix transpose and bit-reverse shuffling steps, converting the algorithm into a 3-step layout-invariant process.

3. Key Contributions

Systematic Characterization: Identified the specific architectural mismatches (precision and granularity) that prevent SoTA GPU HE algorithms from running efficiently on AI ASICs.
BAT (Basis-Aligned Transformation): A novel technique to map high-precision modular integer arithmetic to low-precision dense matrix multiplication, unlocking the MXU for HE.
MAT (Memory-Aligned Transformation): A scheduling technique that embeds data reordering into offline parameter transformation, removing runtime memory overhead.
First AI/HE Co-Acceleration: Demonstrated that existing AI ASICs (TPUs) can accelerate HE operators without hardware modifications, achieving State-of-the-Art (SoTA) energy efficiency.

4. Experimental Results

The authors evaluated CROSS on a real Google TPU v6e (single-host) and compared it against SoTA implementations on CPUs, GPUs (NVIDIA A100, RTX 4090), FPGAs, and HE ASICs.

NTT Throughput:
- Achieved 1.43× higher throughput than WarpDrive (NVIDIA A100).
- Refreshed the SoTA NTT throughput record for practically available devices.
Energy Efficiency (Throughput per Watt):
- 451× higher than OpenFHE (CPU).
- 7.81× higher than WarpDrive (GPU).
- 1.83× higher than FIDESlib (GPU).
- Outperformed specialized HE ASICs (BASALISC, CraterLake) in HE-Add efficiency and narrowed the gap significantly in HE-Mult and Rotate (within 3× of ASICs).
Workload Performance:
- MNIST Inference: 10× speedup over Orion (GPU) with 98% accuracy.
- Logistic Regression: 1.06× throughput per watt compared to Cheddar (RTX 4090).
- Bootstrapping: 1.5× throughput over Cheddar (RTX 4090).

5. Significance and Impact

Paradigm Shift: This work proves that AI accelerators (designed for deep learning) can be the most energy-efficient platform for Homomorphic Encryption, surpassing both general-purpose GPUs and specialized HE ASICs in terms of performance-per-watt for commodity hardware.
No Hardware Cost: Unlike dedicated HE ASICs, CROSS requires zero hardware modifications. It leverages existing, mass-produced AI infrastructure (TPUs) to enable privacy-preserving AI.
Scalability: By enabling efficient HE on AI ASICs, this approach paves the way for "Private AI" at massive scale, allowing cloud providers to perform inference on encrypted data using their existing AI clusters.
Future Directions: The paper highlights that while CROSS is highly efficient, the remaining gap to dedicated ASICs (3-33×) stems from the lack of hardware-specific low-cost shuffling engines and fixed moduli support. Future work may focus on further algorithmic embedding of permutations and optimizing for specific modulus constraints.

In summary, CROSS bridges the gap between the coarse-grained, high-throughput nature of AI ASICs and the fine-grained, high-precision requirements of Homomorphic Encryption, establishing a new standard for energy-efficient privacy-preserving computation.