Quick start

This document covers how to compute on encrypted data homomorphically using the Aegis framework. We will walk you through a complete example step-by-step.

The Aegis framework supports two approaches:

Compiling pretrained ONNX models for FHE inference execution
Enabling users to define Python functions that can be translated to perform encrypted computation and FHE inference.

The following sections detail both approaches.

Using pretrained ONNX models

In this mode, the basic workflow of computation is as follows:

Defining the model.
Save the model as a .onnx file.
Annotate the model's input parameters and compile.
Perform FHE inference.

If a pre-trained ONNX model exists, steps 1 and 2 can be skipped.

Defining the model

import torch.nn as nn

# Define a simple model that performs x + y.
class AddModel(nn.Module):
    def forward(self, x, y):
        return x + y

The above code is a typical example of a PyTorch model definition, and any engineer with a basic understanding of AI development can easily write similar code.

Save the model

import torch
import torch.onnx

# Create model instance and set the model to evaluation mode.
model = AddModel()
model.eval()

# Use placeholders to specify the inputs shape.
dummy_x = torch.zeros(1, dtype=torch.float32) 
dummy_y = torch.zeros(1, dtype=torch.float32) 

# Export the model to ONNX format.
torch.onnx.export(
                    model,
                    (dummy_x, dummy_y),                  # model input
                    "add_model.onnx",                    # saved model file name
                    input_names=['input_x', 'input_y'],  # input parameters name
                    output_names=['output'],             # output name
                 )

The above code exports a PyTorch model named AddModel to the ONNX format. The specific steps include:

Creating an instance of the model and setting it to evaluation (inference) mode.
Defining two input placeholder tensors, dummy_x and dummy_y, to specify the shape and data type of the model's inputs.
Calling the torch.onnx.export method to save the model, along with its input and output information, as an ONNX file named add_model.onnx, and assigning user-friendly names (input_x, input_y, and output) to the inputs and outputs.

Perform FHE inference

Prepare the input data, perform FHE model inference, and decrypt the inference results.

# Generate random input data for testing. in real scenarios, actual data should be used.
input_name = session.get_inputs()[0].name
input_shape = session.get_inputs()[0].shape
input_data = np.random.randn(*input_shape).astype(np.float32)

# Perform model inference.
output_name = session.get_outputs()[0].name
enc_data = session.run([output_name], {input_name: input_data})
clr_data = session.decrypt(enc_data)

Define python functions

In this mode, the basic workflow of computation is as follows:

Define the function you want to compute.
Compile the function into FHE binary code.
Run the binary code for homomorphic evaluation.

Defining the function to compute

Here we define a simple addition function:

def add_1d(X1: list[float, 6], X2: list[float, 6]) -> list[float, 6]:
  for i in range(6):
      X1[i] = X1[i] + X2[i]
  return X1

Compile the function

To compile the target function, first create a Server instance and explicitly specify the function along with its inputs' encryption status when invoking the Server instance's Compiler function for compilation.

server = Server()
compile_result = server.compile(add_1d, {{"X1": "encrypted", "X2": "encrypted"}})

Run the binary code for homomorphic evaluation

Prior to homomorphic evaluation, two prerequisites must be fulfilled: generate cryptographic key pairs and prepare/encrypt input arguments.

Generate key pairs

Generating key pairs requires first create a Client instance and then invoking its keygen method to produce the key pair. as shown:
```
client = Client(compile_result)
client.keygen()
eva_keys = "/tmp/eva_keys.bin"
client.save_eva_keys(eva_keys)
```

encrypt input arguments

Encrypting input arguments requires invoking the Client instance's encrypt method as shown:

private_data_1 = np.array([1, 3, 5, 7, 9, 11])
private_data_2 = np.array([2, 4, 6, 8, 10, 12])
enc_val1, enc_val2 = client.encrypt([private_data_1, private_data_2])

After completing the above steps, execute homomorphic evaluation by invoking the Server instance's run method with the encrypted arguments:

server.load_eva_keys(eva_keys)
enc_output = server.run([enc_val1, enc_val2], compile_result)

The homomorphic evaluation result enc_output is encrypted data. To decrypt, invoke the Client instance's decrypt method with this ciphertext:

pt_output = client.decrypt(enc_output)
print(pt_output)

Using pretrained ONNX models​

Defining the model​

Save the model​

Perform FHE inference​

Define python functions​

Defining the function to compute​

Compile the function​

Run the binary code for homomorphic evaluation​