Use sharded shuffle

akoshelev · Dec 26, 2024 · bf79364 · bf79364
1 parent 3413c3c
commit bf79364
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 11 deletions.
diff --git a/ipa-core/src/protocol/hybrid/breakdown_reveal.rs b/ipa-core/src/protocol/hybrid/breakdown_reveal.rs
@@ -20,7 +20,7 @@ use crate::{
                 AGGREGATE_DEPTH,
             },
             oprf_padding::{apply_dp_padding, PaddingParameters},
-            shuffle::Shuffle,
+            shuffle::ShardedShuffle,
         },
         BooleanProtocols, RecordId,
     },
@@ -66,7 +66,7 @@ pub async fn breakdown_reveal_aggregation<C, BK, V, HV, const B: usize>(
     padding_params: &PaddingParameters,
 ) -> Result<BitDecomposed<Replicated<Boolean, B>>, Error>
 where
-    C: UpgradableContext + Shuffle + ShardedContext,
+    C: UpgradableContext + ShardedShuffle + ShardedContext,
     Boolean: FieldSimd<B>,
     Replicated<Boolean, B>: BooleanProtocols<DZKPUpgraded<C>, B>,
     BK: BooleanArray + U128Conversions,
@@ -94,7 +94,7 @@ where
 
     let attributions = ctx
         .narrow(&Step::Shuffle)
-        .shuffle(attributed_values_padded)
+        .sharded_shuffle(attributed_values_padded)
         .instrument(info_span!("shuffle_attribution_outputs"))
         .await?;
 

diff --git a/ipa-core/src/protocol/hybrid/mod.rs b/ipa-core/src/protocol/hybrid/mod.rs
@@ -33,7 +33,7 @@ use crate::{
         ipa_prf::{
             oprf_padding::{apply_dp_padding, PaddingParameters},
             prf_eval::PrfSharing,
-            shuffle::Shuffle,
+            shuffle::ShardedShuffle,
         },
         prss::FromPrss,
         BooleanProtocols,
@@ -79,7 +79,7 @@ pub async fn hybrid_protocol<'ctx, C, BK, V, HV, const SS_BITS: usize, const B:
 where
     C: UpgradableContext
         + 'ctx
-        + Shuffle
+        + ShardedShuffle
         + ShardedContext
         + FinalizerContext<FinalizingContext = DZKPUpgraded<C>>,
     BK: BreakdownKey<B>,
@@ -121,7 +121,7 @@ where
 
     let shuffled_input_rows = ctx
         .narrow(&Step::InputShuffle)
-        .shuffle(padded_input_rows)
+        .sharded_shuffle(padded_input_rows)
         .instrument(info_span!("shuffle_inputs"))
         .await?;
 

diff --git a/ipa-core/src/protocol/hybrid/oprf.rs b/ipa-core/src/protocol/hybrid/oprf.rs
@@ -181,9 +181,11 @@ where
 
     // reshard reports based on OPRF values. This ensures at the end of this function
     // reports with the same value end up on the same shard.
-    reshard_try_stream(ctx, report_stream, |ctx, _, report| {
-        report.match_key % ctx.shard_count()
-    })
+    reshard_try_stream(
+        ctx.narrow(&HybridStep::ReshardByPrf),
+        report_stream,
+        |ctx, _, report| report.match_key % ctx.shard_count(),
+    )
     .await
 }
 

diff --git a/ipa-core/src/query/runner/hybrid.rs b/ipa-core/src/query/runner/hybrid.rs
@@ -35,7 +35,7 @@ use crate::{
             oprf::{CONV_CHUNK, PRF_CHUNK},
             step::HybridStep,
         },
-        ipa_prf::{oprf_padding::PaddingParameters, prf_eval::PrfSharing, shuffle::Shuffle},
+        ipa_prf::{oprf_padding::PaddingParameters, prf_eval::PrfSharing, shuffle::ShardedShuffle},
         prss::{Endpoint, FromPrss},
         step::ProtocolStep::Hybrid,
         Gate,
@@ -73,7 +73,7 @@ impl<C, HV, R: PrivateKeyRegistry> Query<C, HV, R> {
 impl<C, HV, R> Query<C, HV, R>
 where
     C: UpgradableContext
-        + Shuffle
+        + ShardedShuffle
         + ShardedContext
         + FinalizerContext<FinalizingContext = DZKPUpgraded<C>>,
     HV: BooleanArray + U128Conversions,