microsoft / EVA

Compiler for the SEAL homomorphic encryption library
MIT License
224 stars 58 forks source link

"bad optional access" (naive matmul) #25

Open Wheest opened 2 years ago

Wheest commented 2 years ago

I am looking at GEMM computations in EVA.

EVA uses vectorised computations, though following the paper Secure Outsourced Matrix Computation and Application to Neural Networks, we can run a naive encrypted matmul by having the vector size be 1.

This issue asks if the EVA Extension Library (EXL) will be released, which may already implement this, however it is not current available as far as I know.

I have tried to implement this, however I am getting an error: RuntimeError: bad optional access.

You can see my example before, is there something I am missing?

#!/usr/bin/env python
from eva import EvaProgram, Input, Output, evaluate
from eva.ckks import CKKSCompiler
from eva.seal import generate_keys
from eva.metric import valuation_mse
import numpy as np

def get_gemm(N, K, M):
    gemm = EvaProgram("gemm", vec_size=1)
    with gemm:
        outputs = [[0] * N] * M
        for n in range(N):
            for m in range(M):
                for k in range(K):
                    x = Input(f"x_{n}_{k}")
                    w = Input(f"w_{k}_{m}")
                    outputs[n][m] += x * m
        for n in range(N):
            for m in range(M):
                Output(f"out_{n}_{m}", outputs[n][m])

    gemm.set_input_scales(25)
    gemm.set_output_ranges(10)
    return gemm

def generate_inputs(N, K):
    inputs = dict()
    i = 0
    for n in range(N):
        for k in range(K):
            inputs[f"x_{n}_{k}"] = [i]
            i += 1
    return inputs

def generate_weights(K, M):
    inputs = dict()
    i = 0
    for k in range(K):
        for m in range(M):
            inputs[f"w_{k}_{m}"] = [i]
            i += 1
    return inputs

def main():
    N, K, M = 8, 8, 8
    inputs = generate_inputs(N, K)
    weights = generate_weights(K, M)
    gemm = get_gemm(N, K, M)

    data = {**weights, **inputs}
    print(data)
    for prog in [gemm]:
        print(f"Compiling {prog.name}")

        compiler = CKKSCompiler()
        compiled, params, signature = compiler.compile(prog)
        public_ctx, secret_ctx = generate_keys(params)
        enc_inputs = public_ctx.encrypt(data, signature)
        print("excuting GEMM")
        enc_outputs = public_ctx.execute(compiled, enc_inputs)
        outputs = secret_ctx.decrypt(enc_outputs, signature)

        reference = evaluate(compiled, inputs)

        print("MSE", valuation_mse(outputs, reference))
        print()

if __name__ == "__main__":
    main()
RoryPotter commented 2 years ago

Naively, I would think that the EVA vector size vec_size needs to be at least 2. This may be increased further by EVA for secuirity reasons. If you duplicate each input value, does that fix your problem?

Note that generating your inputs in this way will cause your matrix mutliplications to be more expensive than it could be, which may cause problems with the computations depending on the complexity of your program.

The EVA extension papers explains the optimal algorithms for matrix multications using a diagonal representation of matrices. It also suggests how to create Vector and Matrix input types which would enable basic linear operations.

Wheest commented 2 years ago

Hi, thanks. I'm trying to build the various approaches to matmul, starting from the simplest and going through Halevi and Shoup and beyond, using EVA.

The naive approach is to stand as my baseline. Once I feel comfortable with this, I may look at implementing the Vector and Matrix types if the EVA extensions are not available.

Regarding your suggestion, it seems that I still get the same error.

 def get_gemm(N, K, M):
-    gemm = EvaProgram("gemm", vec_size=1)
+    gemm = EvaProgram("gemm", vec_size=2)
     with gemm:
         outputs = [[0] * N] * M
         for n in range(N):
@@ -30,7 +30,7 @@ def generate_inputs(N, K):
     i = 0
     for n in range(N):
         for k in range(K):
-            inputs[f"x_{n}_{k}"] = [i]
+            inputs[f"x_{n}_{k}"] = [i, i]
             i += 1
     return inputs

@@ -40,7 +40,7 @@ def generate_weights(K, M):
     i = 0
     for k in range(K):
         for m in range(M):
