Tracked Data Structures In Jolt

Tracked Data Types

Before, describing optimisations to polynomial evaluation, we discuss an additional tool to debug numerical algebra libraries. In Jolt sum-check algorithms, almost invariably the bottleneck in computation is number of field multiplications by large scalars. So it might be helpful to be able to count the number of multiplications your algorithm does, and check if we’re under the theoretical limit.

Towards this we introduce TrackedFr data type which is exactly the same as ark_bn254::Fr except for 1 critical difference. Implementation here

Before performing a multiplication operation, we increment a global counter atomically (in a thread safe way). Rust makes this relatively easy to do1. See official docs here

use std::sync::atomic::{AtomicUsize, Ordering};

// A global counter that tracks how many times we've use a*b where a, b are ark_bn254::Fr types
// see jolt-core::field::tracked_fr::TrackedFr for more details
pub static MULT_COUNT: AtomicUsize = AtomicUsize::new(0);

We have the ability to reset this counter at any time as well. So now all we need to do if we wanted to see how many multiplications transpired, is to wrap the function being examined with reset_mult_count() and get_mult_count(). See an example of how to use here.

That’s it! Right now it just tracks multiplications, but you may extend it to any operation ark_bn254::Fr does by simply wrapping the operation as follows:

For all types of multiplication we do w

// crate is jolt-core
use crate::utils::counters::MULT_COUNT;

// Owned * owned
impl Mul<TrackedFr> for TrackedFr {
    type Output = TrackedFr;
    fn mul(self, rhs: TrackedFr) -> Self::Output {
        MULT_COUNT.fetch_add(1, Ordering::Relaxed);
        TrackedFr(self.0 * rhs.0)
    }
}

// Owned * borrowed
impl<'a> Mul<&'a TrackedFr> for TrackedFr {
    type Output = TrackedFr;
    fn mul(self, rhs: &'a TrackedFr) -> Self::Output {
        MULT_COUNT.fetch_add(1, Ordering::Relaxed);
        TrackedFr(self.0 * rhs.0)
    }
}

// Borrowed * owned
impl Mul<TrackedFr> for &TrackedFr {
    type Output = TrackedFr;

    fn mul(self, rhs: TrackedFr) -> Self::Output {
        MULT_COUNT.fetch_add(1, Ordering::Relaxed);
        TrackedFr(self.0 * rhs.0)
    }
}

#[allow(clippy::needless_lifetimes)]
impl<'a, 'b> Mul<&'b TrackedFr> for &'a TrackedFr {
    type Output = TrackedFr;
    fn mul(self, rhs: &'b TrackedFr) -> Self::Output {
        MULT_COUNT.fetch_add(1, Ordering::Relaxed);
        TrackedFr(self.0 * rhs.0)
    }
}

  1. it’s remarkable how far libraries have come that prevent us from writing semaphores and mutexes from scratch↩︎