NumPy fusion¶
An elementwise NumPy expression becomes one loop with no temporaries.
np.sin(a) * 2.0 + np.cos(b) is three NumPy calls and two intermediate
arrays; under ppy run it is a single loop compiled through LLVM.
Run it¶
How the expression is fused¶
The expression tree lowers to the IR's tensor dialect: numpy.sin is
tensor.unary {op = sin}, and * and + are tensor.mul and
tensor.add. The tensor-fusion pass folds the chain into one
tensor.fused region, a computation of one output element from one element
of each input. lower-tensor writes it as one strided loop.
@ppy.pure
@ppy.opt(3)
def normalize(x: np.ndarray) -> np.ndarray:
scale: float = np.sqrt(np.sum(x * x))
return x / scale
A reduction fuses only at the root of a tree. In normalize, x * x feeds
np.sum, so the multiply fuses into the reduction: one pass over x, one
accumulator, no squared array. x / scale is a second loop because it needs
the reduction's result. Nested inside an elementwise expression a reduction
would not be elementwise, and the pass knows the difference.
To see the fused region before it becomes loops:
What the guard checks¶
The kernel takes only what it was compiled for: float64, C-contiguous,
one shape across the operands. The generated boundary checks that on every
call. Anything else (a float32 array, a transposed view, a broadcast) runs
NumPy itself.
Reduction order is preserved bit for bit unless the function is
@ppy.fastmath, so np.sum here gives NumPy's number, not one close to it.
Compared with NumPy, numexpr, Numba, and JAX¶
The two expressions over eight million doubles, in compare/:
fusion_bench.ppy,
fusion_numpy.py,
fusion_numexpr.py,
fusion_numba.py,
fusion_jax.py. Milliseconds, best of five calls,
over five processes.
PPy is the NumPy expression as written, in a function:
NumPy is the same line with no function around it: three calls, two 64 MB temporaries.
numexpr takes the expression as a string and evaluates it in chunks across threads:
Numba is an explicit loop under @njit writing into np.empty_like:
@njit
def blend(a, b):
out = np.empty_like(a)
for i in range(a.shape[0]):
out[i] = math.sin(a[i]) * 2.0 + math.cos(b[i])
return out
JAX is the expression under jax.jit on the CPU, with
jax_enable_x64:
| PPy | NumPy | numexpr | Numba @njit |
JAX jit |
|
|---|---|---|---|---|---|
| normalize | 22.25 ± 0.49 | 21.68 ± 0.54 | 13.03 ± 0.28 | 25.63 ± 0.79 | 8.03 ± 0.14 |
| blend | 73.52 ± 0.72 | 88.02 ± 0.73 | 10.07 ± 0.20 | 84.53 ± 0.84 | 22.84 ± 0.22 |
blend is two transcendentals per element, and on one thread that is what
the time is. PPy, NumPy, and Numba each call sin and cos once per
element, and fusing away NumPy's two temporaries takes off the fifth of the
time that was memory traffic. numexpr and JAX evaluate sin and cos
across vector lanes and across threads, which is where their rows come
from.
normalize is a reduction in NumPy's own order followed by a division
pass, on every tool that keeps the order. PPy's sum is NumPy's sum, so the
row is NumPy's time. @ppy.fastmath is the permission to reassociate it
(Parallel shows what that buys).
Intel Core Ultra 9 386H (16 threads); Numba 0.67.0, numexpr 2.14.2, NumPy 2.5.3 on CPython 3.12.13, JAX 0.11.1 on CPython 3.13.13, PPy on CPython 3.14.5, from a checkout on a native filesystem.
What it prints¶
python numpy_fusion.ppy, ppy numpy_fusion.ppy, ppy run numpy_fusion.ppy
Read on: Plugins: NumPy · Parallel fused kernels · The IR: the tensor dialect
numpy_fusion.ppy is hand-written; there is no .py source and no conversion
step.
05_numpy/numpy_fusion.ppy¶
import numpy as np
import ppy
@ppy.pure
@ppy.opt(3)
def normalize(x: np.ndarray) -> np.ndarray:
scale: float = np.sqrt(np.sum(x * x))
return x / scale
@ppy.pure
def blend(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return np.sin(a) * 2.0 + np.cos(b)
def main() -> None:
values = np.arange(1024, dtype=np.float64)
normalized = normalize(values)
mixed = blend(values, values)
print(round(float(normalized[1]), 8), round(float(mixed[1]), 8), int(values.ndim))
if __name__ == "__main__":
main()
Counterpart programs¶
The programs the comparison above measured, each written the way its tool expects. The PPy one is first.
fusion_bench.ppy (PPy)
"""The two fused expressions of `numpy_fusion.ppy` over eight million doubles, timed: PPY's side."""
import time
import numpy as np
import ppy
@ppy.pure
@ppy.opt(3)
def normalize(x: np.ndarray) -> np.ndarray:
scale: float = np.sqrt(np.sum(x * x))
return x / scale
@ppy.pure
def blend(a: np.ndarray, b: np.ndarray) -> np.ndarray:
return np.sin(a) * 2.0 + np.cos(b)
def main() -> None:
values: np.ndarray = np.linspace(0.0, 10.0, 8_000_000)
other: np.ndarray = np.linspace(1.0, 5.0, 8_000_000)
normalize(values)
blend(values, other)
best = [1e9, 1e9]
for _ in range(5):
started = time.perf_counter()
normalized = normalize(values)
best[0] = min(best[0], (time.perf_counter() - started) * 1000.0)
started = time.perf_counter()
mixed = blend(values, other)
best[1] = min(best[1], (time.perf_counter() - started) * 1000.0)
print(f"# normalize: {best[0]:.2f} ms")
print(f"# blend: {best[1]:.2f} ms")
print(f"{float(normalized[1]):.9f} {float(normalized[-1]):.9f}")
print(f"{float(mixed[1]):.9f} {float(mixed[-1]):.9f}")
main()
fusion_jax.py (Python)
"""The same two expressions under `jax.jit` on the CPU: XLA fuses them."""
import time
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
@jax.jit
def normalize(x):
scale = jnp.sqrt(jnp.sum(x * x))
return x / scale
@jax.jit
def blend(a, b):
return jnp.sin(a) * 2.0 + jnp.cos(b)
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) * 1000.0)
print(f"# {label}: {best:.2f} ms")
return answer
def main():
values = jnp.linspace(0.0, 10.0, 8_000_000)
other = jnp.linspace(1.0, 5.0, 8_000_000)
normalize(values).block_until_ready()
blend(values, other).block_until_ready()
normalized = timed("normalize", lambda: normalize(values).block_until_ready())
mixed = timed("blend", lambda: blend(values, other).block_until_ready())
print(f"{float(normalized[1]):.9f} {float(normalized[-1]):.9f}")
print(f"{float(mixed[1]):.9f} {float(mixed[-1]):.9f}")
main()
fusion_numba.py (Python)
"""The same two expressions under Numba: explicit loops in `@njit`, one pass each."""
import math
import time
import numpy as np
from numba import njit
@njit
def normalize(x):
total = 0.0
for i in range(x.shape[0]):
total += x[i] * x[i]
scale = math.sqrt(total)
out = np.empty_like(x)
for i in range(x.shape[0]):
out[i] = x[i] / scale
return out
@njit
def blend(a, b):
out = np.empty_like(a)
for i in range(a.shape[0]):
out[i] = math.sin(a[i]) * 2.0 + math.cos(b[i])
return out
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) * 1000.0)
print(f"# {label}: {best:.2f} ms")
return answer
def main():
values = np.linspace(0.0, 10.0, 8_000_000)
other = np.linspace(1.0, 5.0, 8_000_000)
normalize(values)
blend(values, other)
normalized = timed("normalize", lambda: normalize(values))
mixed = timed("blend", lambda: blend(values, other))
print(f"{float(normalized[1]):.9f} {float(normalized[-1]):.9f}")
print(f"{float(mixed[1]):.9f} {float(mixed[-1]):.9f}")
main()
fusion_numexpr.py (Python)
"""The same two expressions under numexpr: strings evaluated in chunks across threads."""
import time
import numexpr as ne
import numpy as np
def normalize(x):
scale = float(np.sqrt(ne.evaluate("sum(x * x)", local_dict={"x": x})))
return ne.evaluate("x / scale", local_dict={"x": x, "scale": scale})
def blend(a, b):
return ne.evaluate("sin(a) * 2.0 + cos(b)", local_dict={"a": a, "b": b})
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) * 1000.0)
print(f"# {label}: {best:.2f} ms")
return answer
def main():
values = np.linspace(0.0, 10.0, 8_000_000)
other = np.linspace(1.0, 5.0, 8_000_000)
normalize(values)
blend(values, other)
normalized = timed("normalize", lambda: normalize(values))
mixed = timed("blend", lambda: blend(values, other))
print(f"{float(normalized[1]):.9f} {float(normalized[-1]):.9f}")
print(f"{float(mixed[1]):.9f} {float(mixed[-1]):.9f}")
main()
fusion_numpy.py (Python)
"""The same two expressions as NumPy evaluates them: a temporary per operation."""
import time
import numpy as np
def normalize(x):
scale = np.sqrt(np.sum(x * x))
return x / scale
def blend(a, b):
return np.sin(a) * 2.0 + np.cos(b)
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) * 1000.0)
print(f"# {label}: {best:.2f} ms")
return answer
def main():
values = np.linspace(0.0, 10.0, 8_000_000)
other = np.linspace(1.0, 5.0, 8_000_000)
normalized = timed("normalize", lambda: normalize(values))
mixed = timed("blend", lambda: blend(values, other))
print(f"{float(normalized[1]):.9f} {float(normalized[-1]):.9f}")
print(f"{float(mixed[1]):.9f} {float(mixed[-1]):.9f}")
main()
Source: examples/05_numpy.