FlexHE: A flexible Kernel Generation Framework for Homomorphic Encryption-Based Private Inference

Secure two-party computation (2PC) based on homomorphic encryption (HE) achieves formal data privacy protection and gets increasing adoption for private deep neural network (DNN) inference. As modern HE schemes usually operate on polynomials, existing works rely on manually-designed HE kernels for representative DNN operations. However, this is not only unscalable considering the diverse operator types, shapes, polynomial orders, etc, but also misses important optimization opportunities. In this paper, we introduce FlexHE, a flexible kernel generation framework to enable automatic generation and optimization of HE kernels for 2PC-based private inference. Given a high-level description of DNN operations, FlexHE can systematically define the HE kernel design space considering various optimization dimensions, including loop tiling, reordering, etc. We also analyze the communication and computation impact of different optimization dimensions for design space reduction. To search for the best kernel design, a two-level optimization problem is formulated and iteratively solved with an integer linear programming (ILP) formulation. With extensive experimental results, we not only demonstrate a better coverage of DNN operations including depth-wise Conv3D and dilated Conv3D, but also achieve more than 100×, 7.9×, and 4.2× latency reduction compared to prior-art HElayers, Cheetah, and Falcon, respectively.