用Cython比struct.pack更快?

3

我想做得比 struct.pack 更好。

针对整数打包的特定情况,通过回答此问题,我有以下内容来打包 pack_ints.pyx 中的int列表:

# 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,以使此操作更快?


我怀疑你无法再加快速度了。为什么你认为有可能执行更快的操作呢? - ead
例如,我猜想如果有人通过一些C代码获取到列表对象的指针,他可以打开线程并并行工作,根据核心数量的增加而变得更快。我不知道如何在cython中实现这一点,无法在此函数上使用@nogil - Jay
2
我认为你需要GIL来将Python对象转换为C整数(例如,它可能会引发异常)。 - ead
我建议你做的更有用的事情是修复输入数据。你能否直接生成打包数据,而不是浪费时间生成大量的Python对象,然后再进行转换? - DavidW
1
...以及PyLong_AsLong的实现。您需要调整实现方式,使其不使用GIL(即不能引发异常)。 - DavidW
显示剩余3条评论
2个回答

4
这个答案试图给出一个估计,平行化版本可以获得多少加速。然而,由于这个任务受到内存带宽的限制(Python整数对象至少占用32个字节,并且可能散布在内存中,因此会有许多缓存未命中),我们不应该期望太多。
第一个问题是如何处理错误(元素不是整数或值过大)。我将采用以下策略/简化:当对象
  • 不是整数,
  • 是负整数,
  • 或者整数 >=2^30
时,它将被转换为一个特殊数字(-1),表示发生了错误。只允许非负整数 <2^30 可以使我的工作更容易,因为我必须重新实现PyLong_AsLongAndOverflow而不引发错误,否则检测溢出通常很麻烦(但请参见答案末尾的更复杂方法)。
Python整数对象的内存布局可以在这里找到:
struct _longobject {
    PyObject_VAR_HEAD
    digit ob_digit[1];
};

成员变量ob_size/宏Py_SIZE告诉我们在整数的表示中使用了多少个30位数字(对于负整数,ob_size为负数)。

因此,我的简单规则可以转换为以下C代码(我使用C而不是Cython,因为它是使用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;
}

现在,给定一个列表,我们可以使用以下的C函数并行地将其转换为int缓冲区,该函数使用omp:
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 函数中没有引用计数/竞争条件。

现在,通过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线程和受GIL影响的Python线程是完全不同的东西。
可以(但不一定)在使用具有缓冲区协议的对象时释放GIL以进行omp操作 - 因为这些对象通过缓冲区协议被锁定,不能从不同的Python线程更改。列表没有这样的锁定机制,因此,如果释放了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毫秒。

因此,考虑到约33%的小幅改进,我会选择更健壮的DavidW的方案。


这里是使用略有不同的方式来表示错误值的实现:

  • 非整数对象的结果为-2147483648(即 0x80000000 )- 32位整数可以存储的最小负值。
  • 大于等于2147483647(即 >=0x7fffffff )的整数将被映射/存储为 2147483647 - 32位整数可以存储的最大正数。
  • 小于等于-2147483647(即 <=0x80000001 )的整数将被映射/存储为 -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;
}

谢谢!这段 C 代码能被 Cython 编译吗?另外,如果我使用 tuple() 将列表转换为元组,那么我可以使用 nogil 吗? - Jay
1
@Jay C 代码被合并到 Cython 生成的 c 文件中,并同时编译。将其转换为 tuple 对于 nogil 没有影响,但在这里 nogil 并不重要——循环会像您想要的那样并行运行(由于线程仅进行读取而不是更改,因此没有线程需要 GIL,但其中一个线程应该具有 GIL,以便其他任何事情都不能并行运行和更改——我同意这有点令人困惑...) - DavidW
@Jay 我使用https://cython.readthedocs.io/en/latest/src/userguide/external_C_code.html#including-verbatim-c-code来直接包含C代码。您可以对实现缓冲区协议的对象使用nogil(https://docs.python.org/3/c-api/buffer.html),但元组和列表都不行。然而,正如DavidW所指出的那样,您并不需要这样做——计算已经在并行中进行了。 - ead

2
当我运行原始问题中的代码时,速度提升了约5倍。
当我在这里运行你的代码时,我看到你报告的结果以及一个重要的警告在编译阶段,我认为你忽略了它:
“警告:pack_ints.pyx:13:17:索引应该被标记为更有效的访问。”
我不确定为什么它没有正确地识别类型,但是为了修复它,你应该将“i”的定义改回我最初编写的代码。
cdef int i
# not "i: int"

希望其他人能够尝试更聪明的方法,因为显然这个答案有点荒谬。



我相信当我检查注释时,这行代码是白色的,所以我没有关注警告。如果它跳到X5,我一定会换成cdef。 - Jay
1
@Jay 改变那行代码会使循环在带注释的HTML中变成黄色。而该行代码本身仍然保持为白色。 - DavidW

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接