为何经Numba JIT编译的圆盘均匀点生成函数速度反而更慢？

我编写了一个生成圆盘内均匀分布点的函数，因需频繁运行且处理大数组，原本认为用numba JIT编译能大幅提升速度，但测试后发现编译后的函数反而慢了一倍以上。请问如何找出numba函数的性能瓶颈？

函数代码 #

from numba import njit
import numpy as np

@njit(cache=True)
def generate_points_turbo(centre_point, radius, num_rings, x_axis=np.array([-1, 0, 0]), y_axis=np.array([0, 1, 0])):
    """
    Generate uniformly spaced points inside a circle
    Based on algorithm from:
    http://www.holoborodko.com/pavel/2015/07/23/generating-equidistant-points-on-unit-disk/
    
    Parameters
    ----------
    centre_point : np.ndarray (1, 3)
    radius : float/int
    num_rings : int
    x_axis : np.ndarray
    y_axis : np.ndarray

    Returns
    -------
    points : np.ndarray (n, 3)

    """
    if num_rings > 0:
        delta_R = 1 / num_rings
        ring_radii = np.linspace(delta_R, 1, int(num_rings)) * radius
        k = np.arange(num_rings) + 1
        points_per_ring = np.rint(np.pi / np.arcsin(1 / (2*k))).astype(np.int32)
        num_points = points_per_ring.sum() + 1
        ring_indices = np.zeros(int(num_rings)+1)
        ring_indices[1:] = points_per_ring.cumsum()
        ring_indices += 1
        points = np.zeros((num_points, 3))

        points[0, :] = centre_point

        for indx in range(len(ring_radii)):
            theta = np.linspace(0, 2 * np.pi, points_per_ring[indx]+1)
            points[ring_indices[indx]:ring_indices[indx+1], :] = ((ring_radii[indx] * np.cos(theta[1:]) * x_axis[:, None]).T
                     + (ring_radii[indx] * np.sin(theta[1:]) * y_axis[:, None]).T)
        return points + centre_point

调用方式 #

centre_point = np.array([0,0,0])
radius = 1
num_rings = 15

generate_points_turbo(centre_point, radius, num_rings )

性能测试更新（硬件/数据规模相关性） #

测试发现numba函数的性能转折点因硬件而异：

%timeit generate_points(centre_point, 1, 2)
99.5 µs ± 932 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit generate_points_turbo(centre_point, 1, 2)
213 µs ± 8.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit generate_points(centre_point, 1, 20)
647 µs ± 11.2 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit generate_points_turbo(centre_point, 1, 20)
314 µs ± 8.74 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit generate_points(centre_point, 1, 200)
11.9 ms ± 375 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit generate_points_turbo(centre_point, 1, 200)
7.9 ms ± 243 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

在我的设备上，当num_rings超过12-15后，numba版本函数开始追平甚至超过原函数，但大数据量下的性能提升幅度低于预期，说明函数部分逻辑的性能与数据规模高度相关。

问题解答 #

为什么小数据量下Numba版本更慢？ #

Numba的JIT编译存在启动开销：第一次调用函数时需要编译成机器码，后续调用会复用缓存，但对于极小任务（如num_rings=2），编译开销的占比远大于运行时节省的时间，导致整体耗时更高。当数据规模增大后，编译开销被摊薄，Numba的机器码执行优势才会体现。

如何排查Numba函数的性能瓶颈？ #

1. 用Numba自带的性能分析工具

   • 使用命令行工具：numba --annotate-html profile.html your_script.py，生成的HTML报告会标注每行代码的耗时、是否被编译为机器码，能直观看到瓶颈所在。
   • 在代码中启用 profiling：将装饰器改为@njit(profile=True)，运行函数后会输出详细的性能统计，包括每个函数/代码块的调用次数和耗时。

2. 对比原函数与Numba函数的热点

   • 用Python内置的cProfile分别分析两个函数：

     import cProfile
     cProfile.run('generate_points(centre_point, 1, 200)', sort='cumulative')
     cProfile.run('generate_points_turbo(centre_point, 1, 200)', sort='cumulative')
     

     对比两者的耗时分布，找出Numba处理效率低的环节。

3. 检查Numpy函数的Numba支持度
   并非所有Numpy函数都能被Numba高效编译，部分函数可能会 fallback到Python解释器（导致性能骤降）。比如np.linspace、np.arcsin这类函数，可尝试将其替换为纯Python实现，观察性能变化。

4. 优化内存布局与临时数组

   • 确保输入输出数组为C连续内存（用arr.flags.c_contiguous检查），Numba对连续内存的处理效率更高。
   • 避免在循环内频繁创建临时数组：你的代码中每次循环都生成theta数组，以及广播操作产生的中转数组，这些内存分配和拷贝是潜在瓶颈。可尝试预分配内存，或直接逐元素计算赋值。

5. 调整函数默认参数
   Numba对默认数组参数（如x_axis=np.array([-1,0,0])）的处理可能存在额外开销，建议将默认参数移到函数内部初始化：

   @njit(cache=True)
   def generate_points_turbo(centre_point, radius, num_rings, x_axis=None, y_axis=None):
       if x_axis is None:
           x_axis = np.array([-1, 0, 0])
       if y_axis is None:
           y_axis = np.array([0, 1, 0])
       # 后续逻辑不变
   

针对当前代码的优化建议 #

• 将循环内的np.linspace替换为手动计算theta值，避免临时数组：

  # 替换原theta生成逻辑
  n_points = points_per_ring[indx]
  start_idx = int(ring_indices[indx])
  end_idx = int(ring_indices[indx+1])
  for i in range(n_points):
      theta = 2 * np.pi * i / n_points
      cos_theta = np.cos(theta)
      sin_theta = np.sin(theta)
      points[start_idx + i] = centre_point + ring_radii[indx] * (cos_theta * x_axis + sin_theta * y_axis)
  

• 移除最后的points + centre_point，直接在赋值时加上centre_point，避免全局数组拷贝。

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