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