impl Copy for Field (3% to 5% faster end2end)

Author: TomWambsgans

crates/backend/air/src/symbolic.rs

@@ -4,7 +4,7 @@ use core::fmt::Debug;
 use core::iter::{Product, Sum};
 use core::marker::PhantomData;
 use core::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-use std::rc::Rc;
+use std::cell::RefCell;
 
 use field::{Algebra, Field, InjectiveMonomial, PrimeCharacteristicRing};
 
@@ -66,11 +66,44 @@ pub enum SymbolicOperation {
     Neg,
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub enum SymbolicExpression<F> {
+#[derive(Copy, Clone, Debug)]
+pub struct SymbolicNode<F: Copy> {
+    pub op: SymbolicOperation,
+    pub lhs: SymbolicExpression<F>,
+    pub rhs: SymbolicExpression<F>, // dummy (ZERO) for Neg
+}
+
+// We use an arena as a trick to allow SymbolicExpression to be Copy
+// (ugly trick but fine in practice since SymbolicExpression is only used once at the start of the program)
+thread_local! {
+    static ARENA: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
+}
+
+fn alloc_node<F: Field>(node: SymbolicNode<F>) -> u32 {
+    ARENA.with(|arena| {
+        let mut bytes = arena.borrow_mut();
+        let node_size = std::mem::size_of::<SymbolicNode<F>>();
+        let idx = bytes.len();
+        bytes.resize(idx + node_size, 0);
+        unsafe {
+            std::ptr::write_unaligned(bytes.as_mut_ptr().add(idx) as *mut SymbolicNode<F>, node);
+        }
+        idx as u32
+    })
+}
+
+pub fn get_node<F: Field>(idx: u32) -> SymbolicNode<F> {
+    ARENA.with(|arena| {
+        let bytes = arena.borrow();
+        unsafe { std::ptr::read_unaligned(bytes.as_ptr().add(idx as usize) as *const SymbolicNode<F>) }
+    })
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum SymbolicExpression<F: Copy> {
     Variable(SymbolicVariable<F>),
     Constant(F),
-    Operation(Rc<(SymbolicOperation, Vec<Self>)>),
+    Operation(u32), // index into thread-local arena
 }
 
 impl<F: Field> Default for SymbolicExpression<F> {
@@ -119,7 +152,11 @@ where
     fn add(self, rhs: T) -> Self {
         match (self, rhs.into()) {
             (Self::Constant(lhs), Self::Constant(rhs)) => Self::Constant(lhs + rhs),
-            (lhs, rhs) => Self::Operation(Rc::new((SymbolicOperation::Add, vec![lhs, rhs]))),
+            (lhs, rhs) => Self::Operation(alloc_node(SymbolicNode {
+                op: SymbolicOperation::Add,
+                lhs,
+                rhs,
+            })),
         }
     }
 }
@@ -129,7 +166,7 @@ where
     T: Into<Self>,
 {
     fn add_assign(&mut self, rhs: T) {
-        *self = self.clone() + rhs.into();
+        *self = *self + rhs.into();
     }
 }
 
@@ -148,7 +185,11 @@ impl<F: Field, T: Into<Self>> Sub<T> for SymbolicExpression<F> {
     fn sub(self, rhs: T) -> Self {
         match (self, rhs.into()) {
             (Self::Constant(lhs), Self::Constant(rhs)) => Self::Constant(lhs - rhs),
-            (lhs, rhs) => Self::Operation(Rc::new((SymbolicOperation::Sub, vec![lhs, rhs]))),
+            (lhs, rhs) => Self::Operation(alloc_node(SymbolicNode {
+                op: SymbolicOperation::Sub,
+                lhs,
+                rhs,
+            })),
         }
     }
 }
@@ -158,7 +199,7 @@ where
     T: Into<Self>,
 {
     fn sub_assign(&mut self, rhs: T) {
-        *self = self.clone() - rhs.into();
+        *self = *self - rhs.into();
     }
 }
 
@@ -168,7 +209,11 @@ impl<F: Field> Neg for SymbolicExpression<F> {
     fn neg(self) -> Self {
         match self {
             Self::Constant(c) => Self::Constant(-c),
-            expr => Self::Operation(Rc::new((SymbolicOperation::Neg, vec![expr]))),
+            expr => Self::Operation(alloc_node(SymbolicNode {
+                op: SymbolicOperation::Neg,
+                lhs: expr,
+                rhs: Self::ZERO, // dummy
+            })),
         }
     }
 }
@@ -179,7 +224,11 @@ impl<F: Field, T: Into<Self>> Mul<T> for SymbolicExpression<F> {
     fn mul(self, rhs: T) -> Self {
         match (self, rhs.into()) {
             (Self::Constant(lhs), Self::Constant(rhs)) => Self::Constant(lhs * rhs),
-            (lhs, rhs) => Self::Operation(Rc::new((SymbolicOperation::Mul, vec![lhs, rhs]))),
+            (lhs, rhs) => Self::Operation(alloc_node(SymbolicNode {
+                op: SymbolicOperation::Mul,
+                lhs,
+                rhs,
+            })),
         }
     }
 }
@@ -189,7 +238,7 @@ where
     T: Into<Self>,
 {
     fn mul_assign(&mut self, rhs: T) {
-        *self = self.clone() * rhs.into();
+        *self = *self * rhs.into();
     }
 }
 
@@ -258,7 +307,7 @@ impl<F: Field> AirBuilder for SymbolicAirBuilder<F> {
     fn declare_values(&mut self, values: &[Self::F]) {
         if self.bus_flag_value.is_none() {
             assert_eq!(values.len(), 1);
-            self.bus_flag_value = Some(values[0].clone());
+            self.bus_flag_value = Some(values[0]);
         } else {
             assert!(self.bus_data_values.is_none());
             self.bus_data_values = Some(values.to_vec());
@@ -276,6 +325,9 @@ pub fn get_symbolic_constraints_and_bus_data_values<F: Field, A: Air>(
 where
     A::ExtraData: Default,
 {
+    // Clear the arena before building constraints
+    ARENA.with(|arena| arena.borrow_mut().clear());
+
     let mut builder = SymbolicAirBuilder::<F>::new(air.n_columns(), air.n_down_columns());
     air.eval(&mut builder, &Default::default());
     (

crates/backend/field/src/exponentiation.rs

@@ -17,12 +17,12 @@ pub fn exp_1420470955<R: PrimeCharacteristicRing>(val: R) -> R {
     // Suspect it's possible to improve this with enough effort.
     let p1 = val;
     let p100 = p1.exp_power_of_2(2);
-    let p101 = p100.clone() * p1.clone();
+    let p101 = p100 * p1;
     let p10000 = p100.exp_power_of_2(2);
     let p10101 = p10000 * p101;
     let p10101000000 = p10101.exp_power_of_2(6);
-    let p10101010101 = p10101000000.clone() * p10101.clone();
-    let p101010010101 = p10101000000 * p10101010101.clone();
+    let p10101010101 = p10101000000 * p10101;
+    let p101010010101 = p10101000000 * p10101010101;
     let p101010010101000000000000 = p101010010101.exp_power_of_2(12);
     let p101010010101010101010101 = p101010010101000000000000 * p10101010101;
     let p101010010101010101010101000000 = p101010010101010101010101.exp_power_of_2(6);

crates/backend/field/src/field.rs

@@ -57,7 +57,7 @@ use crate::{Packable, PackedFieldExtension, PackedValue};
 pub trait PrimeCharacteristicRing:
     Sized
     + Default
-    + Clone
+    + Copy
     + Add<Output = Self>
     + AddAssign
     + Sub<Output = Self>
@@ -138,7 +138,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline(always)]
     fn double(&self) -> Self {
-        self.clone() + self.clone()
+        *self + *self
     }
 
     /// The elementary function `halve(a) = a/2`.
@@ -152,7 +152,7 @@ pub trait PrimeCharacteristicRing:
         // is circular when PrimeSubfield = Self. It should also be overwritten by
         // most rings to avoid the multiplication.
         let half = Self::from_prime_subfield(Self::PrimeSubfield::ONE.halve());
-        self.clone() * half
+        *self * half
     }
 
     /// The elementary function `square(a) = a^2`.
@@ -161,7 +161,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline(always)]
     fn square(&self) -> Self {
-        self.clone() * self.clone()
+        *self * *self
     }
 
     /// The elementary function `cube(a) = a^3`.
@@ -170,7 +170,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline(always)]
     fn cube(&self) -> Self {
-        self.square() * self.clone()
+        self.square() * *self
     }
 
     /// Computes the arithmetic generalization of boolean `xor`.
@@ -179,7 +179,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline(always)]
     fn xor(&self, y: &Self) -> Self {
-        self.clone() + y.clone() - self.clone() * y.clone().double()
+        *self + *y - *self * y.double()
     }
 
     /// Computes the arithmetic generalization of a triple `xor`.
@@ -197,7 +197,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline(always)]
     fn andn(&self, y: &Self) -> Self {
-        (Self::ONE - self.clone()) * y.clone()
+        (Self::ONE - *self) * *y
     }
 
     /// The vanishing polynomial for boolean values: `x * (1 - x)`.
@@ -219,12 +219,12 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline]
     fn exp_u64(&self, power: u64) -> Self {
-        let mut current = self.clone();
+        let mut current = *self;
         let mut product = Self::ONE;
 
         for j in 0..bits_u64(power) {
             if (power >> j) & 1 != 0 {
-                product *= current.clone();
+                product *= current;
             }
             current = current.square();
         }
@@ -242,15 +242,15 @@ pub trait PrimeCharacteristicRing:
     fn exp_const_u64<const POWER: u64>(&self) -> Self {
         match POWER {
             0 => Self::ONE,
-            1 => self.clone(),
+            1 => *self,
             2 => self.square(),
             3 => self.cube(),
             4 => self.square().square(),
-            5 => self.square().square() * self.clone(),
+            5 => self.square().square() * *self,
             6 => self.square().cube(),
             7 => {
                 let x2 = self.square();
-                let x3 = x2.clone() * self.clone();
+                let x3 = x2 * *self;
                 let x4 = x2.square();
                 x3 * x4
             }
@@ -264,7 +264,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline]
     fn exp_power_of_2(&self, power_log: usize) -> Self {
-        let mut res = self.clone();
+        let mut res = *self;
         for _ in 0..power_log {
             res = res.square();
         }
@@ -279,7 +279,7 @@ pub trait PrimeCharacteristicRing:
     fn mul_2exp_u64(&self, exp: u64) -> Self {
         // Some rings might want to reimplement this to avoid the
         // exponentiations (and potentially even the multiplication).
-        self.clone() * Self::TWO.exp_u64(exp)
+        *self * Self::TWO.exp_u64(exp)
     }
 
     /// Divide by a given power of two. `div_2exp_u64(a, exp) = a/2^exp`
@@ -291,7 +291,7 @@ pub trait PrimeCharacteristicRing:
     fn div_2exp_u64(&self, exp: u64) -> Self {
         // Some rings might want to reimplement this to avoid the
         // exponentiations (and potentially even the multiplication).
-        self.clone() * Self::from_prime_subfield(Self::PrimeSubfield::ONE.halve().exp_u64(exp))
+        *self * Self::from_prime_subfield(Self::PrimeSubfield::ONE.halve().exp_u64(exp))
     }
 
     /// Construct an iterator which returns powers of `self`: `self^0, self^1, self^2, ...`.
@@ -306,7 +306,7 @@ pub trait PrimeCharacteristicRing:
     #[inline]
     fn shifted_powers(&self, start: Self) -> Powers<Self> {
         Powers {
-            base: self.clone(),
+            base: *self,
             current: start,
         }
     }
@@ -315,7 +315,7 @@ pub trait PrimeCharacteristicRing:
     #[must_use]
     #[inline]
     fn dot_product<const N: usize>(u: &[Self; N], v: &[Self; N]) -> Self {
-        u.iter().zip(v).map(|(x, y)| x.clone() * y.clone()).sum()
+        u.iter().zip(v).map(|(x, y)| *x * *y).sum()
     }
 
     /// Compute the sum of a slice of elements whose length is a compile time constant.
@@ -344,10 +344,10 @@ pub trait PrimeCharacteristicRing:
         // I only tested this on `AVX2` though so there might be a better value for other architectures.
         match N {
             0 => Self::ZERO,
-            1 => input[0].clone(),
-            2 => input[0].clone() + input[1].clone(),
-            3 => input[0].clone() + input[1].clone() + input[2].clone(),
-            4 => (input[0].clone() + input[1].clone()) + (input[2].clone() + input[3].clone()),
+            1 => input[0],
+            2 => input[0] + input[1],
+            3 => input[0] + input[1] + input[2],
+            4 => (input[0] + input[1]) + (input[2] + input[3]),
             5 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<1>(&input[4..]),
             6 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<2>(&input[4..]),
             7 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<3>(&input[4..]),
@@ -447,7 +447,7 @@ pub trait BasedVectorSpace<F: PrimeCharacteristicRing>: Sized {
     #[must_use]
     #[inline]
     fn from_basis_coefficients_slice(slice: &[F]) -> Option<Self> {
-        Self::from_basis_coefficients_iter(slice.iter().cloned())
+        Self::from_basis_coefficients_iter(slice.iter().copied())
     }
 
     /// Fixes a basis for the algebra `A` and uses this to
@@ -958,8 +958,8 @@ impl<R: PrimeCharacteristicRing> Iterator for Powers<R> {
     type Item = R;
 
     fn next(&mut self) -> Option<R> {
-        let result = self.current.clone();
-        self.current *= self.base.clone();
+        let result = self.current;
+        self.current *= self.base;
         Some(result)
     }
 }

crates/backend/koala-bear/src/monty_31/poseidon2_monty.rs

@@ -27,9 +27,9 @@ pub trait InternalLayerBaseParameters<MP: MontyParameters, const WIDTH: usize>:
     #[inline(always)]
     fn generic_internal_linear_layer<R: PrimeCharacteristicRing>(state: &mut [R; WIDTH]) {
         // We mostly delegate to internal_layer_mat_mul but have to handle state[0] separately.
-        let part_sum: R = state[1..].iter().cloned().sum();
-        let full_sum = part_sum.clone() + state[0].clone();
-        state[0] = part_sum - state[0].clone();
+        let part_sum: R = state[1..].iter().copied().sum();
+        let full_sum = part_sum + state[0];
+        state[0] = part_sum - state[0];
         Self::internal_layer_mat_mul(state, full_sum);
     }
 }