TL;DR - Just Tell Us What You Did

What does this mean for end to end performance? – We’ve sped up sum-check, that should speed up everything right? The answer is Yes – in theory, but we cannot yet materialise these improvements in our end to end system. The reason why, has nothing to do with the optimisations we are going to discuss below. A short description of the reason is – our parallelisation code is not perfect – this prevents our traces from always accurately reflecting the impact of low level optimisations (especially optimisations that are in the thick of parallel code). We are in the process of writing a detailed report of this issue. When it’s ready, it’ll be available here.

With that said, here’s are the details of how we came up with this optimisation. First we briefly review the sum-check protocol, which prepares us to understand exactly why speeding up a certain type of multiplication would be beneficial. Then we get to showing how we speed up the multiplication, and end with benches that confirm theory.

Sum-Check Recap

We briefly recap the sum-check protocol¹.

Let \(g \in \mathbb{F}[X_1, \ldots, X_m]\) be a polynomial over \(\mathbb{F}\) with the following structure. \[ g(X_1, \ldots, X_m) = \prod_{j=1}^d q_j(X_1, \ldots, X_m)\] where all \(q_j(X)\)’s are multi-linear polynomials over \(\mathbb{F}\). The sum-check algorithm an interactive protocol² where the prover is trying to convince a verifier that \(\sum_{b_1 \in \{ 0, 1\}} \ldots \sum_{b_m \in \{ 0, 1\}} g(b_1, \ldots, b_m)\) evaluates to \(c \in \mathbb{F}\). In each round \(i\) the prover does the following things - (1) It sends the verifier a univariate polynomial \(p_i(X_i)\). (2) The verifier samples \(r_i \xleftarrow[]{\$}S\) uniformly at random, and the prover computes the polynomial \(g(r_1, \ldots, r_i, X_{i+1}, \ldots, X_m)\). Typically \(S = \mathbb{F}\). We expand more below:

1. Polynomial evaluation The prover computes the following univariate polynomial, using the following steps \[p_i(X_i) = \sum_{b_{i+1} \in \{0,1\}}\ldots\sum_{b_{m} \in \{0,1\}^{}}g(r_1, \ldots, r_{i-1}, X_i, b_{i+1}, b_m)\]

1. If \(g\) has degree \(d\), this requires computing \(p_i(0), \ldots, p_i(d)\), then via standard lagrange interpolation, we’re able to fully define \(p_i(X)\). \[p_i(Z) = \prod_{j=1}^d \Bigg( \sum_{b_{i+1} \in \{0,1\}}\ldots\sum_{b_{m}} q_j\Big(r_1, \ldots, r_{i-1}, Z, b_{i+1}, \ldots, b_m\Big)\Bigg)\] so to get \(p_i(Z)\) we compute for all \(j \in [d]\), \(q_j(r_1, \ldots, r_{i-1}, Z, \vec{b})\) for \(Z \in \{0, \ldots, d\}\) for all \(\vec{b} \in \{0, 1\}^{m-i}\). As \(q_j\) is multi-linear – this can be computed by linear-interpolation.

\[ \begin{aligned} q_j(r_1, \ldots, r_{i-1}, Z, \vec{b}) = & (1-Z) \,q_j(r_1, \ldots, r_{i-1}, 0, b_{i+1}, \ldots, b_m) \\[10pt] & + Z\, q_j(r_1, \ldots, r_{i-1}, 1, b_{i+1}, \ldots, b_m) \end{aligned} \]

As \(d\) is a small constant number, computing the above value requires no multiplication operations. We can compute it via a running sum. See pseudocode below

/// Compute q_j(Z) = (1-Z) * q_j(0) + Z * q_j(1)
/// using only additions
/// when Z is a small integer.
fn compute_q_j_running_sum(z: usize, q_at_0: JoltField, q_at_1: JoltField) -> JoltField {
        let diff = q_at_1 - q_at_0;

        // Compute Z * diff using repeated addition (running sum)
        let mut z_times_diff = JoltField::zero();
        for _ in 0..z {
                z_times_diff += diff;
            }

        q_at_0 + z_times_diff

    }

2. Binding Phase Thee prover needs to binds the multi-linear polynomial \(g\) to \(r_i\), that is to say it computes a new multi-linear polynomial in \(m-i\) variables \(g_i(X_{i+1}, \ldots, X_{m}) = g(r_1, \ldots, r_{i-1}, X_{i+1}, \ldots, X_{m})\). This involves storing all evaluations of \(g_i\) over the hyper cube \(\{0, 1\}^{m-i}\), but unlike the above step \(r_i\) is NOT guaranteed to be a small number. Therefore the running sum described in the pseudocode above is no good to use. We need to use actual multiplication here.

\[ \begin{aligned} q_j(r_1, \ldots, r_{i-1}, \textcolor{orange}{r_i}, \vec{b}) = & (1-\textcolor{orange}{r_i}) \,q_j(r_1, \ldots, r_{i-1}, 0, b_{i+1}, \ldots, b_m) \\[10pt] & + \textcolor{orange}{r_i}\, q_j(r_1, \ldots, r_{i-1}, 1, b_{i+1}, \ldots, b_m) \end{aligned} \]

Now \(q_j(r_1, \ldots, r_{i-1}, 1, b_{i+1}, \ldots, b_m) \in \mathbb{F}\), and traditionally the set \(S\) from which we sample \(r_i\) is also equal to \(\mathbb{F}\). Therefore this step requires \(2^{m-i}\) \(\mathbb{F}\times \mathbb{F}\) multiplications in round \(i\) – a total of \(2^m\) multiplications over \(m\) rounds. It is well known that multiplication for \(\mathbb{F}\) elements is 1000x costlier than addition. So our hope is that we can choose \(S\) cleverly such that this operation can be done faster. Note we cannot make \(S\) too small, as this the soundness probability of any sumcheck protocol is given by \(\frac{d}{|S|}\).

The Challenge Set \(S\)

Define the set \(S \subset \mathbb{F}\) where \(S = \{ x \in \mathbb{F}: \text{the least 2 significant digits of $x$ in Montgomery form are 0} \}\).

cargo bench -p jolt-core --bench challenge_mult_performance

cargo bench -p jolt-core --bench binding
cargo bench -p jolt-core --bench binding --features challenge-254-bit

Some Implementation Details

we use \(256\) bits to represent big integers and our challenge and the modulus for ark-bn254 uses \(254\) bits. Thus, if sampled a 2 limbed \(128\) bit integer \(x\) and used that as the challenge as is x = [x_1, x_0, 0, 0] we are not guaranteed that \(x < p\) which the CIOS algorithm assumes. So we clear the first 3 bits of \(x\) to ensure that \(x\) is ALWAYS less than \(p\).

impl<F: JoltField> MontU128Challenge<F> {
pub fn new(value: u128) -> Self {
    // MontU128 can always be represented by 125 bits.
    // This guarantees that the big integer is never greater than the
    // bn254 modulus
    let val_masked = value & (u128::MAX >> 3);
    let low = val_masked as u64;
    let high = (val_masked >> 64) as u64;
    Self {
        value: [0, 0, low, high],
        _marker: PhantomData,
    }
}

pub fn value(&self) -> [u64; 4] {
    self.value
}

pub fn random<R: Rng + ?Sized>(rng: &mut R) -> Self {
        Self::from(rng.gen::<u128>())
    }
}

The actual multiplication is in the A16z fork of arkworls-algebra crate. Below is a snippet of the optimised CIOS code for the set \(S\) above

Footnotes and References