gh-125022: add support for simple SIMD features detection

Include/internal/pycore_cpuinfo.h

@@ -0,0 +1,179 @@
+/*
+ * @author Bénédikt Tran
+ *
+ * Interface for detecting the different CPUID flags in an opaque manner.
+ * See https://en.wikipedia.org/wiki/CPUID for details on the bit values.
+ *
+ * If a module requires to support SIMD instructions, it should determine
+ * the compiler flags and the instruction sets required for the intrinsics
+ * to work.
+ *
+ * For the headers and expected CPUID bits needed by Intel intrinsics, see
+ * https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html.
+ */
+
+#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 "Python.h"
+#include "pycore_cpuinfo_cpuid_features.h"
+#include "pycore_cpuinfo_xsave_features.h"
+
+typedef struct py_cpuid_features {
+    uint32_t maxleaf;
+    /*
+     * Macro to declare a member flag of 'py_cpuid_features' as a uint8_t.
+     * Whenever this macro is used, do not forget to update the number of
+     * fields and the bitsize of the 'ready' member (see structure end).
+     */
+#define _Py_CPUID_DECL_FLAG(MEMBER_NAME)    uint8_t MEMBER_NAME:1
+    // --- Streaming SIMD Extensions ------------------------------------------
+    _Py_CPUID_DECL_FLAG(sse);
+    _Py_CPUID_DECL_FLAG(sse2);
+    _Py_CPUID_DECL_FLAG(sse3);
+    _Py_CPUID_DECL_FLAG(ssse3); // Supplemental SSE3 instructions
+    _Py_CPUID_DECL_FLAG(sse41); // SSE4.1
+    _Py_CPUID_DECL_FLAG(sse42); // SSE4.2
+
+    // --- Advanced Vector Extensions -----------------------------------------
+    _Py_CPUID_DECL_FLAG(avx);
+    _Py_CPUID_DECL_FLAG(avx_ifma);
+    _Py_CPUID_DECL_FLAG(avx_ne_convert);
+
+    _Py_CPUID_DECL_FLAG(avx_vnni);
+    _Py_CPUID_DECL_FLAG(avx_vnni_int8);
+    _Py_CPUID_DECL_FLAG(avx_vnni_int16);
+
+    // --- Advanced Vector Extensions 2 ---------------------------------------
+    _Py_CPUID_DECL_FLAG(avx2);
+
+    // --- Advanced Vector Extensions (512-bit) -------------------------------
+    /*
+     * AVX-512 instruction set are grouped by the processor generation
+     * that implements them (see https://en.wikipedia.org/wiki/AVX-512).
+     *
+     * We do not include GFNI, VPCLMULQDQ and VAES instructions since
+     * they are not exactly AVX-512 per se, nor do we include BF16 or
+     * FP16 since they operate on bfloat16 and binary16 (half-float).
+     *
+     * See https://en.wikipedia.org/wiki/AVX-512#Instruction_set for
+     * the suffix meanings (for instance 'f' stands for 'Foundation').
+     */
+    _Py_CPUID_DECL_FLAG(avx512_f);
+    _Py_CPUID_DECL_FLAG(avx512_cd);
+
+    _Py_CPUID_DECL_FLAG(avx512_er);
+    _Py_CPUID_DECL_FLAG(avx512_pf);
+
+    _Py_CPUID_DECL_FLAG(avx512_4fmaps);
+    _Py_CPUID_DECL_FLAG(avx512_4vnniw);
+
+    _Py_CPUID_DECL_FLAG(avx512_vpopcntdq);
+
+    _Py_CPUID_DECL_FLAG(avx512_vl);
+    _Py_CPUID_DECL_FLAG(avx512_dq);
+    _Py_CPUID_DECL_FLAG(avx512_bw);
+
+    _Py_CPUID_DECL_FLAG(avx512_ifma);
+    _Py_CPUID_DECL_FLAG(avx512_vbmi);
+
+    _Py_CPUID_DECL_FLAG(avx512_vnni);
+
+    _Py_CPUID_DECL_FLAG(avx512_vbmi2);
+    _Py_CPUID_DECL_FLAG(avx512_bitalg);
+
+    _Py_CPUID_DECL_FLAG(avx512_vp2intersect);
+
+    // --- Instructions -------------------------------------------------------
+    _Py_CPUID_DECL_FLAG(cmov);
+    _Py_CPUID_DECL_FLAG(fma);
+    _Py_CPUID_DECL_FLAG(popcnt);
+    _Py_CPUID_DECL_FLAG(pclmulqdq);
+
+    _Py_CPUID_DECL_FLAG(xsave);   // XSAVE/XRSTOR/XSETBV/XGETBV
+    _Py_CPUID_DECL_FLAG(osxsave); // XSAVE is enabled by the OS
+
+    // --- XCR0 register bits -------------------------------------------------
+    _Py_CPUID_DECL_FLAG(xcr0_sse);
+    // On some Intel CPUs, it is possible for the CPU to support AVX2
+    // instructions even though the underlying OS does not know about
+    // AVX. In particular, only (SSE) XMM registers will be saved and
+    // restored on context-switch, but not (AVX) YMM registers.
+    _Py_CPUID_DECL_FLAG(xcr0_avx);
+    _Py_CPUID_DECL_FLAG(xcr0_avx512_opmask);
+    _Py_CPUID_DECL_FLAG(xcr0_avx512_zmm_hi256);
+    _Py_CPUID_DECL_FLAG(xcr0_avx512_hi16_zmm);
+#undef _Py_CPUID_DECL_FLAG
+    // Whenever a field is added or removed above, update the
+    // number of fields (40) and adjust the bitsize of 'ready'
+    // so that the size of this structure is a multiple of 8.
+    uint8_t ready; // set if the structure is ready for usage
+} py_cpuid_features;
+
+/*
+ * Explicitly initialize all members to zero to guarantee that
+ * we never have an un-initialized attribute at runtime which
+ * could lead to an illegal instruction error.
+ *
+ * This does not mark 'flags' as being ready yet.
+ *
+ * Note: This function does not set any exception and thus never fails.
+ */
+PyAPI_FUNC(void)
+_Py_cpuid_disable_features(py_cpuid_features *flags);
+
+/*
+ * Check whether the structure is ready and flags are inter-compatible,
+ * returning 1 on success and 0 otherwise.
+ *
+ * The caller should disable all CPUID detected features if the check
+ * fails to avoid encountering runtime illegal instruction errors.
+ *
+ * Note: This function does not set any exception and thus never fails.
+ */
+PyAPI_FUNC(int)
+_Py_cpuid_check_features(const py_cpuid_features *flags);
+
+/*
+ * Return 1 if all expected flags are set in 'actual', 0 otherwise.
+ *
+ * If 'actual' or 'expect' are not ready yet, this also returns 0.
+ *
+ * Note: This function does not set any exception and thus never fails.
+ */
+PyAPI_FUNC(int)
+_Py_cpuid_has_features(const py_cpuid_features *actual,
+                       const py_cpuid_features *expect);
+
+/*
+ * Return 1 if 'actual' and 'expect' are identical, 0 otherwise.
+ *
+ * If 'actual' or 'expect' are not ready yet, this also returns 0.
+ *
+ * Note: This function does not set any exception and thus never fails.
+ */
+PyAPI_FUNC(int)
+_Py_cpuid_match_features(const py_cpuid_features *actual,
+                         const py_cpuid_features *expect);
+
+/*
+ * Detect the available features on this machine, storing the result in 'flags'.
+ *
+ * Note: This function does not set any exception and thus never fails.
+ */
+PyAPI_FUNC(void)
+_Py_cpuid_detect_features(py_cpuid_features *flags);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* !Py_INTERNAL_CPUINFO_H */

