1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
|
// NTT (Number Theoretic Transform) implementation for BigDecimal multiplication
#define NTT_PRIMITIVE_ROOT 17
#define NTT_PRIME_BASE1 24
#define NTT_PRIME_BASE2 26
#define NTT_PRIME_BASE3 29
#define NTT_PRIME_SHIFT 27
#define NTT_PRIME1 (((uint32_t)NTT_PRIME_BASE1 << NTT_PRIME_SHIFT) | 1)
#define NTT_PRIME2 (((uint32_t)NTT_PRIME_BASE2 << NTT_PRIME_SHIFT) | 1)
#define NTT_PRIME3 (((uint32_t)NTT_PRIME_BASE3 << NTT_PRIME_SHIFT) | 1)
#define MAX_NTT32_BITS 27
#define NTT_DECDIG_BASE 1000000000
// Calculates base**ex % mod
static uint32_t
mod_pow(uint32_t base, uint32_t ex, uint32_t mod) {
uint32_t res = 1;
uint32_t bit = 1;
while (true) {
if (ex & bit) {
ex ^= bit;
res = ((uint64_t)res * base) % mod;
}
if (!ex) break;
base = ((uint64_t)base * base) % mod;
bit <<= 1;
}
return res;
}
// Recursively performs butterfly operations of NTT
static void
ntt_recursive(int size_bits, uint32_t *input, uint32_t *output, uint32_t *tmp, int depth, uint32_t r, uint32_t prime) {
if (depth > 0) {
ntt_recursive(size_bits, input, tmp, output, depth - 1, ((uint64_t)r * r) % prime, prime);
} else {
tmp = input;
}
uint32_t size_half = (uint32_t)1 << (size_bits - 1);
uint32_t stride = (uint32_t)1 << (size_bits - depth - 1);
uint32_t n = size_half / stride;
uint32_t rn = 1, rm = prime - 1;
for (uint32_t i = 0; i < n; i++) {
uint32_t *aptr = tmp + i * 2 * stride;
uint32_t *bptr = aptr + stride;
uint32_t *out1 = output + stride * i;
uint32_t *out2 = out1 + size_half;
for (uint32_t k = 0; k < stride; k++) {
uint32_t a = aptr[k], b = bptr[k];
out1[k] = (a + (uint64_t)rn * b) % prime;
out2[k] = (a + (uint64_t)rm * b) % prime;
}
rn = ((uint64_t)rn * r) % prime;
rm = ((uint64_t)rm * r) % prime;
}
}
/* Perform NTT on input array.
* base, shift: Represent the prime number as (base << shift | 1)
* r_base: Primitive root of unity modulo prime
* size_bits: log2 of the size of the input array. Should be less or equal to shift
* input: input array of size (1 << size_bits)
*/
static void
ntt(int size_bits, uint32_t *input, uint32_t *output, uint32_t *tmp, int r_base, int base, int shift, int dir) {
uint32_t size = (uint32_t)1 << size_bits;
uint32_t prime = ((uint32_t)base << shift) | 1;
// rmax**(1 << shift) % prime == 1
// r**size % prime == 1
uint32_t rmax = mod_pow((uint32_t)r_base, (uint32_t)base, prime);
uint32_t r = mod_pow(rmax, (uint32_t)1 << (shift - size_bits), prime);
if (dir < 0) r = mod_pow(r, prime - 2, prime);
ntt_recursive(size_bits, input, output, tmp, size_bits - 1, r, prime);
if (dir < 0) {
uint32_t n_inv = mod_pow((uint32_t)size, prime - 2, prime);
for (uint32_t i = 0; i < size; i++) {
output[i] = ((uint64_t)output[i] * n_inv) % prime;
}
}
}
/* Calculate c that satisfies: c % PRIME1 == mod1 && c % PRIME2 == mod2 && c % PRIME3 == mod3
* c = (mod1 * 35002755423056150739595925972 + mod2 * 14584479687667766215746868453 + mod3 * 37919651490985126265126719818) % (PRIME1 * PRIME2 * PRIME3)
* Assume c <= 999999999**2*(1<<27)
*/
static inline void
mod_restore_prime_24_26_29_shift_27(uint32_t mod1, uint32_t mod2, uint32_t mod3, uint32_t *digits) {
// Use mixed radix notation to eliminate modulo by PRIME1 * PRIME2 * PRIME3
// [DIG0, DIG1, DIG2] = DIG0 + DIG1 * PRIME1 + DIG2 * PRIME1 * PRIME2
// DIG0: 0...PRIME1, DIG1: 0...PRIME2, DIG2: 0...PRIME3
// 35002755423056150739595925972 = [1, 3489660916, 3113851359]
// 14584479687667766215746868453 = [0, 13, 1297437912]
// 37919651490985126265126719818 = [0, 0, 3373338954]
uint64_t c0 = mod1;
uint64_t c1 = (uint64_t)mod2 * 13 + (uint64_t)mod1 * 3489660916;
uint64_t c2 = (uint64_t)mod3 * 3373338954 % NTT_PRIME3 + (uint64_t)mod2 * 1297437912 % NTT_PRIME3 + (uint64_t)mod1 * 3113851359 % NTT_PRIME3;
c2 += c1 / NTT_PRIME2;
c1 %= NTT_PRIME2;
c2 %= NTT_PRIME3;
// Base conversion. c fits in 3 digits.
c1 += c2 % NTT_DECDIG_BASE * NTT_PRIME2;
c0 += c1 % NTT_DECDIG_BASE * NTT_PRIME1;
c1 /= NTT_DECDIG_BASE;
digits[0] = c0 % NTT_DECDIG_BASE;
c0 /= NTT_DECDIG_BASE;
c1 += c2 / NTT_DECDIG_BASE % NTT_DECDIG_BASE * NTT_PRIME2;
c0 += c1 % NTT_DECDIG_BASE * NTT_PRIME1;
c1 /= NTT_DECDIG_BASE;
digits[1] = c0 % NTT_DECDIG_BASE;
digits[2] = (uint32_t)(c0 / NTT_DECDIG_BASE + c1 % NTT_DECDIG_BASE * NTT_PRIME1);
}
/*
* NTT multiplication
* Uses three NTTs with mod (24 << 27 | 1), (26 << 27 | 1), and (29 << 27 | 1)
*/
static void
ntt_multiply(size_t a_size, size_t b_size, uint32_t *a, uint32_t *b, uint32_t *c) {
if (a_size < b_size) {
ntt_multiply(b_size, a_size, b, a, c);
return;
}
int ntt_size_bits = (int)bit_length(b_size - 1) + 1;
if (ntt_size_bits > MAX_NTT32_BITS) {
rb_raise(rb_eArgError, "Multiply size too large");
}
// To calculate large_a * small_b faster, split into several batches.
uint32_t ntt_size = (uint32_t)1 << ntt_size_bits;
uint32_t batch_size = ntt_size - (uint32_t)b_size;
uint32_t batch_count = (uint32_t)((a_size + batch_size - 1) / batch_size);
uint32_t *mem = ruby_xcalloc(ntt_size * 9, sizeof(uint32_t));
uint32_t *ntt1 = mem;
uint32_t *ntt2 = mem + ntt_size;
uint32_t *ntt3 = mem + ntt_size * 2;
uint32_t *tmp1 = mem + ntt_size * 3;
uint32_t *tmp2 = mem + ntt_size * 4;
uint32_t *tmp3 = mem + ntt_size * 5;
uint32_t *conv1 = mem + ntt_size * 6;
uint32_t *conv2 = mem + ntt_size * 7;
uint32_t *conv3 = mem + ntt_size * 8;
// Calculate NTT for b in three primes. Result is reused for each batch of a.
memcpy(tmp1, b, b_size * sizeof(uint32_t));
memset(tmp1 + b_size, 0, (ntt_size - b_size) * sizeof(uint32_t));
ntt(ntt_size_bits, tmp1, ntt1, tmp2, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE1, NTT_PRIME_SHIFT, +1);
ntt(ntt_size_bits, tmp1, ntt2, tmp2, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE2, NTT_PRIME_SHIFT, +1);
ntt(ntt_size_bits, tmp1, ntt3, tmp2, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE3, NTT_PRIME_SHIFT, +1);
memset(c, 0, (a_size + b_size) * sizeof(uint32_t));
for (uint32_t idx = 0; idx < batch_count; idx++) {
uint32_t len = idx == batch_count - 1 ? (uint32_t)a_size - idx * batch_size : batch_size;
memcpy(tmp1, a + idx * batch_size, len * sizeof(uint32_t));
memset(tmp1 + len, 0, (ntt_size - len) * sizeof(uint32_t));
// Calculate convolution for this batch in three primes
ntt(ntt_size_bits, tmp1, tmp2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE1, NTT_PRIME_SHIFT, +1);
for (uint32_t i = 0; i < ntt_size; i++) tmp2[i] = ((uint64_t)tmp2[i] * ntt1[i]) % NTT_PRIME1;
ntt(ntt_size_bits, tmp2, conv1, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE1, NTT_PRIME_SHIFT, -1);
ntt(ntt_size_bits, tmp1, tmp2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE2, NTT_PRIME_SHIFT, +1);
for (uint32_t i = 0; i < ntt_size; i++) tmp2[i] = ((uint64_t)tmp2[i] * ntt2[i]) % NTT_PRIME2;
ntt(ntt_size_bits, tmp2, conv2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE2, NTT_PRIME_SHIFT, -1);
ntt(ntt_size_bits, tmp1, tmp2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE3, NTT_PRIME_SHIFT, +1);
for (uint32_t i = 0; i < ntt_size; i++) tmp2[i] = ((uint64_t)tmp2[i] * ntt3[i]) % NTT_PRIME3;
ntt(ntt_size_bits, tmp2, conv3, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE3, NTT_PRIME_SHIFT, -1);
// Restore the original convolution value from three convolutions calculated in three primes.
// Each convolution value is maximum 999999999**2*(1<<27)/2
for (uint32_t i = 0; i < ntt_size; i++) {
uint32_t dig[3];
mod_restore_prime_24_26_29_shift_27(conv1[i], conv2[i], conv3[i], dig);
// Maximum values of dig[0], dig[1], and dig[2] are 999999999, 999999999 and 67108863 respectively
// Maximum overlapped sum (considering overlaps between 2 batches) is less than 4134217722
// so this sum doesn't overflow uint32_t.
for (int j = 0; j < 3; j++) {
// Index check: if dig[j] is non-zero, assign index is within valid range.
if (dig[j]) c[idx * batch_size + i + 1 - (uint32_t)j] += dig[j];
}
}
}
uint32_t carry = 0;
for (int32_t i = (int32_t)(a_size + b_size - 1); i >= 0; i--) {
uint32_t v = c[i] + carry;
c[i] = v % NTT_DECDIG_BASE;
carry = v / NTT_DECDIG_BASE;
}
ruby_xfree(mem);
}
|
