使用 Cythonized Python Wheels 指定精确的 CPU 指令集答案

【问题标题】：Specifying Exact CPU Instruction Set with Cythonized Python Wheels使用 Cythonized Python Wheels 指定精确的 CPU 指令集
【发布时间】：2021-12-26 18:14:27
【问题描述】：

我有一个带有由Cython 编译的本机扩展的 Python 包。由于某些性能需求，编译使用-march=native, -mtune=native 标志完成。这基本上使编译器能够使用任何可用的 ISA 扩展。

此外，我们保留了这个包的非 cythonized 纯 python 版本。它应该用于对性能不太敏感的环境中。

因此，我们总共发布了两个版本：

专为特定平台打造的 Cythonized 车轮
纯蟒蛇轮。

一些包依赖于这个包，并且一些机器与编译包的机器有点不同。由于我们使用了-march=native，因此我们得到SIGILL，因为服务器上缺少某些ISA 扩展。

所以，本质上，如果主机 CPU 与轮子不兼容，我想以某种方式让 pip 忽略本机轮子。

本机轮子确实具有cp37 和平台名称，但我看不到在此处定义更精细的 ISA 要求的方法。我总是可以为 pip 使用--implementation 标志，但我想知道是否有更好的方法让 pip 区分不同的 ISA。

谢谢，

【问题讨论】：

标签： python cython

【解决方案1】：

pip 基础架构不支持这种粒度。

我认为更好的方法是编译两个版本的 Cython 扩展：使用 -march=native 和不使用 -march=native，安装两者并在运行时决定应该加载哪个版本。

这是一个概念证明。

第一个跳转：如何在运行时检查 CPU/OS 组合支持哪些指令。为简单起见，我们将检查 AVX（此 SO-post 有更多详细信息），我仅提供 gcc 特定（另请参阅 this）解决方案 - 称为 impl_picker.pyx：

cdef extern from *:
    """
    int cpu_supports_avx(void){
        return __builtin_cpu_supports("avx");
    }
    """
    int cpu_supports_avx()

def cpu_has_avx_support():
    return cpu_supports_avx() != 0

第二个问题：pyx-file和module必须同名。为避免代码重复，实际代码在 pxi 文件中：

# worker.pxi
cdef extern from *:
    """   
    int compiled_with_avx(void){
        #ifdef __AVX__
            return 1;
        #else
            return 0;
        #endif
    }
    """
    int compiled_with_avx()

def compiled_with_avx_support():
    return compiled_with_avx() != 0

可以看到，函数compiled_with_avx_support 将yield different results，这取决于它是否使用-march=native 编译。

现在我们可以通过包含 *.pxi 文件中的实际代码来定义模块的两个版本。一个叫做worker_native.pyx的模块：

# distutils: extra_compile_args=["-march=native"]

include "worker.pxi"

和worker_fallback.pyx:

include "worker.pxi"

构建一切，例如通过cythonize -i -3 *.pyx，可以这样使用：

from impl_picker import cpu_has_avx_support

# overhead once when imported:
if cpu_has_avx_support():
    import worker_native as worker
else:
    print("using fallback worker")
    import worker_fallback as worker

print("compiled_with_avx_support:", worker.compiled_with_avx_support())

在我的机器上，以上将导致compiled_with_avx_support: True，在旧机器上，将使用“较慢”的worker_fallback，结果将是compiled_with_avx_support: False。

这篇文章的目的不是提供一个有效的setup.py，而只是概述如何实现在运行时选择正确版本的目标。显然，setup.py 可能要复杂得多：例如需要使用不同的编译器设置编译多个 c 文件（参见 SO-post，如何实现）。

【讨论】：