如何使用 Cython 更快地循环数组？

Question

我有一个光学工具，需要在其中计算一系列时间步长和波长的一组复振幅（即波前）。所以本质上我有两个 for 循环，对于每个项目我必须调用下面的方法。 由于计算需要相当长的时间，我尝试使用 Cython 加快一切速度。这就是我现在拥有的：

cimport numpy as cnp
cimport cython
import numpy as np
from libc.math cimport exp, cos, sin

ctypedef cnp.complex128_t DTYPE_t

@cython.cdivision(True)
@cython.boundscheck(False) # compiler directive
@cython.wraparound(False) # compiler directive
cpdef DTYPE_t[:,:,::1] _get_input_complex_amplitudes(
                                  float time,
                                  float wavelength,
                                  double[::1] x_source_sky_coordinates,
                                  double[::1] y_source_sky_coordinates,
                                  double[::1] x_observatory_coordinates,
                                  double[::1] y_observatory_coordinates,
                                  float aperture_radius,
                                  int number_of_inputs,
                                  int grid_size,
                                  DTYPE_t[:,:,::1] out):
    cdef unsigned int ix, iy, index_input
    cdef double arg, cos_val, sin_val
    cdef double factor = 2 * 3.1415926536
    cdef DTYPE_t[:,:,::1] ou = out

    for index_input in range(number_of_inputs):
        for ix in range(grid_size):
            for iy in range(grid_size):
                arg = factor / wavelength * (
                        x_observatory_coordinates[index_input] * x_source_sky_coordinates[ix] +
                        y_observatory_coordinates[index_input] * y_source_sky_coordinates[iy])

                cos_val = cos(arg)
                sin_val = sin(arg)

                ou[index_input,ix, iy] = aperture_radius * (cos_val + 1j * sin_val)
    return ou

但是，使用这段代码，我只能实现 3 倍的加速，并且在我的机器上运行一次该方法仍然需要 0.4 毫秒。我无法想象这是我可以从 Cython 中得到的一切，那么我是否缺少一些东西？通常，时间范围数组的长度为 100，波长数组的长度为 30，因此总共有 3000 个方法调用。

或者，我考虑重写方法来接受数组并同时传递时间和波长范围，而不是循环遍历时间和波长范围。但是，在这种情况下，我需要大量使用

numpy.einsum()

，我听说这不太适合使用 Cython 进行优化。

所以，我的问题是：

1. 我可以做其他事情来提高使用 Cython 的方法的速度吗？
2. 传递 2D 数组而不是循环两个范围在计算上是否更明智？

非常感谢！

Answer 1

sin

和

cos

一般都很贵。原因之一是需要考虑许多极端情况：NaN 值、无穷大、次正规数。另一个原因是精度：据我所知，该函数的误差最多为 1 ULP，这对于双精度来说非常小。如果您知道没有特殊值（NaN、Inf、次正规值）并且性能比准确性更重要，那么您可以启用

-ffast-math

编译器选项。为了提高性能，此选项还假设浮点数学是关联的。但这是相当不安全，所以你应该三思而后行。

您可以根据数学技巧进行一些简单的优化，从而显着提高性能。确实

sin(arg)² + cos(arg)² = 1

，所以

cos(arg) = +/- sqrt(1-sin(arg)²)

。可以很容易地从 a 的值获得符号（特别是如果您知道

arg

在

[X;X+2*pi[

范围内，因为它有助于避免计算昂贵的浮点模数）。

sqrt

比

sin

功能便宜得多，主要是因为主流 CPU 都有专用的单元。另一个优化是预先计算

factor / wavelength

：您可以在循环中使用

factor

并使用

2 * 3.1415926536 / wavelength

进行计算。不过，优化编译器应该已经做到了后者。

sin

和

cos

函数可以使用SIMD指令进行矢量化。这很难手动完成，但有一些库可以做到这一点，例如Intel SVML。一些编译器和 libc 实现也可以使用

-ffast-math

来实现（例如 GCC + 最近的 glibc）。在此代码中，我不期望 GCC 和 Clang 能够对循环进行矢量化。英特尔编译器可能会。如果机器上安装了 Numpy，则可以使用 SVML。在本例中，我们的想法是仅使用 Numpy 进行 sin/cos 计算。最新 x86-64 CPU 上的 SIMD 指令应该使此代码速度提高 2 倍到 3 倍。

另一个解决方案是使用多线程。然而，该函数的执行速度非常快，因此创建线程并等待它们可能会带来显着的性能开销。如果函数调用是独立的，那么最好并行调用而不是并行化该函数。这可以与上述 SIMD 优化相结合。