Skip to content

Commit

Permalink
plaintext_store_db
Browse files Browse the repository at this point in the history
  • Loading branch information
therealyingtong committed Nov 4, 2024
1 parent e525dcf commit 40ac3d5
Show file tree
Hide file tree
Showing 25 changed files with 1,446 additions and 126 deletions.
Empty file added .gitattributes
Empty file.
414 changes: 354 additions & 60 deletions Cargo.lock

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions iris-mpc-common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ ring = "0.17.8"
data-encoding = "2.6.0"
bincode = "1.3.3"
serde-big-array = "0.5.1"
serde_with = "3.11.0"

[dev-dependencies]
float_eq = "1"
Expand Down
6 changes: 4 additions & 2 deletions iris-mpc-common/src/iris_db/iris.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@ use rand::{
distributions::{Bernoulli, Distribution},
Rng,
};
use serde::{Deserialize, Serialize};
use serde_big_array::BigArray;

pub const MATCH_THRESHOLD_RATIO: f64 = 0.375;

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct IrisCodeArray(pub [u64; Self::IRIS_CODE_SIZE_U64]);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
pub struct IrisCodeArray(#[serde(with = "BigArray")] pub [u64; Self::IRIS_CODE_SIZE_U64]);
impl Default for IrisCodeArray {
fn default() -> Self {
Self::ZERO
Expand Down
9 changes: 7 additions & 2 deletions iris-mpc-cpu/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,21 +17,26 @@ bytemuck.workspace = true
dashmap = "6.1.0"
eyre.workspace = true
futures.workspace = true
hawk-pack = { git = "https://github.com/Inversed-Tech/hawk-pack.git", rev = "400db13" }
hawk-pack = { git = "https://github.com/therealyingtong/hawk-pack.git", branch = "new-with-params" }
iris-mpc-common = { path = "../iris-mpc-common" }
itertools.workspace = true
num-traits.workspace = true
rand.workspace = true
rstest = "0.23.0"
serde.workspace = true
sqlx.workspace = true
static_assertions.workspace = true
tokio.workspace = true
tracing.workspace = true
tracing-test = "0.2.5"

[dev-dependencies]
criterion = { version = "0.5.1", features = ["async_tokio"] }
zip = "2.2.0"

[features]
db_dependent = []

[[bench]]
name = "hnsw"
harness = false
harness = false
17 changes: 17 additions & 0 deletions iris-mpc-cpu/benches/assets/.gitattributes
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
hnsw_db_1000000_hawk_graph_links.csv.zip filter=lfs diff=lfs merge=lfs -text
hnsw_db_1000000_hawk_vectors.csv.zip filter=lfs diff=lfs merge=lfs -text
hnsw_db_100000_hawk_graph_links.csv.zip filter=lfs diff=lfs merge=lfs -text
hnsw_db_100000_hawk_vectors.csv.zip filter=lfs diff=lfs merge=lfs -text
hnsw_db_200000_hawk_graph_links.csv.zip filter=lfs diff=lfs merge=lfs -text
hnsw_db_200000_hawk_vectors.csv.zip filter=lfs diff=lfs merge=lfs -text
100K_rust_format_synthetic_data.dat.zip filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_8 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_9 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_3 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_4 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_2 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_5 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_6 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_7 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_0 filter=lfs diff=lfs merge=lfs -text
processed_masked_irises_chunk_1 filter=lfs diff=lfs merge=lfs -text
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_0
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_1
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_2
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_3
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_4
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_5
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_6
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_7
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_8
Git LFS file not shown
3 changes: 3 additions & 0 deletions iris-mpc-cpu/benches/assets/processed_masked_irises_chunk_9
Git LFS file not shown
6 changes: 3 additions & 3 deletions iris-mpc-cpu/benches/hnsw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ fn bench_plaintext_hnsw(c: &mut Criterion) {
group.sample_size(10);
group.sampling_mode(SamplingMode::Flat);

for database_size in [100_usize, 1000, 10000] {
for database_size in [10000] {
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
Expand All @@ -30,7 +30,7 @@ fn bench_plaintext_hnsw(c: &mut Criterion) {

for _ in 0..database_size {
let raw_query = IrisCode::random_rng(&mut rng);
let query = vector.prepare_query(raw_query.clone());
let query = vector.prepare_query(raw_query.clone().into());
let neighbors = searcher
.search_to_insert(&mut vector, &mut graph, &query)
.await;
Expand All @@ -55,7 +55,7 @@ fn bench_plaintext_hnsw(c: &mut Criterion) {
let searcher = HawkSearcher::default();
let mut rng = AesRng::seed_from_u64(0_u64);
let on_the_fly_query = IrisDB::new_random_rng(1, &mut rng).db[0].clone();
let query = db_vectors.prepare_query(on_the_fly_query);
let query = db_vectors.prepare_query(on_the_fly_query.into());
let neighbors = searcher
.search_to_insert(&mut db_vectors, &mut graph, &query)
.await;
Expand Down
1 change: 1 addition & 0 deletions iris-mpc-cpu/migrations/20240909105323_init.down.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
DROP TABLE hawk_vectors;
5 changes: 5 additions & 0 deletions iris-mpc-cpu/migrations/20240909105323_init.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
CREATE TABLE IF NOT EXISTS hawk_vectors (
id integer NOT NULL,
point jsonb NOT NULL,
CONSTRAINT hawk_vectors_pkey PRIMARY KEY (id)
);
53 changes: 26 additions & 27 deletions iris-mpc-cpu/src/hawkers/galois_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,31 +114,6 @@ pub fn setup_local_store_aby3_players() -> eyre::Result<LocalNetAby3NgStoreProto
Ok(LocalNetAby3NgStoreProtocol { runtime, players })
}

impl LocalNetAby3NgStoreProtocol {
pub fn prepare_query(&mut self, code: Vec<GaloisRingSharedIris>) -> PointId {
assert_eq!(code.len(), 3);
assert_eq!(self.players.len(), 3);
let pid0 = self
.players
.get_mut(&Identity::from("alice"))
.unwrap()
.prepare_query(code[0].clone());
let pid1 = self
.players
.get_mut(&Identity::from("bob"))
.unwrap()
.prepare_query(code[1].clone());
let pid2 = self
.players
.get_mut(&Identity::from("charlie"))
.unwrap()
.prepare_query(code[2].clone());
assert_eq!(pid0, pid1);
assert_eq!(pid1, pid2);
pid0
}
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
#[serde(bound = "")]
pub struct DistanceShare<T: IntRing2k> {
Expand Down Expand Up @@ -167,6 +142,30 @@ impl VectorStore for LocalNetAby3NgStoreProtocol {
type QueryRef = PointId; // Vector ID, pending insertion.
type VectorRef = PointId; // Vector ID, inserted.
type DistanceRef = Vec<DistanceShare<u16>>; // Distance represented as shares.
type Data = Vec<GaloisRingSharedIris>;

fn prepare_query(&mut self, code: Vec<GaloisRingSharedIris>) -> PointId {
assert_eq!(code.len(), 3);
assert_eq!(self.players.len(), 3);
let pid0 = self
.players
.get_mut(&Identity::from("alice"))
.unwrap()
.prepare_query(code[0].clone());
let pid1 = self
.players
.get_mut(&Identity::from("bob"))
.unwrap()
.prepare_query(code[1].clone());
let pid2 = self
.players
.get_mut(&Identity::from("charlie"))
.unwrap()
.prepare_query(code[2].clone());
assert_eq!(pid0, pid1);
assert_eq!(pid1, pid2);
pid0
}

async fn insert(&mut self, query: &Self::QueryRef) -> Self::VectorRef {
// The query is now accepted in the store. It keeps the same ID.
Expand Down Expand Up @@ -353,7 +352,7 @@ pub async fn gr_create_ready_made_hawk_searcher<R: RngCore + Clone + CryptoRng>(
let searcher = HawkSearcher::default();

for raw_query in cleartext_database.iter() {
let query = plaintext_vector_store.prepare_query(raw_query.clone());
let query = plaintext_vector_store.prepare_query(raw_query.clone().into());
let neighbors = searcher
.search_to_insert(
&mut plaintext_vector_store,
Expand Down Expand Up @@ -594,7 +593,7 @@ mod tests {
// Now do the work for the plaintext store
let mut plaintext_store = PlaintextStore::default();
let plaintext_preps: Vec<_> = (0..db_dim)
.map(|id| plaintext_store.prepare_query(cleartext_database[id].clone()))
.map(|id| plaintext_store.prepare_query(cleartext_database[id].clone().into()))
.collect();
let mut plaintext_inserts = Vec::new();
for p in plaintext_preps.iter() {
Expand Down
2 changes: 2 additions & 0 deletions iris-mpc-cpu/src/hawkers/mod.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,4 @@
pub mod galois_store;
pub mod plaintext_store;
// #[cfg(feature = "db_dependent")]
pub mod plaintext_store_db;
Loading

0 comments on commit 40ac3d5

Please sign in to comment.