diff --git a/benchmarks/workloads.py b/benchmarks/workloads.py index 37488eb..eebca20 100644 --- a/benchmarks/workloads.py +++ b/benchmarks/workloads.py @@ -243,6 +243,21 @@ def run_convolution(a, b): return fft.convolution(a, b) +# A long array against a short one. This is the other cost regime of +# convolution (and a common one in practice), so it is benchmarked separately +# from the balanced case above. +def build_convolution_lopsided_workload(): + r = _rng(14) + a = [r.randrange(MOD) for _ in range(1 << 16)] + b = [r.randrange(MOD) for _ in range(256)] + return a, b + + +def run_convolution_lopsided(a, b): + fft = convolution_mod.FFT(MOD) + return fft.convolution(a, b) + + # --------------------------------------------------------------------------- # scc: Tarjan's algorithm on a large random directed graph # --------------------------------------------------------------------------- @@ -413,6 +428,11 @@ def run_two_sat(n, clause): ("maxflow", build_maxflow_workload, run_maxflow), ("mincostflow", build_mincostflow_workload, run_mincostflow), ("convolution", build_convolution_workload, run_convolution), + ( + "convolution_lopsided", + build_convolution_lopsided_workload, + run_convolution_lopsided, + ), ("scc", build_scc_workload, run_scc), ("fps", build_fps_workload, run_fps), ("fenwicktree", build_fenwicktree_workload, run_fenwicktree), diff --git a/convolution.py b/convolution.py index 8ef564a..47347f5 100644 --- a/convolution.py +++ b/convolution.py @@ -1,3 +1,6 @@ +import struct + + class FFT: def primitive_root_constexpr(self, m): if m == 2: @@ -86,6 +89,12 @@ def __init__(self, MOD): def butterfly(self, a): n = len(a) h = (n - 1).bit_length() + # Hoisting the instance attributes into locals removes an attribute + # lookup per use from the innermost loops, which dominate the runtime. + mod = self.mod + rate2 = self.rate2 + rate3 = self.rate3 + imag = self.root[2] LEN = 0 while LEN < h: @@ -94,39 +103,78 @@ def butterfly(self, a): rot = 1 for s in range(1 << LEN): offset = s << (h - LEN) - for i in range(p): - l = a[i + offset] - r = a[i + offset + p] * rot - a[i + offset] = (l + r) % self.mod - a[i + offset + p] = (l - r) % self.mod - rot *= self.rate2[(~s & -~s).bit_length() - 1] - rot %= self.mod + end = offset + p + if rot == 1: + for i in range(offset, end): + j = i + p + l = a[i] + r = a[j] + a[i] = (l + r) % mod + a[j] = (l - r) % mod + else: + for i in range(offset, end): + j = i + p + l = a[i] + r = a[j] * rot + a[i] = (l + r) % mod + a[j] = (l - r) % mod + rot = rot * rate2[(~s & -~s).bit_length() - 1] % mod LEN += 1 else: p = 1 << (h - LEN - 2) + p2 = p + p + p3 = p2 + p rot = 1 - imag = self.root[2] for s in range(1 << LEN): - rot2 = (rot * rot) % self.mod - rot3 = (rot2 * rot) % self.mod offset = s << (h - LEN) - for i in range(p): - a0 = a[i + offset] - a1 = a[i + offset + p] * rot - a2 = a[i + offset + 2 * p] * rot2 - a3 = a[i + offset + 3 * p] * rot3 - a1na3imag = (a1 - a3) % self.mod * imag - a[i + offset] = (a0 + a2 + a1 + a3) % self.mod - a[i + offset + p] = (a0 + a2 - a1 - a3) % self.mod - a[i + offset + 2 * p] = (a0 - a2 + a1na3imag) % self.mod - a[i + offset + 3 * p] = (a0 - a2 - a1na3imag) % self.mod - rot *= self.rate3[(~s & -~s).bit_length() - 1] - rot %= self.mod + end = offset + p + if rot == 1: + for i in range(offset, end): + i1 = i + p + i2 = i + p2 + i3 = i + p3 + a0 = a[i] + a1 = a[i1] + a2 = a[i2] + a3 = a[i3] + a02 = a0 + a2 + a13 = a1 + a3 + a0n2 = a0 - a2 + a1na3imag = (a1 - a3) % mod * imag + a[i] = (a02 + a13) % mod + a[i1] = (a02 - a13) % mod + a[i2] = (a0n2 + a1na3imag) % mod + a[i3] = (a0n2 - a1na3imag) % mod + else: + rot2 = rot * rot % mod + rot3 = rot2 * rot % mod + for i in range(offset, end): + i1 = i + p + i2 = i + p2 + i3 = i + p3 + a0 = a[i] + a1 = a[i1] * rot + a2 = a[i2] * rot2 + a3 = a[i3] * rot3 + a02 = a0 + a2 + a13 = a1 + a3 + a0n2 = a0 - a2 + a1na3imag = (a1 - a3) % mod * imag + a[i] = (a02 + a13) % mod + a[i1] = (a02 - a13) % mod + a[i2] = (a0n2 + a1na3imag) % mod + a[i3] = (a0n2 - a1na3imag) % mod + rot = rot * rate3[(~s & -~s).bit_length() - 1] % mod LEN += 2 def butterfly_inv(self, a): n = len(a) h = (n - 1).bit_length() + mod = self.mod + irate2 = self.irate2 + irate3 = self.irate3 + iimag = self.iroot[2] + LEN = h while LEN: if LEN == 1: @@ -134,58 +182,186 @@ def butterfly_inv(self, a): irot = 1 for s in range(1 << (LEN - 1)): offset = s << (h - LEN + 1) - for i in range(p): - l = a[i + offset] - r = a[i + offset + p] - a[i + offset] = (l + r) % self.mod - a[i + offset + p] = (l - r) * irot % self.mod - irot *= self.irate2[(~s & -~s).bit_length() - 1] - irot %= self.mod + end = offset + p + if irot == 1: + for i in range(offset, end): + j = i + p + l = a[i] + r = a[j] + a[i] = (l + r) % mod + a[j] = (l - r) % mod + else: + for i in range(offset, end): + j = i + p + l = a[i] + r = a[j] + a[i] = (l + r) % mod + a[j] = (l - r) * irot % mod + irot = irot * irate2[(~s & -~s).bit_length() - 1] % mod LEN -= 1 else: p = 1 << (h - LEN) + p2 = p + p + p3 = p2 + p irot = 1 - iimag = self.iroot[2] for s in range(1 << (LEN - 2)): - irot2 = (irot * irot) % self.mod - irot3 = (irot * irot2) % self.mod offset = s << (h - LEN + 2) - for i in range(p): - a0 = a[i + offset] - a1 = a[i + offset + p] - a2 = a[i + offset + 2 * p] - a3 = a[i + offset + 3 * p] - a2na3iimag = (a2 - a3) * iimag % self.mod - a[i + offset] = (a0 + a1 + a2 + a3) % self.mod - a[i + offset + p] = (a0 - a1 + a2na3iimag) * irot % self.mod - a[i + offset + 2 * p] = (a0 + a1 - a2 - a3) * irot2 % self.mod - a[i + offset + 3 * p] = ( - (a0 - a1 - a2na3iimag) * irot3 % self.mod - ) - irot *= self.irate3[(~s & -~s).bit_length() - 1] - irot %= self.mod + end = offset + p + if irot == 1: + for i in range(offset, end): + i1 = i + p + i2 = i + p2 + i3 = i + p3 + a0 = a[i] + a1 = a[i1] + a2 = a[i2] + a3 = a[i3] + a01 = a0 + a1 + a23 = a2 + a3 + a0n1 = a0 - a1 + a2na3iimag = (a2 - a3) * iimag % mod + a[i] = (a01 + a23) % mod + a[i1] = (a0n1 + a2na3iimag) % mod + a[i2] = (a01 - a23) % mod + a[i3] = (a0n1 - a2na3iimag) % mod + else: + irot2 = irot * irot % mod + irot3 = irot * irot2 % mod + for i in range(offset, end): + i1 = i + p + i2 = i + p2 + i3 = i + p3 + a0 = a[i] + a1 = a[i1] + a2 = a[i2] + a3 = a[i3] + a01 = a0 + a1 + a23 = a2 + a3 + a0n1 = a0 - a1 + a2na3iimag = (a2 - a3) * iimag % mod + a[i] = (a01 + a23) % mod + a[i1] = (a0n1 + a2na3iimag) * irot % mod + a[i2] = (a01 - a23) * irot2 % mod + a[i3] = (a0n1 - a2na3iimag) * irot3 % mod + irot = irot * irate3[(~s & -~s).bit_length() - 1] % mod LEN -= 2 + # struct codes for the fixed-width integers usable as packing slots + _code = {1: "B", 2: "H", 4: "I", 8: "Q"} + # slot widths (in bytes) that struct can pack/unpack with at most two fields + _widths = (1, 2, 4, 8, 9, 10, 12, 16) + _packers = {} + _CHUNK = 256 + + def _pack(self, a, nb, w): + """Encode a as sum(a[i] << (8 * nb * i)), one nb-byte slot per element. + + w is the struct field width in bytes (None to fall back to to_bytes). + """ + if w is None: + return int.from_bytes( + b"".join(x.to_bytes(nb, "little") for x in a), "little" + ) + entry = self._packers.get((nb, w)) + if entry is None: + unit = self._code[w] + ("%dx" % (nb - w) if nb > w else "") + entry = (struct.Struct("<" + unit * self._CHUNK).pack, unit) + self._packers[(nb, w)] = entry + pack_chunk, unit = entry + n = len(a) + chunk = self._CHUNK + if n <= chunk: + return int.from_bytes(struct.pack("<" + unit * n, *a), "little") + tail = n % chunk + end = n - tail + parts = [pack_chunk(*a[i : i + chunk]) for i in range(0, end, chunk)] + if tail: + parts.append(struct.pack("<" + unit * tail, *a[end:])) + return int.from_bytes(b"".join(parts), "little") + + def _unpack(self, buf, nb, length, mod): + """Inverse of _pack: read length slots of nb bytes, reduced mod mod.""" + code = self._code.get(nb) if nb <= 8 else self._code.get(nb - 8) + if code is not None: + if nb <= 8: + res = [x % mod for (x,) in struct.iter_unpack("<" + code, buf)] + else: + res = [ + (lo | hi << 64) % mod + for lo, hi in struct.iter_unpack("> 3, (bmax.bit_length() + 7) >> 3, 1) + nb = max(((min(n, m) * amax * bmax).bit_length() + 7) >> 3, vb) + w = None + if nb <= 16 and vb <= 8: + for width in self._widths: + if width >= nb: + nb = width + break + w = 1 + while w < vb: + w <<= 1 + prod = self._pack(a, nb, w) * self._pack(b, nb, w) + return self._unpack(prod.to_bytes(nb * (n + m), "little"), nb, n + m - 1, mod) + def convolution(self, a, b): n = len(a) m = len(b) if not (a) or not (b): return [] - if min(n, m) <= 40: + mod = self.mod + if n * m <= 40: res = [0] * (n + m - 1) - for i in range(n): - for j in range(m): - res[i + j] += a[i] * b[j] - res[i + j] %= self.mod + for i, ai in enumerate(a): + if ai: + ai %= mod + for j, bj in enumerate(b, i): + res[j] = (res[j] + ai * bj) % mod return res + # Both remaining paths want representatives in [0, mod). + amax = max(a) + if min(a) < 0 or amax >= mod: + a = [x % mod for x in a] + amax = max(a) + bmax = max(b) + if min(b) < 0 or bmax >= mod: + b = [x % mod for x in b] + bmax = max(b) z = 1 << ((n + m - 2).bit_length()) + # Cost model, fitted on CPython 3.11. Kronecker substitution costs + # about 8e-10 * (hi / lo) * (lo * nb) ** 1.585 seconds (Karatsuba on + # lo * nb bytes, repeated hi / lo times for a lopsided product), while + # the NTT costs about 2.7e-7 * z * log2(z). Their ratio gives the + # constant below. The NTT only wins once both operands are large: + # small, lopsided or small-coefficient inputs stay on the integer path. + lo, hi = (n, m) if n < m else (m, n) + nb = max(((lo * amax * bmax).bit_length() + 7) >> 3, 1) + if hi * (lo * nb) ** 1.585 < 337.0 * lo * z * (z.bit_length() - 1): + return self._convolution_int(a, b, amax, bmax) + # butterfly_inv is unnormalized, so the result needs scaling by 1/z. + # Folding that into b before its transform costs m multiplications + # instead of a separate pass (and an extra list) over the z outputs. + iz = pow(z, mod - 2, mod) a = a + [0] * (z - n) - b = b + [0] * (z - m) + b = [x * iz % mod for x in b] + [0] * (z - m) self.butterfly(a) self.butterfly(b) - c = [(a[i] * b[i]) % self.mod for i in range(z)] + c = [x * y % mod for x, y in zip(a, b)] self.butterfly_inv(c) - iz = pow(z, self.mod - 2, self.mod) - for i in range(n + m - 1): - c[i] = (c[i] * iz) % self.mod - return c[: n + m - 1] + del c[n + m - 1 :] + return c diff --git a/tests/test_convolution.py b/tests/test_convolution.py index 796673a..91d710f 100644 --- a/tests/test_convolution.py +++ b/tests/test_convolution.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 +import random import sys import os import unittest @@ -8,6 +9,19 @@ import convolution +MODS = [998244353, 167772161, 469762049, 754974721] + + +def naive_convolution(a, b, mod): + """O(nm) reference implementation.""" + if not a or not b: + return [] + res = [0] * (len(a) + len(b) - 1) + for i, x in enumerate(a): + for j, y in enumerate(b): + res[i + j] = (res[i + j] + x * y) % mod + return res + class TestConvolution(unittest.TestCase): """Test cases for convolution module""" @@ -22,6 +36,81 @@ def test_practice2_f(self): ) self.practice2_f(1, 1, [10000000], [10000000], [871938225]) + def test_empty(self): + conv = convolution.FFT(998244353) + self.assertEqual(conv.convolution([], [1, 2, 3]), []) + self.assertEqual(conv.convolution([1, 2, 3], []), []) + self.assertEqual(conv.convolution([], []), []) + + def test_against_naive_random(self): + """Random small cases against the O(nm) reference, over several mods. + + Sizes and coefficient ranges are varied deliberately: convolution + dispatches between a schoolbook, an integer-packing and an NTT path + depending on both, so a single size would only cover one of them. + """ + rng = random.Random(998244353) + for mod in MODS: + conv = convolution.FFT(mod) + for _ in range(60): + n = rng.randint(1, 200) + m = rng.randint(1, 200) + vmax = rng.choice([2, 256, mod]) + a = [rng.randrange(vmax) for _ in range(n)] + b = [rng.randrange(vmax) for _ in range(m)] + self.assertEqual( + conv.convolution(a, b), naive_convolution(a, b, mod), (mod, n, m) + ) + + def test_unnormalized_input(self): + """Negative coefficients and coefficients >= mod are reduced.""" + mod = 998244353 + conv = convolution.FFT(mod) + rng = random.Random(7) + for n, m in [(3, 4), (60, 60), (200, 5)]: + a = [rng.randrange(-3 * mod, 3 * mod) for _ in range(n)] + b = [rng.randrange(-3 * mod, 3 * mod) for _ in range(m)] + self.assertEqual(conv.convolution(a, b), naive_convolution(a, b, mod)) + + def test_zero_coefficients(self): + mod = 998244353 + conv = convolution.FFT(mod) + n = 100 + self.assertEqual(conv.convolution([0] * n, [0] * n), [0] * (2 * n - 1)) + self.assertEqual(conv.convolution([0] * n, [1] * n), [0] * (2 * n - 1)) + + def test_large_and_lopsided(self): + """Cases big enough to reach the NTT and integer-packing paths.""" + mod = 998244353 + conv = convolution.FFT(mod) + rng = random.Random(11) + for n, m in [(1, 3000), (3000, 1), (2000, 17), (1000, 1000), (4096, 4096)]: + a = [rng.randrange(mod) for _ in range(n)] + b = [rng.randrange(mod) for _ in range(m)] + c = conv.convolution(a, b) + self.assertEqual(len(c), n + m - 1) + # Spot-check individual coefficients rather than the whole + # O(nm) reference, which would be too slow at these sizes. + for k in (0, (n + m - 1) // 2, n + m - 2): + expect = sum( + a[i] * b[k - i] for i in range(max(0, k - m + 1), min(n - 1, k) + 1) + ) + self.assertEqual(c[k], expect % mod, (n, m, k)) + + def test_butterfly_round_trip(self): + """butterfly followed by butterfly_inv is the identity times n.""" + mod = 998244353 + conv = convolution.FFT(mod) + rng = random.Random(3) + for h in (1, 2, 3, 8): + n = 1 << h + a = [rng.randrange(mod) for _ in range(n)] + work = a[:] + conv.butterfly(work) + conv.butterfly_inv(work) + inv_n = pow(n, mod - 2, mod) + self.assertEqual([x * inv_n % mod for x in work], a) + if __name__ == "__main__": unittest.main()