小NumPy数组基础算术运算性能优化问询

我在处理6-10个元素的一维np.float32数组时，执行一组简单的基础数学运算，发现NumPy的运算开销远高于Fortran实现。由于需要执行数百万次这类运算，希望找到优化方案，而非仅接受“这是固有开销”。该性能差异在多台设备上均可复现（环境：Windows 10 64位、Python 3.9.10 64位、NumPy 1.21.5 MKL）。

NumPy实现与测试 #

核心代码 #

def NumPyFunc(array1, array2, float1, float2, float3):
    output1 = (array2 - array1) / (float2 - float1)
    output2 = array1 + output1 * (float3 - float1)
    return output1, output2

输入示例 #

import numpy

sz = 6
array1 = 3000.0 * numpy.random.uniform(size=(sz,)).astype(numpy.float32)
array2 = 2222.0 * numpy.random.uniform(size=(sz,)).astype(numpy.float32)
float1 = float(numpy.random.uniform(100000, 1e7))
float2 = float(numpy.random.uniform(100000, 1e7))
float3 = float(numpy.random.uniform(100000, 1e7))

性能测试结果 #

• 设备1：

  %timeit NumPyFunc(array1, array2, float1, float2, float3)
  3.33 µs ± 18 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
  

• 设备2：

  %timeit NumPyFunc(array1, array2, float1, float2, float3)
  1.5 µs ± 19.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
  

Fortran实现与测试 #

为对比性能，我用f2py封装了Fortran子程序，测试结果显示其速度比NumPy快5倍以上。

Fortran核心子程序 #

pure subroutine f90_small_arrays(n, array1, array2, float1, float2, float3, output1, output2)
    implicit none
    integer, intent(in) :: n
    real(4), intent(in), dimension(n) :: array1, array2
    real(4), intent(in) :: float1, float2, float3
    real(4), intent(out), dimension(n) :: output1, output2

    output1 = (array2 - array1) / (float2 - float1)
    output2 = array1 + output1 * (float3 - float1)
end subroutine f90_small_arrays

Python调用代码 #

from f90_small_arrays import f90_small_arrays

def FortranFunc(array1, array2, float1, float2, float3):
    output1, output2 = f90_small_arrays(array1, array2, float1, float2, float3)
    return output1, output2

性能测试结果 #

• 设备1：

  %timeit FortranFunc(array1, array2, float1, float2, float3)
  654 ns ± 0.869 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
  

• 设备2：

  %timeit FortranFunc(array1, array2, float1, float2, float3)
  286 ns ± 5.92 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
  

现状与优化需求 #

• 升级到NumPy 1.26.3后，性能反而比1.21.5版本下降15%。
• 虽可用Fortran替代，但核心运算代码独立于Python文件，降低了整体可读性。
• 考虑过Numba JIT优化，但py2exe分发含Numba的应用流程繁琐，暂不采用。

希望得到针对NumPy的具体优化建议，尽可能缩小与Fortran的性能差距，而非仅接受“这是NumPy的固有开销”。

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