Merge pull request private-attribution#814 from richajaindce/oneshot

Oneshot for OPRF + enabling compact-gate for it

benches/oneshot/ipa.rs

@@ -9,7 +9,7 @@ use ipa::{
     ff::Fp32BitPrime,
     helpers::{query::IpaQueryConfig, GatewayConfig},
     test_fixture::{
-        ipa::{ipa_in_the_clear, test_ipa, IpaSecurityModel},
+        ipa::{ipa_in_the_clear, test_ipa, test_oprf_ipa, CappingOrder, IpaSecurityModel},
         EventGenerator, EventGeneratorConfig, TestWorld, TestWorldConfig,
     },
 };
@@ -70,6 +70,8 @@ struct Args {
     /// Needed for benches.
     #[arg(long, hide = true)]
     bench: bool,
+    #[arg(short = 'o', long)]
+    oprf: bool,
 }
 
 impl Args {
@@ -121,25 +123,35 @@ async fn run(args: Args) -> Result<(), Error> {
     .take(args.query_size)
     .collect::<Vec<_>>();
 
+    let order = if args.oprf {
+        CappingOrder::CapMostRecentFirst
+    } else {
+        CappingOrder::CapOldestFirst
+    };
     let expected_results = ipa_in_the_clear(
         &raw_data,
         args.per_user_cap,
         args.attribution_window(),
         args.breakdown_keys,
+        &order,
     );
 
     let world = TestWorld::new_with(config.clone());
     tracing::trace!("Preparation complete in {:?}", _prep_time.elapsed());
 
     let _protocol_time = Instant::now();
-    test_ipa::<BenchField>(
-        &world,
-        &raw_data,
-        &expected_results,
-        args.config(),
-        args.mode,
-    )
-    .await;
+    if args.oprf {
+        test_oprf_ipa::<BenchField>(&world, raw_data, &expected_results, args.config()).await;
+    } else {
+        test_ipa::<BenchField>(
+            &world,
+            &raw_data,
+            &expected_results,
+            args.config(),
+            args.mode,
+        )
+        .await;
+    }
     tracing::trace!(
         "{m:?} IPA for {q} records took {t:?}",
         m = args.mode,

pre-commit

@@ -99,6 +99,9 @@ check "Concurrency tests" \
 check "IPA benchmark" \
     cargo bench --bench oneshot_ipa --no-default-features --features="enable-benches descriptive-gate" -- -n 62
 
+check "IPA OPRF benchmark" \
+    cargo bench --bench oneshot_ipa --no-default-features --features="enable-benches descriptive-gate" -- -n 62 --oprf
+
 check "Arithmetic circuit benchmark" \
     cargo bench --bench oneshot_arithmetic --no-default-features --features "enable-benches descriptive-gate"
 

scripts/collect_steps.py

@@ -149,9 +149,7 @@ def extract_intermediate_steps(steps):
 
     return steps
 
-
-if __name__ == "__main__":
-    steps = set()
+def ipa_steps(steps):
     for c in PER_USER_CAP:
         for w in ATTRIBUTION_WINDOW:
             for b in BREAKDOWN_KEYS:
@@ -171,6 +169,40 @@ def extract_intermediate_steps(steps):
                     print(" ".join(args), file=sys.stderr)
                     steps.update(collect_steps(args))
 
+OPRF_BREAKDOWN_KEYS = [256]
+OPRF_USER_CAP = [16, 64, 128]
+OPRF_SECURITY_MODEL = ["semi-honest"]
+OPRF_TRIGGER_VALUE = [7]
+
+def oprf_steps(steps):
+    for c in OPRF_USER_CAP:
+        for w in ATTRIBUTION_WINDOW:
+            for b in OPRF_BREAKDOWN_KEYS:
+                for m in OPRF_SECURITY_MODEL:
+                    for tv in OPRF_TRIGGER_VALUE:
+                        args = ARGS + [
+                            "-n",
+                            str(QUERY_SIZE),
+                            "-c",
+                            str(c),
+                            "-w",
+                            str(w),
+                            "-b",
+                            str(b),
+                            "-m",
+                            m,
+                            "-t",
+                            str(tv),
+                            "-o"
+                        ]
+                    print(" ".join(args), file=sys.stderr)
+                    steps.update(collect_steps(args))
+
+if __name__ == "__main__":
+    steps = set()
+    ipa_steps(steps)
+    oprf_steps(steps)
+
     full_steps = extract_intermediate_steps(steps)
     sorted_steps = sorted(full_steps)
 

src/bin/report_collector.rs

@@ -26,7 +26,7 @@ use ipa::{
     protocol::{BreakdownKey, MatchKey},
     report::{KeyIdentifier, DEFAULT_KEY_ID},
     test_fixture::{
-        ipa::{ipa_in_the_clear, IpaSecurityModel, TestRawDataRecord},
+        ipa::{ipa_in_the_clear, CappingOrder, IpaSecurityModel, TestRawDataRecord},
         EventGenerator, EventGeneratorConfig,
     },
 };
@@ -247,6 +247,7 @@ async fn ipa(
             ipa_query_config.per_user_credit_cap,
             ipa_query_config.attribution_window_seconds,
             ipa_query_config.max_breakdown_key,
+            &CappingOrder::CapOldestFirst,
         );
 
         // pad the output vector to the max breakdown key, to make sure it is aligned with the MPC results

src/protocol/ipa/mod.rs

@@ -464,7 +464,7 @@ pub mod tests {
         test_executor::{run, run_with},
         test_fixture::{
             input::GenericReportTestInput,
-            ipa::{ipa_in_the_clear, test_ipa, IpaSecurityModel},
+            ipa::{ipa_in_the_clear, test_ipa, CappingOrder, IpaSecurityModel},
             logging, EventGenerator, EventGeneratorConfig, Reconstruct, Runner, TestWorld,
             TestWorldConfig,
         },
@@ -815,6 +815,7 @@ pub mod tests {
                 per_user_cap,
                 ATTRIBUTION_WINDOW_SECONDS,
                 MAX_BREAKDOWN_KEY,
+                &CappingOrder::CapOldestFirst,
             );
 
             let config = TestWorldConfig {

src/protocol/mod.rs

@@ -6,7 +6,6 @@ pub mod context;
 pub mod dp;
 pub mod ipa;
 pub mod modulus_conversion;
-#[cfg(feature = "descriptive-gate")]
 pub mod prf_sharding;
 pub mod prss;
 pub mod sort;
@@ -22,11 +21,13 @@ pub use basics::BasicProtocols;
 
 use crate::{
     error::Error,
-    ff::{Gf40Bit, Gf8Bit},
+    ff::{Gf20Bit, Gf3Bit, Gf40Bit, Gf8Bit},
 };
 
 pub type MatchKey = Gf40Bit;
 pub type BreakdownKey = Gf8Bit;
+pub type TriggerValue = Gf3Bit;
+pub type Timestamp = Gf20Bit;
 
 /// Unique identifier of the MPC query requested by report collectors
 /// TODO(615): Generating this unique id may be tricky as it may involve communication between helpers and

src/protocol/prf_sharding/mod.rs

@@ -37,11 +37,11 @@ pub mod bucket;
 pub mod feature_label_dot_product;
 
 pub struct PrfShardedIpaInputRow<BK: GaloisField, TV: GaloisField, TS: GaloisField> {
-    prf_of_match_key: u64,
-    is_trigger_bit: Replicated<Gf2>,
-    breakdown_key: Replicated<BK>,
-    trigger_value: Replicated<TV>,
-    timestamp: Replicated<TS>,
+    pub prf_of_match_key: u64,
+    pub is_trigger_bit: Replicated<Gf2>,
+    pub breakdown_key: Replicated<BK>,
+    pub trigger_value: Replicated<TV>,
+    pub timestamp: Replicated<TS>,
 }
 
 impl<BK: GaloisField, TV: GaloisField, TS: GaloisField> PrfShardedIpaInputRow<BK, TV, TS> {