Concrete — Zama's Fully Homomorphic Encryption Compiler

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

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

The importance of compilers for homomorphic encryption has been clear since the rise of practical Fully Homomorphic Encryption (FHE) schemes. Managing noise, choosing appropriate crypto parameters, and finding the best set and order of operations for a specific computation are problems for cryptographers. A homomorphic encryption compiler for developers is supposed to take all that complexity away. 

At Zama, we started prototyping compilers a while back using different techniques and technologies. HNP was one of our prototypes that was made available to the community, but it didn’t solve major FHE compilation challenges. So we kept digging.

Following the announcement of Concrete v1.0.0, we released the compiler Zama have been working on for almost two years. This blog post explains how you can use it. A later blog post will be dedicated to the Compiler’s internals.

A First Simple Example

The first thing you need to know is what the input to the Compiler looks like. The Compiler expects an input program in MLIR. Below is a simple example doing an addition of two encrypted unsigned integers.

func.func @main(%arg0: !FHE.eint<6>, %arg1: !FHE.eint<6>) -> !FHE.eint<6> {
   %result = "FHE.add_eint"(%arg0, %arg1): (!FHE.eint<6>, !FHE.eint<6>) -> (!FHE.eint<6>)
   return %result: !FHE.eint<6>
}

It’s pretty easy to read. We define the function [.c-inline-code]main[.c-inline-code] that gets two arguments of type [.c-inline-code]FHE.eint<6>[.c-inline-code], which then defines an encrypted unsigned integer of 6 bits and returns a value of the same type. The function itself makes a call to the [.c-inline-code]FHE.add_eint[.c-inline-code] operation, which will add the two operands.

You can use the Compiler via different APIs and tools. Currently available APIs are Python, Cpp, and C. You also have a CLI tool that is quite handy for debugging the compilation pipeline at different stages and intermediate representations.

The Python API is the easiest to set up, so we can begin to compile and run the previous MLIR example. Start with installing [.c-inline-code]concrete-python[.c-inline-code] via [.c-inline-code]pip install concrete-python[.c-inline-code]. This doesn’t just install the Compiler, but the Python frontend as well (we will talk about frontends later in this post). You should be able to import the Compiler module if installation was successful.

import concrete.compiler as compiler

One of the main entry points to the Compiler is the [.c-inline-code]LibrarySupport[.c-inline-code] class, which allows you to compile and run FHE programs while storing artifacts on disk. A similar class ([.c-inline-code]JITSupport[.c-inline-code]) is available which keeps artifacts in memory.

mlir_input = """
   func.func @main(%arg0: !FHE.eint<6>, %arg1: !FHE.eint<6>) -> !FHE.eint<6> {
       %result = "FHE.add_eint"(%arg0, %arg1): (!FHE.eint<6>, !FHE.eint<6>) -> (!FHE.eint<6>)
       return %result: !FHE.eint<6>
   }"""

# create a new compiler engine that will write artifacts to the `out` directory
engine = compiler.LibrarySupport.new("./out")

# create the default compilation options, you can modify those at your convenience via the `set_*` methods
options = compiler.CompilationOptions.new()

# use the compiler engine to compile the program with the options you defined
compilation_result = engine.compile(mlir_input, options)

At this point, you have completed the compilation, and should find both the compiled library [.c-inline-code]sharedlib.so[.c-inline-code], as well as the [.c-inline-code]client_parameters.concrete.params.json[.c-inline-code] under the [.c-inline-code]out[.c-inline-code] directory. The shared library contains the executable function, which can be loaded and executed. The json file is a description of the inputs and outputs, as well as crypto parameters for the compiled function. 

These parameters allow you to generate the right key set, including both secret keys and evaluation keys. Secret keys have the capability of decrypting ciphertexts, and thus should only be accessible to parties doing decryption. Evaluation keys are public material that can be sent to a server in order to run an encrypted computation. Crypto parameters are the keystone for FHE computation to ensure security, correctness, and efficiency. FHE libraries require these parameters to be chosen by the user, but the crypto parameters are an output of the compilation in the Compiler, so you don’t need to provide them.

Now you can make use of the compiled function by generating keys, encrypting the arguments, calling the compiled function, and finally decrypting the result.

client_parameters = engine.load_client_parameters(compilation_result)
key_set = compiler.ClientSupport.key_set(client_parameters)
args = [1, 3]
public_arguments = compiler.ClientSupport.encrypt_arguments(client_parameters, key_set, args)


