The Architecture of Concrete, Zama's Fully Homomorphic Encryption Compiler Leveraging MLIR

> Part I: Concrete, Zama's Fully Homomorphic Compiler‍

> Part II: The Architecture of Concrete, Zama's Fully Homomorphic Encryption Compiler Leveraging MLIR

In a previous blog post, we introduced Concrete, Zama’s Fully Homomorphic Encryption compiler, and discussed the rationale for using MLIR. In this blog post, we will take a deeper look into the architecture of our compiler.

Zama’s Approach: A Modular Framework

When we started to work on our compiler in early 2020, one of our main objectives was to release it under an open source license. We were disappointed that, due to the lack of such tools in the open, FHE had only gained little traction outside a small group of cryptography experts. We wanted FHE and all of its applications to become accessible to everyone, including developers without a background in cryptography. We admire open-source frameworks for deep learning like PyTorch and Keras, as they made deep learning available to an audience beyond machine learning specialists. Similarly, user-friendly and open-source tools for FHE are necessary to encourage widespread adoption and help unleash the full potential of this technology. An FHE compiler is an integral part of a comprehensive toolchain that makes it easy for developers to integrate FHE into their applications. This is why we are putting so much effort into the design, implementation and architecture of Concrete.

One of the design goals of Concrete was to provide a modular stack allowing for extensions and the customization or replacement of individual components. At its core is Concrete Compiler, which translates the MLIR-based intermediate representation of a plaintext program into an FHE program. Being a part of the modular Concrete framework, this compiler can be replaced by any other compiler that supports the same MLIR input dialect (more about this in the next section). On top of Concrete compiler, we added our first frontend, called Concrete Python, which processes programs directly written in Python and produces the MLIR-based intermediate representation required by the compiler. The frontend plays an important role for the adoption of the compiler: application developers are used to writing Python and should not be burdened with the task of writing MLIR directly, which is not only unnecessarily complicated for application development, but also tedious and error prone. We expect more pluggable frontends for other languages to be created in the near future, which will hopefully come from our open-source community. Being able to include the community in the development was a major motivation to develop Concrete as an open-source project, and we are looking forward to community contributions. Finally, on top of Concrete Python, we created our application package Concrete ML, which adds privacy on top of classical ML frameworks, such as scikit-learn, XGBoost, Torch, Keras and ONNX. Other application frameworks will come in the future and we are excited to mention that BioPython developed by Tekkare already uses Concrete ML today.

If you want to learn more about Concrete ML, check out its documentation here.

Low-level translation to boolean circuits vs. high-level short integer encodings in Concrete

Other TFHE-based compilers (e.g., Cingulata or Google’s FHE transpiler) traditionally take a low-level approach, transpiling arithmetic circuits into their optimized boolean equivalent and using TFHE boolean gates to homomorphically evaluate the circuit. Our compiler takes a high-level approach, translating the arithmetic circuit into a TFHE equivalent before determining the most efficient encoding and cryptographic parameters depending on the constraints imposed by the circuit.

Both approaches come with their benefits and drawbacks. On one hand, the boolean encoding approach allows for the representation of arbitrary-sized integers (i.e. without any restriction on the maximum bitwidth), the arithmetic operations need to be transpiled into their boolean equivalent, and it makes it necessary to evaluate all the boolean gates in FHE. This is very costly, as each boolean gate is implemented with a bootstrapping operation, the most computationally intensive operation in FHE. On the other hand, the short-integer encoding approach used by Concrete restricts bootstrapping to values with a maximum width of 8 bits by default and a maximum of 16 bits using our current extension for CRT encoding, and leverages very fast leveled operators (e.g., additions and cleartext multiplications) that can be implemented without bootstrapping.

In the future, we would like to implement a more hybrid strategy into Concrete that selects the approach depending on the input program. The ongoing standardization of MLIR-based FHE dialects in the HEIR project by Zama, Google and other industry leaders will allow Google and Zama tools to be integrated on an open-source basis and represents a possible path for upstream contributions of FHE tools directly in the MLIR project.

Compiler overview

As previously explained, we heavily rely on MLIR as infrastructure for our compiler. The choice of MLIR was motivated by the engineering cost efficiency of using already-existing dialects (e.g., tensor, memref, linalg, scf, omp), compilation and optimization passes (canonicalization, linalg generalization, dead code elimination, bufferization) and the infrastructure for testing.

