Parallel fused kernels¶
@ppy.parallel splits a fused NumPy loop across the worker pool, and the
output is bit-identical to the serial kernel and to NumPy. An elementwise
loop has no order to lose. A reduction does, and the compiler will not split
one that reassociates unless the function says so.
Run it¶
The decorator asks, the analysis answers¶
@ppy.pure
@ppy.parallel
@ppy.opt(3)
def parallel(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
This is eight million elements and five operations in one fused loop,
chunked across [tool.ppy.parallel] threads workers and joined. serial
is the same expression without the decorator.
The program checks np.array_equal across serial, parallel, and NumPy, and
it prints bit-identical: True on every path. A loop the analysis cannot
prove splittable stays serial and says why in an optimization remark.
A sum keeps its order unless you let it go¶
@ppy.pure
@ppy.fastmath
def relaxed_total(a: np.ndarray) -> float:
return np.sum(a * a)
@ppy.pure
def strict_total(a: np.ndarray) -> float:
return np.sum(a * a)
Splitting a floating-point sum changes where the rounding happens, so
strict_total is left in NumPy's order and equals NumPy's number exactly.
@ppy.fastmath permits the reassociation: relaxed_total vectorizes and
splits, and lands within 1e-3 of the strict answer on eight million squares.
You make the choice per function, and the default is the exact one.
parallel.range for other loops¶
parallel.range is the spelling for a loop the program itself declares
splittable, including pointers, buffers, and reductions
(parallel range). @ppy.parallel stays the
switch for fused NumPy loops like these.
Compared with NumPy, numexpr, Numba, and JAX¶
The fused expression and the sum of squares over eight million doubles, in
compare/: fused_bench.ppy,
fused_numpy.py, fused_numexpr.py,
fused_numba.py, fused_jax.py.
Milliseconds, best of five calls, over five processes.
PPy is the NumPy expression in a function, with @ppy.parallel to
split it. The serial row is the same expression fused into one loop with no
decorator, and strict_total is NumPy's sum in NumPy's order:
@ppy.pure
@ppy.parallel
@ppy.opt(3)
def parallel(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
NumPy is the expression as written: five operations, four 64 MB temporaries, one thread.
numexpr is the expression as a string, compiled to its own virtual machine and evaluated in chunks across threads:
Numba is an explicit loop under @njit(parallel=True) writing into
np.empty_like:
@njit(parallel=True)
def fused(a, b):
out = np.empty_like(a)
for i in prange(a.shape[0]):
out[i] = (a[i] * b[i] + a[i]) * (b[i] - a[i]) + a[i] * 0.5 - b[i] * 0.25
return out
JAX is the expression under jax.jit on the CPU, fused by XLA and run
on its thread pool, with jax_enable_x64:
| PPy | NumPy | numexpr | Numba prange |
JAX jit |
|
|---|---|---|---|---|---|
| fused, serial | 12.09 ± 0.41 | — | — | — | — |
| fused | 4.10 ± 0.18 | 56.26 ± 0.39 | 6.80 ± 0.31 | 3.37 ± 0.41 | 5.91 ± 0.46 |
| sum of squares | 13.35 ± 0.50 | 12.35 ± 0.27 | 6.90 ± 0.13 | 3.99 ± 0.08 | 0.65 ± 0.01 |
| sum of squares, relaxed | 4.47 ± 0.21 | — | — | — | — |
Fusing the expression is what removes NumPy's four temporaries. The serial fused loop reads the two inputs once and writes the output once. Splitting that loop across the cores is a memory-bandwidth problem that PPy, Numba, JAX, and numexpr solve the same way, within a couple of milliseconds of each other.
The fused kernel also checks its own result in the same loop: one add-reduction of non-finite elements the vectorizer keeps in a register, and one guard after. That is how it keeps NumPy's floating-point reporting without a second pass over 64 MB.
The ordered sum is NumPy's order and NumPy's time. The relaxed one vectorizes on one thread. JAX's reduction, reassociated across its pool, is the row that shows what the same permission buys with threads.
Intel Core Ultra 9 386H (16 threads), threads backend; NumPy 2.5.3, numexpr 2.14.2, Numba 0.67.0 on CPython 3.12.13, JAX 0.11.1 on CPython 3.13.13, PPy on CPython 3.13.13, from a checkout on a native filesystem.
What it prints¶
python parallel.ppy
fused serial 114.8 ms sample=-0.249979000059
fused parallel 123.4 ms sample=-0.249979000059
numpy 121.4 ms sample=-0.249979000059
bit-identical: True
strict == numpy: True
relaxed close : True
ppy parallel.ppy
fused serial 113.8 ms sample=-0.249979000059
fused parallel 127.6 ms sample=-0.249979000059
numpy 132.2 ms sample=-0.249979000059
bit-identical: True
strict == numpy: True
relaxed close : True
ppy run parallel.ppy
fused serial 14.9 ms sample=-0.249979000059
fused parallel 5.6 ms sample=-0.249979000059
numpy 18.3 ms sample=-0.249979000059
bit-identical: True
strict == numpy: True
relaxed close : True
Read on: Parallel loops · NumPy fusion
parallel.ppy is hand-written; there is no .py source and no conversion step.
07_parallel/parallel.ppy¶
import time
import numpy as np
import ppy
@ppy.pure
@ppy.opt(3)
def serial(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
@ppy.pure
@ppy.parallel
@ppy.opt(3)
def parallel(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
@ppy.pure
@ppy.fastmath
def relaxed_total(a: np.ndarray) -> float:
return np.sum(a * a)
@ppy.pure
def strict_total(a: np.ndarray) -> float:
return np.sum(a * a)
def report(label: str, elapsed: float, sample: float) -> None:
print(f"{label:16s} {elapsed * 1000.0:8.1f} ms sample={sample:.12f}")
def main() -> None:
x: np.ndarray = np.linspace(0.0, 10.0, 8000000)
y: np.ndarray = np.linspace(1.0, 5.0, 8000000)
serial(x, y)
parallel(x, y)
start: float = time.perf_counter()
a = serial(x, y)
report("fused serial", time.perf_counter() - start, float(a[7]))
start = time.perf_counter()
b = parallel(x, y)
report("fused parallel", time.perf_counter() - start, float(b[7]))
start = time.perf_counter()
c = (x * y + x) * (y - x) + x * 0.5 - y * 0.25
report("numpy", time.perf_counter() - start, float(c[7]))
print("bit-identical:", bool(np.array_equal(a, c)) and bool(np.array_equal(b, c)))
print("strict == numpy:", strict_total(x) == float(np.sum(x * x)))
print("relaxed close :", abs(relaxed_total(x) - float(np.sum(x * x))) < 1e-3)
if __name__ == "__main__":
main()
Counterpart programs¶
The programs the comparison above measured, each written the way its tool expects. The PPy one is first.
fused_bench.ppy (PPy)
"""The fused kernels of `parallel.ppy` over eight million doubles, timed: the PPY side."""
import time
import numpy as np
import ppy
@ppy.pure
@ppy.opt(3)
def serial(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
@ppy.pure
@ppy.parallel
@ppy.opt(3)
def parallel(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
@ppy.pure
def strict_total(a: np.ndarray) -> float:
return np.sum(a * a)
@ppy.pure
@ppy.fastmath
def relaxed_total(a: np.ndarray) -> float:
return np.sum(a * a)
def main() -> None:
x: np.ndarray = np.linspace(0.0, 10.0, 8000000)
y: np.ndarray = np.linspace(1.0, 5.0, 8000000)
serial(x, y)
parallel(x, y)
strict_total(x)
relaxed_total(x)
best = [1e9, 1e9, 1e9, 1e9]
for _ in range(5):
started = time.perf_counter()
out = serial(x, y)
best[0] = min(best[0], (time.perf_counter() - started) * 1000.0)
started = time.perf_counter()
out = parallel(x, y)
best[1] = min(best[1], (time.perf_counter() - started) * 1000.0)
started = time.perf_counter()
strict_total(x)
best[2] = min(best[2], (time.perf_counter() - started) * 1000.0)
started = time.perf_counter()
relaxed_total(x)
best[3] = min(best[3], (time.perf_counter() - started) * 1000.0)
labels = ["fused, serial", "fused", "sum of squares", "sum of squares, relaxed"]
for label, took in zip(labels, best, strict=True):
print(f"# {label}: {took:.2f} ms")
out = parallel(x, y)
print(f"{float(out[7]):.12f} {float(out[-1]):.12f}")
print(f"{strict_total(x):.3f}")
if abs(relaxed_total(x) - strict_total(x)) > 1e-3:
print("the relaxed sum drifted")
main()
fused_jax.py (Python)
"""The same expression under JAX on the CPU: `jax.jit` fuses it through XLA."""
import time
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
@jax.jit
def fused(a, b):
return (a * b + a) * (b - a) + a * 0.5 - b * 0.25
@jax.jit
def sum_of_squares(a):
return jnp.sum(a * a)
def timed(label, run):
best = 1e9
answer = None
for _ in range(5):
started = time.perf_counter()
answer = run()
best = min(best, time.perf_counter() - started)
print(f"# {label}: {best * 1000:.2f} ms")
return answer
def main():
x = jnp.linspace(0.0, 10.0, 8_000_000)
y = jnp.linspace(1.0, 5.0, 8_000_000)
fused(x, y).block_until_ready()
sum_of_squares(x).block_until_ready()
out = timed("fused", lambda: fused(x, y).block_until_ready())
print(f"{float(out[7]):.12f} {float(out[-1]):.12f}")
print(f"{timed('sum of squares', lambda: float(sum_of_squares(x).block_until_ready())):.3f}")
main()
fused_numba.py (Python)
"""The same expression under Numba: an explicit loop in `@njit(parallel=True)` with `prange`."""
import time
import numpy as np
from numba import njit, prange
@njit(parallel=True)
def fused(a, b):
out = np.empty_like(a)
for i in prange(a.shape[0]):
out[i] = (a[i] * b[i] + a[i]) * (b[i] - a[i]) + a[i] * 0.5 - b[i] * 0.25
return out
@njit
def sum_of_squares(a):
total = 0.0
for i in range(a.shape[0]):
total += a[i] * a[i]
return total
def timed(label, run):
best = 1e9
answer = None
for _ in range(5):
started = time.perf_counter()
answer = run()
best = min(best, time.perf_counter() - started)
print(f"# {label}: {best * 1000:.2f} ms")
return answer
def main():
x = np.linspace(0.0, 10.0, 8_000_000)
y = np.linspace(1.0, 5.0, 8_000_000)
fused(x, y)
sum_of_squares(x)
out = timed("fused", lambda: fused(x, y))
print(f"{float(out[7]):.12f} {float(out[-1]):.12f}")
print(f"{timed('sum of squares', lambda: sum_of_squares(x)):.3f}")
main()
fused_numexpr.py (Python)
"""The same expression under numexpr: one string, evaluated in chunks across threads."""
import time
import numexpr as ne
import numpy as np
def timed(label, run):
best = 1e9
answer = None
for _ in range(5):
started = time.perf_counter()
answer = run()
best = min(best, time.perf_counter() - started)
print(f"# {label}: {best * 1000:.2f} ms")
return answer
def main():
x = np.linspace(0.0, 10.0, 8_000_000)
y = np.linspace(1.0, 5.0, 8_000_000)
names = {"x": x, "y": y}
expression = "(x * y + x) * (y - x) + x * 0.5 - y * 0.25"
ne.evaluate(expression, local_dict=names)
out = timed("fused", lambda: ne.evaluate(expression, local_dict=names))
print(f"{float(out[7]):.12f} {float(out[-1]):.12f}")
total = lambda: float(ne.evaluate("sum(x * x)", local_dict=names)) # noqa: E731
print(f"{timed('sum of squares', total):.3f}")
main()
fused_numpy.py (Python)
"""The fused expression as NumPy evaluates it: five array operations, four temporaries."""
import time
import numpy as np
def timed(label, run):
best = 1e9
answer = None
for _ in range(5):
started = time.perf_counter()
answer = run()
best = min(best, time.perf_counter() - started)
print(f"# {label}: {best * 1000:.2f} ms")
return answer
def main():
x = np.linspace(0.0, 10.0, 8_000_000)
y = np.linspace(1.0, 5.0, 8_000_000)
out = timed("fused", lambda: (x * y + x) * (y - x) + x * 0.5 - y * 0.25)
print(f"{float(out[7]):.12f} {float(out[-1]):.12f}")
print(f"{timed('sum of squares', lambda: float(np.sum(x * x))):.3f}")
main()
Source: examples/07_parallel.