💻 computer science

KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

KeyMemRT is an MLIR-based compiler and runtime framework that utilizes dataflow analysis to automatically manage the lifetimes of FHE rotation keys, significantly reducing memory consumption and improving execution speed compared to existing state-of-the-art compilers.

Original authors: Eymen Ünay, Björn Franke, Jackson Woodruff

Published 2026-01-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Eymen Ünay, Björn Franke, Jackson Woodruff

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 are trying to solve a massive, complex puzzle, but you are only allowed to look at the pieces while wearing a blindfold. You have to manipulate the pieces without seeing them, using a special set of rules that keep the puzzle hidden from prying eyes. This is essentially what Fully Homomorphic Encryption (FHE) does: it allows computers to perform calculations on data while that data remains completely encrypted.

However, there's a huge catch. To do this "blindfolded" math, the computer needs a giant library of special "magic keys" (called rotation keys).

The Problem: The "Key Hoarder"

Think of these magic keys like a massive collection of physical keys for a hotel with thousands of rooms.

The Old Way (ANT-ACE): Imagine a hotel manager who, before a guest even arrives, grabs every single key in the building and dumps them all onto the front desk. Even if the guest only needs to visit Room 101, the manager keeps the keys for Rooms 102 through 5,000 right there on the desk.
- The Result: The desk (memory) gets cluttered and overflowing. If the hotel is huge, the desk becomes so full of keys that there's no room left for anything else. The system crashes or slows down because it's trying to manage a mountain of unnecessary keys.
The "Slow" Way (Fhelipe): Another manager tries to save space by only keeping a few "master keys." To open Room 101, they don't have the direct key, so they have to use a master key to open Room 1, then another to open Room 2, and so on, chaining them together until they reach 101.
- The Result: The desk is tidy, but the guest has to wait a long time while the manager fumbles through a long chain of keys. The process is slow.

The Solution: KeyMemRT

The authors of this paper built a new system called KeyMemRT. Think of it as a super-smart, automated concierge service that manages the keys perfectly.

It Knows Exactly What You Need: Instead of grabbing all the keys at once, KeyMemRT analyzes the guest's itinerary (the program's code) to see exactly which rooms they will visit and in what order.
Just-in-Time Delivery: It keeps the desk clear. It only brings out the specific key for Room 101 right before the guest needs it.
Instant Cleanup: As soon as the guest leaves Room 101, the concierge immediately takes that key back and puts it away, making space for the next key.
The "Prefetch" Trick: To make sure the guest never has to wait, the concierge starts bringing the next key over while the guest is still in the current room. This happens in the background, so the process feels fast and smooth.

What Did They Achieve?

The researchers tested this new system against the old methods using various complex tasks (like recognizing images in medical scans or financial data).

Memory Savings: Compared to the "Key Hoarder" method (ANT-ACE), KeyMemRT reduced the amount of memory needed by 1.74 times. It's like fitting a whole hotel's worth of keys into a single briefcase instead of a warehouse.
Speed: Compared to the "Slow Chain" method (Fhelipe), KeyMemRT was 1.73 times faster. It didn't just save space; it also made the work happen much quicker because it didn't waste time chaining keys together.

Why Does This Matter?

Currently, FHE is hard to use because it requires computers with enormous amounts of memory (hundreds of gigabytes) just to hold these keys. This makes it expensive and impractical for many real-world uses.

KeyMemRT acts like a translator and a traffic controller. It takes complex, messy code and optimizes it so that it can run on standard computers without needing a supercomputer just to hold the keys. It makes privacy-preserving computing (like analyzing sensitive medical or financial data without ever seeing the raw numbers) much more scalable and efficient.

In short: They built a smart system that stops FHE programs from hoarding useless keys and stops them from wasting time making keys from scratch, allowing private computing to run faster and on smaller, cheaper machines.

Technical Summary: KeyMemRT Compiler and Runtime

Problem Statement

Fully Homomorphic Encryption (FHE) enables privacy-preserving computation on encrypted data but faces significant barriers to real-world adoption due to excessive computational and memory requirements. A primary bottleneck is the management of rotation keys. In batched FHE schemes (e.g., CKKS, BGV, BFV), rotation operations are essential for intra-vector computations and bootstrapping. However, these schemes require a unique rotation key for every specific rotation index.

As FHE programs scale, the number of required rotation indices grows, often necessitating hundreds of gigabytes or even terabytes of memory to store all keys simultaneously. Existing FHE compilers generally treat rotation keys as a single, static entity, loading all keys into memory for the duration of program execution. This approach lacks memory scalability. Conversely, existing strategies that attempt to reduce memory, such as Fhelipe, rely on decomposing arbitrary rotation indices into chains of powers-of-two. While this reduces the number of stored keys, it significantly increases computational latency due to the chaining overhead. There is a lack of compilers that offer fine-grained, automatic key management to balance memory usage and execution speed.

Methodology

The authors present KeyMemRT, a compiler and runtime framework designed to manage rotation key lifetimes individually. The system is built on the MLIR (Multi-Level Intermediate Representation) infrastructure and introduces a novel dialect and runtime architecture.

