support simple SIMD detection

python · picnixz · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 6, 2024
commit 9a6ccb54c1052aa7fd2cd8d486e703c8da848d9d
diff --git a/Include/internal/pycore_cpuinfo.h b/Include/internal/pycore_cpuinfo.h
@@ -0,0 +1,26 @@
+#ifndef Py_INTERNAL_CPUINFO_H
+#define Py_INTERNAL_CPUINFO_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#include <stdbool.h>
+
+typedef struct {
+    bool sse, sse2, sse3, sse41, sse42, avx, avx2, avx512vbmi;
+    bool done;
+} cpu_simd_flags;
+
+extern void
+detect_cpu_simd_features(cpu_simd_flags *flags);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* !Py_INTERNAL_CPUINFO_H */
@@ -434,6 +434,7 @@ PYTHON_OBJS=	\
 		Python/context.o \
 		Python/critical_section.o \
 		Python/crossinterp.o \
+		Python/cpuinfo.o \
 		Python/dynamic_annotations.o \
 		Python/errors.o \
 		Python/flowgraph.o \
@@ -1191,6 +1192,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_complexobject.h \
 		$(srcdir)/Include/internal/pycore_condvar.h \
 		$(srcdir)/Include/internal/pycore_context.h \
+		$(srcdir)/Include/internal/pycore_cpuinfo.h \
 		$(srcdir)/Include/internal/pycore_critical_section.h \
 		$(srcdir)/Include/internal/pycore_crossinterp.h \
 		$(srcdir)/Include/internal/pycore_descrobject.h \

@@ -225,6 +225,7 @@
     <ClInclude Include="..\Include\internal\pycore_complexobject.h" />
     <ClInclude Include="..\Include\internal\pycore_condvar.h" />
     <ClInclude Include="..\Include\internal\pycore_context.h" />
+    <ClInclude Include="..\Include\internal\pycore_cpuinfo.h" />
     <ClInclude Include="..\Include\internal\pycore_critical_section.h" />
     <ClInclude Include="..\Include\internal\pycore_crossinterp.h" />
     <ClInclude Include="..\Include\internal\pycore_descrobject.h" />
@@ -584,6 +585,7 @@
     <ClCompile Include="..\Python\codegen.c" />
     <ClCompile Include="..\Python\compile.c" />
     <ClCompile Include="..\Python\context.c" />
+    <ClCompile Include="..\Python\cpuinfo.c" />
     <ClCompile Include="..\Python\critical_section.c" />
     <ClCompile Include="..\Python\crossinterp.c" />
     <ClCompile Include="..\Python\dynamic_annotations.c" />

@@ -597,6 +597,9 @@
     <ClInclude Include="..\Include\internal\pycore_context.h">
       <Filter>Include\internal</Filter>
     </ClInclude>
+    <ClInclude Include="..\Include\internal\pycore_cpuinfo.h">
+      <Filter>Include\cpython</Filter>
+    </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_critical_section.h">
       <Filter>Include\internal</Filter>
     </ClInclude>
@@ -1304,6 +1307,9 @@
     <ClCompile Include="..\Python\compile.c">
       <Filter>Python</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\cpuinfo.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="..\Python\critical_section.c">
       <Filter>Python</Filter>
     </ClCompile>

