使用 Cython 优化简单的 CPU 绑定循环并替换列表

Question

我正在尝试评估一些方法，但我遇到了性能方面的障碍。

为什么我的 cython 代码这么慢？？我的期望是代码会运行比毫秒快得多（对于只有 256 ** 2 个条目的 2d 循环可能是纳秒）。

这是我的测试结果：

$ python setup.py build_ext --inplace; python test.py
running build_ext
        counter: 0.00236220359802 sec
       pycounter: 0.00323309898376 sec
      percentage: 73.1 %

我的初始代码如下所示：

#!/usr/bin/env python
# encoding: utf-8
# filename: loop_testing.py

def generate_coords(dim, length):
    """Generates a list of coordinates from dimensions and size
    provided.

    Parameters:
        dim -- dimension
        length -- size of each dimension

    Returns:
        A list of coordinates based on dim and length
    """
    values = []
    if dim == 2:
        for x in xrange(length):
            for y in xrange(length):
                values.append((x, y))

    if dim == 3:
        for x in xrange(length):
            for y in xrange(length):
                for z in xrange(length):
                    values.append((x, y, z))

    return values

这可以满足我的需要，但速度很慢。对于给定的暗淡，长度 = (2, 256)，我看到 iPython 上的时间约为 2.3 毫秒。

为了加快速度，我开发了一个 cython 等效项（我认为它是等效项）。

#!/usr/bin/env python
# encoding: utf-8
# filename: loop_testing.pyx
# cython: boundscheck=False
# cython: wraparound=False

cimport cython
from cython.parallel cimport prange

import numpy as np
cimport numpy as np


ctypedef int DTYPE

# 2D point updater
cpdef inline void _counter_2d(DTYPE[:, :] narr, int val) nogil:
    cdef:
        DTYPE count = 0
        DTYPE index = 0
        DTYPE x, y

    for x in range(val):
        for y in range(val):
            narr[index][0] = x
            narr[index][1] = y
            index += 1

cpdef DTYPE[:, :] counter(dim=2, val=256):
    narr = np.zeros((val**dim, dim), dtype=np.dtype('i4'))
    _counter_2d(narr, val)
    return narr

def pycounter(dim=2, val=256):
    vals = []
    for x in xrange(val):
        for y in xrange(val):
            vals.append((x, y))
    return vals

及调用时机：

#!/usr/bin/env python
# filename: test.py
"""
Usage:
    test.py [options]
    test.py [options] <val>
    test.py [options] <dim> <val>

Options:
    -h --help       This Message
    -n              Number of loops [default: 10]
"""

if __name__ == "__main__":
    from docopt import docopt
    from timeit import Timer

    args = docopt(__doc__)
    dim = args.get("<dim>") or 2
    val = args.get("<val>") or 256
    n = args.get("-n") or 10
    dim = int(dim)
    val = int(val)
    n = int(n)

    tests = ['counter', 'pycounter']
    timing = {}
    for test in tests:
        code = "{}(dim=dim, val=val)".format(test)
        variables = "dim, val = ({}, {})".format(dim, val)
        setup = "from loop_testing import {}; {}".format(test, variables)
        t = Timer(code, setup=setup)
        timing[test] = t.timeit(n) / n

    for test, val in timing.iteritems():
        print "{:>20}: {} sec".format(test, val)
    print "{:>20}: {:>.3} %".format("percentage", timing['counter'] / timing['pycounter'] * 100)

为了参考，setup.py 构建 cython 代码：

from distutils.core import setup
from Cython.Build import cythonize
import numpy

include_path = [numpy.get_include()]

setup(
    name="looping",
    ext_modules=cythonize('loop_testing.pyx'),  # accepts a glob pattern
    include_dirs=include_path,
)

编辑： Link 到工作版本：https://github.com/brianbruggeman/cython_experimentation

Answer 1

看起来您的 Cython 代码对 numpy 数组做了一些奇怪的事情，并没有真正利用 C 编译。要检查生成的代码，运行

cython -a loop_testing.pyx

如果您避免使用 numpy 部分并对 Python 函数进行简单的 Cython 转换会怎样？

编辑：看起来您可以完全避免使用 Cython 以获得相当不错的加速。（~30x 在我的机器上）

def npcounter(dim=2, val=256):
  return np.indices((val,)*dim).reshape((dim,-1)).T

Answer 2

这个 Cython 代码很慢，因为 narr[index][0] = x 分配严重依赖于 Python C-API。相反，使用 narr[index, 0] = x 被翻译成纯 C，并解决了这个问题。

正如@perimosocordiae 指出的那样，使用带有注释的 cythonize 绝对是调试此类问题的方法。

在某些情况下，在 setup.py 中为 gcc 显式指定编译标志也是值得的，

setup(
   [...]
   extra_compile_args=['-O2', '-march=native'],
   extra_link_args=['-O2', '-march=native'])

假设合理的默认编译标志，这应该不是必需的。但是，例如，在我的 Linux 系统上，默认情况下似乎根本没有优化，添加上述标志会显着提高性能。

使用 Cython 优化简单的 CPU 绑定循环并替换列表

Optimizing simple CPU-bound loops using Cython and replacing a list

performance

loops

numpy

cython