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

我有一个带有本机扩展的 Python 包，由Cython https://cython.org/。由于一些性能需要，编译是用-march=native, -mtune=native旗帜。这基本上使编译器能够使用任何可用的 ISA 扩展。

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

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

• Cythonized 轮是为非常特定的平台构建的
• 纯蟒蛇轮。

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

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

原生轮子确实有cp37和平台名称，但我在这里没有看到定义更细粒度 ISA 要求的方法。我可以随时使用--implementationpip 的标志，但我想知道 pip 是否有更好的方法来区分不同的 ISA。

Thanks,

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

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

这是一个概念证明。

第一个要跳的环节：如何在运行时检查 CPU/OS 组合支持哪些指令。为了简单起见，我们将检查 AVX（这SO-post https://stackoverflow.com/q/6121792/5769463有更多详细信息）并且我仅提供特定于 gcc 的（另请参阅this https://gcc.gnu.org/bugzilla/show_bug.cgi?id=85100) 解决方案 - 称为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 will 产生不同的结果 https://stackoverflow.com/q/28939652/5769463，取决于它是否是用编译的-march=native or not.

现在我们只需包含 *.pxi 文件中的实际代码即可定义该模块的两个版本。一个模块称为worker_native.pyx:

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

include "worker.pxi"

and 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 https://stackoverflow.com/a/59364990/5769463，如何实现这一点）。