With modules

This document explains how to compile Fully Homomorphic Encryption (FHE) modules containing multiple functions using Concrete.

Deploying a server that contains many compatible functions is important for some use cases. With Concrete, you can compile FHE modules containing as many functions as needed.

These modules support the composition of different functions, meaning that the encrypted result of one function can be used as the input for another function without needing to decrypt it first. Additionally, a module is deployed in a single artifact, making it as simple to use as a single-function project.

Single inputs / outputs

The following example demonstrates how to create an FHE module:

from concrete import fhe

@fhe.module()
class Counter:
    @fhe.function({"x": "encrypted"})
    def inc(x):
        return x + 1 % 20

    @fhe.function({"x": "encrypted"})
    def dec(x):
        return x - 1 % 20

Then, to compile the Counter module, use the compile method with a dictionary of input-sets for each function:

inputset = list(range(20))
CounterFhe = Counter.compile({"inc": inputset, "dec": inputset})

After the module is compiled, you can encrypt and call the different functions as follows:

x = 5
x_enc = CounterFhe.inc.encrypt(x)
x_inc_enc = CounterFhe.inc.run(x_enc)
x_inc = CounterFhe.inc.decrypt(x_inc_enc)
assert x_inc == 6

x_inc_dec_enc = CounterFhe.dec.run(x_inc_enc)
x_inc_dec = CounterFhe.dec.decrypt(x_inc_dec_enc)
assert x_inc_dec == 5

for _ in range(10):
    x_enc = CounterFhe.inc.run(x_enc)
x_dec = CounterFhe.inc.decrypt(x_enc)
assert x_dec == 15

You can generate the keyset beforehand by calling keygen() method on the compiled module:

CounterFhe.keygen()

Multi inputs / outputs

Composition is not limited to single input / single output. Here is an example that computes the 10 first elements of the Fibonacci sequence in FHE:

from concrete import fhe

def noise_reset(x):
   return fhe.univariate(lambda x: x)(x)

@fhe.module()
class Fibonacci:

    @fhe.function({"n1th": "encrypted", "nth": "encrypted"})
    def fib(n1th, nth):
       return noise_reset(nth), noise_reset(n1th + nth)

print("Compiling `Fibonacci` module ...")
inputset = list(zip(range(0, 100), range(0, 100)))
FibonacciFhe = Fibonacci.compile({"fib": inputset})

print("Generating keyset ...")
FibonacciFhe.keygen()

print("Encrypting initial values")
n1th = 1
nth = 2
(n1th_enc, nth_enc) = FibonacciFhe.fib.encrypt(n1th, nth)

print(f"|           ||        (n-1)-th       |         n-th          |")
print(f"| iteration || decrypted | cleartext | decrypted | cleartext |")
for i in range(10):
   (n1th_enc, nth_enc) = FibonacciFhe.fib.run(n1th_enc, nth_enc)
   (n1th, nth) = Fibonacci.fib(n1th, nth)

    # For demo purpose; no decryption is needed.
   (n1th_dec, nth_dec) = FibonacciFhe.fib.decrypt(n1th_enc, nth_enc)
   print(f"|     {i}     || {n1th_dec:<9} | {n1th:<9} | {nth_dec:<9} | {nth:<9} |")

Executing this script will provide the following output:

Compiling `Fibonacci` module ...
Generating keyset ...
Encrypting initial values
|           ||        (n-1)-th       |         n-th          |
| iteration || decrypted | cleartext | decrypted | cleartext |
|     0     || 2         | 2         | 3         | 3         |
|     1     || 3         | 3         | 5         | 5         |
|     2     || 5         | 5         | 8         | 8         |
|     3     || 8         | 8         | 13        | 13        |
|     4     || 13        | 13        | 21        | 21        |
|     5     || 21        | 21        | 34        | 34        |
|     6     || 34        | 34        | 55        | 55        |
|     7     || 55        | 55        | 89        | 89        |
|     8     || 89        | 89        | 144       | 144       |
|     9     || 144       | 144       | 233       | 233       |

Iterations

Modules support iteration with cleartext iterands to some extent, particularly for loops structured like this:

for i in some_cleartext_constant_range:
    # Do something in FHE in the loop body, implemented as an FHE function.

Unbounded loops or complex dynamic conditions are also supported, as long as these conditions are computed in pure cleartext in Python. The following example computes the Collatz sequence:

from concrete import fhe

@fhe.module()
class Collatz:

    @fhe.function({"x": "encrypted"})
    def collatz(x):

        y = x // 2
        z = 3 * x + 1

        is_x_odd = fhe.bits(x)[0]

        # In a fast way, compute ans = is_x_odd * (z - y) + y
        ans = fhe.multivariate(lambda b, x: b * x)(is_x_odd, z - y) + y

        is_one = ans == 1

        return ans, is_one


print("Compiling `Collatz` module ...")
inputset = [i for i in range(63)]
CollatzFhe = Collatz.compile({"collatz": inputset})

print("Generating keyset ...")
CollatzFhe.keygen()

print("Encrypting initial value")
x = 19
x_enc = CollatzFhe.collatz.encrypt(x)
is_one_enc = None

print(f"| decrypted | cleartext |")
while is_one_enc is None or not CollatzFhe.collatz.decrypt(is_one_enc):
    x_enc, is_one_enc = CollatzFhe.collatz.run(x_enc)
    x, is_one = Collatz.collatz(x)

    # For demo purpose; no decryption is needed.
    x_dec = CollatzFhe.collatz.decrypt(x_enc)
    print(f"| {x_dec:<9} | {x:<9} |")

This script prints the following output:

