arkworks-rs / algebra

Libraries for finite field, elliptic curve, and polynomial arithmetic
https://arkworks.rs
Apache License 2.0
624 stars 243 forks source link

bounded multi-scalar multiplications. #685

Open mmaker opened 1 year ago

mmaker commented 1 year ago

Here's a much faster Pippegenger algorithm for MSM's with u8 integers.

We are talking about > 3x faster for sizes of ~1k elements. It would be nice to have a generic interface for MSM that depending on the vector of scalars does something smarter than just convert to field elements and invoke the generic MSM algorithm.

/// Simpler pippenger for u8 scalars.
pub fn u8msm<G: CurveGroup>(bases: &[G::Affine], scalars: &[u8]) -> G {
    let mut buckets = [G::zero(); 8];
    for (base, scalar) in bases.iter().zip(scalars.iter()) {
        buckets.iter_mut().enumerate().for_each(|(i, bucket)| {
            if scalar & (1u8 << i) != 0 {
                bucket.add_assign(base);
            }
        });
    }

    for b in (1..8).rev() {
        buckets[b].double_in_place();
        buckets[b - 1].add_assign(buckets[b]);
    }
    buckets[0]
}
mmagician commented 1 year ago

For an example where this approach is already used, see e.g. the Lasso implementation, where a very simple solution reportedly speeds up MSMs by ~30% on 2^24 inputs - in their instance the BigInts are known to be small. https://github.com/a16z/Lasso/blob/e52d160c44b701f538de39b29502546f85af4ac6/src/msm/mod.rs#L95. Maybe we could automatically derive implementations for {u8, ..., u64}msm with a macro? Tagging Lasso implementers @moodlezoup @sragss

mmaker commented 1 year ago

I'd rather go with a trait VariableBaseMSM<Scalar> that gets implemented as VariableBaseMSM<G::ScalarField>, VariableBaseMSM<u8>, etc.. That'd spare us the conversion to BigInts and back, and seems more ergonomic as the user doesn't even have to convert these values. What do you think?

mmagician commented 1 year ago

That indeed sounds optimal

mmagician commented 1 year ago

Actually, when checking the downstream usages of VariableBaseMSM in e.g. arkworks spartan it's not obvious that VariableBaseMSM<u8/u64/etc> will actually be directly useful (even though more performant in isolation). Typically you'd already have the inputs as some Field, so in order to call MSM the implemented for u64, you would anyway need to perform the conversions.

mmaker commented 1 year ago

they could implement domain-specific functions or do something like the following I think?

type GG = u64;

pub trait VariableBaseMSM<T>: Sized {

    fn msm_unchecked(bases: &[Self], scalars: &[T]) -> Self;
}

impl VariableBaseMSM<u8> for GG { 
    fn msm_unchecked(bases: &[Self], scalars: &[u8]) -> Self {
        bases.iter().zip(scalars.iter()).map(|(&x, &y)| x*(y as GG)).sum()   
    }
}

impl VariableBaseMSM<u16> for GG { 
    fn msm_unchecked(bases: &[Self], scalars: &[u16]) -> Self {
        bases.iter().zip(scalars.iter()).map(|(&x, &y)| x*(y as GG)).sum()   
    }
}

fn commit<T>(ck: &[GG], scalars: &[T]) -> GG
    where GG: VariableBaseMSM<T> {
    GG::msm_unchecked(ck, scalars)
}

We do something similar with msm_bigint: we could just have it as a type of scalar

Pratyush commented 1 year ago

A slightly simpler solution would be to just put a NUM_BITS variable on the msm method. This would allow us to avoid the proliferation of impl VariableBaseMSM<T> for each T.

const FULL: u16 = u16::MAX;
pub trait VariableBaseMSM: Sized {
    // by default we assume the scalars are full-sized field elements
    fn msm_unchecked<const NUM_BITS: usize = FULL>(bases: &[Self], scalars: &[F]) -> Self {
        // Some logic based on `NUM_BITS`
    }
}

fn u64_commit<G: VariableBaseMSM>(ck: &[G], scalars: &[G::Scalar])  {
    GG::msm_unchecked<64>(ck, scalars)
}

We could also offer a "simple" method which internally figures out a maximum size S of the input scalars, and then invokes msm_unchecked<S>

mmaker commented 11 months ago

Agreed -- but when you have bounded elements they are often not as a Field type, so we'd have to come up with something to make the API more ergonomic

burdges commented 11 months ago

Is the issue the delegation to SWCurveConfig::msm and similar? It's likely some of that interface should be improved anyways, but they enable things like FPGAs or native calls from WASM, but they space of trade offs gets subtle.

In unsafe code, we could literally enforce a type equivalent to &[ #[repr(align(runtime_limb_size))] [[u8; runtime_limb_size]; runtime_limbs_count] ]