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)
}
}
it’s remarkable how far libraries have come that prevent us from writing semaphores and mutexes from scratch↩︎