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