Concrete is built to make a developer’s life as easy as possible. To this end, we provide Concrete Python, an easy to use frontend for the development of FHE applications. You can learn more about Concrete Python with this written and video tutorial.

Concrete Python has its own internal pipeline (tracing, fusing, MLIR code generation) to transpile the Python code into a high-level IR that can be compiled by the Compiler. This high-level IR is a scheme-independent and crypto-free representation of an arithmetic circuit that simply specifies which values are to be processed as encrypted values and which values should be processed as plaintext values. It also defines a set of operators that can be executed in FHE.

Concrete Compiler provides bindings to allow the Concrete Python transpiler to build the IR and invoke the Concrete Compiler engine and tools. These bindings can easily be used by other frontends built by us or by the community.

We have defined several dialects that are used by the compiler to transform a high level crypto-free representation of the computation DAG on encrypted integers into a binary executable. The dialects, in decreasing order of their level of abstraction are:

The FHE dialect defines crypto-free FHE types and high level operators on encrypted scalars,
The FHELinalg dialect defines high level operators on encrypted tensors,
The TFHE dialect introduces crypto-system dependent parameters and operators,
The Concrete dialect models backend operators in preparation of code generation,

Note: The current pipeline architecture may evolve in the future, in particular to introduce new dialects that are currently defined as part of the HEIR project.

Concrete defines the TFHE optimization problem by collecting constraints on the high level representation of the arithmetic circuit and uses the Concrete Optimizer to solve it. The solution of the optimization problem is a set of guidelines used by the compiler to convert the FHE circuit into its optimized TFHE equivalent. Boolean encoding generally relies on static parameters and a static topology of boolean gates. By contrast, in our approach, the accuracy and efficiency depend on the choice of parameters and the topological transformations which can be optimally decided by the constraint analysis on the circuit to compile. The details of the optimization problems will be discussed in a future blog post.

Finally, the TFHE operators are lowered into the Concrete dialect, which allows us to prepare the code generation to LLVM IR, essentially by replacing operators with calls to the backends exposing TFHE crypto primitives.

Auto parallelism and runtime

One main advantage of the compiler is its capability to automatically extract parallelism in order to exploit the hardware. To achieve this, we exploit two levels of parallelism: loop parallelism for shared memory systems, and dataflow parallelism for accelerators and distributed systems.

To extract loop parallelism, we rely on the existing MLIR infrastructure. Thanks to the transformations available in MLIR’s Linalg Dialect, high-level FHELinalg operations can be parallelized by first lowering them to linalg.generic operations, then lowering to parallel loops using built-in Linalg transformations and then generating IR that translates to calls to the OpenMP runtime.

We have chosen to introduce a new runtime and a new dialect to model the dataflow parallelism into our compiler. This is due to the fact that the MLIR project doesn’t offer a builtin dialect for dataflow parallelism. Our runtime is itself based on the HPX C++ library, which already provides great infrastructure for distributed computing and scheduling of tasks within and across hosts. The tasks are created by the compiler and lowered to API calls to our dataflow runtime thanks to a set of passes operating on the RT dialect which allow us to model tasks and the runtime abstractions.

Conclusion

In this blog post, we presented an overview of the internals of the Concrete compiler. One of the major aspects of optimizing an integer TFHE DAG is to find accurate and optimal crypto parameters and topology, which is done by the Concrete Optimizer, and will be the subject of another blog post. If you are interested in learning more, you can take a look at some of our technical papers that explain the mathematical theory behind this tool.

We are proud to develop everything in the open, thanks to existing and available libraries, and equally giving access to our own tools to the community. We are happy to see our community of external contributors grow, especially with the Zama Bounty Program and we are grateful for the positive feedback from the community (And near 900 GitHub stars on our repository).

In a nutshell, the Concrete Framework contains:

an easy-to-use frontend, Concrete Python,
a reusable compiler infrastructure,
multiple backends (CPU, GPU, and later, hardware accelerators),
a specific TFHE optimizer,
a  compilation pipeline and runtime designed to scale.

If you want to try it check out the examples in the Concrete ML documentation. We are excited to get feedback from the community and see what you’ll be building on top of Concrete