1. The `kmrt` MLIR Dialect

KeyMemRT introduces a new dialect called kmrt to represent key management directly within the IR.

New Types: It defines a rot_key type parameterized by the rotation index, allowing the compiler to distinguish between keys for different indices.
New Operations: The dialect includes operations for explicit key management:
- kmrt.load_key: Loads a specific rotation key.
- kmrt.clear_key: Removes a key from memory.
- kmrt.prefetch_key: Provides hints to the runtime for asynchronous loading.
- kmrt.use_key / kmrt.assume_key: Manage key scope and liveness.
Integration: The compiler lowers high-level operations (e.g., from the ckks dialect via HEIR) into the openfhe dialect while inserting kmrt operations to manage key lifecycles.

2. Compiler Optimization Pipeline

The KeyMemRT compiler performs a series of lowering and optimization passes to reduce memory footprint without sacrificing performance:

CKKS to OpenFHE Lowering: Associates every rotation operation with a specific rot_key value, inverting the problem from creating liveness ranges to merging them.
Baby-Step Giant-Step (BSGS) Decomposition: Reduces the number of required keys from $N$ to $\approx 2\sqrt{N}$ for consecutive rotation indices by chaining rotations. This is particularly effective for matrix-vector multiplication patterns.
Merge Rotation Keys: Uses dataflow analysis to detect sequences where a key is cleared and immediately reloaded with the same index, removing redundant load/clear pairs.
Rotation Hoisting: Identifies rotation patterns applied to common inputs and transforms them into faster alternatives supported by the OpenFHE library.
Bootstrap Key Management: Materializes the internal rotation keys used by bootstrap operations (which refresh ciphertext levels) into the IR. This allows the compiler to manage bootstrap keys with the same fine-grained control as explicit rotation keys, a feature previously unaddressed.
Common Subexpression Elimination (CSE): Removes dead code resulting from key optimizations.
Hint Placement: Inserts prefetch_key hints to enable the runtime to overlap key loading with computation.

3. KeyMemRT Runtime

The runtime acts as an intermediary between the compiled C++ code and the FHE library (OpenFHE). It manages the dynamic loading and unloading of keys from disk to RAM based on compiler-inserted hints.

Low-Memory Mode: Executes directives sequentially, loading keys only when strictly required to minimize peak memory usage.
Balanced Mode: Utilizes a prefetching mechanism to load keys in parallel with computation. It maintains an internal queue, loading keys ahead of time up to a user-defined limit and clearing them when space is needed, thereby reducing latency overhead.

Key Contributions

New MLIR Dialect (kmrt): The first dialect to expose rotation keys as IR elements with specific types and memory management operations, enabling fine-grained analysis.
Automatic Fine-Grained Key Management: The first FHE compiler to automatically track key liveness, optimize bootstrap keys, and manage key lifetimes individually rather than as a global set.
Runtime Framework: A runtime that supports on-demand key loading and clearing, utilizing compiler hints to overlap I/O with computation.
Integration: A post-optimizing compiler toolchain that can be targeted by existing FHE frontends (demonstrated with Orion and HEIR).

Results

The authors evaluated KeyMemRT against two state-of-the-art approaches: ANT-ACE (which keeps all keys in memory) and Fhelipe (which uses decomposition to reduce keys but increases latency). Benchmarks included various ML models (MLP, LoLa, LeNet, ResNet variants, AlexNet) with FHE parameters set to a ring dimension of $2^{16}$ and multiplicative depth of 30.

Memory Reduction:
- Compared to ANT-ACE, KeyMemRT achieved a geometric mean memory reduction of 1.74x. In specific cases like LeNet, memory usage was reduced by 2.59x.
- Compared to Fhelipe, KeyMemRT achieved a memory reduction of 1.16x.
Latency/Speedup:
- Compared to ANT-ACE, KeyMemRT (Balanced mode) achieved a speedup of 1.20x.
- Compared to Fhelipe, KeyMemRT achieved a significant speedup of 1.73x (geometric mean), with improvements up to 2.5x on AlexNet. This is because KeyMemRT avoids the computational overhead of chaining rotations required by Fhelipe.
Scalability: KeyMemRT demonstrated the ability to run large programs (e.g., ResNet-50) that exceeded 100 GB of memory requirements under ANT-ACE, by keeping memory usage significantly lower.

Significance and Claims

The paper claims that KeyMemRT is the first FHE compiler to provide automatic, fine-grained key management and bootstrap key management. The authors assert that this work unlocks memory scalability for FHE, allowing arbitrarily large programs to run without the prohibitive memory burden of storing all rotation keys simultaneously.

By decoupling memory usage from the total number of keys and instead managing keys based on their liveness, KeyMemRT enables FHE programs to scale better on existing hardware. The authors emphasize that their approach does not compromise performance; unlike decomposition-based methods that trade memory for speed, KeyMemRT achieves both memory reduction and latency improvements by loading only the necessary keys at the precise moment they are needed. This addresses a critical barrier to FHE adoption in resource-constrained or large-scale environments.