使用 cython 比 struct.pack 快

Question

我正在努力做得比 struct.pack 更好。

通过对的回答，以打包整数的具体情况为例，我有以下内容来打包 pack_ints.pyx 中的整数列表：

# cython: language_level=3, boundscheck=False
import cython

@cython.boundscheck(False)
@cython.wraparound(False)
def pack_ints(int_col):

    int_buf = bytearray(4*len(int_col))
    cdef int[::1] buf_view = memoryview(int_buf).cast('i')

    idx: int = 0
    for idx in range(len(int_col)):
        buf_view[idx] = int_col[idx]


    return int_buf

在ipython中使用此测试代码：

from struct import pack 
import pyximport; pyximport.install(language_level=3) 
import pack_ints 

amount = 10**7 
ints = list(range(amount)) 

res1 = pack(f'{amount}i', *ints) 
res2 = pack_ints.pack_ints(ints) 
assert(res1 == res2) 

%timeit pack(f'{amount}i', *ints)  
%timeit pack_ints.pack_ints(ints)

我得到：

304 ms ± 2.18 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
212 ms ± 6.54 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我尝试将 int_buf 键入为 array('b')，但没有看到任何改进。

是否有任何其他方法可以对此进行改进，或者以不同的方式使用 cython 来加快此操作？

Answer 1

当我运行来自原始问题的代码时，我得到了大约 5 倍的加速。

当我运行你的代码在这里时，我看到你报告的结果 加上编译阶段的重要警告 我认为你忽略了：

warning: pack_ints.pyx:13:17: Index should be typed for more efficient access

我不确定为什么它没有正确选择类型，但要修复它，您应该将 i 的定义改回我最初编写的代码：

cdef int i
# not "i: int"

希望其他人会过来尝试一些更聪明的东西，因为这显然有点荒谬，这是一个答案。

Answer 2

这个答案试图给出一个估计，并行化版本可以产生多少加速。然而，因为这个任务是内存带宽限制的（Python-integer-objects 至少需要 32 个字节并且可以分散在内存中，所以会有很多缓存未命中）我们不应该期望太多。

第一个问题，如何处理错误（元素不是整数或值太大）。我将关注 strategy/simplification：当对象

不是整数，
为负整数，
或整数 >=2^30

它将被转换为一个特殊的数字 (-1)，表示出现问题。只允许非负整数 <2^30 让我的生活更轻松，因为我必须重新实现 PyLong_AsLongAndOverflow 而不会引发错误，否则检测溢出通常很麻烦（但是，请参阅答案末尾的版本更复杂的方法）。

可以找到 Python 的整数对象的内存布局 here:

struct _longobject {
    PyObject_VAR_HEAD
    digit ob_digit[1];
};

Member ob_size/macro Py_SIZE告诉我们在表示整数时使用了多少个30位数字（ob_size表示负整数）。

我的简单规则因此转化为以下 C 代码（我使用 C 而不是 Cython，因为它是 simpler/more 使用 Python 的 C-API):

#include <Python.h>

// returns -1 if vv is not an integer,
//            negative, or > 2**30-1
int to_int(PyObject *vv){ 
   if (PyLong_Check(vv)) {
       PyLongObject * v = (PyLongObject *)vv;
       Py_ssize_t i = Py_SIZE(v);
       if(i==0){
           return 0;
       }
       if(i==1){//small enought for a digit
           return v->ob_digit[0];
       }
       //negative (i<0) or too big (i>1)
       return -1;
   }
   return -1;
}

现在给定一个列表，我们可以将其转换为 int-buffer 并与以下使用 omp 的 C 函数并行：

void convert_list(PyListObject *lst, int *output){
    Py_ssize_t n = Py_SIZE(lst);
    PyObject **data = lst->ob_item;
    #pragma omp parallel for
    for(Py_ssize_t i=0; i<n; ++i){
        output[i] = to_int(data[i]);
    }
}

没什么好说的- PyListObject-API用于并行访问列表的元素。可以做到，因为to_int-function.

中没有ref counting/racing条件

现在，将它与 Cython 捆绑在一起：

%%cython -c=-fopenmp --link-args=-fopenmp
import cython

cdef extern from *:
    """
    #include <Python.h>

    int to_int(PyObject *vv){ 
       ... code
    }

    void convert_list(PyListObject *lst, int *output){
        ... code
    }
    """
    void convert_list(list lst, int *output)

@cython.boundscheck(False)
@cython.wraparound(False)
def pack_ints_ead(list int_col):
    cdef char[::1] int_buf = bytearray(4*len(int_col))
    convert_list(int_col, <int*>(&int_buf[0]))
    return int_buf.base

一个重要的细节是：convert_list 一定不能是 nogil（因为它不是）！ Omp 线程和 Python-threads（受 GIL 影响）是完全不同的东西。

可以（但不是必须）在使用具有 buffer-protocol 的对象时为 omp 操作释放 GIL - 因为这些对象通过缓冲区协议被锁定并且不能从不同的 Python 更改-线程。 list 没有这样的锁定机制，因此，如果 GIL 被释放，列表可能会在另一个线程中更改，我们所有的指针都可能失效。

现在开始计时（列表稍大）：

amount = 5*10**7 
ints = list(range(amount)) 


%timeit pack(f'{amount}i', *ints)  
# 1.51 s ± 38.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit pack_ints_DavidW(ints) 
# 284 ms ± 3.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit pack_ints_ead(ints) 
# 177 ms ± 11.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

顺便说一下，关闭 pack_ints_ead 的并行化会导致运行 209 毫秒的时间。

所以考虑到 ca 的适度改善。 33%，我会选择更强大的 DavidW 解决方案。

这里是用一种稍微不同的方式来表示错误值的实现：

不是整数对象会导致 -2147483648（即 0x80000000）- 32 位整数可以存储的最小负值。
整数 >=2147483647（即 >=0x7fffffff）将被映射 to/stored 为 2147483647 - 32 位整数可以存储的最大正数。
整数 <=-2147483647（即 <=0x80000001）将被映射 to/stored 为 -2147483647
所有其他整数都映射到它们的正确值。

主要优点是，它适用于更大范围的整数值。该算法产生的时间与第一个简单版本几乎相同运行（可能慢 2-3%）：

int to_int(PyObject *vv){ 
   if (PyLong_Check(vv)) {
       PyLongObject * v = (PyLongObject *)vv;
       Py_ssize_t i = Py_SIZE(v);
       int sign = i<0 ? -1 : 1;
       i = abs(i);
       if(i==0){
           return 0;
       }
       if(i==1){//small enought for a digit
           return sign*v->ob_digit[0];
       }
       if(i==2 && (v->ob_digit[1]>>1)==0){
           int add = (v->ob_digit[1]&1) << 30;
           return sign*(v->ob_digit[0]+add);
       }
       return sign * 0x7fffffff;
   }
   return 0x80000000;
}

使用 cython 比 struct.pack 快

Going faster than struct.pack with cython

python

struct

cython

pack