使用numba加速大型numpy数组向量化的性能优化探究

环境配置 #

from timeit import timeit

import numba as nb
import numpy as np
import numpy.typing as npt

TIMESTAMPS = 3
REGIONS = 150
HEIGHT = 240
WIDTH = 320

TRIALS = 100

np.random.seed(0)

FRAME_MASKS = np.ascontiguousarray(
    np.random.randint(0, 2, (TIMESTAMPS, HEIGHT, WIDTH), dtype=bool)
)
REGION_MASKS = np.ascontiguousarray(
    np.random.randint(0, 2, (REGIONS, HEIGHT, WIDTH), dtype=bool)
)
REGION_AREAS = np.sum(REGION_MASKS, axis=(-1, -2))

# OLFR = "Overlap of Frames with Regions"

# 预运行编译代码
olfr_native_for_loops(*test_params)
olfr_fully_vectorized(*test_params)
olfr_with_numba(*test_params)
olfr_all_flags_numba(*test_params)

# 计时测试
test_params = (FRAME_MASKS, REGION_MASKS, REGION_AREAS)
for prefix, func in [
    ["for_loops", lambda: olfr_native_for_loops(*test_params)],
    ["fully_vectorized", lambda: olfr_fully_vectorized(*test_params)],
    ["with_numba", lambda: olfr_with_numba(*test_params)],
    ["all_flags_numba", lambda: olfr_all_flags_numba(*test_params)],
]:
    speed = timeit(func, number=TRIALS)
    print(f"Done. {prefix} speed={speed:.4f}")

1. 原生for循环 #

def olfr_native_for_loops(
    frame_masks: npt.ArrayLike,  # T x H x W
    region_masks: npt.ArrayLike,  # R x H x W
    region_areas: npt.ArrayLike,  # R
):
    ratios = np.zeros((TIMESTAMPS, REGIONS))

    for t in range(TIMESTAMPS):
        for r in range(REGIONS):
            raw_intersection = frame_masks[t] & region_masks[r]
            intersection_area = np.sum(raw_intersection, axis=(-1, -2))
            ratios[t, r] = intersection_area / region_areas[r]

    return ratios

运行结果：

Done. for_loops speed=2.0005

该方案内部已具备部分向量化特性，但外层for循环仍有优化空间，因此尝试下一步向量化改造。

2. 完全向量化 #

def olfr_fully_vectorized(
    frame_masks: npt.ArrayLike,  # T x H x W
    region_masks: npt.ArrayLike,  # R x H x W
    region_areas: npt.ArrayLike,  # R
):
    raw_intersection = frame_masks[:, np.newaxis, ...] & region_masks[np.newaxis]
    intersection_area = np.sum(raw_intersection, axis=(-1, -2))
    ratios = intersection_area / region_areas
    return ratios

运行结果：

Done. fully_vectorized speed=2.4903

代码更简洁，但性能不升反降。原因是向量化后生成的中间数组过大，触发内存缓存驱逐。生产环境中该方案CPU使用率仅为原生方案的一半，推测BLAS未能调度足够线程实现并行化，因此转向Numba尝试优化。

3. 常规Numba #

@nb.njit
def olfr_with_numba(
    frame_masks: npt.ArrayLike,  # T x H x W
    region_masks: npt.ArrayLike,  # R x H x W
    region_areas: npt.ArrayLike,  # R
):
    raw_intersection = frame_masks.reshape(TIMESTAMPS, 1, -1) & region_masks.reshape(
        1, REGIONS, -1
    )
    intersection_area = np.sum(raw_intersection, axis=-1)
    result = intersection_area / region_areas
    return result

运行结果：

Done. with_numba speed=14.7773

性能进一步下降，查阅资料得知Numba对大数组的默认处理效率不佳，因此尝试启用Numba的并行、快速数学等参数。

4. 带参数的Numba #

@nb.njit(parallel=True, fastmath=True, nogil=True, cache=True)
def olfr_all_flags_numba(
    frame_masks: npt.ArrayLike,  # T x H x W
    region_masks: npt.ArrayLike,  # R x H x W
    region_areas: npt.ArrayLike,  # R
):
    raw_intersection = frame_masks.reshape(TIMESTAMPS, 1, -1) & region_masks.reshape(
        1, REGIONS, -1
    )
    intersection_area = np.sum(raw_intersection, axis=-1)
    result = intersection_area / region_areas
    return result

运行结果：

Done. all_flags_numba speed=7.6477

性能有所提升，但仍不及原生for循环方案。

核心疑问 #

生产代码中通过ThreadPool手动实现多线程可显著提速，为何BLAS/Numba无法自动完成此类并行调度？

附：AWS c5.4xlarge CPU参数 #

CPU(s):              16
On-line CPU(s) list: 0-15
Thread(s) per core:  2
Core(s) per socket:  8
Socket(s):           1
NUMA node(s):        1
Vendor ID:           GenuineIntel
CPU family:          6
Model:               85
Model name:          Intel(R) Xeon(R) Platinum 8124M CPU @ 3.00GHz
Stepping:            4
CPU MHz:             3400.948
L1d cache:           32K
L1i cache:           32K
L2 cache:            1024K
L3 cache:            25344K

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