server_lambda = engine.load_server_lambda(compilation_result)
evaluation_keys = key_set.get_evaluation_keys()
public_result = engine.server_call(server_lambda, public_arguments, evaluation_keys)


result = compiler.ClientSupport.decrypt_result(key_set, public_result)

The client parameters are first loaded from the compilation output directory and keys are generated accordingly (different functions may require a different set of keys). The [.c-inline-code]server_lambda[.c-inline-code] refers to the entry-point of the compiled library. It will be loaded from the [.c-inline-code]sharedlib.so[.c-inline-code]. You get the set of evaluation keys and call the compiled function. The result of the addition is still encrypted, so you should decrypt it at the end.

While the previous code does compile, encrypt, run, and decrypt on the same machine, encrypted inputs and outputs and keys are serializable. This means you can implement a use case (see Typical Workflow here) where the client holds the keys and does all the encryption/decryption, while the server’s job is to run the encrypted computation.

Some Linear Algebra

The previous FHE program is quite simple, but you can still use the same Python code to compile a more advanced example that works on tensors of encrypted integers. You only have to change the [.c-inline-code]mlir_input[.c-inline-code] and the input arguments.

mlir_input = """
func.func @main(%arg0: tensor<4x4x!FHE.eint<6>>, %arg1: tensor<4x2xi7>) -> tensor<4x2x!FHE.eint<6>> {
   %0 = "FHELinalg.matmul_eint_int"(%arg0, %arg1): (tensor<4x4x!FHE.eint<6>>, tensor<4x2xi7>) -> (tensor<4x2x!FHE.eint<6>>)
   %tlu = arith.constant dense<[40, 13, 20, 62, 47, 41, 46, 30, 59, 58, 17, 4, 34, 44, 49, 5, 10, 63, 18, 21, 33, 45, 7, 14, 24, 53, 56, 3, 22, 29, 1, 39, 48, 32, 38, 28, 15, 12, 52, 35, 42, 11, 6, 43, 0, 16, 27, 9, 31, 51, 36, 37, 55, 57, 54, 2, 8, 25, 50, 23, 61, 60, 26, 19]> : tensor<64xi64>
   %result = "FHELinalg.apply_lookup_table"(%0, %tlu): (tensor<4x2x!FHE.eint<6>>, tensor<64xi64>) -> (tensor<4x2x!FHE.eint<6>>)
   return %result: tensor<4x2x!FHE.eint<6>>
}"""

import numpy as np