-            inputs[f"w_{k}_{m}"] = [i]
+            inputs[f"w_{k}_{m}"] = [i, i]
             i += 1
     return inputs
RoryPotter commented 2 years ago

The naive approach is in some ways harder than trying to implement the algorithms, like the horizontal sum which is in the EVA library. It also means you lose the parallelism from SIMD batching.

The code in https://github.com/microsoft/EVA/issues/8#issuecomment-831469050 follows the same similar approach to yours, but uses numpy in a clever way to do the multiplications. Note that doing it this way can lead to very large EVA programs and take a long time.

RoryPotter commented 2 years ago

The code in https://github.com/microsoft/EVA/issues/4#issuecomment-799023880 also provides some useful ideas on how to create EVA programs and work with vectors/matrices.

Wheest commented 2 years ago

Thank you, this has been very helpful. I have got a naive GEMM working, see below. I now need to try and understand how the algorithms for matmul using a vectorised approach works.

I note that the convolution example has rotations and similar, is there any documentation that can give more intuition into what the syntax means there?

#!/usr/bin/env python
from eva import EvaProgram, Input, Output, evaluate
from eva import evaluate
from eva.ckks import CKKSCompiler
from eva.seal import generate_keys
from eva.metric import valuation_mse
from random import uniform
import numpy as np
import unittest
import math
import time

def linear(x, y):
    return np.matmul(x, y)

def generate_inputs(N, K):
    inputs = dict()
    inputs_np = np.zeros((N, K))
    i = 0
    for n in range(N):
        for k in range(K):
            inputs[f"x_{n}_{k}"] = [i]
            inputs_np[n, k] = i
            i += 1
    return inputs, inputs_np

def generate_weights(K, M):
    inputs = dict()
    inputs_np = np.zeros((K, M))
    i = 0
    for k in range(K):
        for m in range(M):
            inputs[f"w_{k}_{m}"] = [i]
            inputs_np[k, m] = i
            i += 1
    return inputs, inputs_np

def generate_matmul(N, K, M):
    gemm = EvaProgram("gemm", vec_size=1)
    with gemm:
        a = np.array([Input(f"x_{n}_{k}") for n in range(N) for k in range(K)]).reshape(
            N, K
        )
        b = np.array([Input(f"w_{k}_{m}") for k in range(K) for m in range(M)]).reshape(
            K, M
        )

        out = linear(a, b)

        for n in range(out.shape[0]):
            for m in range(out.shape[1]):
                Output(f"y_{n}_{m}", out[n][m])

    gemm.set_input_scales(32)
    gemm.set_output_ranges(32)
    return gemm

def benchmark_matmul(N, K, M):
    inputs, inputs_np = generate_inputs(N, K)
    weights, weights_np = generate_weights(K, M)

    matmul = generate_matmul(N, K, M)

    data = {**weights, **inputs}
    compiler = CKKSCompiler(config={"security_level": "128", "warn_vec_size": "false"})
    compiled, params, signature = compiler.compile(matmul)
    public_ctx, secret_ctx = generate_keys(params)
    enc_inputs = public_ctx.encrypt(data, signature)
    start = time.time()
    enc_outputs = public_ctx.execute(compiled, enc_inputs)
    end = time.time()
    run_time = end - start

    outputs = secret_ctx.decrypt(enc_outputs, signature)

    y = np.array([outputs[f"y_{n}_{m}"] for n in range(N) for m in range(M)]).reshape(
        N, M
    )

    start = time.time()
    true_y = linear(inputs_np, weights_np)
    end = time.time()
    plain_run_time = end - start

    correct = np.allclose(y, true_y)
    if not correct:
        raise ValueError(f"We were wrong for size {N}")
    return run_time, plain_run_time

if __name__ == "__main__":
    results_cipher = dict()
    results_plain = dict()
    for size in [4, 8, 16, 32]:
        N, K, M = size, size, size
        time_cipher, time_plain = benchmark_matmul(N, K, M)
        results_cipher[f"{N}_{K}_{M}"] = time_cipher
        results_plain[f"{N}_{K}_{M}"] = time_plain
        print(results_cipher)
        print(results_plain)
        print()