Bring back the base field arithmetic gate (#343)

* Bring back the base field arithmetic gate

* fix

src/bin/bench_recursion.rs

@@ -26,6 +26,7 @@ fn bench_prove<F: RichField + Extendable<D>, const D: usize>() -> Result<()> {
         num_wires: 126,
         num_routed_wires: 33,
         constant_gate_size: 6,
+        use_base_arithmetic_gate: false,
         security_bits: 128,
         rate_bits: 3,
         num_challenges: 3,

src/gadgets/arithmetic.rs

@@ -1,8 +1,8 @@
 use std::borrow::Borrow;
 
 use crate::field::extension_field::Extendable;
-use crate::field::field_types::RichField;
-use crate::gates::arithmetic::ArithmeticExtensionGate;
+use crate::field::field_types::{PrimeField, RichField};
+use crate::gates::arithmetic_base::ArithmeticGate;
 use crate::gates::exponentiation::ExponentiationGate;
 use crate::iop::target::{BoolTarget, Target};
 use crate::plonk::circuit_builder::CircuitBuilder;
@@ -33,18 +33,117 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         multiplicand_1: Target,
         addend: Target,
     ) -> Target {
-        let multiplicand_0_ext = self.convert_to_ext(multiplicand_0);
-        let multiplicand_1_ext = self.convert_to_ext(multiplicand_1);
-        let addend_ext = self.convert_to_ext(addend);
+        // If we're not configured to use the base arithmetic gate, just call arithmetic_extension.
+        if !self.config.use_base_arithmetic_gate {
+            let multiplicand_0_ext = self.convert_to_ext(multiplicand_0);
+            let multiplicand_1_ext = self.convert_to_ext(multiplicand_1);
+            let addend_ext = self.convert_to_ext(addend);
 
-        self.arithmetic_extension(
+            return self
+                .arithmetic_extension(
+                    const_0,
+                    const_1,
+                    multiplicand_0_ext,
+                    multiplicand_1_ext,
+                    addend_ext,
+                )
+                .0[0];
+        }
+
+        // See if we can determine the result without adding an `ArithmeticGate`.
+        if let Some(result) =
+            self.arithmetic_special_cases(const_0, const_1, multiplicand_0, multiplicand_1, addend)
+        {
+            return result;
+        }
+
+        // See if we've already computed the same operation.
+        let operation = BaseArithmeticOperation {
             const_0,
             const_1,
-            multiplicand_0_ext,
-            multiplicand_1_ext,
-            addend_ext,
-        )
-        .0[0]
+            multiplicand_0,
+            multiplicand_1,
+            addend,
+        };
+        if let Some(&result) = self.base_arithmetic_results.get(&operation) {
+            return result;
+        }
+
+        // Otherwise, we must actually perform the operation using an ArithmeticExtensionGate slot.
+        let result = self.add_base_arithmetic_operation(operation);
+        self.base_arithmetic_results.insert(operation, result);
+        result
+    }
+
+    fn add_base_arithmetic_operation(&mut self, operation: BaseArithmeticOperation<F>) -> Target {
+        let (gate, i) = self.find_base_arithmetic_gate(operation.const_0, operation.const_1);
+        let wires_multiplicand_0 = Target::wire(gate, ArithmeticGate::wire_ith_multiplicand_0(i));
+        let wires_multiplicand_1 = Target::wire(gate, ArithmeticGate::wire_ith_multiplicand_1(i));
+        let wires_addend = Target::wire(gate, ArithmeticGate::wire_ith_addend(i));
+
+        self.connect(operation.multiplicand_0, wires_multiplicand_0);
+        self.connect(operation.multiplicand_1, wires_multiplicand_1);
+        self.connect(operation.addend, wires_addend);
+
+        Target::wire(gate, ArithmeticGate::wire_ith_output(i))
+    }
+
+    /// Checks for special cases where the value of
+    /// `const_0 * multiplicand_0 * multiplicand_1 + const_1 * addend`
+    /// can be determined without adding an `ArithmeticGate`.
+    fn arithmetic_special_cases(
+        &mut self,
+        const_0: F,
+        const_1: F,
+        multiplicand_0: Target,
+        multiplicand_1: Target,
+        addend: Target,
+    ) -> Option<Target> {
+        let zero = self.zero();
+
+        let mul_0_const = self.target_as_constant(multiplicand_0);
+        let mul_1_const = self.target_as_constant(multiplicand_1);
+        let addend_const = self.target_as_constant(addend);
+
+        let first_term_zero =
+            const_0 == F::ZERO || multiplicand_0 == zero || multiplicand_1 == zero;
+        let second_term_zero = const_1 == F::ZERO || addend == zero;
+
+        // If both terms are constant, return their (constant) sum.
+        let first_term_const = if first_term_zero {
+            Some(F::ZERO)
+        } else if let (Some(x), Some(y)) = (mul_0_const, mul_1_const) {
+            Some(x * y * const_0)
+        } else {
+            None
+        };
+        let second_term_const = if second_term_zero {
+            Some(F::ZERO)
+        } else {
+            addend_const.map(|x| x * const_1)
+        };
+        if let (Some(x), Some(y)) = (first_term_const, second_term_const) {
+            return Some(self.constant(x + y));
+        }
+
+        if first_term_zero && const_1.is_one() {
+            return Some(addend);
+        }
+
+        if second_term_zero {
+            if let Some(x) = mul_0_const {
+                if (x * const_0).is_one() {
+                    return Some(multiplicand_1);
+                }
+            }
+            if let Some(x) = mul_1_const {
+                if (x * const_0).is_one() {
+                    return Some(multiplicand_0);
+                }
+            }
+        }
+
+        None
     }
 
     /// Computes `x * y + z`.
@@ -116,7 +215,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
     /// Exponentiate `base` to the power of `2^power_log`.
     pub fn exp_power_of_2(&mut self, base: Target, power_log: usize) -> Target {
-        if power_log > ArithmeticExtensionGate::<D>::new_from_config(&self.config).num_ops {
+        if power_log > ArithmeticGate::new_from_config(&self.config).num_ops {
             // Cheaper to just use `ExponentiateGate`.
             return self.exp_u64(base, 1 << power_log);
         }
@@ -170,8 +269,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let base_t = self.constant(base);
         let exponent_bits: Vec<_> = exponent_bits.into_iter().map(|b| *b.borrow()).collect();
 
-        if exponent_bits.len() > ArithmeticExtensionGate::<D>::new_from_config(&self.config).num_ops
-        {
+        if exponent_bits.len() > ArithmeticGate::new_from_config(&self.config).num_ops {
             // Cheaper to just use `ExponentiateGate`.
             return self.exp_from_bits(base_t, exponent_bits);
         }
@@ -221,3 +319,13 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         self.inverse_extension(x_ext).0[0]
     }
 }
+
+/// Represents a base arithmetic operation in the circuit. Used to memoize results.
+#[derive(Copy, Clone, Eq, PartialEq, Hash)]
+pub(crate) struct BaseArithmeticOperation<F: PrimeField> {
+    const_0: F,
+    const_1: F,
+    multiplicand_0: Target,
+    multiplicand_1: Target,
+    addend: Target,
+}

src/gadgets/arithmetic_extension.rs

@@ -4,7 +4,7 @@ use crate::field::extension_field::target::{ExtensionAlgebraTarget, ExtensionTar
 use crate::field::extension_field::FieldExtension;
 use crate::field::extension_field::{Extendable, OEF};
 use crate::field::field_types::{Field, PrimeField, RichField};
-use crate::gates::arithmetic::ArithmeticExtensionGate;
+use crate::gates::arithmetic_extension::ArithmeticExtensionGate;
 use crate::iop::generator::{GeneratedValues, SimpleGenerator};
 use crate::iop::target::Target;
 use crate::iop::witness::{PartitionWitness, Witness};
@@ -32,7 +32,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         }
 
         // See if we've already computed the same operation.
-        let operation = ArithmeticOperation {
+        let operation = ExtensionArithmeticOperation {
             const_0,
             const_1,
             multiplicand_0,
@@ -51,7 +51,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
     fn add_arithmetic_extension_operation(
         &mut self,
-        operation: ArithmeticOperation<F, D>,
+        operation: ExtensionArithmeticOperation<F, D>,
     ) -> ExtensionTarget<D> {
         let (gate, i) = self.find_arithmetic_gate(operation.const_0, operation.const_1);
         let wires_multiplicand_0 = ExtensionTarget::from_range(
@@ -519,9 +519,9 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     }
 }
 
-/// Represents an arithmetic operation in the circuit. Used to memoize results.
+/// Represents an extension arithmetic operation in the circuit. Used to memoize results.
 #[derive(Copy, Clone, Eq, PartialEq, Hash)]
-pub(crate) struct ArithmeticOperation<F: PrimeField + Extendable<D>, const D: usize> {
+pub(crate) struct ExtensionArithmeticOperation<F: PrimeField + Extendable<D>, const D: usize> {
     const_0: F,
     const_1: F,
     multiplicand_0: ExtensionTarget<D>,

src/gates/arithmetic_base.rs

@@ -0,0 +1,212 @@
+use crate::field::extension_field::target::ExtensionTarget;
+use crate::field::extension_field::Extendable;
+use crate::field::field_types::RichField;
+use crate::gates::gate::Gate;
+use crate::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
+use crate::iop::target::Target;
+use crate::iop::witness::{PartitionWitness, Witness};
+use crate::plonk::circuit_builder::CircuitBuilder;
+use crate::plonk::circuit_data::CircuitConfig;
+use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
+
+/// A gate which can perform a weighted multiply-add, i.e. `result = c0 x y + c1 z`. If the config
+/// supports enough routed wires, it can support several such operations in one gate.
+#[derive(Debug)]
+pub struct ArithmeticGate {
+    /// Number of arithmetic operations performed by an arithmetic gate.
+    pub num_ops: usize,
+}
+
+impl ArithmeticGate {
+    pub fn new_from_config(config: &CircuitConfig) -> Self {
+        Self {
+            num_ops: Self::num_ops(config),
+        }
+    }
+
+    /// Determine the maximum number of operations that can fit in one gate for the given config.
+    pub(crate) fn num_ops(config: &CircuitConfig) -> usize {
+        let wires_per_op = 4;
+        config.num_routed_wires / wires_per_op
+    }
+
+    pub fn wire_ith_multiplicand_0(i: usize) -> usize {
+        4 * i
+    }
+    pub fn wire_ith_multiplicand_1(i: usize) -> usize {
+        4 * i + 1
+    }
+    pub fn wire_ith_addend(i: usize) -> usize {
+        4 * i + 2
+    }
+    pub fn wire_ith_output(i: usize) -> usize {
+        4 * i + 3
+    }
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticGate {
+    fn id(&self) -> String {
+        format!("{:?}", self)
+    }
+
+    fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
+        let const_0 = vars.local_constants[0];
+        let const_1 = vars.local_constants[1];
+
+        let mut constraints = Vec::new();
+        for i in 0..self.num_ops {
+            let multiplicand_0 = vars.local_wires[Self::wire_ith_multiplicand_0(i)];
+            let multiplicand_1 = vars.local_wires[Self::wire_ith_multiplicand_1(i)];
+            let addend = vars.local_wires[Self::wire_ith_addend(i)];
+            let output = vars.local_wires[Self::wire_ith_output(i)];
+            let computed_output = multiplicand_0 * multiplicand_1 * const_0 + addend * const_1;
+
+            constraints.push(output - computed_output);
+        }
+
+        constraints
+    }
+
+    fn eval_unfiltered_base(&self, vars: EvaluationVarsBase<F>) -> Vec<F> {
+        let const_0 = vars.local_constants[0];
+        let const_1 = vars.local_constants[1];
+
+        let mut constraints = Vec::new();
+        for i in 0..self.num_ops {
+            let multiplicand_0 = vars.local_wires[Self::wire_ith_multiplicand_0(i)];
+            let multiplicand_1 = vars.local_wires[Self::wire_ith_multiplicand_1(i)];
+            let addend = vars.local_wires[Self::wire_ith_addend(i)];
+            let output = vars.local_wires[Self::wire_ith_output(i)];
+            let computed_output = multiplicand_0 * multiplicand_1 * const_0 + addend * const_1;
+
+            constraints.push(output - computed_output);
+        }
+
+        constraints
+    }
+
+    fn eval_unfiltered_recursively(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        vars: EvaluationTargets<D>,
+    ) -> Vec<ExtensionTarget<D>> {
+        let const_0 = vars.local_constants[0];
+        let const_1 = vars.local_constants[1];
+
+        let mut constraints = Vec::new();
+        for i in 0..self.num_ops {
+            let multiplicand_0 = vars.local_wires[Self::wire_ith_multiplicand_0(i)];
+            let multiplicand_1 = vars.local_wires[Self::wire_ith_multiplicand_1(i)];
+            let addend = vars.local_wires[Self::wire_ith_addend(i)];
+            let output = vars.local_wires[Self::wire_ith_output(i)];
+            let computed_output = {
+                let scaled_mul =
+                    builder.mul_many_extension(&[const_0, multiplicand_0, multiplicand_1]);
+                builder.mul_add_extension(const_1, addend, scaled_mul)
+            };
+
+            let diff = builder.sub_extension(output, computed_output);
+            constraints.push(diff);
+        }
+
+        constraints
+    }
+
+    fn generators(
+        &self,
+        gate_index: usize,
+        local_constants: &[F],
+    ) -> Vec<Box<dyn WitnessGenerator<F>>> {
+        (0..self.num_ops)
+            .map(|i| {
+                let g: Box<dyn WitnessGenerator<F>> = Box::new(
+                    ArithmeticBaseGenerator {
+                        gate_index,
+                        const_0: local_constants[0],
+                        const_1: local_constants[1],
+                        i,
+                    }
+                    .adapter(),
+                );
+                g
+            })
+            .collect::<Vec<_>>()
+    }
+
+    fn num_wires(&self) -> usize {
+        self.num_ops * 4
+    }
+
+    fn num_constants(&self) -> usize {
+        2
+    }
+
+    fn degree(&self) -> usize {
+        3
+    }
+
+    fn num_constraints(&self) -> usize {
+        self.num_ops
+    }
+}
+
+#[derive(Clone, Debug)]
+struct ArithmeticBaseGenerator<F: RichField + Extendable<D>, const D: usize> {
+    gate_index: usize,
+    const_0: F,
+    const_1: F,
+    i: usize,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
+    for ArithmeticBaseGenerator<F, D>
+{
+    fn dependencies(&self) -> Vec<Target> {
+        [
+            ArithmeticGate::wire_ith_multiplicand_0(self.i),
+            ArithmeticGate::wire_ith_multiplicand_1(self.i),
+            ArithmeticGate::wire_ith_addend(self.i),
+        ]
+        .iter()
+        .map(|&i| Target::wire(self.gate_index, i))
+        .collect()
+    }
+
+    fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
+        let get_wire =
+            |wire: usize| -> F { witness.get_target(Target::wire(self.gate_index, wire)) };
+
+        let multiplicand_0 = get_wire(ArithmeticGate::wire_ith_multiplicand_0(self.i));
+        let multiplicand_1 = get_wire(ArithmeticGate::wire_ith_multiplicand_1(self.i));
+        let addend = get_wire(ArithmeticGate::wire_ith_addend(self.i));
+
+        let output_target = Target::wire(self.gate_index, ArithmeticGate::wire_ith_output(self.i));
+
+        let computed_output =
+            multiplicand_0 * multiplicand_1 * self.const_0 + addend * self.const_1;
+
+        out_buffer.set_target(output_target, computed_output)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+
+    use crate::field::goldilocks_field::GoldilocksField;
+    use crate::gates::arithmetic_base::ArithmeticGate;
+    use crate::gates::gate_testing::{test_eval_fns, test_low_degree};
+    use crate::plonk::circuit_data::CircuitConfig;
+
+    #[test]
+    fn low_degree() {
+        let gate = ArithmeticGate::new_from_config(&CircuitConfig::standard_recursion_config());
+        test_low_degree::<GoldilocksField, _, 4>(gate);
+    }
+
+    #[test]
+    fn eval_fns() -> Result<()> {
+        let gate = ArithmeticGate::new_from_config(&CircuitConfig::standard_recursion_config());
+        test_eval_fns::<GoldilocksField, _, 4>(gate)
+    }
+}

src/gates/arithmetic_extension.rs

@@ -12,7 +12,8 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 use crate::plonk::circuit_data::CircuitConfig;
 use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
 
-/// A gate which can a linear combination `c0*x*y+c1*z` twice with the same `x`.
+/// A gate which can perform a weighted multiply-add, i.e. `result = c0 x y + c1 z`. If the config
+/// supports enough routed wires, it can support several such operations in one gate.
 #[derive(Debug)]
 pub struct ArithmeticExtensionGate<const D: usize> {
     /// Number of arithmetic operations performed by an arithmetic gate.
@@ -206,7 +207,7 @@ mod tests {
     use anyhow::Result;
 
     use crate::field::goldilocks_field::GoldilocksField;
-    use crate::gates::arithmetic::ArithmeticExtensionGate;
+    use crate::gates::arithmetic_extension::ArithmeticExtensionGate;
     use crate::gates::gate_testing::{test_eval_fns, test_low_degree};
     use crate::plonk::circuit_data::CircuitConfig;
 

src/gates/gate_tree.rs

@@ -223,7 +223,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Tree<GateRef<F, D>> {
 mod tests {
     use super::*;
     use crate::field::goldilocks_field::GoldilocksField;
-    use crate::gates::arithmetic::ArithmeticExtensionGate;
+    use crate::gates::arithmetic_extension::ArithmeticExtensionGate;
     use crate::gates::base_sum::BaseSumGate;
     use crate::gates::constant::ConstantGate;
     use crate::gates::gmimc::GMiMCGate;

src/gates/mod.rs

@@ -1,7 +1,8 @@
 // Gates have `new` methods that return `GateRef`s.
 #![allow(clippy::new_ret_no_self)]
 
-pub mod arithmetic;
+pub mod arithmetic_base;
+pub mod arithmetic_extension;
 pub mod arithmetic_u32;
 pub mod assert_le;
 pub mod base_sum;

src/iop/generator.rs

@@ -87,7 +87,7 @@ pub(crate) fn generate_partial_witness<'a, F: RichField + Extendable<D>, const D
     assert_eq!(
         remaining_generators, 0,
         "{} generators weren't run",
-        remaining_generators
+        remaining_generators,
     );
 
     witness

src/plonk/circuit_builder.rs

@@ -12,9 +12,11 @@ use crate::field::fft::fft_root_table;
 use crate::field::field_types::{Field, RichField};
 use crate::fri::commitment::PolynomialBatchCommitment;
 use crate::fri::{FriConfig, FriParams};
-use crate::gadgets::arithmetic_extension::ArithmeticOperation;
+use crate::gadgets::arithmetic::BaseArithmeticOperation;
+use crate::gadgets::arithmetic_extension::ExtensionArithmeticOperation;
 use crate::gadgets::arithmetic_u32::U32Target;
-use crate::gates::arithmetic::ArithmeticExtensionGate;
+use crate::gates::arithmetic_base::ArithmeticGate;
+use crate::gates::arithmetic_extension::ArithmeticExtensionGate;
 use crate::gates::arithmetic_u32::{U32ArithmeticGate, NUM_U32_ARITHMETIC_OPS};
 use crate::gates::constant::ConstantGate;
 use crate::gates::gate::{Gate, GateInstance, GateRef, PrefixedGate};
@@ -74,8 +76,11 @@ pub struct CircuitBuilder<F: RichField + Extendable<D>, const D: usize> {
     constants_to_targets: HashMap<F, Target>,
     targets_to_constants: HashMap<Target, F>,
 
+    /// Memoized results of `arithmetic` calls.
+    pub(crate) base_arithmetic_results: HashMap<BaseArithmeticOperation<F>, Target>,
+
     /// Memoized results of `arithmetic_extension` calls.
-    pub(crate) arithmetic_results: HashMap<ArithmeticOperation<F, D>, ExtensionTarget<D>>,
+    pub(crate) arithmetic_results: HashMap<ExtensionArithmeticOperation<F, D>, ExtensionTarget<D>>,
 
     batched_gates: BatchedGates<F, D>,
 }
@@ -93,6 +98,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             marked_targets: Vec::new(),
             generators: Vec::new(),
             constants_to_targets: HashMap::new(),
+            base_arithmetic_results: HashMap::new(),
             arithmetic_results: HashMap::new(),
             targets_to_constants: HashMap::new(),
             batched_gates: BatchedGates::new(),
@@ -742,11 +748,13 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     }
 }
 
-/// Various gate types can contain multiple copies in a single Gate. This helper struct lets a CircuitBuilder track such gates that are currently being "filled up."
+/// Various gate types can contain multiple copies in a single Gate. This helper struct lets a
+/// CircuitBuilder track such gates that are currently being "filled up."
 pub struct BatchedGates<F: RichField + Extendable<D>, const D: usize> {
     /// A map `(c0, c1) -> (g, i)` from constants `(c0,c1)` to an available arithmetic gate using
     /// these constants with gate index `g` and already using `i` arithmetic operations.
     pub(crate) free_arithmetic: HashMap<(F, F), (usize, usize)>,
+    pub(crate) free_base_arithmetic: HashMap<(F, F), (usize, usize)>,
 
     /// A map `(c0, c1) -> (g, i)` from constants `vec_size` to an available arithmetic gate using
     /// these constants with gate index `g` and already using `i` random accesses.
@@ -771,6 +779,7 @@ impl<F: RichField + Extendable<D>, const D: usize> BatchedGates<F, D> {
     pub fn new() -> Self {
         Self {
             free_arithmetic: HashMap::new(),
+            free_base_arithmetic: HashMap::new(),
             free_random_access: HashMap::new(),
             current_switch_gates: Vec::new(),
             current_u32_arithmetic_gate: None,
@@ -781,6 +790,37 @@ impl<F: RichField + Extendable<D>, const D: usize> BatchedGates<F, D> {
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
+    /// Finds the last available arithmetic gate with the given constants or add one if there aren't any.
+    /// Returns `(g,i)` such that there is an arithmetic gate with the given constants at index
+    /// `g` and the gate's `i`-th operation is available.
+    pub(crate) fn find_base_arithmetic_gate(&mut self, const_0: F, const_1: F) -> (usize, usize) {
+        let (gate, i) = self
+            .batched_gates
+            .free_base_arithmetic
+            .get(&(const_0, const_1))
+            .copied()
+            .unwrap_or_else(|| {
+                let gate = self.add_gate(
+                    ArithmeticGate::new_from_config(&self.config),
+                    vec![const_0, const_1],
+                );
+                (gate, 0)
+            });
+
+        // Update `free_arithmetic` with new values.
+        if i < ArithmeticGate::num_ops(&self.config) - 1 {
+            self.batched_gates
+                .free_base_arithmetic
+                .insert((const_0, const_1), (gate, i + 1));
+        } else {
+            self.batched_gates
+                .free_base_arithmetic
+                .remove(&(const_0, const_1));
+        }
+
+        (gate, i)
+    }
+
     /// Finds the last available arithmetic gate with the given constants or add one if there aren't any.
     /// Returns `(g,i)` such that there is an arithmetic gate with the given constants at index
     /// `g` and the gate's `i`-th operation is available.
@@ -941,36 +981,36 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         (gate, instance)
     }
 
+    /// Fill the remaining unused arithmetic operations with zeros, so that all
+    /// `ArithmeticGate` are run.
+    fn fill_base_arithmetic_gates(&mut self) {
+        let zero = self.zero();
+        for ((c0, c1), (_gate, i)) in self.batched_gates.free_base_arithmetic.clone() {
+            for _ in i..ArithmeticGate::num_ops(&self.config) {
+                // If we directly wire in zero, an optimization will skip doing anything and return
+                // zero. So we pass in a virtual target and connect it to zero afterward.
+                let dummy = self.add_virtual_target();
+                self.arithmetic(c0, c1, dummy, dummy, dummy);
+                self.connect(dummy, zero);
+            }
+        }
+        assert!(self.batched_gates.free_base_arithmetic.is_empty());
+    }
+
     /// Fill the remaining unused arithmetic operations with zeros, so that all
     /// `ArithmeticExtensionGenerator`s are run.
     fn fill_arithmetic_gates(&mut self) {
         let zero = self.zero_extension();
-        let remaining_arithmetic_gates = self
-            .batched_gates
-            .free_arithmetic
-            .values()
-            .copied()
-            .collect::<Vec<_>>();
-        for (gate, i) in remaining_arithmetic_gates {
-            for j in i..ArithmeticExtensionGate::<D>::num_ops(&self.config) {
-                let wires_multiplicand_0 = ExtensionTarget::from_range(
-                    gate,
-                    ArithmeticExtensionGate::<D>::wires_ith_multiplicand_0(j),
-                );
-                let wires_multiplicand_1 = ExtensionTarget::from_range(
-                    gate,
-                    ArithmeticExtensionGate::<D>::wires_ith_multiplicand_1(j),
-                );
-                let wires_addend = ExtensionTarget::from_range(
-                    gate,
-                    ArithmeticExtensionGate::<D>::wires_ith_addend(j),
-                );
-
-                self.connect_extension(zero, wires_multiplicand_0);
-                self.connect_extension(zero, wires_multiplicand_1);
-                self.connect_extension(zero, wires_addend);
+        for ((c0, c1), (_gate, i)) in self.batched_gates.free_arithmetic.clone() {
+            for _ in i..ArithmeticExtensionGate::<D>::num_ops(&self.config) {
+                // If we directly wire in zero, an optimization will skip doing anything and return
+                // zero. So we pass in a virtual target and connect it to zero afterward.
+                let dummy = self.add_virtual_extension_target();
+                self.arithmetic_extension(c0, c1, dummy, dummy, dummy);
+                self.connect_extension(dummy, zero);
             }
         }
+        assert!(self.batched_gates.free_arithmetic.is_empty());
     }
 
     /// Fill the remaining unused random access operations with zeros, so that all
@@ -1064,6 +1104,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
     fn fill_batched_gates(&mut self) {
         self.fill_arithmetic_gates();
+        self.fill_base_arithmetic_gates();
         self.fill_random_access_gates();
         self.fill_switch_gates();
         self.fill_u32_arithmetic_gates();

src/plonk/circuit_data.rs

@@ -26,6 +26,9 @@ pub struct CircuitConfig {
     pub num_wires: usize,
     pub num_routed_wires: usize,
     pub constant_gate_size: usize,
+    /// Whether to use a dedicated gate for base field arithmetic, rather than using a single gate
+    /// for both base field and extension field arithmetic.
+    pub use_base_arithmetic_gate: bool,
     pub security_bits: usize,
     pub rate_bits: usize,
     /// The number of challenge points to generate, for IOPs that have soundness errors of (roughly)
@@ -55,6 +58,7 @@ impl CircuitConfig {
             num_wires: 143,
             num_routed_wires: 25,
             constant_gate_size: 6,
+            use_base_arithmetic_gate: true,
             security_bits: 100,
             rate_bits: 3,
             num_challenges: 2,