Use 32 bit histogram values

akoshelev · Jun 14, 2024 · a9b346a · a9b346a
1 parent 743f8de
commit a9b346a
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 18 deletions.
diff --git a/ipa-core/src/bin/report_collector.rs b/ipa-core/src/bin/report_collector.rs
@@ -19,7 +19,7 @@ use ipa_core::{
         CsvSerializer, IpaQueryResult, Verbosity,
     },
     config::NetworkConfig,
-    ff::{boolean_array::BA16, FieldType},
+    ff::{boolean_array::BA32, FieldType},
     helpers::query::{IpaQueryConfig, QueryConfig, QuerySize, QueryType},
     hpke::{KeyRegistry, PublicKeyOnly},
     net::MpcHelperClient,
@@ -256,7 +256,7 @@ async fn ipa(
     let mut key_registries = KeyRegistries::default();
     let actual = match query_style {
         IpaQueryStyle::Oprf => {
-            playbook_oprf_ipa::<BA16, _>(
+            playbook_oprf_ipa::<BA32, _>(
                 input_rows,
                 helper_clients,
                 query_id,

diff --git a/ipa-core/src/protocol/ipa_prf/aggregation/mod.rs b/ipa-core/src/protocol/ipa_prf/aggregation/mod.rs
@@ -1,4 +1,5 @@
 use std::{
+    any::type_name,
     convert::Infallible,
     iter::{self, repeat},
     pin::Pin,
@@ -16,7 +17,7 @@ use crate::{
     },
     protocol::{
         basics::{BooleanArrayMul, BooleanProtocols, SecureMul},
-        boolean::{step::SixteenBitStep, NBitStep},
+        boolean::{step::ThirtyTwoBitStep, NBitStep},
         context::Context,
         ipa_prf::{
             aggregation::step::{AggregateValuesStep, AggregationStep as Step},
@@ -240,6 +241,9 @@ where
     Boolean: FieldSimd<B>,
     Replicated<Boolean, B>: BooleanProtocols<C, B>,
 {
+    // Step used to add trigger values.
+    type AdditionStep = ThirtyTwoBitStep;
+
     let mut depth = 0;
     while num_rows > 1 {
         // We reduce pairwise, passing through the odd record at the end if there is one, so the
@@ -267,16 +271,18 @@ where
                             }
                             Ok(mut chunk_pair) => {
                                 assert_eq!(chunk_pair.len(), 2);
+                                assert!(
+                                    OV::BITS <= AdditionStep::BITS,
+                                    "{} not large enough to accommodate the sum of {} bit values",
+                                    type_name::<AdditionStep>(),
+                                    OV::BITS,
+                                );
                                 let b = chunk_pair.pop().unwrap();
                                 let a = chunk_pair.pop().unwrap();
                                 let record_id = RecordId::from(i);
                                 if a.len() < usize::try_from(OV::BITS).unwrap() {
-                                    assert!(
-                                        OV::BITS <= SixteenBitStep::BITS,
-                                        "SixteenBitStep not large enough to accomodate this sum"
-                                    );
                                     // If we have enough output bits, add and keep the carry.
-                                    let (mut sum, carry) = integer_add::<_, SixteenBitStep, B>(
+                                    let (mut sum, carry) = integer_add::<_, AdditionStep, B>(
                                         ctx.narrow(&AggregateValuesStep::Add),
                                         record_id,
                                         &a,
@@ -286,11 +292,7 @@ where
                                     sum.push(carry);
                                     Ok(sum)
                                 } else {
-                                    assert!(
-                                        OV::BITS <= SixteenBitStep::BITS,
-                                        "SixteenBitStep not large enough to accommodate this sum"
-                                    );
-                                    integer_sat_add::<C, SixteenBitStep, B>(
+                                    integer_sat_add::<C, AdditionStep, B>(
                                         ctx.narrow(&AggregateValuesStep::SaturatingAdd),
                                         record_id,
                                         &a,

diff --git a/ipa-core/src/protocol/ipa_prf/aggregation/step.rs b/ipa-core/src/protocol/ipa_prf/aggregation/step.rs
@@ -22,7 +22,7 @@ impl From<usize> for BucketStep {
 
 #[derive(CompactStep)]
 pub(crate) enum AggregateValuesStep {
-    #[step(child = crate::protocol::boolean::step::SixteenBitStep)]
+    #[step(child = crate::protocol::boolean::step::ThirtyTwoBitStep)]
     Add,
     #[step(child = crate::protocol::ipa_prf::boolean_ops::step::SaturatedAdditionStep)]
     SaturatingAdd,

diff --git a/ipa-core/src/query/executor.rs b/ipa-core/src/query/executor.rs
@@ -22,7 +22,7 @@ use typenum::Unsigned;
 #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
 use crate::{ff::Fp32BitPrime, query::runner::execute_test_multiply};
 use crate::{
-    ff::{boolean_array::BA16, FieldType, Serializable},
+    ff::{boolean_array::BA32, FieldType, Serializable},
     helpers::{
         negotiate_prss,
         query::{QueryConfig, QueryType},
@@ -90,13 +90,14 @@ pub fn execute(
             move |prss, gateway, config, input| {
                 let ctx = SemiHonestContext::new(prss, gateway);
                 Box::pin(
-                    OprfIpaQuery::<BA16>::new(ipa_config, key_registry)
+                    OprfIpaQuery::<BA32>::new(ipa_config, key_registry)
                         .execute(ctx, config.size, input)
                         .then(|res| ready(res.map(|out| Box::new(out) as Box<dyn Result>))),
                 )
             },
         ),
-        // TODO(953): This is not doing anything differently than the Fp32BitPrime case.
+        // TODO(953): This is not doing anything differently than the Fp32BitPrime case, except
+        // using 16 bits for histogram values
         #[cfg(any(test, feature = "weak-field"))]
         (QueryType::OprfIpa(ipa_config), FieldType::Fp31) => do_query(
             config,
@@ -105,7 +106,7 @@ pub fn execute(
             move |prss, gateway, config, input| {
                 let ctx = SemiHonestContext::new(prss, gateway);
                 Box::pin(
-                    OprfIpaQuery::<BA16>::new(ipa_config, key_registry)
+                    OprfIpaQuery::<crate::ff::boolean_array::BA16>::new(ipa_config, key_registry)
                         .execute(ctx, config.size, input)
                         .then(|res| ready(res.map(|out| Box::new(out) as Box<dyn Result>))),
                 )

diff --git a/ipa-core/src/secret_sharing/vector/transpose.rs b/ipa-core/src/secret_sharing/vector/transpose.rs
@@ -609,6 +609,7 @@ impl_transpose_shares_bool_to_ba!(BA256, 256, 256, test_transpose_shares_bool_to
 impl_transpose_shares_bool_to_ba_small!(BA8, 8, 256, test_transpose_shares_bool_to_ba_8x256);
 impl_transpose_shares_bool_to_ba!(BA16, 16, 256, test_transpose_shares_bool_to_ba_16x256);
 impl_transpose_shares_bool_to_ba!(BA16, 16, 32, test_transpose_shares_bool_to_ba_16x32);
+impl_transpose_shares_bool_to_ba!(BA32, 32, 256, test_transpose_shares_bool_to_ba_32x256);
 impl_transpose_shares_bool_to_ba_small!(BA8, 8, 32, test_transpose_shares_bool_to_ba_8x32);
 
 // Usage: Aggregation output tests