Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

Eymen Ünay, Björn Franke, Jackson Woodruff

TL;DR

This paper tackles the memory bottleneck imposed by rotation and bootstrap keys in CKKS-based FHE by introducing KeyMemRT, an MLIR-based compiler and runtime with a novel kmrt dialect for fine-grained key management. It automates per-key lifetimes, key moves, and bootstrap handling, enabling memory-scalable FHE and supporting arbitrary rotation indices without memory blow-up. The compiler lowers CKKS programs to OpenFHE while inserting key-management operations, and the runtime loads/clears keys on demand with optional prefetching. Empirical results show memory reductions up to 1.74x and latency improvements up to 1.73x against strong baselines (ANT-ACE, Fhelipe) across a range of ML models, demonstrating practical impact for scalable privacy-preserving computing.

Abstract

Fully Homomorphic Encryption (FHE) enables privacy preserving computation but it suffers from high latency and memory consumption. The computations are secured with special keys called rotation keys which often take up the majority of memory. In complex FHE applications, these rotation keys can cause a large memory bottleneck limiting program throughput. Existing compilers make little effort to solve this problem, instead relying on systems with massive memory availability. This resource requirement is a barrier to FHE uptake because optimizing FHE programs by hand is challenging due to their scale, complexity and expertise required. In this work, we present KeyMemRT; an MLIR based compiler and runtime framework that individually manages rotation key lifetimes to lower memory utilization and to allow arbitrary number of rotation indices to be supported without memory bloating. KeyMemRT relies on dataflow analysis to determine key lifetimes and is the first FHE compiler to provide automatic key management, handle fine-grained key-mangement and manage boostrap keys. We implement frontends for Orion and HEIR and show improvements over state-of-the-art FHE compilers. KeyMemRT achieves memory reduction of 1.74x and a speedup of 1.20x over ANT-ACE, and memory reduction of 1.16x and a speedup of 1.73x over memory-optimized compiler Fhelipe. We provide KeyMemRT as a post-optimizing compiler that can be targeted by any FHE compiler.

KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

TL;DR

This paper tackles the memory bottleneck imposed by rotation and bootstrap keys in CKKS-based FHE by introducing KeyMemRT, an MLIR-based compiler and runtime with a novel kmrt dialect for fine-grained key management. It automates per-key lifetimes, key moves, and bootstrap handling, enabling memory-scalable FHE and supporting arbitrary rotation indices without memory blow-up. The compiler lowers CKKS programs to OpenFHE while inserting key-management operations, and the runtime loads/clears keys on demand with optional prefetching. Empirical results show memory reductions up to 1.74x and latency improvements up to 1.73x against strong baselines (ANT-ACE, Fhelipe) across a range of ML models, demonstrating practical impact for scalable privacy-preserving computing.

Abstract

Fully Homomorphic Encryption (FHE) enables privacy preserving computation but it suffers from high latency and memory consumption. The computations are secured with special keys called rotation keys which often take up the majority of memory. In complex FHE applications, these rotation keys can cause a large memory bottleneck limiting program throughput. Existing compilers make little effort to solve this problem, instead relying on systems with massive memory availability. This resource requirement is a barrier to FHE uptake because optimizing FHE programs by hand is challenging due to their scale, complexity and expertise required. In this work, we present KeyMemRT; an MLIR based compiler and runtime framework that individually manages rotation key lifetimes to lower memory utilization and to allow arbitrary number of rotation indices to be supported without memory bloating. KeyMemRT relies on dataflow analysis to determine key lifetimes and is the first FHE compiler to provide automatic key management, handle fine-grained key-mangement and manage boostrap keys. We implement frontends for Orion and HEIR and show improvements over state-of-the-art FHE compilers. KeyMemRT achieves memory reduction of 1.74x and a speedup of 1.20x over ANT-ACE, and memory reduction of 1.16x and a speedup of 1.73x over memory-optimized compiler Fhelipe. We provide KeyMemRT as a post-optimizing compiler that can be targeted by any FHE compiler.
Paper Structure (41 sections, 1 equation, 10 figures, 3 tables)

This paper contains 41 sections, 1 equation, 10 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Motivating Example
  3. What is Memory-Scalable FHE?
  4. ANT-ACE: No Key Management
  5. Fhelipe: Binary Decomposition
  6. KeyMemRT: Fine-Grained Key Liveness Analysis
  7. Background
  8. Rotation Keys
  9. Scalability of Rotation Keys
  10. Bootstrap Keys
  11. Existing Memory Management Strategies
  12. Rotation Chaining
  13. KeyMemRT: System Overview
  14. KeyMemRT Compiler
  15. MLIR based kmrt Dialect
  16. ...and 26 more sections

Figures (10)

  • Figure 1: Currently the memory requirements from larger numbers of keys is not scalable. KeyMemRT provides lower memory growth without overhead allowing FHE programs to scale better.
  • Figure 2: The KeyMemRT compiler takes the ckks dialect as input and optimizes the code to keep keys out of memory as much as possible by unoverlapping key livenesses. The KeyMemRT runtime manages keys, loading them in and clearing them as required to keep memory utilization down. Fhelipe uses decomposition to reduce the number of keys required, but this results in long runtimes. ANT-ACE loads all the keys in, which results in fast computation but wasted memory.
  • Figure 3: Design of KeyMemRT. The MLIR-based compiler works with the openfhe dialect and inserts key-management operations in our kmrt dialect. The runtime uses these indications to manage keys.
  • Figure 4: The full pipeline of the KeyMemRT compiler. CKKS to OpenFHE (\ref{['Sec.Comp.CtO']}) and Lower Linear Transform (\ref{['Sec.Comp.LLT']}) passes perform lowering and generate the kmrt operations. Baby-Step Giant-Step (BSGS) Decomposition (\ref{['Sec.Comp.BSGS']}) reduces the number of keys in the program as explained in Section \ref{['Sec.Background.RC']}. First Merge Rotation Keys (\ref{['Sec.Comp.MRK']}) pass manages rotation keys of explicit rotations by the user and the decomposed rotations from BSGS. Rotation Hoisting (\ref{['Sec.Comp.RH']}) speeds up rotations applied to common inputs. After compiling and profiling at this stage, the FHE profile is used in Bootstrap Removal (\ref{['Sec.Comp.BR']}) to speed upthe program. Bootstrap Key Management (\ref{['Sec.Comp.BKM']}) materializes the keys used opaquely by the library to the IR level. Second Merge Rotation Keys (\ref{['Sec.Comp.MRK']}) pass merges key livenesses of the bootstrap related rotation keys with the explicit rotations. Common Subexpression Elimination (\ref{['Sec.Comp.CSE']}) removes cleartext dead code resulted from the kmrt optimizations. Insert Clear Ops (\ref{['Sec.Comp.ICO']}) reduces memory consumption by clearing plaintexts and ciphertexts after last use. Lastly, Key Prefetching (\ref{['Sec.Comp.KP']}) places hint operations that will inform the runtime. The final generated code is a C++ program with OpenFHE API and KeyMemRT runtime calls.
  • Figure 5: Lowering from the target independent ckks dialect to kmrt and openfhe dialects. Highlighted lines show the added operations for loading and clearing rotation keys. This pass is described in Section \ref{['Sec.Comp.CtO']}.
  • ...and 5 more figures