Compiling `Collatz` module ...
Generating keyset ...
Encrypting initial value
| decrypted | cleartext |
| 58        | 58        |
| 29        | 29        |
| 88        | 88        |
| 44        | 44        |
| 22        | 22        |
| 11        | 11        |
| 34        | 34        |
| 17        | 17        |
| 52        | 52        |
| 26        | 26        |
| 13        | 13        |
| 40        | 40        |
| 20        | 20        |
| 10        | 10        |
| 5         | 5         |
| 16        | 16        |
| 8         | 8         |
| 4         | 4         |
| 2         | 2         |
| 1         | 1         |

In this example, a while loop iterates until the decrypted value equals 1. The loop body is implemented in FHE, but the iteration control must be in cleartext.

Runtime optimization

By default, when using modules, all inputs and outputs of every function are compatible, sharing the same precision and crypto-parameters. This approach applies the crypto-parameters of the most costly code path to all code paths. This simplicity may be costly and unnecessary for some use cases.

To optimize runtime, we provide finer-grained control over the composition policy via the composition module attribute. Here is an example:

from concrete import fhe

@fhe.module()
class Collatz:

    @fhe.function({"x": "encrypted"})
    def collatz(x):
        y = x // 2
        z = 3 * x + 1
        is_x_odd = fhe.bits(x)[0]
        ans = fhe.multivariate(lambda b, x: b * x)(is_x_odd, z - y) + y
        is_one = ans == 1
        return ans, is_one

    composition = fhe.AllComposable()

You have 3 options for the composition attribute:

  1. fhe.AllComposable (default): This policy ensures that all ciphertexts used in the module are compatible. It is the least restrictive policy but the most costly in terms of performance.

  2. fhe.NotComposable: This policy is the most restrictive but the least costly. It is suitable when you do not need any composition and only want to pack multiple functions in a single artifact.

  3. fhe.Wired: This policy allows you to define custom composition rules. You can specify which outputs of a function can be forwarded to which inputs of another function.

Note that, in case of complex composition logic another option is to rely on [[composing_functions_with_modules#Automatic module tracing]] to automatically derive the composition from examples.

Here is an example:
from concrete import fhe
from fhe import Wired, Wire, Output, Input

@fhe.module()
class Collatz:

    @fhe.function({"x": "encrypted"})
    def collatz(x):
        y = x // 2
        z = 3 * x + 1
        is_x_odd = fhe.bits(x)[0]
        ans = fhe.multivariate(lambda b, x: b * x)(is_x_odd, z - y) + y
        is_one = ans == 1
        return ans, is_one

    composition = Wired(
        {
            Wire(Output(collatz, 0), Input(collatz, 0)
        }
    )

In this case, the policy states that the first output of the collatz function can be forwarded to the first input of collatz, but not the second output (which is decrypted every time, and used for control flow).

You can use the fhe.Wire between any two functions. It is also possible to define wires with fhe.AllInputs and fhe.AllOutputs ends. For instance, in the previous example:

    composition = Wired(
        {
            Wire(AllOutputs(collatz), AllInputs(collatz))
        }
    )

This policy would be equivalent to using the fhe.AllComposable policy.

Automatic module tracing

When a module's composition logic is static and straightforward, declaratively defining a Wired policy is usually the simplest approach. However, in cases where modules have more complex or dynamic composition logic, deriving an accurate list of Wire components to be used in the policy can become challenging.

Another related problem is defining different function input-sets. When the composition logic is simple, these can be provided manually. But as the composition gets more convoluted, computing a consistent ensemble of inputsets for a module may become intractable.

For those advanced cases, you can derive the composition rules and the input-sets automatically from user-provided examples. Consider the following module:

from concrete import fhe
from fhe import Wired

@fhe.module()
class MyModule:
    @fhe.function({"x": "encrypted"})
    def increment(x):
        return (x + 1) % 100

    @fhe.function({"x": "encrypted"})
    def decrement(x):
        return (x - 1) % 100

    @fhe.function({"x": "encrypted"})
    def decimate(x):
        return (x / 10) % 100

    composition = fhe.Wired()

You can use the wire_pipeline context manager to activate the module tracing functionality:

# A single inputset used during tracing is defined
inputset = [np.random.randint(1, 100, size=()) for _ in range(100)]

# The inputset is passed to the `wire_pipeline` method, which itself returns an iterator over the inputset samples.
with MyModule.wire_pipeline(inputset) as samples_iter:

    # The inputset is iterated over
    for s in samples_iter:

        # Here we provide an example of how we expect the module functions to be used at runtime in fhe.
        MyModule.increment(MyModule.decimate(MyModule.decrement(s)))

# It is not needed to provide any inputsets to the `compile` method after tracing the wires, since those were already computed automatically during the module tracing.
module = MyModule.compile(
    p_error=0.01,
)

Note that any dynamic branching is possible during module tracing. However, for complex runtime logic, ensure that the input set provides sufficient examples to cover all potential code paths.

Current Limitations

Depending on the functions, composition may add a significant overhead compared to a non-composable version.

To be composable, a function must meet the following condition: every output that can be forwarded as input (according to the composition policy) must contain a noise-refreshing operation. Since adding a noise refresh has a noticeable impact on performance, Concrete does not automatically include it.

For instance, to implement a function that doubles an encrypted value, you might write:

@fhe.module()
class Doubler:
    @fhe.compiler({"counter": "encrypted"})
    def double(counter):
       return counter * 2

This function is valid with the fhe.NotComposable policy. However, if compiled with the fhe.AllComposable policy, it will raise a RuntimeError: Program cannot be composed: ..., indicating that an extra Programmable Bootstrapping (PBS) step must be added.

To resolve this and make the circuit valid, add a PBS at the end of the circuit:

@fhe.module()
class Doubler:
    @fhe.compiler({"counter": "encrypted"})
    def double(counter):
       return fhe.refresh(counter * 2)

Last updated