KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

TL;DR

KeyMemRT is an MLIR-based compiler and runtime that automatically manages rotation keys in FHE, significantly reducing memory usage and improving performance for complex encrypted computations.

Contribution

It introduces automatic, fine-grained rotation key management in FHE compilers, a novel approach that reduces memory and enhances efficiency without requiring extensive manual optimization.

Findings

Memory reduction of 1.74x over ANT-ACE

Speedup of 1.20x over ANT-ACE

Memory reduction of 1.16x and 1.73x speedup over Fhelipe

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.