From 806641d13f57b7cc98f7258f6a0df8e06f580aa9 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Fri, 22 Oct 2021 19:11:05 -0700 Subject: [PATCH] Small optimizations (#319) * Small optimizations * Small optimizations * feedback * inline * feedback * fix unused import --- src/field/crandall_field.rs | 7 +++++++ src/field/field_types.rs | 25 ++++++++++++++++++++++ src/field/goldilocks_field.rs | 39 ++++++++++++++++++++++++++++------- src/gates/base_sum.rs | 2 +- src/gates/gmimc.rs | 2 +- src/gates/poseidon.rs | 4 ++-- src/hash/poseidon.rs | 4 ++-- src/plonk/vanishing_poly.rs | 4 ++-- 8 files changed, 71 insertions(+), 16 deletions(-) diff --git a/src/field/crandall_field.rs b/src/field/crandall_field.rs index 1d55ee3b..648830a4 100644 --- a/src/field/crandall_field.rs +++ b/src/field/crandall_field.rs @@ -86,6 +86,7 @@ impl Field for CrandallField { #[inline] fn from_canonical_u64(n: u64) -> Self { + debug_assert!(n < Self::ORDER); Self(n) } @@ -219,6 +220,12 @@ impl PrimeField for CrandallField { let (sum, over) = self.0.overflowing_add(rhs); Self(sum.wrapping_sub((over as u64) * Self::ORDER)) } + + #[inline] + unsafe fn sub_canonical_u64(&self, rhs: u64) -> Self { + let (sum, under) = self.0.overflowing_sub(rhs); + Self(sum.wrapping_add((under as u64) * Self::ORDER)) + } } impl Neg for CrandallField { diff --git a/src/field/field_types.rs b/src/field/field_types.rs index 8fda4fc6..4fe10b17 100644 --- a/src/field/field_types.rs +++ b/src/field/field_types.rs @@ -62,26 +62,32 @@ pub trait Field: fn order() -> BigUint; + #[inline] fn is_zero(&self) -> bool { *self == Self::ZERO } + #[inline] fn is_nonzero(&self) -> bool { *self != Self::ZERO } + #[inline] fn is_one(&self) -> bool { *self == Self::ONE } + #[inline] fn double(&self) -> Self { *self + *self } + #[inline] fn square(&self) -> Self { *self * *self } + #[inline] fn cube(&self) -> Self { self.square() * *self } @@ -340,6 +346,16 @@ pub trait PrimeField: Field { fn from_noncanonical_u64(n: u64) -> Self; + #[inline] + fn add_one(&self) -> Self { + unsafe { self.add_canonical_u64(1) } + } + + #[inline] + fn sub_one(&self) -> Self { + unsafe { self.sub_canonical_u64(1) } + } + /// Equivalent to *self + Self::from_canonical_u64(rhs), but may be cheaper. The caller must /// ensure that 0 <= rhs < Self::ORDER. The function may return incorrect results if this /// precondition is not met. It is marked unsafe for this reason. @@ -348,6 +364,15 @@ pub trait PrimeField: Field { // Default implementation. *self + Self::from_canonical_u64(rhs) } + + /// Equivalent to *self - Self::from_canonical_u64(rhs), but may be cheaper. The caller must + /// ensure that 0 <= rhs < Self::ORDER. The function may return incorrect results if this + /// precondition is not met. It is marked unsafe for this reason. + #[inline] + unsafe fn sub_canonical_u64(&self, rhs: u64) -> Self { + // Default implementation. + *self - Self::from_canonical_u64(rhs) + } } /// An iterator over the powers of a certain base element `b`: `b^0, b^1, b^2, ...`. diff --git a/src/field/goldilocks_field.rs b/src/field/goldilocks_field.rs index a7d46363..6b13dd80 100644 --- a/src/field/goldilocks_field.rs +++ b/src/field/goldilocks_field.rs @@ -13,7 +13,6 @@ use crate::field::extension_field::quartic::QuarticExtension; use crate::field::extension_field::Extendable; use crate::field::field_types::{Field, PrimeField, RichField}; use crate::field::inversion::try_inverse_u64; -use crate::util::assume; const EPSILON: u64 = (1 << 32) - 1; @@ -93,6 +92,7 @@ impl Field for GoldilocksField { #[inline] fn from_canonical_u64(n: u64) -> Self { + debug_assert!(n < Self::ORDER); Self(n) } @@ -103,6 +103,12 @@ impl Field for GoldilocksField { fn rand_from_rng(rng: &mut R) -> Self { Self::from_canonical_u64(rng.gen_range(0..Self::ORDER)) } + + #[inline] + fn multiply_accumulate(&self, x: Self, y: Self) -> Self { + // u64 + u64 * u64 cannot overflow. + reduce128((self.0 as u128) + (x.0 as u128) * (y.0 as u128)) + } } impl PrimeField for GoldilocksField { @@ -126,6 +132,18 @@ impl PrimeField for GoldilocksField { fn from_noncanonical_u64(n: u64) -> Self { Self(n) } + + #[inline] + unsafe fn add_canonical_u64(&self, rhs: u64) -> Self { + let (res_wrapped, carry) = self.0.overflowing_add(rhs); + Self(res_wrapped.wrapping_add(EPSILON * (carry as u64))) + } + + #[inline] + unsafe fn sub_canonical_u64(&self, rhs: u64) -> Self { + let (res_wrapped, borrow) = self.0.overflowing_sub(rhs); + Self(res_wrapped.wrapping_sub(EPSILON * (borrow as u64))) + } } impl Neg for GoldilocksField { @@ -153,6 +171,7 @@ impl Add for GoldilocksField { } impl AddAssign for GoldilocksField { + #[inline] fn add_assign(&mut self, rhs: Self) { *self = *self + rhs; } @@ -263,7 +282,9 @@ impl RichField for GoldilocksField {} /// the registers, so its use is not recommended when either input will be used again. #[inline(always)] #[cfg(target_arch = "x86_64")] -unsafe fn add_with_wraparound(x: u64, y: u64) -> u64 { +unsafe fn add_no_canonicalize_trashing_input(x: u64, y: u64) -> u64 { + use crate::util::assume; + let res_wrapped: u64; let adjustment: u64; asm!( @@ -288,7 +309,7 @@ unsafe fn add_with_wraparound(x: u64, y: u64) -> u64 { #[inline(always)] #[cfg(not(target_arch = "x86_64"))] -unsafe fn add_with_wraparound(x: u64, y: u64) -> u64 { +unsafe fn add_no_canonicalize_trashing_input(x: u64, y: u64) -> u64 { let (res_wrapped, carry) = x.overflowing_add(y); res_wrapped.wrapping_add(EPSILON * (carry as u64)) } @@ -300,12 +321,14 @@ unsafe fn add_with_wraparound(x: u64, y: u64) -> u64 { /// the registers, so its use is not recommended when either input will be used again. #[inline(always)] #[cfg(target_arch = "x86_64")] -unsafe fn sub_with_wraparound(x: u64, y: u64) -> u64 { +unsafe fn sub_no_canonicalize_trashing_input(x: u64, y: u64) -> u64 { + use crate::util::assume; + let res_wrapped: u64; let adjustment: u64; asm!( "sub {0}, {1}", - "sbb {1:e}, {1:e}", // See add_with_wraparound. + "sbb {1:e}, {1:e}", // See add_no_canonicalize_trashing_input. inlateout(reg) x => res_wrapped, inlateout(reg) y => adjustment, options(pure, nomem, nostack), @@ -316,7 +339,7 @@ unsafe fn sub_with_wraparound(x: u64, y: u64) -> u64 { #[inline(always)] #[cfg(not(target_arch = "x86_64"))] -unsafe fn sub_with_wraparound(x: u64, y: u64) -> u64 { +unsafe fn sub_no_canonicalize_trashing_input(x: u64, y: u64) -> u64 { let (res_wrapped, borrow) = x.overflowing_sub(y); res_wrapped.wrapping_sub(EPSILON * (borrow as u64)) } @@ -329,9 +352,9 @@ fn reduce128(x: u128) -> GoldilocksField { let x_hi_hi = x_hi >> 32; let x_hi_lo = x_hi & EPSILON; - let t0 = unsafe { sub_with_wraparound(x_lo, x_hi_hi) }; + let t0 = unsafe { sub_no_canonicalize_trashing_input(x_lo, x_hi_hi) }; let t1 = x_hi_lo * EPSILON; - let t2 = unsafe { add_with_wraparound(t0, t1) }; + let t2 = unsafe { add_no_canonicalize_trashing_input(t0, t1) }; GoldilocksField(t2) } diff --git a/src/gates/base_sum.rs b/src/gates/base_sum.rs index 2b1e021a..3249821b 100644 --- a/src/gates/base_sum.rs +++ b/src/gates/base_sum.rs @@ -68,7 +68,7 @@ impl, const D: usize, const B: usize> Gate fo let constraints_iter = limbs.iter().map(|&limb| { (0..B) - .map(|i| limb - F::from_canonical_usize(i)) + .map(|i| unsafe { limb.sub_canonical_u64(i as u64) }) .product::() }); constraints.extend(constraints_iter); diff --git a/src/gates/gmimc.rs b/src/gates/gmimc.rs index 4ef87be6..716da5bd 100644 --- a/src/gates/gmimc.rs +++ b/src/gates/gmimc.rs @@ -118,7 +118,7 @@ impl + GMiMC, const D: usize, const WIDTH: u // Assert that `swap` is binary. let swap = vars.local_wires[Self::WIRE_SWAP]; - constraints.push(swap * (swap - F::ONE)); + constraints.push(swap * swap.sub_one()); let mut state = Vec::with_capacity(12); for i in 0..4 { diff --git a/src/gates/poseidon.rs b/src/gates/poseidon.rs index 279ffbbd..493bd263 100644 --- a/src/gates/poseidon.rs +++ b/src/gates/poseidon.rs @@ -3,7 +3,7 @@ use std::marker::PhantomData; use crate::field::extension_field::target::ExtensionTarget; use crate::field::extension_field::Extendable; -use crate::field::field_types::{Field, RichField}; +use crate::field::field_types::{Field, PrimeField, RichField}; use crate::gates::gate::Gate; use crate::hash::poseidon; use crate::hash::poseidon::Poseidon; @@ -180,7 +180,7 @@ where // Assert that `swap` is binary. let swap = vars.local_wires[Self::WIRE_SWAP]; - constraints.push(swap * (swap - F::ONE)); + constraints.push(swap * swap.sub_one()); let mut state = Vec::with_capacity(WIDTH); for i in 0..4 { diff --git a/src/hash/poseidon.rs b/src/hash/poseidon.rs index 5d34c64d..bb3e8a75 100644 --- a/src/hash/poseidon.rs +++ b/src/hash/poseidon.rs @@ -477,8 +477,8 @@ where #[inline(always)] fn sbox_monomial, const D: usize>(x: F) -> F { // x |--> x^7 - let x2 = x * x; - let x4 = x2 * x2; + let x2 = x.square(); + let x4 = x2.square(); let x3 = x * x2; x3 * x4 } diff --git a/src/plonk/vanishing_poly.rs b/src/plonk/vanishing_poly.rs index 0bb03b50..ae79d2dc 100644 --- a/src/plonk/vanishing_poly.rs +++ b/src/plonk/vanishing_poly.rs @@ -1,6 +1,6 @@ use crate::field::extension_field::target::ExtensionTarget; use crate::field::extension_field::{Extendable, FieldExtension}; -use crate::field::field_types::{Field, RichField}; +use crate::field::field_types::{Field, PrimeField, RichField}; use crate::gates::gate::PrefixedGate; use crate::iop::target::Target; use crate::plonk::circuit_builder::CircuitBuilder; @@ -140,7 +140,7 @@ pub(crate) fn eval_vanishing_poly_base, const D: us for i in 0..num_challenges { let z_x = local_zs[i]; let z_gz = next_zs[i]; - vanishing_z_1_terms.push(l1_x * (z_x - F::ONE)); + vanishing_z_1_terms.push(l1_x * z_x.sub_one()); numerator_values.extend((0..num_routed_wires).map(|j| { let wire_value = vars.local_wires[j];