Skip to content

Optimize convolution with Kronecker substitution and performance tuning #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
shakayami wants to merge 1 commit into from
+342 −57
Closed

Optimize convolution with Kronecker substitution and performance tuning #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions benchmarks/workloads.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -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),
Expand Down
290 changes: 233 additions & 57 deletions convolution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
import struct


class FFT:
def primitive_root_constexpr(self, m):
if m == 2:
Expand Down Expand Up @@ -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:
Expand All @@ -94,98 +103,265 @@ 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:
p = 1 << (h - LEN)
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("<Q" + code, buf)
]
del res[length:]
return res
frm = int.from_bytes
return [frm(buf[i : i + nb], "little") % mod for i in range(0, nb * length, nb)]

def _convolution_int(self, a, b, amax, bmax):
"""Kronecker substitution: multiply the two operands as one big int.

Each coefficient gets its own zero-padded slot, wide enough that no
coefficient of the product can carry into the next slot, so CPython's
(C-level, Karatsuba) integer multiply does the whole convolution.
"""
n = len(a)
m = len(b)
mod = self.mod
vb = max((amax.bit_length() + 7) >> 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
Loading
Loading