websnark/src/build_fft.js
src/build_fft.jsRaw
1/*
2 Copyright 2019 0KIMS association.
3
4 This file is part of websnark (Web Assembly zkSnark Prover).
5
6 websnark is a free software: you can redistribute it and/or modify it
7 under the terms of the GNU General Public License as published by
8 the Free Software Foundation, either version 3 of the License, or
9 (at your option) any later version.
10
11 websnark is distributed in the hope that it will be useful, but WITHOUT
12 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
13 or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
14 License for more details.
15
16 You should have received a copy of the GNU General Public License
17 along with websnark. If not, see <https://www.gnu.org/licenses/>.
18*/
19
20const bigInt = require("big-integer");
21const utils = require("./utils.js");
22
23module.exports = function buildFFT(module, prefix, f1mPrefix) {
24
25 const n64 = module.modules[f1mPrefix].n64;
26 const n8 = n64*8;
27 const q = module.modules[f1mPrefix].q;
28
29 let rem = q.minus(bigInt(1));
30 let maxBits = 0;
31 while (!rem.isOdd()) {
32 maxBits ++;
33 rem = rem.shiftRight(1);
34 }
35
36 let nr = bigInt(2);
37
38 while ( nr.modPow(q.shiftRight(1), q).equals(1) ) nr = nr.add(1);
39
40 const w = new Array(maxBits+1);
41 w[maxBits] = nr.modPow(rem, q);
42
43 let n=maxBits-1;
44 while (n>=0) {
45 w[n] = w[n+1].modPow(2, q);
46 n--;
47 }
48
49 const bytes = [];
50 const R = bigInt(1).shiftLeft(n8*8).mod(q);
51
52 for (let i=0; i<w.length; i++) {
53 const m = w[i].times(R).mod(q);
54 bytes.push(...utils.bigInt2BytesLE(m, n8));
55 }
56
57 const ROOTs = module.alloc(bytes);
58
59 const i2 = new Array(maxBits+1);
60 i2[0] = bigInt(1);
61
62 for (let i=1; i<=maxBits; i++) {
63 i2[i] = i2[i-1].times(2);
64 }
65
66 const bytesi2 =[];
67 for (let i=0; i<=maxBits; i++) {
68 const m = i2[i].modInv(q).times(R).mod(q);
69 bytesi2.push(...utils.bigInt2BytesLE(m, n8));
70 }
71
72 const INV2 = module.alloc(bytesi2);
73
74 function rev(x) {
75 let r=0;
76 for (let i=0; i<8; i++) {
77 if (x & (1 << i)) {
78 r = r | (0x80 >> i);
79 }
80 }
81 return r;
82 }
83
84 const rtable = Array(256);
85 for (let i=0; i<256; i++) {
86 rtable[i] = rev(i);
87 }
88
89 const REVTABLE = module.alloc(rtable);
90
91
92 function buildLog2() {
93 const f = module.addFunction(prefix+"__log2");
94 f.addParam("n", "i32");
95 f.setReturnType("i32");
96 f.addLocal("bits", "i32");
97 f.addLocal("aux", "i32");
98
99 const c = f.getCodeBuilder();
100
101 f.addCode(
102 c.setLocal(
103 "aux",
104 c.i32_shr_u(
105 c.getLocal("n"),
106 c.i32_const(1)
107 )
108 )
109 );
110 f.addCode(c.setLocal("bits", c.i32_const(0)));
111
112 f.addCode(c.block(c.loop(
113 c.br_if(
114 1,
115 c.i32_eqz(c.getLocal("aux"))
116 ),
117
118 c.setLocal(
119 "aux",
120 c.i32_shr_u(
121 c.getLocal("aux"),
122 c.i32_const(1)
123 )
124 ),
125
126 c.setLocal(
127 "bits",
128 c.i32_add(
129 c.getLocal("bits"),
130 c.i32_const(1)
131 )
132 ),
133
134 c.br(0)
135 )));
136
137 f.addCode(c.if(
138 c.i32_ne(
139 c.getLocal("n"),
140 c.i32_shl(
141 c.i32_const(1),
142 c.getLocal("bits")
143 )
144 ),
145 c.unreachable()
146 ));
147
148 f.addCode(c.if(
149 c.i32_gt_u(
150 c.getLocal("bits"),
151 c.i32_const(maxBits)
152 ),
153 c.unreachable()
154 ));
155
156 f.addCode(c.getLocal("bits"));
157 }
158
159 function buildFFT() {
160 const f = module.addFunction(prefix+"_fft");
161 f.addParam("px", "i32");
162 f.addParam("n", "i32");
163 f.addParam("odd", "i32");
164
165 f.addLocal("bits", "i32");
166
167 const c = f.getCodeBuilder();
168
169 f.addCode(
170 c.setLocal(
171 "bits",
172 c.call(
173 prefix + "__log2",
174 c.getLocal("n")
175 )
176 )
177 );
178
179
180 f.addCode(c.call(
181 prefix+"__rawfft",
182 c.getLocal("px"),
183 c.getLocal("bits"),
184 c.getLocal("odd"),
185 ));
186
187 }
188
189 function buildIFFT() {
190 const f = module.addFunction(prefix+"_ifft");
191 f.addParam("px", "i32");
192 f.addParam("n", "i32");
193 f.addParam("odd", "i32");
194 f.addLocal("bits", "i32");
195
196 const c = f.getCodeBuilder();
197
198 f.addCode(
199 c.setLocal(
200 "bits",
201 c.call(
202 prefix + "__log2",
203 c.getLocal("n")
204 )
205 )
206 );
207
208 f.addCode(c.call(
209 prefix+"__rawfft",
210 c.getLocal("px"),
211 c.getLocal("bits"),
212 c.getLocal("odd")
213 ));
214
215 f.addCode(c.call(
216 prefix+"__finalInverse",
217 c.getLocal("px"),
218 c.getLocal("bits"),
219 ));
220
221 }
222
223 function buildRawFFT() {
224 const f = module.addFunction(prefix+"__rawfft");
225 f.addParam("px", "i32");
226 f.addParam("bits", "i32");
227 f.addParam("odd", "i32");
228 f.addLocal("s", "i32");
229 f.addLocal("k", "i32");
230 f.addLocal("j", "i32");
231 f.addLocal("m", "i32");
232 f.addLocal("mdiv2", "i32");
233 f.addLocal("n", "i32");
234 f.addLocal("pwm", "i32");
235 f.addLocal("idx1", "i32");
236 f.addLocal("idx2", "i32");
237
238 const c = f.getCodeBuilder();
239
240 const W = c.i32_const(module.alloc(n8));
241 const T = c.i32_const(module.alloc(n8));
242 const U = c.i32_const(module.alloc(n8));
243
244 f.addCode(
245 c.call(prefix + "__reversePermutation", c.getLocal("px"), c.getLocal("bits")),
246 c.setLocal("n", c.i32_shl(c.i32_const(1), c.getLocal("bits"))),
247 c.setLocal("s", c.i32_const(1)),
248 c.block(c.loop(
249 c.br_if(
250 1,
251 c.i32_gt_u(
252 c.getLocal("s"),
253 c.getLocal("bits")
254 )
255 ),
256 c.setLocal("m", c.i32_shl(c.i32_const(1), c.getLocal("s"))),
257 c.setLocal("pwm",
258 c.i32_add(
259 c.i32_const(ROOTs),
260 c.i32_mul(
261 c.getLocal("s"),
262 c.i32_const(n8)
263 )
264 )
265 ),
266 c.setLocal("k", c.i32_const(0)),
267 c.block(c.loop(
268 c.br_if(
269 1,
270 c.i32_ge_u(
271 c.getLocal("k"),
272 c.getLocal("n")
273 )
274 ),
275
276 c.if(
277 c.getLocal("odd"),
278 c.call(
279 f1mPrefix + "_copy",
280 c.i32_add(
281 c.getLocal("pwm"),
282 c.i32_const(n8)
283 ),
284 W
285 ),
286 c.call(f1mPrefix + "_one", W)
287 ),
288
289
290 c.setLocal("mdiv2", c.i32_shr_u(c.getLocal("m"), c.i32_const(1)) ),
291 c.setLocal("j", c.i32_const(0)),
292 c.block(c.loop(
293 c.br_if(
294 1,
295 c.i32_ge_u(
296 c.getLocal("j"),
297 c.getLocal("mdiv2")
298 )
299 ),
300
301 c.setLocal(
302 "idx1",
303 c.i32_add(
304 c.getLocal("px"),
305 c.i32_mul(
306 c.i32_add(
307 c.getLocal("k"),
308 c.getLocal("j")
309 ),
310 c.i32_const(n8)
311 )
312 )
313 ),
314
315 c.setLocal(
316 "idx2",
317 c.i32_add(
318 c.getLocal("idx1"),
319 c.i32_mul(
320 c.getLocal("mdiv2"),
321 c.i32_const(n8)
322 )
323 )
324 ),
325
326 c.call(
327 f1mPrefix + "_mul",
328 W,
329 c.getLocal("idx2"),
330 T
331 ),
332
333 c.call(
334 f1mPrefix + "_copy",
335 c.getLocal("idx1"),
336 U
337 ),
338
339 c.call(
340 f1mPrefix + "_add",
341 U,
342 T,
343 c.getLocal("idx1"),
344 ),
345
346 c.call(
347 f1mPrefix + "_sub",
348 U,
349 T,
350 c.getLocal("idx2"),
351 ),
352
353 c.call(
354 f1mPrefix + "_mul",
355 W,
356 c.getLocal("pwm"),
357 W,
358 ),
359
360 c.setLocal("j", c.i32_add(c.getLocal("j"), c.i32_const(1))),
361 c.br(0)
362 )),
363
364 c.setLocal("k", c.i32_add(c.getLocal("k"), c.getLocal("m"))),
365 c.br(0)
366 )),
367
368 c.setLocal("s", c.i32_add(c.getLocal("s"), c.i32_const(1))),
369 c.br(0)
370 ))
371 );
372 }
373
374 function buildCopyInterleaved() {
375 const f = module.addFunction(prefix+"_copyNInterleaved");
376 f.addParam("ps", "i32");
377 f.addParam("pd", "i32");
378 f.addParam("n", "i32");
379 f.addLocal("pi", "i32");
380 f.addLocal("po", "i32");
381 f.addLocal("pn", "i32");
382
383 const c = f.getCodeBuilder();
384
385 f.addCode(
386 c.setLocal("pi", c.getLocal("ps")),
387 c.setLocal("po", c.getLocal("pd")),
388 c.setLocal(
389 "pn",
390 c.i32_add(
391 c.getLocal("ps"),
392 c.i32_mul(
393 c.getLocal("n"),
394 c.i32_const(n8)
395 )
396 )
397 ),
398 c.block(c.loop(
399 c.br_if(
400 1,
401 c.i32_eq(
402 c.getLocal("pi"),
403 c.getLocal("pn")
404 )
405 ),
406
407 c.call(f1mPrefix + "_copy", c.getLocal("pi"), c.getLocal("po")),
408
409 c.setLocal("pi", c.i32_add(c.getLocal("pi"), c.i32_const(n8))),
410 c.setLocal("po", c.i32_add(c.getLocal("po"), c.i32_const(n8*2))),
411 c.br(0)
412 ))
413 );
414 }
415
416
417
418 function buildToMontgomery() {
419 const f = module.addFunction(prefix+"_toMontgomeryN");
420 f.addParam("ps", "i32");
421 f.addParam("pd", "i32");
422 f.addParam("n", "i32");
423 f.addLocal("pi", "i32");
424 f.addLocal("po", "i32");
425 f.addLocal("pn", "i32");
426
427 const c = f.getCodeBuilder();
428
429 f.addCode(
430 c.setLocal("pi", c.getLocal("ps")),
431 c.setLocal("po", c.getLocal("pd")),
432 c.setLocal(
433 "pn",
434 c.i32_add(
435 c.getLocal("ps"),
436 c.i32_mul(
437 c.getLocal("n"),
438 c.i32_const(n8)
439 )
440 )
441 ),
442 c.block(c.loop(
443 c.br_if(
444 1,
445 c.i32_eq(
446 c.getLocal("pi"),
447 c.getLocal("pn")
448 )
449 ),
450
451 c.call(f1mPrefix + "_toMontgomery", c.getLocal("pi"), c.getLocal("po")),
452
453 c.setLocal("pi", c.i32_add(c.getLocal("pi"), c.i32_const(n8))),
454 c.setLocal("po", c.i32_add(c.getLocal("po"), c.i32_const(n8))),
455 c.br(0)
456 ))
457 );
458 }
459
460
461 function buildMulN() {
462 const f = module.addFunction(prefix+"_mulN");
463 f.addParam("px", "i32");
464 f.addParam("py", "i32");
465 f.addParam("n", "i32");
466 f.addParam("pd", "i32");
467 f.addLocal("pix", "i32");
468 f.addLocal("piy", "i32");
469 f.addLocal("po", "i32");
470 f.addLocal("lastpix", "i32");
471
472 const c = f.getCodeBuilder();
473
474 f.addCode(
475 c.setLocal("pix", c.getLocal("px")),
476 c.setLocal("piy", c.getLocal("py")),
477 c.setLocal("po", c.getLocal("pd")),
478 c.setLocal(
479 "lastpix",
480 c.i32_add(
481 c.getLocal("px"),
482 c.i32_mul(
483 c.getLocal("n"),
484 c.i32_const(n8)
485 )
486 )
487 ),
488 c.block(c.loop(
489 c.br_if(
490 1,
491 c.i32_eq(
492 c.getLocal("pix"),
493 c.getLocal("lastpix")
494 )
495 ),
496
497 c.call(f1mPrefix + "_mul", c.getLocal("pix"), c.getLocal("piy"), c.getLocal("po")),
498
499 c.setLocal("pix", c.i32_add(c.getLocal("pix"), c.i32_const(n8))),
500 c.setLocal("piy", c.i32_add(c.getLocal("piy"), c.i32_const(n8))),
501 c.setLocal("po", c.i32_add(c.getLocal("po"), c.i32_const(n8))),
502 c.br(0)
503 ))
504 );
505 }
506
507 function buildFromMontgomery() {
508 const f = module.addFunction(prefix+"_fromMontgomeryN");
509 f.addParam("ps", "i32");
510 f.addParam("pd", "i32");
511 f.addParam("n", "i32");
512 f.addLocal("pi", "i32");
513 f.addLocal("po", "i32");
514 f.addLocal("pn", "i32");
515
516 const c = f.getCodeBuilder();
517
518 f.addCode(
519 c.setLocal("pi", c.getLocal("ps")),
520 c.setLocal("po", c.getLocal("pd")),
521 c.setLocal(
522 "pn",
523 c.i32_add(
524 c.getLocal("ps"),
525 c.i32_mul(
526 c.getLocal("n"),
527 c.i32_const(n8)
528 )
529 )
530 ),
531 c.block(c.loop(
532 c.br_if(
533 1,
534 c.i32_eq(
535 c.getLocal("pi"),
536 c.getLocal("pn")
537 )
538 ),
539
540 c.call(f1mPrefix + "_fromMontgomery", c.getLocal("pi"), c.getLocal("po")),
541
542 c.setLocal("pi", c.i32_add(c.getLocal("pi"), c.i32_const(n8))),
543 c.setLocal("po", c.i32_add(c.getLocal("po"), c.i32_const(n8))),
544 c.br(0)
545 ))
546 );
547 }
548
549
550 function buildFinalInverse() {
551 const f = module.addFunction(prefix+"__finalInverse");
552 f.addParam("px", "i32");
553 f.addParam("bits", "i32");
554 f.addLocal("n", "i32");
555 f.addLocal("ndiv2", "i32");
556 f.addLocal("pInv2", "i32");
557 f.addLocal("i", "i32");
558 f.addLocal("mask", "i32");
559 f.addLocal("idx1", "i32");
560 f.addLocal("idx2", "i32");
561
562 const c = f.getCodeBuilder();
563
564 const T = c.i32_const(module.alloc(n8));
565
566 f.addCode(
567 c.setLocal("n", c.i32_shl( c.i32_const(1), c.getLocal("bits"))),
568
569 c.setLocal(
570 "pInv2",
571 c.i32_add(
572 c.i32_const(INV2),
573 c.i32_mul(
574 c.getLocal("bits"),
575 c.i32_const(n8)
576 )
577 )
578 ),
579
580 c.setLocal("mask", c.i32_sub( c.getLocal("n") , c.i32_const(1))),
581 c.setLocal("i", c.i32_const(1)),
582 c.setLocal(
583 "ndiv2",
584 c.i32_shr_u(
585 c.getLocal("n"),
586 c.i32_const(1)
587 )
588 ),
589 c.block(c.loop(
590 c.br_if(
591 1,
592 c.i32_eq(
593 c.getLocal("i"),
594 c.getLocal("ndiv2")
595 )
596 ),
597
598 c.setLocal("idx1",
599 c.i32_add(
600 c.getLocal("px"),
601 c.i32_mul(
602 c.getLocal("i"),
603 c.i32_const(n8)
604 )
605 )
606 ),
607
608 c.setLocal("idx2",
609 c.i32_add(
610 c.getLocal("px"),
611 c.i32_mul(
612 c.i32_sub(
613 c.getLocal("n"),
614 c.getLocal("i")
615 ),
616 c.i32_const(n8)
617 )
618 )
619 ),
620
621 c.call(f1mPrefix + "_copy", c.getLocal("idx1"), T),
622 c.call(f1mPrefix + "_mul", c.getLocal("idx2") , c.getLocal("pInv2"), c.getLocal("idx1") ),
623 c.call(f1mPrefix + "_mul", T , c.getLocal("pInv2"), c.getLocal("idx2")),
624
625// c.call(f1mPrefix + "_mul", c.getLocal("idx1") , c.getLocal("pInv2"), c.getLocal("idx1") ),
626// c.call(f1mPrefix + "_mul", c.getLocal("idx2") , c.getLocal("pInv2"), c.getLocal("idx1") ),
627
628 c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
629
630 c.br(0)
631 )),
632
633 c.call(f1mPrefix + "_mul", c.getLocal("px") , c.getLocal("pInv2"), c.getLocal("px")),
634
635 c.setLocal("idx2",
636 c.i32_add(
637 c.getLocal("px"),
638 c.i32_mul(
639 c.getLocal("ndiv2"),
640 c.i32_const(n8)
641 )
642 )
643 ),
644
645 c.call(f1mPrefix + "_mul", c.getLocal("idx2"),c.getLocal("pInv2"), c.getLocal("idx2"))
646
647 );
648 }
649
650 function buildReversePermutation() {
651 const f = module.addFunction(prefix+"__reversePermutation");
652 f.addParam("px", "i32");
653 f.addParam("bits", "i32");
654 f.addLocal("n", "i32");
655 f.addLocal("i", "i32");
656 f.addLocal("ri", "i32");
657 f.addLocal("idx1", "i32");
658 f.addLocal("idx2", "i32");
659
660 const c = f.getCodeBuilder();
661
662 const T = c.i32_const(module.alloc(n8));
663
664 f.addCode(
665 c.setLocal("n", c.i32_shl( c.i32_const(1), c.getLocal("bits"))),
666 c.setLocal("i", c.i32_const(0)),
667 c.block(c.loop(
668 c.br_if(
669 1,
670 c.i32_eq(
671 c.getLocal("i"),
672 c.getLocal("n")
673 )
674 ),
675
676 c.setLocal("idx1",
677 c.i32_add(
678 c.getLocal("px"),
679 c.i32_mul(
680 c.getLocal("i"),
681 c.i32_const(n8)
682 )
683 )
684 ),
685
686 c.setLocal("ri", c.call(prefix + "__rev", c.getLocal("i"), c.getLocal("bits"))),
687
688 c.setLocal("idx2",
689 c.i32_add(
690 c.getLocal("px"),
691 c.i32_mul(
692 c.getLocal("ri"),
693 c.i32_const(n8)
694 )
695 )
696 ),
697
698 c.if(
699 c.i32_lt_u(
700 c.getLocal("i"),
701 c.getLocal("ri")
702 ),
703 [
704 ...c.call(f1mPrefix + "_copy", c.getLocal("idx1"), T),
705 ...c.call(f1mPrefix + "_copy", c.getLocal("idx2") , c.getLocal("idx1")),
706 ...c.call(f1mPrefix + "_copy", T , c.getLocal("idx2"))
707 ]
708 ),
709
710 c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
711
712 c.br(0)
713 ))
714 );
715 }
716
717 function buildRev() {
718 const f = module.addFunction(prefix+"__rev");
719 f.addParam("x", "i32");
720 f.addParam("bits", "i32");
721 f.setReturnType("i32");
722
723 const c = f.getCodeBuilder();
724
725 f.addCode(
726 c.i32_rotl(
727 c.i32_add(
728 c.i32_add(
729 c.i32_shl(
730 c.i32_load8_u(
731 c.i32_and(
732 c.getLocal("x"),
733 c.i32_const(0xFF)
734 ),
735 REVTABLE,
736 0
737 ),
738 c.i32_const(24)
739 ),
740 c.i32_shl(
741 c.i32_load8_u(
742 c.i32_and(
743 c.i32_shr_u(
744 c.getLocal("x"),
745 c.i32_const(8)
746 ),
747 c.i32_const(0xFF)
748 ),
749 REVTABLE,
750 0
751 ),
752 c.i32_const(16)
753 ),
754 ),
755 c.i32_add(
756 c.i32_shl(
757 c.i32_load8_u(
758 c.i32_and(
759 c.i32_shr_u(
760 c.getLocal("x"),
761 c.i32_const(16)
762 ),
763 c.i32_const(0xFF)
764 ),
765 REVTABLE,
766 0
767 ),
768 c.i32_const(8)
769 ),
770 c.i32_load8_u(
771 c.i32_and(
772 c.i32_shr_u(
773 c.getLocal("x"),
774 c.i32_const(24)
775 ),
776 c.i32_const(0xFF)
777 ),
778 REVTABLE,
779 0
780 ),
781 )
782 ),
783 c.getLocal("bits")
784 )
785 );
786 }
787
788
789 buildRev();
790 buildReversePermutation();
791 buildRawFFT();
792 buildCopyInterleaved();
793 buildFromMontgomery();
794 buildToMontgomery();
795 buildFinalInverse();
796 buildLog2();
797 buildFFT();
798 buildIFFT();
799 buildMulN();
800
801 module.exportFunction(prefix+"_fft");
802 module.exportFunction(prefix+"_ifft");
803 module.exportFunction(prefix+"_toMontgomeryN");
804 module.exportFunction(prefix+"_fromMontgomeryN");
805 module.exportFunction(prefix+"_copyNInterleaved");
806 module.exportFunction(prefix+"_mulN");
807
808};