plonky2-ecdsa/src/curve/curve_msm.rs
src/curve/curve_msm.rsRaw
1use alloc::vec::Vec;
2
3use itertools::Itertools;
4use plonky2::field::types::{Field, PrimeField};
5use plonky2_maybe_rayon::*;
6
7use crate::curve::curve_summation::affine_multisummation_best;
8use crate::curve::curve_types::{AffinePoint, Curve, ProjectivePoint};
9
10/// In Yao's method, we compute an affine summation for each digit. In a parallel setting, it would
11/// be easiest to assign individual summations to threads, but this would be sub-optimal because
12/// multi-summations can be more efficient than repeating individual summations (see
13/// `affine_multisummation_best`). Thus we divide digits into large chunks, and assign chunks of
14/// digits to threads. Note that there is a delicate balance here, as large chunks can result in
15/// uneven distributions of work among threads.
16const DIGITS_PER_CHUNK: usize = 80;
17
18#[derive(Clone, Debug)]
19pub struct MsmPrecomputation<C: Curve> {
20 /// For each generator (in the order they were passed to `msm_precompute`), contains a vector
21 /// of powers, i.e. [(2^w)^i] for i < DIGITS.
22 // TODO: Use compressed coordinates here.
23 powers_per_generator: Vec<Vec<AffinePoint<C>>>,
24
25 /// The window size.
26 w: usize,
27}
28
29pub fn msm_precompute<C: Curve>(
30 generators: &[ProjectivePoint<C>],
31 w: usize,
32) -> MsmPrecomputation<C> {
33 MsmPrecomputation {
34 powers_per_generator: generators
35 .into_par_iter()
36 .map(|&g| precompute_single_generator(g, w))
37 .collect(),
38 w,
39 }
40}
41
42fn precompute_single_generator<C: Curve>(g: ProjectivePoint<C>, w: usize) -> Vec<AffinePoint<C>> {
43 let digits = (C::ScalarField::BITS + w - 1) / w;
44 let mut powers: Vec<ProjectivePoint<C>> = Vec::with_capacity(digits);
45 powers.push(g);
46 for i in 1..digits {
47 let mut power_i_proj = powers[i - 1];
48 for _j in 0..w {
49 power_i_proj = power_i_proj.double();
50 }
51 powers.push(power_i_proj);
52 }
53 ProjectivePoint::batch_to_affine(&powers)
54}
55
56pub fn msm_parallel<C: Curve>(
57 scalars: &[C::ScalarField],
58 generators: &[ProjectivePoint<C>],
59 w: usize,
60) -> ProjectivePoint<C> {
61 let precomputation = msm_precompute(generators, w);
62 msm_execute_parallel(&precomputation, scalars)
63}
64
65pub fn msm_execute<C: Curve>(
66 precomputation: &MsmPrecomputation<C>,
67 scalars: &[C::ScalarField],
68) -> ProjectivePoint<C> {
69 assert_eq!(precomputation.powers_per_generator.len(), scalars.len());
70 let w = precomputation.w;
71 let digits = (C::ScalarField::BITS + w - 1) / w;
72 let base = 1 << w;
73
74 // This is a variant of Yao's method, adapted to the multi-scalar setting. Because we use
75 // extremely large windows, the repeated scans in Yao's method could be more expensive than the
76 // actual group operations. To avoid this, we store a multimap from each possible digit to the
77 // positions in which that digit occurs in the scalars. These positions have the form (i, j),
78 // where i is the index of the generator and j is an index into the digits of the scalar
79 // associated with that generator.
80 let mut digit_occurrences: Vec<Vec<(usize, usize)>> = Vec::with_capacity(digits);
81 for _i in 0..base {
82 digit_occurrences.push(Vec::new());
83 }
84 for (i, scalar) in scalars.iter().enumerate() {
85 let digits = to_digits::<C>(scalar, w);
86 for (j, &digit) in digits.iter().enumerate() {
87 digit_occurrences[digit].push((i, j));
88 }
89 }
90
91 let mut y = ProjectivePoint::ZERO;
92 let mut u = ProjectivePoint::ZERO;
93
94 for digit in (1..base).rev() {
95 for &(i, j) in &digit_occurrences[digit] {
96 u = u + precomputation.powers_per_generator[i][j];
97 }
98 y = y + u;
99 }
100
101 y
102}
103
104pub fn msm_execute_parallel<C: Curve>(
105 precomputation: &MsmPrecomputation<C>,
106 scalars: &[C::ScalarField],
107) -> ProjectivePoint<C> {
108 assert_eq!(precomputation.powers_per_generator.len(), scalars.len());
109 let w = precomputation.w;
110 let digits = (C::ScalarField::BITS + w - 1) / w;
111 let base = 1 << w;
112
113 // This is a variant of Yao's method, adapted to the multi-scalar setting. Because we use
114 // extremely large windows, the repeated scans in Yao's method could be more expensive than the
115 // actual group operations. To avoid this, we store a multimap from each possible digit to the
116 // positions in which that digit occurs in the scalars. These positions have the form (i, j),
117 // where i is the index of the generator and j is an index into the digits of the scalar
118 // associated with that generator.
119 let mut digit_occurrences: Vec<Vec<(usize, usize)>> = Vec::with_capacity(digits);
120 for _i in 0..base {
121 digit_occurrences.push(Vec::new());
122 }
123 for (i, scalar) in scalars.iter().enumerate() {
124 let digits = to_digits::<C>(scalar, w);
125 for (j, &digit) in digits.iter().enumerate() {
126 digit_occurrences[digit].push((i, j));
127 }
128 }
129
130 // For each digit, we add up the powers associated with all occurrences that digit.
131 let digits: Vec<usize> = (0..base).collect();
132 let digit_acc: Vec<ProjectivePoint<C>> = digits
133 .par_chunks(DIGITS_PER_CHUNK)
134 .flat_map(|chunk| {
135 let summations: Vec<Vec<AffinePoint<C>>> = chunk
136 .iter()
137 .map(|&digit| {
138 digit_occurrences[digit]
139 .iter()
140 .map(|&(i, j)| precomputation.powers_per_generator[i][j])
141 .collect()
142 })
143 .collect();
144 affine_multisummation_best(summations)
145 })
146 .collect();
147 // println!("Computing the per-digit summations (in parallel) took {}s", start.elapsed().as_secs_f64());
148
149 let mut y = ProjectivePoint::ZERO;
150 let mut u = ProjectivePoint::ZERO;
151 for digit in (1..base).rev() {
152 u = u + digit_acc[digit];
153 y = y + u;
154 }
155 // println!("Final summation (sequential) {}s", start.elapsed().as_secs_f64());
156 y
157}
158
159pub(crate) fn to_digits<C: Curve>(x: &C::ScalarField, w: usize) -> Vec<usize> {
160 let scalar_bits = C::ScalarField::BITS;
161 let num_digits = (scalar_bits + w - 1) / w;
162
163 // Convert x to a bool array.
164 let x_canonical: Vec<_> = x
165 .to_canonical_biguint()
166 .to_u64_digits()
167 .iter()
168 .cloned()
169 .pad_using(scalar_bits / 64, |_| 0)
170 .collect();
171 let mut x_bits = Vec::with_capacity(scalar_bits);
172 for i in 0..scalar_bits {
173 x_bits.push((x_canonical[i / 64] >> (i as u64 % 64) & 1) != 0);
174 }
175
176 let mut digits = Vec::with_capacity(num_digits);
177 for i in 0..num_digits {
178 let mut digit = 0;
179 for j in ((i * w)..((i + 1) * w).min(scalar_bits)).rev() {
180 digit <<= 1;
181 digit |= x_bits[j] as usize;
182 }
183 digits.push(digit);
184 }
185 digits
186}
187
188#[cfg(test)]
189mod tests {
190 use alloc::vec;
191
192 use num::BigUint;
193 use plonky2::field::secp256k1_scalar::Secp256K1Scalar;
194
195 use super::*;
196 use crate::curve::secp256k1::Secp256K1;
197
198 #[test]
199 fn test_to_digits() {
200 let x_canonical = [
201 0b10101010101010101010101010101010,
202 0b10101010101010101010101010101010,
203 0b11001100110011001100110011001100,
204 0b11001100110011001100110011001100,
205 0b11110000111100001111000011110000,
206 0b11110000111100001111000011110000,
207 0b00001111111111111111111111111111,
208 0b11111111111111111111111111111111,
209 ];
210 let x = Secp256K1Scalar::from_noncanonical_biguint(BigUint::from_slice(&x_canonical));
211 assert_eq!(x.to_canonical_biguint().to_u32_digits(), x_canonical);
212 assert_eq!(
213 to_digits::<Secp256K1>(&x, 17),
214 vec![
215 0b01010101010101010,
216 0b10101010101010101,
217 0b01010101010101010,
218 0b11001010101010101,
219 0b01100110011001100,
220 0b00110011001100110,
221 0b10011001100110011,
222 0b11110000110011001,
223 0b01111000011110000,
224 0b00111100001111000,
225 0b00011110000111100,
226 0b11111111111111110,
227 0b01111111111111111,
228 0b11111111111111000,
229 0b11111111111111111,
230 0b1,
231 ]
232 );
233 }
234
235 #[test]
236 fn test_msm() {
237 let w = 5;
238
239 let generator_1 = Secp256K1::GENERATOR_PROJECTIVE;
240 let generator_2 = generator_1 + generator_1;
241 let generator_3 = generator_1 + generator_2;
242
243 let scalar_1 = Secp256K1Scalar::from_noncanonical_biguint(BigUint::from_slice(&[
244 11111111, 22222222, 33333333, 44444444,
245 ]));
246 let scalar_2 = Secp256K1Scalar::from_noncanonical_biguint(BigUint::from_slice(&[
247 22222222, 22222222, 33333333, 44444444,
248 ]));
249 let scalar_3 = Secp256K1Scalar::from_noncanonical_biguint(BigUint::from_slice(&[
250 33333333, 22222222, 33333333, 44444444,
251 ]));
252
253 let generators = vec![generator_1, generator_2, generator_3];
254 let scalars = vec![scalar_1, scalar_2, scalar_3];
255
256 let precomputation = msm_precompute(&generators, w);
257 let result_msm = msm_execute(&precomputation, &scalars);
258
259 let result_naive = Secp256K1::convert(scalar_1) * generator_1
260 + Secp256K1::convert(scalar_2) * generator_2
261 + Secp256K1::convert(scalar_3) * generator_3;
262
263 assert_eq!(result_msm, result_naive);
264 }
265}