Draft OPA implementation (#15)

This commit is contained in:
Alex Ozdemir
2021-08-24 08:29:54 -07:00
committed by GitHub
parent 3691bb91b8
commit 2226ad901f
7 changed files with 618 additions and 5 deletions

30
Cargo.lock generated
View File

@@ -235,6 +235,7 @@ dependencies = [
"rand",
"rsmt2",
"rug",
"serde_json",
"structopt",
"thiserror",
"typed-arena",
@@ -573,6 +574,12 @@ dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "0.4.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b71991ff56294aa922b450139ee08b3bfc70982c6b2c7562771375cf73542dd4"
[[package]]
name = "lazy_static"
version = "1.4.0"
@@ -940,12 +947,35 @@ version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e3bad0ee36814ca07d7968269dd4b7ec89ec2da10c4bb613928d3077083c232"
[[package]]
name = "ryu"
version = "1.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "71d301d4193d031abdd79ff7e3dd721168a9572ef3fe51a1517aba235bd8f86e"
[[package]]
name = "scopeguard"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
[[package]]
name = "serde"
version = "1.0.129"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d1f72836d2aa753853178eda473a3b9d8e4eefdaf20523b919677e6de489f8f1"
[[package]]
name = "serde_json"
version = "1.0.66"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "336b10da19a12ad094b59d870ebde26a45402e5b470add4b5fd03c5048a32127"
dependencies = [
"itoa",
"ryu",
"serde",
]
[[package]]
name = "sha-1"
version = "0.8.2"

View File

@@ -28,6 +28,7 @@ funty = "=1.1"
ahash = "0.7"
good_lp = { version = "1.1", features = ["lp-solvers", "coin_cbc"], default-features = false }
lp-solvers = "0.0.4"
serde_json = "1.0"
[dev-dependencies]
quickcheck = "1"

View File

@@ -0,0 +1,393 @@
//! ILP-based sharing assignment
//!
//! Loosely based on ["Efficient MPC via Program Analysis: A Framework for Efficient Optimal
//! Mixing"](https://dl.acm.org/doi/pdf/10.1145/3319535.3339818) by Ishaq, Muhammad and Milanova,
//! Ana L. and Zikas, Vassilis.
//!
//! Our actual ILP is as follows:
//!
//! Let `s`, `t` denote terms, and `a`, `b` denote protocols.
//!
//! Let `T[t, a]` be a binary variable indicating whether term `t` is evaluated using protocol `a`.
//! Let `C[t, a, b]` be a binary variable indicating whether term `t` needs to be converted from
//! `a` to `b`.
//!
//! Since each term is evaluated using one protocol,
//!
//! `forall t. 1 = \sum_a T[t, a] (1)`
//!
//! Sometimes conversions are needed
//!
//! `forall t a b. forall s in Uses(t). C[t, a, b] >= T[t, a] + T[s, b] - 1 (2)`
//!
//! The constraint (2) is intendend to encode
//!
//! `forall t a b. C[t, a, b] = OR_(s in Uses(t)) T[t, a] AND T[s, b]`
//!
//! It does this well because (a) the system is SAT and (b) our objective is a linear combination
//! of all variables (term and conversion) scaled by their cost. In trying to minimize that, `C`
//! will be set to the smallest value possible (0) if either of the variables on the right of (2)
//! are 0. If they are both 1 (for ANY `s`), then it must be 1.
use ahash::{AHashMap, AHashSet};
use serde_json::Value;
use super::{ShareType, SharingMap, SHARE_TYPES};
use crate::ir::term::*;
use crate::target::ilp::{variable, Expression, Ilp, Variable};
use std::{fs::File, path::Path, env::var};
/// A cost model for ABY operations and share conversions
#[derive(Debug)]
pub struct CostModel {
/// Conversion costs: maps (from, to) pairs to cost
conversions: AHashMap<(ShareType, ShareType), f64>,
/// Operator costs: maps (op, type) to cost
ops: AHashMap<Op, AHashMap<ShareType, f64>>,
}
impl CostModel {
/// Create a cost model from an OPA json file, like [this](https://github.com/ishaq/OPA/blob/d613c15ff715fa62c03e37b673548f94c16bfe0d/solver/sample-costs.json)
pub fn from_opa_cost_file(p: &impl AsRef<Path>) -> CostModel {
use ShareType::*;
let get_cost_opt =
|op_name: &str, obj: &serde_json::map::Map<String, Value>| -> Option<f64> {
let o = obj.get(op_name)?;
Some(
o.get("1")
.unwrap_or_else(|| panic!("Missing op '1' entry in {:#?}", o))
.as_f64()
.expect("not a number"),
)
};
let get_cost = |op_name: &str, obj: &serde_json::map::Map<String, Value>| -> f64 {
get_cost_opt(op_name, obj).unwrap()
};
let mut conversions = AHashMap::new();
let mut ops = AHashMap::new();
let f = File::open(p).expect("Missing file");
let json: Value = serde_json::from_reader(f).expect("Bad JSON");
let obj = json.as_object().unwrap();
for (_width, json) in obj {
//let w = u32::from_str(width).expect("bad width");
let obj = json.as_object().unwrap();
// conversions
conversions.insert((Arithmetic, Boolean), get_cost("a2b", obj));
conversions.insert((Boolean, Arithmetic), get_cost("b2a", obj));
conversions.insert((Yao, Boolean), get_cost("y2b", obj));
conversions.insert((Boolean, Yao), get_cost("b2y", obj));
conversions.insert((Yao, Arithmetic), get_cost("y2a", obj));
conversions.insert((Arithmetic, Yao), get_cost("a2y", obj));
let ops_from_name = |name: &str| {
match name {
// assume comparisions are unsigned
"ge" => vec![BV_ULE],
"le" => vec![BV_SLE],
"gt" => vec![BV_ULT],
"lt" => vec![BV_SLT],
// assume n-ary ops apply to BVs
"add" => vec![BV_ADD],
"mul" => vec![BV_MUL],
"and" => vec![BV_AND],
"or" => vec![BV_OR],
"xor" => vec![BV_XOR],
// assume eq applies to BVs
"eq" => vec![Op::Eq],
"shl" => vec![BV_SHL],
// assume shr is logical, not arithmetic
"shr" => vec![BV_LSHR],
"sub" => vec![BV_SUB],
"mux" => vec![ITE],
"ne" => vec![],
_ => panic!("Unknown operator name: {}", name),
}
};
for (op_name, json) in obj {
// HACK: assumes the presence of 2 partitions names into conversion and otherwise.
if !op_name.contains("2") {
for op in ops_from_name(op_name) {
let obj = json.as_object().unwrap();
for (share_type, share_name) in
&[(Arithmetic, "a"), (Boolean, "b"), (Yao, "y")]
{
if let Some(cost) = get_cost_opt(share_name, obj) {
ops.entry(op.clone())
.or_insert_with(|| AHashMap::new())
.insert(*share_type, cost);
}
}
}
}
}
}
CostModel { conversions, ops }
}
}
/// Uses an ILP to assign...
pub fn assign(c: &Computation) -> SharingMap {
let p = format!(
"{}/third_party/opa/sample_costs.json",
var("CARGO_MANIFEST_DIR").expect("Could not find env var CARGO_MANIFEST_DIR")
);
let costs = CostModel::from_opa_cost_file(&p);
build_ilp(c, &costs)
}
fn build_ilp(c: &Computation, costs: &CostModel) -> SharingMap {
let mut terms: TermSet = TermSet::new();
let mut def_uses: AHashSet<(Term, Term)> = AHashSet::new();
for o in &c.outputs {
for t in PostOrderIter::new(o.clone()) {
terms.insert(t.clone());
for c in &t.cs {
def_uses.insert((c.clone(), t.clone()));
}
}
}
let terms: AHashMap<Term, usize> = terms.into_iter().enumerate().map(|(i, t)| (t, i)).collect();
let mut term_vars: AHashMap<(Term, ShareType), (Variable, f64, String)> = AHashMap::new();
let mut conv_vars: AHashMap<(Term, ShareType, ShareType), (Variable, f64)> = AHashMap::new();
let mut ilp = Ilp::new();
// build variables for all term assignments
for (t, i) in terms.iter() {
let mut vars = vec![];
if let Op::Var(_, _) = &t.op {
for ty in &SHARE_TYPES {
let name = format!("t_{}_{}", i, ty.char());
let v = ilp.new_variable(variable().binary(), name.clone());
term_vars.insert((t.clone(), *ty), (v, 0.0, name));
vars.push(v);
}
} else if let Some(costs) = costs.ops.get(&t.op) {
for (ty, cost) in costs {
let name = format!("t_{}_{}", i, ty.char());
let v = ilp.new_variable(variable().binary(), name.clone());
term_vars.insert((t.clone(), *ty), (v, *cost, name));
vars.push(v);
}
} else {
panic!("No cost for op {}", &t.op)
}
// Sum of assignments is at least 1.
ilp.new_constraint(
vars.into_iter()
.fold((0.0).into(), |acc: Expression, v| acc + v)
>> 1.0,
);
}
// build variables for all conversions assignments
for (def, use_) in &def_uses {
let def_i = terms.get(def).unwrap();
for from_ty in &SHARE_TYPES {
for to_ty in &SHARE_TYPES {
// if def can be from_ty, and use can be to_ty
if term_vars.contains_key(&(def.clone(), *from_ty))
&& term_vars.contains_key(&(use_.clone(), *to_ty))
&& from_ty != to_ty
{
let v = ilp.new_variable(
variable().binary(),
format!("c_{}_{}2{}", def_i, from_ty.char(), to_ty.char()),
);
dbg!((from_ty, to_ty));
conv_vars.insert(
(def.clone(), *from_ty, *to_ty),
(v, *costs.conversions.get(&(*from_ty, *to_ty)).unwrap()),
);
}
}
}
}
let def_uses: AHashMap<Term, Vec<Term>> = {
let mut t = AHashMap::new();
for (d, u) in def_uses {
t.entry(d).or_insert_with(|| Vec::new()).push(u);
}
t
};
for (def, uses) in def_uses {
for use_ in uses {
for from_ty in &SHARE_TYPES {
for to_ty in &SHARE_TYPES {
conv_vars.get(&(def.clone(), *from_ty, *to_ty)).map(|c| {
term_vars.get(&(def.clone(), *from_ty)).map(|t_from| {
// c[term i from pi to pi'] >= t[term j with pi'] + t[term i with pi] - 1
term_vars
.get(&(use_.clone(), *to_ty))
.map(|t_to| ilp.new_constraint(c.0 >> t_from.0 + t_to.0 - 1.0))
})
});
}
}
}
}
ilp.maximize(
-conv_vars
.values()
.map(|(a, b)| (a, b))
.chain(term_vars.values().map(|(a, b, _)| (a, b)))
.fold(0.0.into(), |acc: Expression, (v, cost)| {
acc + v.clone() * *cost
}),
);
println!("ILP: {:#?}", ilp);
let (_opt, solution) = ilp.default_solve().unwrap();
println!("Solution: {:#?}", solution);
let mut assignment = TermMap::new();
for ((term, ty), (_, _, var_name)) in &term_vars {
if solution.get(var_name).unwrap() == &1.0 {
assignment.insert(term.clone(), *ty);
}
}
assignment
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_cost_model() {
let p = format!(
"{}/third_party/opa/sample_costs.json",
var("CARGO_MANIFEST_DIR").expect("Could not find env var CARGO_MANIFEST_DIR")
);
let c = CostModel::from_opa_cost_file(&p);
// random checks from the file...
assert_eq!(
&1127.0,
c.ops.get(&BV_MUL).unwrap().get(&ShareType::Yao).unwrap()
);
assert_eq!(
&1731.0,
c.ops
.get(&BV_MUL)
.unwrap()
.get(&ShareType::Boolean)
.unwrap()
);
assert_eq!(
&7.0,
c.ops
.get(&BV_XOR)
.unwrap()
.get(&ShareType::Boolean)
.unwrap()
);
}
#[test]
fn mul1_bv_opt() {
let p = format!(
"{}/third_party/opa/sample_costs.json",
var("CARGO_MANIFEST_DIR").expect("Could not find env var CARGO_MANIFEST_DIR")
);
let costs = CostModel::from_opa_cost_file(&p);
let cs = Computation {
outputs: vec![term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
leaf_term(Op::Var("b".to_owned(), Sort::BitVector(32)))
]],
metadata: ComputationMetadata::default(),
values: None,
};
let assignment = build_ilp(&cs, &costs);
dbg!(&assignment);
}
#[test]
fn huge_mul_then_eq() {
let p = format!(
"{}/third_party/opa/sample_costs.json",
var("CARGO_MANIFEST_DIR").expect("Could not find env var CARGO_MANIFEST_DIR")
);
let costs = CostModel::from_opa_cost_file(&p);
let cs = Computation {
outputs: vec![term![Op::Eq;
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32)))
]
]
]
]
]
]
],
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32)))
]],
metadata: ComputationMetadata::default(),
values: None,
};
let assignment = build_ilp(&cs, &costs);
// Big enough to do the math with arith
assert_eq!(
&ShareType::Arithmetic,
assignment.get(&cs.outputs[0].cs[0]).unwrap()
);
// Then convert to boolean
assert_eq!(&ShareType::Boolean, assignment.get(&cs.outputs[0]).unwrap());
}
#[test]
fn big_mul_then_eq() {
let p = format!(
"{}/third_party/opa/sample_costs.json",
var("CARGO_MANIFEST_DIR").expect("Could not find env var CARGO_MANIFEST_DIR")
);
let costs = CostModel::from_opa_cost_file(&p);
let cs = Computation {
outputs: vec![term![Op::Eq;
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
term![BV_MUL;
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32))),
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32)))
]
]
]
],
leaf_term(Op::Var("a".to_owned(), Sort::BitVector(32)))
]],
metadata: ComputationMetadata::default(),
values: None,
};
let assignment = build_ilp(&cs, &costs);
// All yao
assert_eq!(
&ShareType::Yao,
assignment.get(&cs.outputs[0].cs[0]).unwrap()
);
assert_eq!(&ShareType::Yao, assignment.get(&cs.outputs[0]).unwrap());
}
}