Include/internal/pycore_cpuinfo_cpuid_features.h

@@ -0,0 +1,102 @@
+/*
+ * @author Bénédikt Tran
+ * @seealso Tools/cpuinfo/cpuid_features_gen.py
+ *
+ * The enumeration describes masks to apply on CPUID output registers.
+ *
+ * Member names are Py_CPUID_MASK_<REGISTER>_L<LEAF>[S<SUBLEAF>]_<FEATURE>,
+ * where <> (resp. []) denotes a required (resp. optional) group and:
+ *
+ * - REGISTER is EAX, EBX, ECX or EDX,
+ * - LEAF is the initial value of the EAX register (1 or 7),
+ * - SUBLEAF is the initial value of the ECX register (omitted if 0), and
+ * - FEATURE is a SIMD feature (with one or more specialized instructions).
+ *
+ * For maintainability, the flags are ordered by registers, leafs, subleafs,
+ * and bits. See https://en.wikipedia.org/wiki/CPUID for the values.
+ *
+ * Note 1: The LEAF is also called the 'page' or the 'level'.
+ * Note 2: The SUBLEAF is also referred to as the 'count'.
+ *
+ * The LEAF value should only 1 or 7 as other values may have different
+ * meanings depending on the underlying architecture.
+ */
+
+#ifndef Py_INTERNAL_CPUINFO_CPUID_FEATURES_H
+#define Py_INTERNAL_CPUINFO_CPUID_FEATURES_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#include "Python.h"
+
+// fmt: off
+/*[python input]
+import importlib
+import os
+import sys
+
+ROOT = os.getcwd()
+TOOL = os.path.join(ROOT, "Tools/cpuinfo/cpuid_features_gen.py")
+TOOL = os.path.realpath(TOOL)
+
+if not os.path.exists(TOOL):
+    raise FileNotFoundError(TOOL)
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(TOOL)))
+module = importlib.import_module("cpuinfo.cpuid_features_gen")
+print(module.generate_cpuid_features_enum("py_cpuid_feature_mask"))
+[python start generated code]*/
+typedef enum py_cpuid_feature_mask {
+    /* CPUID (LEAF=1, SUBLEAF=0) [ECX] */
+    Py_CPUID_MASK_ECX_L1_SSE3                   = 0x00000001,    // bit = 0
+    Py_CPUID_MASK_ECX_L1_PCLMULQDQ              = 0x00000002,    // bit = 1
+    Py_CPUID_MASK_ECX_L1_SSSE3                  = 0x00000200,    // bit = 9
+    Py_CPUID_MASK_ECX_L1_FMA                    = 0x00001000,    // bit = 12
+    Py_CPUID_MASK_ECX_L1_SSE4_1                 = 0x00080000,    // bit = 19
+    Py_CPUID_MASK_ECX_L1_SSE4_2                 = 0x00100000,    // bit = 20
+    Py_CPUID_MASK_ECX_L1_POPCNT                 = 0x00800000,    // bit = 23
+    Py_CPUID_MASK_ECX_L1_XSAVE                  = 0x04000000,    // bit = 26
+    Py_CPUID_MASK_ECX_L1_OSXSAVE                = 0x08000000,    // bit = 27
+    Py_CPUID_MASK_ECX_L1_AVX                    = 0x10000000,    // bit = 28
+    /* CPUID (LEAF=1, SUBLEAF=0) [EDX] */
+    Py_CPUID_MASK_EDX_L1_CMOV                   = 0x00008000,    // bit = 15
+    Py_CPUID_MASK_EDX_L1_SSE                    = 0x02000000,    // bit = 25
+    Py_CPUID_MASK_EDX_L1_SSE2                   = 0x04000000,    // bit = 26
+    /* CPUID (LEAF=7, SUBLEAF=0) [EBX] */
+    Py_CPUID_MASK_EBX_L7_AVX2                   = 0x00000020,    // bit = 5
+    Py_CPUID_MASK_EBX_L7_AVX512_F               = 0x00010000,    // bit = 16
+    Py_CPUID_MASK_EBX_L7_AVX512_DQ              = 0x00020000,    // bit = 17
+    Py_CPUID_MASK_EBX_L7_AVX512_IFMA            = 0x00200000,    // bit = 21
+    Py_CPUID_MASK_EBX_L7_AVX512_PF              = 0x04000000,    // bit = 26
+    Py_CPUID_MASK_EBX_L7_AVX512_ER              = 0x08000000,    // bit = 27
+    Py_CPUID_MASK_EBX_L7_AVX512_CD              = 0x10000000,    // bit = 28
+    Py_CPUID_MASK_EBX_L7_AVX512_BW              = 0x40000000,    // bit = 30
+    Py_CPUID_MASK_EBX_L7_AVX512_VL              = 0x80000000,    // bit = 31
+    /* CPUID (LEAF=7, SUBLEAF=0) [ECX] */
+    Py_CPUID_MASK_ECX_L7_AVX512_VBMI            = 0x00000002,    // bit = 1
+    Py_CPUID_MASK_ECX_L7_AVX512_VBMI2           = 0x00000040,    // bit = 6
+    Py_CPUID_MASK_ECX_L7_AVX512_VNNI            = 0x00000800,    // bit = 11
+    Py_CPUID_MASK_ECX_L7_AVX512_BITALG          = 0x00001000,    // bit = 12
+    Py_CPUID_MASK_ECX_L7_AVX512_VPOPCNTDQ       = 0x00004000,    // bit = 14
+    /* CPUID (LEAF=7, SUBLEAF=0) [EDX] */
+    Py_CPUID_MASK_EDX_L7_AVX512_4VNNIW          = 0x00000004,    // bit = 2
+    Py_CPUID_MASK_EDX_L7_AVX512_4FMAPS          = 0x00000008,    // bit = 3
+    Py_CPUID_MASK_EDX_L7_AVX512_VP2INTERSECT    = 0x00000100,    // bit = 8
+    /* CPUID (LEAF=7, SUBLEAF=1) [EAX] */
+    Py_CPUID_MASK_EAX_L7S1_AVX_VNNI             = 0x00000010,    // bit = 4
+    Py_CPUID_MASK_EAX_L7S1_AVX_IFMA             = 0x00800000,    // bit = 23
+    /* CPUID (LEAF=7, SUBLEAF=1) [EDX] */
+    Py_CPUID_MASK_EDX_L7S1_AVX_VNNI_INT8        = 0x00000010,    // bit = 4
+    Py_CPUID_MASK_EDX_L7S1_AVX_NE_CONVERT       = 0x00000020,    // bit = 5
+    Py_CPUID_MASK_EDX_L7S1_AVX_VNNI_INT16       = 0x00000400,    // bit = 10
+} py_cpuid_feature_mask;
+/*[python end generated code: output=c4460242e465fa91 input=61d2b5f1bc368b94]*/
+// fmt: on
+
+#endif // !Py_INTERNAL_CPUINFO_CPUID_FEATURES_H

