plonky2-u32/src/gates/add_many_u32.rs
src/gates/add_many_u32.rsRaw
1use alloc::format;
2use alloc::string::{String, ToString};
3use alloc::vec::Vec;
4use plonky2::util::serialization::{IoResult, Buffer, Write, Read};
5use core::marker::PhantomData;
6
7use itertools::unfold;
8use plonky2::field::extension::Extendable;
9use plonky2::field::types::Field;
10use plonky2::gates::gate::Gate;
11use plonky2::gates::util::StridedConstraintConsumer;
12use plonky2::hash::hash_types::RichField;
13use plonky2::iop::ext_target::ExtensionTarget;
14use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
15use plonky2::iop::target::Target;
16use plonky2::iop::wire::Wire;
17use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
18use plonky2::plonk::circuit_builder::CircuitBuilder;
19use plonky2::plonk::circuit_data::CircuitConfig;
20use plonky2::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
21use plonky2::util::ceil_div_usize;
22
23const LOG2_MAX_NUM_ADDENDS: usize = 4;
24const MAX_NUM_ADDENDS: usize = 16;
25
26/// A gate to perform addition on `num_addends` different 32-bit values, plus a small carry
27#[derive(Copy, Clone, Debug)]
28pub struct U32AddManyGate<F: RichField + Extendable<D>, const D: usize> {
29 pub num_addends: usize,
30 pub num_ops: usize,
31 _phantom: PhantomData<F>,
32}
33
34impl<F: RichField + Extendable<D>, const D: usize> U32AddManyGate<F, D> {
35 pub fn new_from_config(config: &CircuitConfig, num_addends: usize) -> Self {
36 Self {
37 num_addends,
38 num_ops: Self::num_ops(num_addends, config),
39 _phantom: PhantomData,
40 }
41 }
42
43 pub(crate) fn num_ops(num_addends: usize, config: &CircuitConfig) -> usize {
44 debug_assert!(num_addends <= MAX_NUM_ADDENDS);
45 let wires_per_op = (num_addends + 3) + Self::num_limbs();
46 let routed_wires_per_op = num_addends + 3;
47 (config.num_wires / wires_per_op).min(config.num_routed_wires / routed_wires_per_op)
48 }
49
50 pub fn wire_ith_op_jth_addend(&self, i: usize, j: usize) -> usize {
51 debug_assert!(i < self.num_ops);
52 debug_assert!(j < self.num_addends);
53 (self.num_addends + 3) * i + j
54 }
55 pub fn wire_ith_carry(&self, i: usize) -> usize {
56 debug_assert!(i < self.num_ops);
57 (self.num_addends + 3) * i + self.num_addends
58 }
59
60 pub fn wire_ith_output_result(&self, i: usize) -> usize {
61 debug_assert!(i < self.num_ops);
62 (self.num_addends + 3) * i + self.num_addends + 1
63 }
64 pub fn wire_ith_output_carry(&self, i: usize) -> usize {
65 debug_assert!(i < self.num_ops);
66 (self.num_addends + 3) * i + self.num_addends + 2
67 }
68
69 pub fn limb_bits() -> usize {
70 2
71 }
72 pub fn num_result_limbs() -> usize {
73 ceil_div_usize(32, Self::limb_bits())
74 }
75 pub fn num_carry_limbs() -> usize {
76 ceil_div_usize(LOG2_MAX_NUM_ADDENDS, Self::limb_bits())
77 }
78 pub fn num_limbs() -> usize {
79 Self::num_result_limbs() + Self::num_carry_limbs()
80 }
81
82 pub fn wire_ith_output_jth_limb(&self, i: usize, j: usize) -> usize {
83 debug_assert!(i < self.num_ops);
84 debug_assert!(j < Self::num_limbs());
85 (self.num_addends + 3) * self.num_ops + Self::num_limbs() * i + j
86 }
87}
88
89impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for U32AddManyGate<F, D> {
90 fn id(&self) -> String {
91 format!("{self:?}")
92 }
93
94 fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
95 dst.write_usize(self.num_addends);
96 dst.write_usize(self.num_ops)
97 }
98
99 fn deserialize(src: &mut Buffer) -> IoResult<Self> {
100 let num_addends = src.read_usize()?;
101 let num_ops = src.read_usize()?;
102 Ok(Self { num_addends, num_ops, _phantom: PhantomData })
103 }
104
105 fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
106 let mut constraints = Vec::with_capacity(self.num_constraints());
107 for i in 0..self.num_ops {
108 let addends: Vec<F::Extension> = (0..self.num_addends)
109 .map(|j| vars.local_wires[self.wire_ith_op_jth_addend(i, j)])
110 .collect();
111 let carry = vars.local_wires[self.wire_ith_carry(i)];
112
113 let computed_output = addends.iter().fold(F::Extension::ZERO, |x, &y| x + y) + carry;
114
115 let output_result = vars.local_wires[self.wire_ith_output_result(i)];
116 let output_carry = vars.local_wires[self.wire_ith_output_carry(i)];
117
118 let base = F::Extension::from_canonical_u64(1 << 32u64);
119 let combined_output = output_carry * base + output_result;
120
121 constraints.push(combined_output - computed_output);
122
123 let mut combined_result_limbs = F::Extension::ZERO;
124 let mut combined_carry_limbs = F::Extension::ZERO;
125 let base = F::Extension::from_canonical_u64(1u64 << Self::limb_bits());
126 for j in (0..Self::num_limbs()).rev() {
127 let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
128 let max_limb = 1 << Self::limb_bits();
129 let product = (0..max_limb)
130 .map(|x| this_limb - F::Extension::from_canonical_usize(x))
131 .product();
132 constraints.push(product);
133
134 if j < Self::num_result_limbs() {
135 combined_result_limbs = base * combined_result_limbs + this_limb;
136 } else {
137 combined_carry_limbs = base * combined_carry_limbs + this_limb;
138 }
139 }
140 constraints.push(combined_result_limbs - output_result);
141 constraints.push(combined_carry_limbs - output_carry);
142 }
143
144 constraints
145 }
146
147 fn eval_unfiltered_base_one(
148 &self,
149 vars: EvaluationVarsBase<F>,
150 mut yield_constr: StridedConstraintConsumer<F>,
151 ) {
152 for i in 0..self.num_ops {
153 let addends: Vec<F> = (0..self.num_addends)
154 .map(|j| vars.local_wires[self.wire_ith_op_jth_addend(i, j)])
155 .collect();
156 let carry = vars.local_wires[self.wire_ith_carry(i)];
157
158 let computed_output = addends.iter().fold(F::ZERO, |x, &y| x + y) + carry;
159
160 let output_result = vars.local_wires[self.wire_ith_output_result(i)];
161 let output_carry = vars.local_wires[self.wire_ith_output_carry(i)];
162
163 let base = F::from_canonical_u64(1 << 32u64);
164 let combined_output = output_carry * base + output_result;
165
166 yield_constr.one(combined_output - computed_output);
167
168 let mut combined_result_limbs = F::ZERO;
169 let mut combined_carry_limbs = F::ZERO;
170 let base = F::from_canonical_u64(1u64 << Self::limb_bits());
171 for j in (0..Self::num_limbs()).rev() {
172 let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
173 let max_limb = 1 << Self::limb_bits();
174 let product = (0..max_limb)
175 .map(|x| this_limb - F::from_canonical_usize(x))
176 .product();
177 yield_constr.one(product);
178
179 if j < Self::num_result_limbs() {
180 combined_result_limbs = base * combined_result_limbs + this_limb;
181 } else {
182 combined_carry_limbs = base * combined_carry_limbs + this_limb;
183 }
184 }
185 yield_constr.one(combined_result_limbs - output_result);
186 yield_constr.one(combined_carry_limbs - output_carry);
187 }
188 }
189
190 fn eval_unfiltered_circuit(
191 &self,
192 builder: &mut CircuitBuilder<F, D>,
193 vars: EvaluationTargets<D>,
194 ) -> Vec<ExtensionTarget<D>> {
195 let mut constraints = Vec::with_capacity(self.num_constraints());
196
197 for i in 0..self.num_ops {
198 let addends: Vec<ExtensionTarget<D>> = (0..self.num_addends)
199 .map(|j| vars.local_wires[self.wire_ith_op_jth_addend(i, j)])
200 .collect();
201 let carry = vars.local_wires[self.wire_ith_carry(i)];
202
203 let mut computed_output = carry;
204 for addend in addends {
205 computed_output = builder.add_extension(computed_output, addend);
206 }
207
208 let output_result = vars.local_wires[self.wire_ith_output_result(i)];
209 let output_carry = vars.local_wires[self.wire_ith_output_carry(i)];
210
211 let base: F::Extension = F::from_canonical_u64(1 << 32u64).into();
212 let base_target = builder.constant_extension(base);
213 let combined_output =
214 builder.mul_add_extension(output_carry, base_target, output_result);
215
216 constraints.push(builder.sub_extension(combined_output, computed_output));
217
218 let mut combined_result_limbs = builder.zero_extension();
219 let mut combined_carry_limbs = builder.zero_extension();
220 let base = builder
221 .constant_extension(F::Extension::from_canonical_u64(1u64 << Self::limb_bits()));
222 for j in (0..Self::num_limbs()).rev() {
223 let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
224 let max_limb = 1 << Self::limb_bits();
225
226 let mut product = builder.one_extension();
227 for x in 0..max_limb {
228 let x_target =
229 builder.constant_extension(F::Extension::from_canonical_usize(x));
230 let diff = builder.sub_extension(this_limb, x_target);
231 product = builder.mul_extension(product, diff);
232 }
233 constraints.push(product);
234
235 if j < Self::num_result_limbs() {
236 combined_result_limbs =
237 builder.mul_add_extension(base, combined_result_limbs, this_limb);
238 } else {
239 combined_carry_limbs =
240 builder.mul_add_extension(base, combined_carry_limbs, this_limb);
241 }
242 }
243 constraints.push(builder.sub_extension(combined_result_limbs, output_result));
244 constraints.push(builder.sub_extension(combined_carry_limbs, output_carry));
245 }
246
247 constraints
248 }
249
250 fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
251 (0..self.num_ops)
252 .map(|i| {
253 WitnessGeneratorRef::new(
254 U32AddManyGenerator {
255 gate: *self,
256 row,
257 i,
258 _phantom: PhantomData,
259 }
260 .adapter(),
261 )
262 })
263 .collect()
264 }
265
266 fn num_wires(&self) -> usize {
267 (self.num_addends + 3) * self.num_ops + Self::num_limbs() * self.num_ops
268 }
269
270 fn num_constants(&self) -> usize {
271 0
272 }
273
274 fn degree(&self) -> usize {
275 1 << Self::limb_bits()
276 }
277
278 fn num_constraints(&self) -> usize {
279 self.num_ops * (3 + Self::num_limbs())
280 }
281}
282
283#[derive(Clone, Debug)]
284struct U32AddManyGenerator<F: RichField + Extendable<D>, const D: usize> {
285 gate: U32AddManyGate<F, D>,
286 row: usize,
287 i: usize,
288 _phantom: PhantomData<F>,
289}
290
291impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
292 for U32AddManyGenerator<F, D>
293{
294 fn id(&self) -> String {
295 "U32AddManyGenerator".to_string()
296 }
297
298 fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
299 self.gate.serialize(dst)?;
300 dst.write_usize(self.row)?;
301 dst.write_usize(self.i)
302 }
303
304 fn deserialize(src: &mut Buffer) -> IoResult<Self> {
305 let gate = U32AddManyGate::deserialize(src)?;
306 let row = src.read_usize()?;
307 let i = src.read_usize()?;
308 Ok(Self { gate, row, i, _phantom: PhantomData })
309 }
310
311 fn dependencies(&self) -> Vec<Target> {
312 let local_target = |column| Target::wire(self.row, column);
313
314 (0..self.gate.num_addends)
315 .map(|j| local_target(self.gate.wire_ith_op_jth_addend(self.i, j)))
316 .chain([local_target(self.gate.wire_ith_carry(self.i))])
317 .collect()
318 }
319
320 fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
321 let local_wire = |column| Wire {
322 row: self.row,
323 column,
324 };
325
326 let get_local_wire = |column| witness.get_wire(local_wire(column));
327
328 let addends: Vec<_> = (0..self.gate.num_addends)
329 .map(|j| get_local_wire(self.gate.wire_ith_op_jth_addend(self.i, j)))
330 .collect();
331 let carry = get_local_wire(self.gate.wire_ith_carry(self.i));
332
333 let output = addends.iter().fold(F::ZERO, |x, &y| x + y) + carry;
334 let output_u64 = output.to_canonical_u64();
335
336 let output_carry_u64 = output_u64 >> 32;
337 let output_result_u64 = output_u64 & ((1 << 32) - 1);
338
339 let output_carry = F::from_canonical_u64(output_carry_u64);
340 let output_result = F::from_canonical_u64(output_result_u64);
341
342 let output_carry_wire = local_wire(self.gate.wire_ith_output_carry(self.i));
343 let output_result_wire = local_wire(self.gate.wire_ith_output_result(self.i));
344
345 out_buffer.set_wire(output_carry_wire, output_carry);
346 out_buffer.set_wire(output_result_wire, output_result);
347
348 let num_result_limbs = U32AddManyGate::<F, D>::num_result_limbs();
349 let num_carry_limbs = U32AddManyGate::<F, D>::num_carry_limbs();
350 let limb_base = 1 << U32AddManyGate::<F, D>::limb_bits();
351
352 let split_to_limbs = |mut val, num| {
353 unfold((), move |_| {
354 let ret = val % limb_base;
355 val /= limb_base;
356 Some(ret)
357 })
358 .take(num)
359 .map(F::from_canonical_u64)
360 };
361
362 let result_limbs = split_to_limbs(output_result_u64, num_result_limbs);
363 let carry_limbs = split_to_limbs(output_carry_u64, num_carry_limbs);
364
365 for (j, limb) in result_limbs.chain(carry_limbs).enumerate() {
366 let wire = local_wire(self.gate.wire_ith_output_jth_limb(self.i, j));
367 out_buffer.set_wire(wire, limb);
368 }
369 }
370}
371
372#[cfg(test)]
373mod tests {
374 use anyhow::Result;
375 use plonky2::field::extension::quartic::QuarticExtension;
376 use plonky2::field::goldilocks_field::GoldilocksField;
377 use plonky2::field::types::Sample;
378 use plonky2::gates::gate_testing::{test_eval_fns, test_low_degree};
379 use plonky2::hash::hash_types::HashOut;
380 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
381 use rand::rngs::OsRng;
382 use rand::Rng;
383
384 use super::*;
385
386 #[test]
387 fn low_degree() {
388 test_low_degree::<GoldilocksField, _, 4>(U32AddManyGate::<GoldilocksField, 4> {
389 num_addends: 4,
390 num_ops: 3,
391 _phantom: PhantomData,
392 })
393 }
394
395 #[test]
396 fn eval_fns() -> Result<()> {
397 const D: usize = 2;
398 type C = PoseidonGoldilocksConfig;
399 type F = <C as GenericConfig<D>>::F;
400 test_eval_fns::<F, C, _, D>(U32AddManyGate::<GoldilocksField, D> {
401 num_addends: 4,
402 num_ops: 3,
403 _phantom: PhantomData,
404 })
405 }
406
407 #[test]
408 fn test_gate_constraint() {
409 type F = GoldilocksField;
410 type FF = QuarticExtension<GoldilocksField>;
411 const D: usize = 4;
412 const NUM_ADDENDS: usize = 10;
413 const NUM_U32_ADD_MANY_OPS: usize = 3;
414
415 fn get_wires(addends: Vec<Vec<u64>>, carries: Vec<u64>) -> Vec<FF> {
416 let mut v0 = Vec::new();
417 let mut v1 = Vec::new();
418
419 let num_result_limbs = U32AddManyGate::<F, D>::num_result_limbs();
420 let num_carry_limbs = U32AddManyGate::<F, D>::num_carry_limbs();
421 let limb_base = 1 << U32AddManyGate::<F, D>::limb_bits();
422 for op in 0..NUM_U32_ADD_MANY_OPS {
423 let adds = &addends[op];
424 let ca = carries[op];
425
426 let output = adds.iter().sum::<u64>() + ca;
427 let output_result = output & ((1 << 32) - 1);
428 let output_carry = output >> 32;
429
430 let split_to_limbs = |mut val, num| {
431 unfold((), move |_| {
432 let ret = val % limb_base;
433 val /= limb_base;
434 Some(ret)
435 })
436 .take(num)
437 .map(F::from_canonical_u64)
438 };
439
440 let mut result_limbs: Vec<_> =
441 split_to_limbs(output_result, num_result_limbs).collect();
442 let mut carry_limbs: Vec<_> =
443 split_to_limbs(output_carry, num_carry_limbs).collect();
444
445 for a in adds {
446 v0.push(F::from_canonical_u64(*a));
447 }
448 v0.push(F::from_canonical_u64(ca));
449 v0.push(F::from_canonical_u64(output_result));
450 v0.push(F::from_canonical_u64(output_carry));
451 v1.append(&mut result_limbs);
452 v1.append(&mut carry_limbs);
453 }
454
455 v0.iter().chain(v1.iter()).map(|&x| x.into()).collect()
456 }
457
458 let mut rng = OsRng;
459 let addends: Vec<Vec<_>> = (0..NUM_U32_ADD_MANY_OPS)
460 .map(|_| (0..NUM_ADDENDS).map(|_| rng.gen::<u32>() as u64).collect())
461 .collect();
462 let carries: Vec<_> = (0..NUM_U32_ADD_MANY_OPS)
463 .map(|_| rng.gen::<u32>() as u64)
464 .collect();
465
466 let gate = U32AddManyGate::<F, D> {
467 num_addends: NUM_ADDENDS,
468 num_ops: NUM_U32_ADD_MANY_OPS,
469 _phantom: PhantomData,
470 };
471
472 let vars = EvaluationVars {
473 local_constants: &[],
474 local_wires: &get_wires(addends, carries),
475 public_inputs_hash: &HashOut::rand(),
476 };
477
478 assert!(
479 gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()),
480 "Gate constraints are not satisfied."
481 );
482 }
483}