View File

@@ -2,8 +2,11 @@
use crate::ir::term::{Computation, PostOrderIter, TermMap};
pub mod ilp;
/// The sharing scheme used for an operation
pub enum SharingType {
#[derive(Debug,PartialEq,Eq,Hash,Clone,Copy)]
pub enum ShareType {
/// Arithmetic sharing (additive mod `Z_(2^l)`)
Arithmetic,
/// Boolean sharing (additive mod `Z_2`)
@@ -12,15 +15,29 @@ pub enum SharingType {
Yao,
}
/// List of share types.
pub const SHARE_TYPES: [ShareType; 3] = [ShareType::Arithmetic, ShareType::Boolean, ShareType::Yao];
impl ShareType {
fn char(&self) -> char {
match self {
&ShareType::Arithmetic => 'a',
&ShareType::Yao => 'y',
&ShareType::Boolean => 'b',
}
}
}
/// A map from terms (operations or inputs) to sharing schemes they use
pub type SharingMap = TermMap<SharingType>;
pub type SharingMap = TermMap<ShareType>;
/// Assigns boolean sharing to all terms
pub fn all_boolean_sharing(c: &Computation) -> SharingMap {
c.outputs
.iter()
.flat_map(|output| {
PostOrderIter::new(output.clone()).map(|term| (term.clone(), SharingType::Boolean))
PostOrderIter::new(output.clone()).map(|term| (term.clone(), ShareType::Boolean))
})
.collect()
}

View File

@@ -3,11 +3,12 @@
pub mod trans;
use ahash::AHashMap as HashMap;
use good_lp::{
pub(crate) use good_lp::{
Constraint, Expression, ProblemVariables, ResolutionError, Solution, Solver, SolverModel,
Variable, VariableDefinition,
Variable, VariableDefinition, variable
};
use log::debug;
use std::fmt::{self, Formatter, Debug};
/// An integer linear program
pub struct Ilp {
@@ -21,6 +22,16 @@ pub struct Ilp {
maximize: Expression,
}
impl Debug for Ilp {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_struct("Ilp")
.field("var_names", &self.var_names)
.field("constraints", &self.constraints)
.field("maximize", &self.maximize)
.finish_non_exhaustive()
}
}
impl Ilp {
/// Create an empty ILP
pub fn new() -> Self {
@@ -80,6 +91,12 @@ impl Ilp {
Err(e) => panic!("Error in solving: {}", e),
}
}
/// Solve, using the default solver of [good_lp].
pub fn default_solve(
self,
) -> Result<(f64, HashMap<String, f64>), IlpUnsat> {
self.solve(good_lp::default_solver)
}
}
/// Why the ILP could not be solved

3
third_party/opa/README.md vendored Normal file
View File

@@ -0,0 +1,3 @@
# From the OPA distribution
[Source](https://github.com/ishaq/OPA)

152
third_party/opa/sample_costs.json vendored Normal file
View File

@@ -0,0 +1,152 @@
{
"32": {
"a2b": {
"1": 2596.4
},
"a2y": {
"1": 2665.2
},
"add": {
"b": {
"1": 160
},
"y": {
"1": 48
},
"a": {
"1": 1
}
},
"and": {
"b": {
"1": 117
},
"y": {
"1": 32
}
},
"b2a": {
"1": 1868.3999999999999
},
"b2y": {
"1": 2293
},
"eq": {
"b": {
"1": 489
},
"y": {
"1": 39
}
},
"ge": {
"b": {
"1": 733
},
"y": {
"1": 60
}
},
"gt": {
"b": {
"1": 573
},
"y": {
"1": 40
}
},
"le": {
"b": {
"1": 618
},
"y": {
"1": 41
}
},
"lt": {
"b": {
"1": 739
},
"y": {
"1": 60
}
},
"mul": {
"b": {
"1": 1731
},
"y": {
"1": 1127
},
"a": {
"1": 104
}
},
"mux": {
"b": {
"1": 108
},
"y": {
"1": 37
}
},
"ne": {
"b": {
"1": 484
},
"y": {
"1": 38
}
},
"or": {
"b": {
"1": 123
},
"y": {
"1": 40
}
},
"shl": {
"b": {
"1": 981
},
"y": {
"1": 224
}
},
"shr": {
"b": {
"1": 1015
},
"y": {
"1": 224
}
},
"sub": {
"b": {
"1": 52
},
"y": {
"1": 49
},
"a": {
"1": 1
}
},
"xor": {
"b": {
"1": 7
},
"y": {
"1": 23
}
},
"y2a": {
"1": 3207
},
"y2b": {
"1": 2040.2
}
}
}