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

Question

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

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

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

为特定平台构建的 Cythonized wheel
纯-python轮.

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

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

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

谢谢，

Answer 1

pip 基础设施不支持这种粒度。

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

这是一个概念证明。

第一个要跳的箍：如何在运行时检查CPU/OS组合支持哪些指令。为简单起见，我们将检查 AVX（此 SO-post has more details) and I offer only a gcc-specific (see also 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文件和模块必须同名。为避免代码重复，实际代码在 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 将，这取决于它是否使用 -march=native 编译。

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

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

include "worker.pxi"

和worker_fallback.pyx：

include "worker.pxi"

构建一切，例如via 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。

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

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

Specifying Exact CPU Instruction Set with Cythonized Python Wheels

python

cython