為什么 Cython 在這里比 Python+numpy 慢？

Question

我想實現一些快速凸分析操作 - 近端運算符等。 我是 Cython 的新手，我認為這將是適合這項工作的工具。 我在純 Python 和 Cython （下面的mwe_py.py和mwe_c.pyx ）中都有實現。 但是，當我比較它們時，Python + Numpy 版本明顯快於 Cython 版本。 為什么是這樣？ 我嘗試過使用內存視圖，它應該允許更快的索引/操作； 但是，性能差異非常明顯！ 任何有關如何在下面修復mwe_c.pyx以接近“最佳” Cython 代碼的建議將不勝感激。

import pyximport; pyximport.install(language_level=3)

import mwe_c
import mwe_py
import numpy as np
from time import time

n = 100000
nreps = 10000
x = np.random.randn(n)
z = np.random.randn(n)
tau = 1.0

t0 = time()
for _ in range(nreps):
    out = mwe_c.prox_translation(mwe_c.prox_two_norm, x, z, tau)
t1 = time()
print(t1 - t0)


t0 = time()
for _ in range(nreps):
    out = mwe_py.prox_translation(mwe_py.prox_two_norm, x, z, tau)
t1 = time()
print(t1 - t0)

分別給出輸出：

10.76103401184082  # (seconds)
5.988733291625977  # (seconds)

下面是mwe_py.py ：

import numpy.linalg as la

def proj_two_norm(x):
    """projection onto l2 unit ball"""
    X = la.norm(x)
    if X <= 1:
        return x
    return x / X

def prox_two_norm(x, tau):
    """proximal mapping of l2 norm with parameter tau"""
    return x - proj_two_norm(x / tau)

def prox_translation(prox_func, x, z, tau=None):
  """Compute prox(f(. - z))(x) where prox_func(x, tau) is prox(tau * f)(x)."""
  if tau is None:
      tau = 1.0
  return z + prox_func(x - z, tau)

最后，這里是mwe_c.pyx ：

import numpy as np
cimport numpy as np


cdef double [::1] aasubtract(double [::1] x, double [::1] z):
    cdef unsigned int i, m = len(x), n = len(z);
    assert m == n, f"vectors must have the same length"
    cdef double [::1] out = np.copy(x);
    for i in range(n):
        out[i] -= z[i]
    return out


cdef double [::1] vsdivide(double [::1] x, double tau):
    """Divide an array by a scalar element-wise."""
    cdef:
        unsigned int i, n = len(x);
        double [::1] out = np.copy(x);
    for i in range(n):
        out[i] /= tau
    return out


cdef double two_norm(double [::1] x):
    cdef:
        double out = 0.0;
        unsigned int i, n=len(x);
    for i in range(n):
        out = out + x[i]**2
    out = out **.5
    return out


cdef double [::1] proj_two_norm(double [::1] x):
    """project x onto the unit two ball."""
    cdef double x_norm = two_norm(x);
    cdef unsigned int i, n = len(x);
    cdef double [::1] p = np.copy(x);
    if x_norm <= 1:
        return p
    for i in range(n):
        p[i] = p[i] / x_norm
    return p


cpdef double [::1] prox_two_norm(double [::1] x, double tau):
    """double [::1] prox_two_norm(double [::1] x, double tau)"""
    cdef unsigned int i, n = len(x);
    cdef double [::1] out = x.copy(), Px = x.copy();
    Px = proj_two_norm(vsdivide(Px, tau));
    for i in range(n):
        out[i] = out[i] - Px[i]
    return out


cpdef prox_translation(
    prox_func,
    double [::1] x,
    double [::1] z,
    double tau=1.0
):
    cdef:
        unsigned int i, n = len(x);
        double [::1] out = prox_func(aasubtract(x, z), tau);
    for i in range(n):
        out[i] += z[i];
    return out

Answer 1

主要問題是您將優化的 Numpy 代碼與優化程度較低的 Cython 代碼進行比較。 事實上，Numpy 利用SIMD 指令（如 x86-64 處理器上的 SSE 和 AVX/AVX2）能夠連續計算許多項目。 Cython 默認使用默認的-O2優化級別，它不會啟用任何自動矢量化策略，從而導致標量代碼變慢（除非您使用最新版本的 GCC）。 您可以使用-O3告訴大多數編譯器（例如舊的 GCC 和 Clang）來啟用自動矢量化。 請注意，這不足以生成非常快速的代碼。 實際上，為了兼容性，編譯器僅在 x86-64 處理器上使用舊版 SIMD 指令。 -mavx和-mavx2啟用 AVX/AVX-2 指令集，以便生成更快的代碼，前提是您的機器支持它（否則它會簡單地崩潰）。 -mfma也可能有幫助。 -march=native也可以用於 select 目標平台上可用的最佳指令集。 請注意，Numpy 在運行時（部分）執行此檢查（感謝 GCC 特定的 C 功能）。

第二個主要問題是out = out + x[i]**2導致編譯器無法在不破壞 IEEE-754 標准的情況下優化的大循環依賴鏈。 實際上，要執行的加法鏈很長，處理器執行此操作的速度不能比使用當前代碼串行執行每條加法指令的速度更快。 問題是添加兩個浮點數具有相當大的延遲（在非常現代的 x86-64 處理器上通常為 3 到 4 個周期）。 這意味着處理器無法流水線化指令。 事實上，現代處理器通常可以並行執行兩個加法（每個內核），但當前循環阻止了這種情況。 最后，這個循環是完全受延遲限制的。 您可以通過手動展開循環來解決此問題。

使用-ffast-math可以幫助編譯器進行此類優化，但會以破壞 IEEE-754 標准為代價。 如果使用此選項，則無需使用特殊值（例如 NaN 數字）或某些操作。 欲了解更多信息，請閱讀gcc 的 ffast-math 究竟做了什么？ .

此外，請注意數組副本很昂貴，我不確定是否需要所有副本。 您可以創建一個新的空數組並對其進行填充，而不是對數組的副本進行操作。 這會更快，尤其是對於大 arrays。

最后，分裂很慢。 請考慮乘以逆。 由於 IEEE-754 標准，編譯器無法進行這種優化，但您可以輕松做到。 話雖如此，您需要確保這在您的情況下沒問題，因為它可能會稍微改變結果。 使用-ffast-math也應該會自動解決這個問題。

請注意，Numpy 許多開發人員都知道編譯器和處理器是如何工作的，因此他們已經進行了手動優化以生成快速代碼（就像我多次做過的那樣）。 在處理大型 arrays 時，您幾乎無法擊敗 Numpy，除非您合並循環或使用多線程。 實際上，與計算單元相比，如今的 RAM 相當慢，並且 Numpy 創建了許多臨時 arrays 。 Cython 可用於避免創建大多數臨時 arrays。