Include/internal/pycore_cpuinfo_xsave_features.h

@@ -0,0 +1,49 @@
+/*
+ * @author Bénédikt Tran
+ * @seealso Tools/cpuinfo/xsave_features_gen.py
+ *
+ * XSAVE state components (XCR0 control register).
+ *
+ * See https://en.wikipedia.org/wiki/Control_register#XCR0_and_XSS.
+ */
+#ifndef Py_INTERNAL_CPUINFO_XSAVE_FEATURES_H
+#define Py_INTERNAL_CPUINFO_XSAVE_FEATURES_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#include "Python.h"
+
+// fmt: off
+/*[python input]
+import importlib
+import os
+import sys
+
+ROOT = os.getcwd()
+TOOL = os.path.join(ROOT, "Tools/cpuinfo/xsave_features_gen.py")
+TOOL = os.path.realpath(TOOL)
+
+if not os.path.exists(TOOL):
+    raise FileNotFoundError(TOOL)
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(TOOL)))
+module = importlib.import_module("cpuinfo.xsave_features_gen")
+print(module.generate_xsave_features_enum("py_xsave_feature_mask"))
+[python start generated code]*/
+typedef enum py_xsave_feature_mask {
+    Py_XSAVE_MASK_XCR0_SSE              = 0x00000002,    // bit = 1
+    Py_XSAVE_MASK_XCR0_AVX              = 0x00000004,    // bit = 2
+    Py_XSAVE_MASK_XCR0_AVX512_OPMASK    = 0x00000020,    // bit = 5
+    Py_XSAVE_MASK_XCR0_AVX512_ZMM_HI256 = 0x00000040,    // bit = 6
+    Py_XSAVE_MASK_XCR0_AVX512_HI16_ZMM  = 0x00000080,    // bit = 7
+} py_xsave_feature_mask;
+/*[python end generated code: output=9a476ed0abbc617b input=41f35058299c0118]*/
+// fmt: on
+
+#endif // !Py_INTERNAL_CPUINFO_XSAVE_FEATURES_H

Makefile.pre.in

@@ -435,6 +435,7 @@ PYTHON_OBJS=	\
 		Python/codegen.o \
 		Python/compile.o \
 		Python/context.o \
+		Python/cpuinfo.o \
 		Python/critical_section.o \
 		Python/crossinterp.o \
 		Python/dynamic_annotations.o \
@@ -1286,6 +1287,9 @@ 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_cpuinfo_cpuid_features.h \
+		$(srcdir)/Include/internal/pycore_cpuinfo_xsave_features.h \
 		$(srcdir)/Include/internal/pycore_critical_section.h \
 		$(srcdir)/Include/internal/pycore_crossinterp.h \
 		$(srcdir)/Include/internal/pycore_crossinterp_data_registry.h \