Limit the number of bits for bool_and operation to 9 (#1058)

Continues the work started in #1056 to reduce the total number of step transitions. Confirmed with @benjaminsavage that in WALR we don't need more than 9 bits and in OPRF IPA it is the same because the number of breakdowns is currently limited to 512

ipa-core/src/protocol/boolean/and.rs

@@ -1,13 +1,29 @@
 use std::iter::zip;
 
+use ipa_macros::Step;
+
 use crate::{
     error::Error,
     ff::boolean::Boolean,
-    protocol::{basics::SecureMul, context::Context, step::BitOpStep, RecordId},
+    protocol::{basics::SecureMul, context::Context, RecordId},
     secret_sharing::{replicated::semi_honest::AdditiveShare, BitDecomposed, FieldSimd},
 };
 
-/// Matrix bitwise AND for use with vectors of bit-decomposed values
+const MAX_BITS: usize = 9;
+
+#[derive(Step)]
+pub(crate) enum BoolAndStep {
+    #[dynamic(9)] // keep in sync with MAX_BITS
+    Bit(usize),
+}
+
+/// Matrix bitwise AND for use with vectors of bit-decomposed values. Supports up to 9 bits of input
+/// that is enough to support both WALR and PRF IPA use cases. IPA currently supports up to
+/// 512 breakdowns (see [`MAX_BREAKDOWN`] limitation) and WALR does not need more than that.
+/// Limiting the number of bits helps with our static compact gate compilation, so we want this
+/// number to be as small as possible.
+///
+/// [`MAX_BREAKDOWN`]: crate::protocol::ipa_prf::aggregation::bucket::move_single_value_to_bucket
 ///
 /// ## Errors
 /// Propagates errors from the multiplication protocol.
@@ -16,7 +32,7 @@ use crate::{
 //
 // Supplying an iterator saves constructing a complete copy of the argument
 // in memory when it is a uniform constant.
-pub async fn bool_and<'a, C, BI, const N: usize>(
+pub async fn bool_and_9_bit<'a, C, BI, const N: usize>(
     ctx: C,
     record_id: RecordId,
     a: &BitDecomposed<AdditiveShare<Boolean, N>>,
@@ -31,10 +47,15 @@ where
 {
     let b = b.into_iter();
     assert_eq!(a.len(), b.len());
+    assert!(
+        a.len() <= MAX_BITS,
+        "Up to {MAX_BITS} values are supported, but was given a value of {len} bits",
+        len = a.len()
+    );
 
     BitDecomposed::try_from(
         ctx.parallel_join(zip(a.iter(), b).enumerate().map(|(i, (a, b))| {
-            let ctx = ctx.narrow(&BitOpStep::Bit(i));
+            let ctx = ctx.narrow(&BoolAndStep::Bit(i));
             a.multiply(b, ctx, record_id)
         }))
         .await?,

ipa-core/src/protocol/ipa_prf/aggregation/bucket.rs

@@ -6,7 +6,7 @@ use crate::{
     ff::boolean::Boolean,
     helpers::repeat_n,
     protocol::{
-        basics::SecureMul, boolean::and::bool_and, context::Context,
+        basics::SecureMul, boolean::and::bool_and_9_bit, context::Context,
         ipa_prf::prf_sharding::BinaryTreeDepthStep, RecordId,
     },
     secret_sharing::{replicated::semi_honest::AdditiveShare, BitDecomposed, FieldSimd},
@@ -124,7 +124,7 @@ where
                     let index_contribution = &row_contribution[tree_index];
 
                     (robust || tree_index + span < breakdown_count).then(|| {
-                        bool_and(
+                        bool_and_9_bit(
                             bucket_c,
                             record_id,
                             index_contribution,

ipa-core/src/protocol/ipa_prf/prf_sharding/feature_label_dot_product.rs

@@ -13,7 +13,7 @@ use crate::{
     helpers::{repeat_n, stream::TryFlattenItersExt},
     protocol::{
         basics::{SecureMul, ShareKnownValue},
-        boolean::{and::bool_and, or::or},
+        boolean::{and::bool_and_9_bit, or::or},
         context::{Context, UpgradedSemiHonestContext},
         ipa_prf::aggregation::aggregate_values,
         BooleanProtocols, RecordId,
@@ -109,7 +109,7 @@ impl InputsRequiredFromPrevRow {
         bit_decomposed_output
             .transpose_from(&input_row.feature_vector)
             .unwrap_infallible();
-        let capped_attributed_feature_vector = bool_and(
+        let capped_attributed_feature_vector = bool_and_9_bit(
             ctx,
             record_id,
             &bit_decomposed_output,