From ef72299c9f1663bedf1ebd600c289f080e5c75c7 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 15:31:14 +0000 Subject: [PATCH] Speed up convolution with Kronecker substitution and tighter NTT loops Profiling the existing NTT showed the cost is almost entirely CPython interpreter overhead in the butterfly inner loops, and that the NTT is only the right algorithm once both operands are large. Three changes: 1. Add an integer-packing (Kronecker substitution) path. Each coefficient is placed in its own zero-padded slot of a big integer, so a single CPython multiply performs the whole convolution in C. This is faster than the NTT whenever the smaller operand is small, the operands are lopsided, or the coefficients are small; a fitted cost model picks between the two. Slot width is derived from the actual maximum coefficient rather than the modulus, and struct is used for packing/unpacking instead of per-element to_bytes. 2. Optimize the butterfly loops: hoist instance attributes into locals, precompute the p/2p/3p index offsets, share common subexpressions between the four radix-4 outputs, and skip the twiddle multiplications for the groups where the rotation is 1. 3. Fold the 1/z normalization into b before its transform, so the inverse transform's output needs no extra scaling pass or extra list. The schoolbook path now triggers on n*m instead of min(n,m), since integer packing beats it above roughly 40 coefficient products. Measured on CPython 3.11 with mod 998244353 (old -> new): n=m=1000 7.8ms -> 1.6ms (4.9x) n=m=8192 80.3ms -> 46.8ms (1.7x) n=m=65536 1040.3ms -> 690.0ms (1.5x) n=65536, m=256 773.5ms -> 50.1ms (15.4x) n=262144, m=256 4121.7ms -> 200.6ms (20.6x) n=m=8192, 0/1 76.9ms -> 5.8ms (13.2x) Peak memory is unchanged to 75% lower. Tests cover both dispatch paths, unnormalized/negative/zero inputs and several moduli, and the benchmark suite gains a lopsided convolution workload alongside the balanced one. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_015N4uDXqbEPEaCbPoE4Utz4 --- benchmarks/workloads.py | 20 +++ convolution.py | 290 ++++++++++++++++++++++++++++++-------- tests/test_convolution.py | 89 ++++++++++++ 3 files changed, 342 insertions(+), 57 deletions(-) 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()