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