diff --git a/Cargo.lock b/Cargo.lock index 29dd6944..769992c6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/Cargo.toml b/Cargo.toml index 3567e214..0345292d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/src/target/aby/assignment/ilp.rs b/src/target/aby/assignment/ilp.rs new file mode 100644 index 00000000..fb84bbd5 --- /dev/null +++ b/src/target/aby/assignment/ilp.rs @@ -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>, +} + +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) -> CostModel { + use ShareType::*; + let get_cost_opt = + |op_name: &str, obj: &serde_json::map::Map| -> Option { + 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| -> 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 = 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> = { + 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()); + } +} diff --git a/src/target/aby/assignment/mod.rs b/src/target/aby/assignment/mod.rs index d5151e2b..e11fd59e 100644 --- a/src/target/aby/assignment/mod.rs +++ b/src/target/aby/assignment/mod.rs @@ -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; +pub type SharingMap = TermMap; /// 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() } diff --git a/src/target/ilp/mod.rs b/src/target/ilp/mod.rs index 2929eead..c9e0b187 100644 --- a/src/target/ilp/mod.rs +++ b/src/target/ilp/mod.rs @@ -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), IlpUnsat> { + self.solve(good_lp::default_solver) + } } /// Why the ILP could not be solved diff --git a/third_party/opa/README.md b/third_party/opa/README.md new file mode 100644 index 00000000..46e8fded --- /dev/null +++ b/third_party/opa/README.md @@ -0,0 +1,3 @@ +# From the OPA distribution + +[Source](https://github.com/ishaq/OPA) diff --git a/third_party/opa/sample_costs.json b/third_party/opa/sample_costs.json new file mode 100644 index 00000000..f0b85d8a --- /dev/null +++ b/third_party/opa/sample_costs.json @@ -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 + } + } +} +