x = np.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], dtype=np.uint8)
y = np.array([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=np.uint8)
args = [x, y]

The above MLIR function defines a matrix multiplication between a tensor of encrypted values, and a tensor of clear values. It then applies a lookup table representing a nonlinear random function.

If you are not familiar with MLIR, [.c-inline-code]FHE[.c-inline-code] and [.c-inline-code]FHELinalg[.c-inline-code] are two dialects— a grouping of functionality which can be used to extend the MLIR system—which define a set of operations. Using operations from these dialects is how you would define an FHE computation. You can read more about the available operations here and here.

There is also a CLI tool that allows you to compile and debug FHE programs. You can build it from here. We will consider [.c-inline-code]fhe.mlir[.c-inline-code] to contain our first example program.

$ concretecompiler --action=compile -o output_directory fhe.mlir

This is the same as compiling it using the Python API. 

Now the CLI provides some more options for debugging intermediate representations. You can try to run it with the [.c-inline-code]–help[.c-inline-code] option to get all the possible options, but interesting options to check intermediate representations are listed below:

--action=                                   - output mode
    =roundtrip                                       -   Parse input module and regenerate textual representation
    =dump-fhe                                        -   Dump FHE module
    =dump-fhe-no-linalg                              -   Lower FHELinalg to FHE and dump result
    =dump-tfhe                                       -   Lower to TFHE and dump result
    =dump-concrete                                   -   Lower to Concrete and dump result
    =dump-sdfg                                       -   Lower to SDFG operations annd dump result
    =dump-std                                        -   Lower to std and dump result
    =dump-llvm-dialect                               -   Lower to LLVM dialect and dump result
    =dump-llvm-ir                                    -   Lower to LLVM-IR and dump result
    =dump-optimized-llvm-ir                          -   Lower to LLVM-IR, optimize and dump result

The dump-[dialect_name] action allows you to see what the IR looks like at a certain level of the compilation, making sure lowering from one dialect to the other works as expected. This is interesting, especially if you want to understand the compilation pipeline. For example, you can see how our previous example program looks in the Concrete dialect:

$ concretecompiler --action=dump-concrete fhe.mlir
module {
  func.func @main(%arg0: tensor<769xi64>, %arg1: tensor<769xi64>) -> tensor<769xi64> {
    %0 = "Concrete.add_lwe_tensor"(%arg0, %arg1) : (tensor<769xi64>, tensor<769xi64>) -> tensor<769xi64>
    return %0 : tensor<769xi64>
  }
}

At this level, the Compiler has already figured out crypto parameters, and especially the size of LWE ciphertexts, which is apparent in the IR.

The second example program is more verbose as it makes use of more complex operations that work on tensors. [.c-inline-code]fhe2.mlir[.c-inline-code] should contain the second example program:

$ concretecompiler –action=dump-concrete fhe2.mlir
module {
  func.func @main(%arg0: tensor<4x4x4097xi64>, %arg1: tensor<4x2xi7>) -> tensor<4x2x4097xi64> {
    %c0 = arith.constant 0 : index
    %c4 = arith.constant 4 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %cst = arith.constant dense<[40, 13, 20, 62, 47, 41, 46, 30, 59, 58, 17, 4, 34, 44, 49, 5, 10, 63, 18, 21, 33, 45, 7, 14, 24, 53, 56, 3, 22, 29, 1, 39, 48, 32, 38, 28, 15, 12, 52, 35, 42, 11, 6, 43, 0, 16, 27, 9, 31, 51, 36, 37, 55, 57, 54, 2, 8, 25, 50, 23, 61, 60, 26, 19]> : tensor<64xi64>
    %generated = tensor.generate  {
    ^bb0(%arg2: index, %arg3: index, %arg4: index):
      %c0_i64 = arith.constant 0 : i64
      tensor.yield %c0_i64 : i64
    } : tensor<4x2x4097xi64>
    %0 = scf.for %arg2 = %c0 to %c4 step %c1 iter_args(%arg3 = %generated) -> (tensor<4x2x4097xi64>) {
      %3 = scf.for %arg4 = %c0 to %c2 step %c1 iter_args(%arg5 = %arg3) -> (tensor<4x2x4097xi64>) {
        %4 = scf.for %arg6 = %c0 to %c4 step %c1 iter_args(%arg7 = %arg5) -> (tensor<4x2x4097xi64>) {
          %extracted_slice = tensor.extract_slice %arg0[%arg2, %arg6, 0] [1, 1, 4097] [1, 1, 1] : tensor<4x4x4097xi64> to tensor<4097xi64>
          %extracted = tensor.extract %arg1[%arg6, %arg4] : tensor<4x2xi7>
          %extracted_slice_0 = tensor.extract_slice %arg7[%arg2, %arg4, 0] [1, 1, 4097] [1, 1, 1] : tensor<4x2x4097xi64> to tensor<4097xi64>
          %5 = arith.extsi %extracted : i7 to i64
          %6 = "Concrete.mul_cleartext_lwe_tensor"(%extracted_slice, %5) : (tensor<4097xi64>, i64) -> tensor<4097xi64>
          %7 = "Concrete.add_lwe_tensor"(%extracted_slice_0, %6) : (tensor<4097xi64>, tensor<4097xi64>) -> tensor<4097xi64>
          %inserted_slice = tensor.insert_slice %7 into %arg7[%arg2, %arg4, 0] [1, 1, 4097] [1, 1, 1] : tensor<4097xi64> into tensor<4x2x4097xi64>
          scf.yield %inserted_slice : tensor<4x2x4097xi64>
        }
        scf.yield %4 : tensor<4x2x4097xi64>
      }
      scf.yield %3 : tensor<4x2x4097xi64>
    }
    %1 = bufferization.alloc_tensor() : tensor<4x2x4097xi64>
    %2 = scf.for %arg2 = %c0 to %c4 step %c1 iter_args(%arg3 = %1) -> (tensor<4x2x4097xi64>) {
      %3 = scf.for %arg4 = %c0 to %c2 step %c1 iter_args(%arg5 = %arg3) -> (tensor<4x2x4097xi64>) {
        %extracted_slice = tensor.extract_slice %0[%arg2, %arg4, 0] [1, 1, 4097] [1, 1, 1] : tensor<4x2x4097xi64> to tensor<4097xi64>
        %4 = "Concrete.encode_expand_lut_for_bootstrap_tensor"(%cst) {isSigned = false, outputBits = 6 : i32, polySize = 4096 : i32} : (tensor<64xi64>) -> tensor<4096xi64>
        %5 = "Concrete.keyswitch_lwe_tensor"(%extracted_slice) {baseLog = 3 : i32, level = 6 : i32, lwe_dim_in = 4096 : i32, lwe_dim_out = 868 : i32} : (tensor<4097xi64>) -> tensor<869xi64>
        %6 = "Concrete.bootstrap_lwe_tensor"(%5, %4) {baseLog = 22 : i32, glweDimension = 1 : i32, inputLweDim = 868 : i32, level = 1 : i32, polySize = 4096 : i32} : (tensor<869xi64>, tensor<4096xi64>) -> tensor<4097xi64>
        %inserted_slice = tensor.insert_slice %6 into %arg5[%arg2, %arg4, 0] [1, 1, 4097] [1, 1, 1] : tensor<4097xi64> into tensor<4x2x4097xi64>
        scf.yield %inserted_slice : tensor<4x2x4097xi64>
      }
      scf.yield %3 : tensor<4x2x4097xi64>
    }
    return %2 : tensor<4x2x4097xi64>
  }
}

One last set of options to try is [.c-inline-code]--debug[.c-inline-code] (of course!) and [.c-inline-code]--verbose[.c-inline-code]. This tells you a lot about what’s happening during compilation, especially lowering passes.

Compiler Frontends

Seeing your first program, you might wonder if you have to write the MLIR for any circuit you want to run. Thankfully, that’s not the case. You could do it, but it would be more user friendly to build a frontend layer on top of the Compiler, just like Concrete Python. The goal of the frontend would be to go from any representation of your choice to the equivalent MLIR representation from where the Compiler could take over. An extreme example could be to build a drag and drop tool that lets you build a computation graph visually, which then translates it to MLIR. As of today, Concrete Python is the only frontend built on top of the Compiler, but with Concrete being completely open-source, external users are encouraged to build custom frontends for their needs.

There are a CAPI, Python Bindings, and Rust Bindings that let you build an MLIR representation, so you could integrate them with pretty much everything. Below, we show a Python script that generates an example MLIR:

import concrete.lang as concretelang
from concrete.lang.dialects.fhe import EncryptedIntegerType, SubEintOp
from mlir.dialects import func
from mlir.ir import (
   Context,
   InsertionPoint,
   Location,
   Module,
)

with Context() as ctx, Location.unknown():
   concretelang.register_dialects(ctx)
   module = Module.create()
   with InsertionPoint(module.body):
       eint8_type = EncryptedIntegerType.get(ctx, 8)
       parameters = [eint8_type, eint8_type]
       @func.FuncOp.from_py_func(*parameters)
       def main(*args):
           return SubEintOp(eint8_type, args[0], args[1]).result

print(module)

MLIR Python bindings are also packaged as part of Concrete Python. We can combine them with custom FHE dialects to generate a function that takes as input two encrypted integers of 8 bits, subtract one from the other, and return the result. If you run this script, you should get the following MLIR:

module {
  func.func @main(%arg0: !FHE.eint<8>, %arg1: !FHE.eint<8>) -> !FHE.eint<8> {
    %0 = "FHE.sub_eint"(%arg0, %arg1) : (!FHE.eint<8>, !FHE.eint<8>) -> !FHE.eint<8>
    return %0 : !FHE.eint<8>
  }
}

The CAPI differs from the Python API in terms of MLIR code generation, but can be used with much more than just Python. It opens up the ability to build new and creative frontends on top of the Concrete Compiler.

Conclusion

You can use the Concrete Compiler to compile an encrypted computation, but it's also possible to build any frontend you like, as long as it can generate an MLIR program at the end. Using the examples provided above, you are able to use the Compiler to suit your needs and complete your builds without a thorough understanding of cryptography.

Subscribe to the Zama newsletter if you liked this blog post, we will soon share more about the internals of the Zama Concrete compiler and how the parameters are being chosen.

Additional links

• Star the Concrete Github repository to endorse our work.
• Review the Concrete documentation.
• Get support on our community channels.
• Help advance the FHE space with the Zama Bounty Program.