diff --git a/Python/cpuinfo.c b/Python/cpuinfo.c
@@ -0,0 +1,110 @@
+/*
+ * Naive CPU SIMD features detection.
+ *
+ * See Modules/black2module.c.
+ */
+
+#include "Python.h"
+#include "pycore_cpuinfo.h"
+
+#include <stdbool.h>
+
+#if defined(__x86_64__) 
F438
&& defined(__GNUC__)
+#include <cpuid.h>
+#elif defined(_M_X64)
+#include <intrin.h>
+#endif
+
+// AVX2 cannot be compiled on macOS ARM64 (yet it can be compiled on x86_64).
+// However, since autoconf incorrectly assumes so when compiling a universal2
+// binary, we disable all AVX-related instructions.
+#if defined(__APPLE__) && defined(__arm64__)
+#  undef CAN_COMPILE_SIMD_AVX_INSTRUCTIONS
+#  undef CAN_COMPILE_SIMD_AVX2_INSTRUCTIONS
+#  undef CAN_COMPILE_SIMD_AVX512_VBMI_INSTRUCTIONS
+#endif
+
+#define EDX1_SSE            (1 << 25)   // sse, EDX, page 1, bit 25
+#define EDX1_SSE2           (1 << 26)   // sse2, EDX, page 1, bit 26
+#define ECX1_SSE3           (1 << 9)    // sse3, ECX, page 1, bit 0
+#define ECX1_SSE4_1         (1 << 19)   // sse4.1, ECX, page 1, bit 19
+#define ECX1_SSE4_2         (1 << 20)   // sse4.2, ECX, page 1, bit 20
+#define ECX1_AVX            (1 << 28)   // avx, ECX, page 1, bit 28
+#define EBX7_AVX2           (1 << 5)    // avx2, EBX, page 7, bit 5
+#define ECX7_AVX512_VBMI    (1 << 1)    // avx512-vbmi, ECX, page 7, bit 1
+
+void
+detect_cpu_simd_features(cpu_simd_flags *flags)
+{
+    if (flags->done) {
+        return;
+    }
+
+    int eax1 = 0, ebx1 = 0, ecx1 = 0, edx1 = 0;
+    int eax7 = 0, ebx7 = 0, ecx7 = 0, edx7 = 0;
+#if defined(__x86_64__) && defined(__GNUC__)
+    __cpuid_count(1, 0, eax1, ebx1, ecx1, edx1);
+    __cpuid_count(7, 0, eax7, ebx7, ecx7, edx7);
+#elif defined(_M_X64)
+    int info1[4] = {0};
+    __cpuidex(info1, 1, 0);
+    eax1 = info1[0];
+    ebx1 = info1[1];
+    ecx1 = info1[2];
+    edx1 = info1[3];
+
+    int info7[4] = {0};
+    __cpuidex(info7, 7, 0);
+    eax7 = info7[0];
+    ebx7 = info7[1];
+    ecx7 = info7[2];
+    edx7 = info7[3];
+#else
+    // use (void) expressions to avoid warnings
+    (void) eax1; (void) ebx1; (void) ecx1; (void) edx1;
+    (void) eax7; (void) ebx7; (void) ecx7; (void) edx7;
+#endif
+
+#ifdef CAN_COMPILE_SIMD_SSE_INSTRUCTIONS
+    flags->sse = (edx1 & EDX1_SSE) != 0;
+#else
+    flags->sse = false;
+#endif
+#ifdef CAN_COMPILE_SIMD_SSE2_INSTRUCTIONS
+    flags->sse2 = (edx1 & EDX1_SSE2) != 0;
+#else
+    flags->sse2 = false;
+#endif
+#ifdef CAN_COMPILE_SIMD_SSE3_INSTRUCTIONS
+    flags->sse3 = (ecx1 & ECX1_SSE3) != 0;
+    #else
+#endif
+    flags->sse3 = false;
+#ifdef CAN_COMPILE_SIMD_SSE4_1_INSTRUCTIONS
+    flags->sse41 = (ecx1 & ECX1_SSE4_1) != 0;
+#else
+    flags->sse41 = false;
+#endif
+#ifdef CAN_COMPILE_SIMD_SSE4_2_INSTRUCTIONS
+    flags->sse42 = (ecx1 & ECX1_SSE4_2) != 0;
+#else
+    flags->sse42 = false;
+#endif
+#ifdef CAN_COMPILE_SIMD_AVX_INSTRUCTIONS
+    flags->avx = (ecx1 & ECX1_AVX) != 0;
+#else
+    flags->avx = false;
+#endif
+#ifdef CAN_COMPILE_SIMD_AVX2_INSTRUCTIONS
+    flags->avx2 = (ebx7 & EBX7_AVX2) != 0;
+#else
+    flags->avx2 = false;
+#endif
+#ifdef CAN_COMPILE_SIMD_AVX512_VBMI_INSTRUCTIONS
+    flags->avx512vbmi = (ecx7 & ECX7_AVX512_VBMI) != 0;
+#else
+    flags->avx512vbmi = false;
+#endif
+
+    flags->done = true;
+}