为何显式代数运算比Numpy的np.prod和np.all更快？

实验代码 #

import numpy as np
import time

N_ROWS = int(2e5)
N_ITER = int(1e3)

np.random.seed(123)

mat = np.random.rand(N_ROWS, 2)
mat_bool = mat > .5

# np.prod
start = time.time()
for i in range(N_ITER):
    _ = np.prod(mat, axis=1)
duration = time.time() - start
print(f"np.prod took {duration}s")

start = time.time()
for i in range(N_ITER):
    _ = mat[:, 0] * mat[:, 1]
duration = time.time() - start
print(f"Manual prod took {duration}s")

# np.all
start = time.time()
for i in range(N_ITER):
    _ = np.all(mat_bool, axis=1)
duration = time.time() - start
print(f"np.all took {duration}s")

start = time.time()
for i in range(N_ITER):
    _ = mat_bool[:, 0] * mat_bool[:, 1]
duration = time.time() - start
print(f"manual all took {duration}s")

实验结果 #

np.prod took 2.5077707767486572s
Manual prod took 0.0815896987915039s
np.all took 2.831434488296509s
manual all took 0.11392521858215332s

（Numpy版本：'1.24.2'）

核心原因 #

• 通用函数的额外适配开销：np.prod和np.all是Numpy的通用型函数，设计目标是支持任意维度、任意轴数、多种数据类型的输入。这意味着它们内部要做大量前置工作：比如校验输入维度合法性、归一化轴参数（处理负索引）、适配不同数据类型的计算逻辑。这些步骤都会产生额外的运行时间，而手动运算完全针对固定2列的场景，跳过了所有通用适配逻辑。

• 计算路径的复杂度差异：手动的列相乘是直接调用Numpy最基础的元素级乘法ufunc，这是底层优化最彻底的向量化操作。而np.prod(axis=1)需要对每一行的元素做累积乘法，np.all(axis=1)需要对每一行做累积逻辑与，这类累积操作的计算路径比直接两元素相乘更长，当数组行数达到20万、循环次数1000次时，累积操作的开销会被显著放大。

• 内存缓存利用效率：对于2列的数组，手动操作是对两个一维数组做元素级运算，即使Numpy默认是行优先存储，两个列的访问模式完全一致，CPU缓存的命中率较高。而np.prod/np.all需要逐行遍历处理每个元素，内存访问模式相对零散，缓存利用效率不如直接的列运算。

内容的提问来源于stack exchange，提问作者Xavier