MPC and HEIR

[j2kun]

There's been some chatter internally at Google around MPC applications and compilers. I have mentioned in a few talks recently that we want the work in HEIR to also be useful in these settings, and so this thread is intended to be a place to discuss some of the details of what could be possible, either by extending the secret dialect in a way that is useful for both MPC and FHE, or else by having a new dialect with passes that can apply to both.

For some context (and please correct me if I'm wrong, I'm only learning about MPC's details right now), MPC seems to be based on Shamir secret sharing over a finite ring of coefficients. Each participating party gets one share [x]_i of a secret value x, uses the partial homomorphism of Shamir's scheme to locally compute things like [x+y]_i from [x]_i and [y]_i, but requires complex communication protocols (and/or pre-computed random shares like "Beaver triples") to compute other things.

The main areas of MPC protocol research appears to be

• Designing new MPC primitives
• Reducing the communication complexity of MPC primitives as much as possible, possibly trading off for additional precomputation.
• Figuring out how to represent unsupported operations in terms of MPC primitives (e.g., representing a conv2d in terms of a matrix multiplication by repeating and reordering the matrices involved)

I don't think HEIR will contribute much to the first one, but the second and third are plausible. While some of these could be peephole optimizations, it seems more likely that one would need to do some sort of optimization/synthesis, wherein a global analysis of the protocol runs with a cost model representing communication complexity, though there are some primitives like the RELU function in SecureNN 1 that this could apply to at a smaller scale.

[j2kun]

I'm going to sketch out some ideas here as MLIR pseudocode

We probably want to have different representations of the protocol as a whole vs the operations done by each participant (if/whether MPC is always symmetric I'm not sure)

Maybe a protocol could start with secret representing what data is secret and what ops need to be computed on secret data. (not sure if the type below should be !secret.secret<tensor<Nxi32>> or tensor<Nx!secret.secret<i32>>

func.func @mpc_computation(%input : !secret.secret<tensor<1000xi32>>) {
  %linear_layer = secret.generic
    ins(%weights, %input : tensor<1000x256xi32>, !secret.secret<tensor<1000xi32>>) {
    ^bb0(%plain_weights: tensor<1000x256xi32>, %plain_input: tensor<1000xi32>) :
      %0 = linalg.matmul %plain_input, %weights : tensor<256xi32>
      secret.yield %0 : tensor<256xi32>
    }
    %relu_activation =  secret.generic
    ins(%layer1 : !secret.secret<tensor<256xi32>>) {
    ^bb0(%plain_layer1: tensor<256xi32>) :
      %0 = some_dialect.relu %plain_layer1 : tensor<256xi32>
      secret.yield %0 : tensor<256xi32>
    }
    func.return %relu_activation
}

Secret here can be used to express that all servers involved are unaware the contents of the secret, i.e., they receive only shares. Then a pass could lower this to something like this, where !secret.secret<i32, K> denotes that the secret is known to entity K.

module attributes {mpc.party_count = 3} {
  func.func @mpc_computation_server_0(%input : tensor<1000x!secret.secret<i32, 0>>) {
    <some ops...>
  }
  
  func.func @mpc_computation_server_1(%input : tensor<1000x!secret.secret<i32, 1>>) {
    <some ops...>
  }

  // Maybe this server needs no input because its role in the protocol is to generate beaver triples
  // Main point here is that the contents of each server's role can be asymmetric
  func.func @mpc_computation_server_2() {
    <some ops...>
  }
}

Here "secret" is semantically interpreted as a single share, private to one server unless they choose to reveal it. The ops inside these functions would need to be able to support low-level MPC ops like addition of two shares, addition and multiplication by a static constant and "opening" a share to its plaintext value, given enough shares. So these would be something like:

shamir.add %0, %1 : (!secret.secret<i32, 0>, !secret.secret<i32, 0>) -> !secret.secret<i32, 0>
shamir.add_plaintext %0, %1 : (!secret.secret<i32, 0>, i32) -> !secret.secret<i32, 0>
shamir.reconstruct %2, %3, %4 (!secret.secret<i32, 0>, !secret.secret<i32, 1>, !secret.secret<i32, 2>) -> i32

In the last one, we're assuming that the threshold for recovering a secret from its shares is 3, and that an op verifier would require that the K in !secret.secret<i32, K> is distinct across all options. The op would probably need to be variadic, and have some way to know the chosen threshold for recovery.

There would also need to be a way to construct shares, especially shares consisting of random numbers as used in beaver triples:

%0 = arith.constant 5 : i32
%1 = shamir.split %0 {shares=3} : i32 ->  (!secret.secret<i32, 0>, !secret.secret<i32, 1>, !secret.secret<i32, 2>)

One would also need a mechanism to communicate, like

mpc.broadcast %0 : !secret.secret<i32, 0>
mpc.send %0 {recipient=1} : !secret.secret<i32, 0> -> ()
mpc.receive {recipient=1} : () -> !secret.secret<i32, 1>

I think there's a bit of delicacy here in reusing secret.secret. Because the value x is secret, for which we have shares [x]_0, [x]_1, [x]_2, but each share is also secret to each server until they decide to send them to other servers in the protocol. So an op like mpc.receive {recipient=1} : () -> !secret.secret<i32, 1> is a bit misleading: once server 0 receives it, that share is no longer secret to server 1. But it's still annotated with 1 because it was the share originally constructed for server 1, and we need to track whether we have distinct shares in order to reconstruct the secret. There are also situations in which, say, server 2 is responsible for generating shares that are distributed to servers 0 and 1, so we might have something like this running on server 2:

%0 = random.rand_int : i32
%s0, %s1 = shamir.split %0 {shares=2} : i32 ->  (!secret.secret<i32, 0>, !secret.secret<i32, 1>)
mpc.send %s0 {recipient=0}: !secret.secret<i32, 0> -> ()
mpc.send %s1 {recipient=1}: !secret.secret<i32, 0> -> ()

The optimizations could come in:

• At the step of the high-level secret ops, say, reordering the network architecture greedily to reduce later communication costs (as done in https://eprint.iacr.org/2019/1049 for RELU/MaxPool)
• While lowering high-level generic ops to their implementation in shamir, effectively picking the best representation of the high level ops in MPC primitives
• After all the primitives have been lowered to shamir + mpc, at which point optimization patterns/dataflow could apply across primitives, and the goal would be to minimize communication costs.
• After splitting into multiple servers, in trying to re-balance the work across servers to reduce latency and/or communication. One of the optimizations in https://eprint.iacr.org/2019/1049 for RELU seems to do this, effectively recognizing what secrets are effectively random and "load balancing" the sending of those shares from server 2 (the RNG server) to the other two servers.

The next step IMO is to pick a layer at which we want to demonstrate some interesting optimization can occur.

[j2kun]

Coincidentally, there was an ODM just yesterday on MPI in MLIR, which might cover some of the send/receive parts of what we want to do. https://www.youtube.com/watch?v=VB6Ec3RnMEE

[AlexanderViand-Intel]

This is a great topic! Combining FHE+MPC(+ZKP) exponentially increases the range of applications for which we can achieve practical solutions, but it also exponentially increases the complexity of designing these solutions. I think there's more than enough here to have an entire meeting/WG session just dedicated to this topic.

Note that secret sharing makes up only one half of the MPC world, with the other being Garbled Circuit-based approaches. Kind of like the LWE/RLWE split in FHE, one's better at arithmetic (secret sharing) and one's better at arbitrary binary things (GC). Also note, that Shamir secret sharing is only one approach to secret-sharing-based MPC, and there are others (e.g., additive secret sharing).

There's a lot of prior art in this space and I might provide a more complete/curated reading list at some point, but for now I'll just throw out some keywords/things to look at:

• Marcella's SoK on MPC (SP'19) gives a great overview over the different frameworks:
  https://ieeexplore.ieee.org/document/8835312 https://github.com/MPC-SoK/frameworks

• Note that these frameworks often contain "compilers", but they're more like basic DSL->Circuit translators and generally not super advanced.

• In practice, I find MP-SPDZ to be by far the most actively developed/user-friendly framework that still offers cutting edge research things
  https://eprint.iacr.org/2020/521 https://github.com/data61/MP-SPDZ

• There are, however, also "real" compilers for MPC that do a lot more than just basic translation. These come in a variety of flavours, but the most important ones are the hybird protocol compilers, which automatically partition a program and switch between different types of MPC (kind of like FHE scheme switching). Looking at ABY and it's follow-ups should be a good start. However, there are also MPC compilers that are more generally focussed, e.g., Viaduct or CirC

[AlexanderViand-Intel]

Oh, and a lot of MPC protocols internally use FHE (mostly BGV, iirc). Afaik, this is mostly in the offline-phase (which, is "offline" in the "offline algorithm" sense, not in the networking sense - there's usually lots of data sent around!), to generate Beaver triples.

These generally use their own homebrew implementation of the scheme (usually no keyswitching needed, etc) and I've always wondered if they could be outperformed by the state of the art libraries or by pre-compiling the circuit with something like HEIR

[asraa]

Like Alex mentioned above, even describing secret sharing has other axes (shamir vs additive) and also threshold values that can make the reconstruction parameters different. Obviously starting with one style of lowering is better.

There's also a lot of prior MPC + compiler work, one of which is https://ieeexplore.ieee.org/document/9797347 (developing cost models for MPC) and Silph (https://eprint.iacr.org/2023/060).

[j2kun]

I came across this interesting paper that characterizes the inefficiencies in PI schemes (hybrid or not): https://arxiv.org/abs/2207.07177

Something to note is that the networks are much larger (Res-Net 18, with 150 layers), but come with even higher bandwidth and storage costs than FHE (~50G / inference) and, depending on the request rate (e.g., 1 request per 15 minutes), the inference latency can take up to 30 minutes.

[j2kun]

I was also looking around at possible MPC backends: https://github.com/rdragos/awesome-mpc#frameworks

This interesting project has implementations in a bunch of these libraries: https://github.com/MPC-SoK/frameworks

[j2kun]

Ah, this one is the same paper that Alex posted above: "Marcella's SoK on MPC"

[j2kun]

Looks like this paper from Usenix 23 is based on MLIR, has an HLO front-end (from JAX), and lowers to XLA

• paper: https://www.usenix.org/system/files/atc23-ma.pdf
• tablegen: https://github.com/secretflow/spu/tree/main/libspu/dialect

[asraa]

That SPU paper cites and uses Cheetah, a 2PC protocol for neural net inference: https://eprint.iacr.org/2022/207.pdf

[j2kun]

They also have this comment:

Coding practice and conventions in this repository follow the MLIR Developer Guide in this repo as part of the intent to act as an incubator for technology to upstream.

[AlexanderViand-Intel]

That's a great line in general, and "incubator for technology to upstream" is a nice description!

[j2kun]

Just pointing out a few things as I go:

• It seems that all the tutorial examples in the ABY library are broken, and open GH issues complain about memory leaks and other instabilities and broken things. Seems like a dud to me.
• The MP-SPDZ library executes its protocols in a virtual machine, so it seems unsuitable as a backend target.

[j2kun]

A small update: I was able to get TinyGarble to work with some of our existing boolean circuits, resulting in an inference on a small neural network (350k boolean gates, roughly half of them nands) in about 0.1s on my local machine.

[j2kun]

This means that we could use our existing compiler flow, use heir-translate --emit-verilog, and then continue on through TinyGarble to get a working baseline MPC implementation. It does require a few minor changes to the generated verilog API (needs its inputs in a specific order/name)

[j2kun]

TIL that Intel Labs has TinyGarble2: https://github.com/IntelLabs/TinyGarble2.0

@AlexanderViand-Intel maybe there's interest in integrating it? Interestingly, the paper for this project has a Google contributor who is on my current team: Baiyu Li

[AlexanderViand-Intel]

While there's not a lot of interest at Intel (I did talk to Ro about TinyGarble2 but he's of course much more focused on FHE at the moment) and I guess interest from your side has also been limited, I do have an update on "MPC and HEIR":

I'm likely going to be supervising an ETH Zurich MSc thesis this fall that's going to be aimed at building an MLIR-based MPC compiler (of course based on HEIR / secret, and maybe also reusing some BGV/(R)LWE stuff for a SPDZ triple generation implementation)