TMP: usage proof-of-concept

Author: picnixz

Modules/blake2module.c

@@ -16,29 +16,11 @@
 #include "pyconfig.h"
 #include "Python.h"
 #include "hashlib.h"
-#include "pycore_strhex.h"       // _Py_strhex()
+#include "pycore_cpuinfo.h"         // py_cpuid_features
+#include "pycore_strhex.h"          // _Py_strhex()
 #include "pycore_typeobject.h"
 #include "pycore_moduleobject.h"
 
-// QUICK CPU AUTODETECTION
-//
-// See https://github.com/python/cpython/pull/119316 -- we only enable
-// vectorized versions for Intel CPUs, even though HACL*'s "vec128" modules also
-// run on ARM NEON. (We could enable them on POWER -- but I don't have access to
-// a test machine to see if that speeds anything up.)
-//
-// Note that configure.ac and the rest of the build are written in such a way
-// that if the configure script finds suitable flags to compile HACL's SIMD128
-// (resp. SIMD256) files, then Hacl_Hash_Blake2b_Simd128.c (resp. ...) will be
-// pulled into the build automatically, and then only the CPU autodetection will
-// need to be updated here.
-
-#if defined(__x86_64__) && defined(__GNUC__)
-#include <cpuid.h>
-#elif defined(_M_X64)
-#include <intrin.h>
-#endif
-
 #include <stdbool.h>
 
 // SIMD256 can't be compiled on macOS ARM64, and performance of SIMD128 isn't
@@ -51,83 +33,6 @@
 #  undef HACL_CAN_COMPILE_SIMD256
 #endif
 
-// ECX
-#define ECX_SSE3 (1 << 0)
-#define ECX_SSSE3 (1 << 9)
-#define ECX_SSE4_1 (1 << 19)
-#define ECX_SSE4_2 (1 << 20)
-#define ECX_AVX (1 << 28)
-
-// EBX
-#define EBX_AVX2 (1 << 5)
-
-// EDX
-#define EDX_SSE (1 << 25)
-#define EDX_SSE2 (1 << 26)
-#define EDX_CMOV (1 << 15)
-
-// zero-initialized by default
-typedef struct {
-    bool sse, sse2, sse3, sse41, sse42, cmov, avx, avx2;
-    bool done;
-} cpu_flags;
-
-void detect_cpu_features(cpu_flags *flags) {
-  if (!flags->done) {
-    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 };
-    int info7[4] = { 0 };
-    __cpuidex(info1, 1, 0);
-    __cpuidex(info7, 7, 0);
-    eax1 = info1[0];
-    ebx1 = info1[1];
-    ecx1 = info1[2];
-    edx1 = info1[3];
-    eax7 = info7[0];
-    ebx7 = info7[1];
-    ecx7 = info7[2];
-    edx7 = info7[3];
-#endif
-    (void) eax1; (void) ebx1; (void) ecx1; (void) edx1;
-    (void) eax7; (void) ebx7; (void) ecx7; (void) edx7;
-
-
-    flags->avx = (ecx1 & ECX_AVX) != 0;
-
-    flags->avx2 = (ebx7 & EBX_AVX2) != 0;
-
-    flags->sse = (edx1 & EDX_SSE) != 0;
-    flags->sse2 = (edx1 & EDX_SSE2) != 0;
-    flags->cmov = (edx1 & EDX_CMOV) != 0;
-
-    flags->sse3 = (ecx1 & ECX_SSE3) != 0;
-    /* ssse3 = (ecx1 & ECX_SSSE3) != 0; */
-    flags->sse41 = (ecx1 & ECX_SSE4_1) != 0;
-    flags->sse42 = (ecx1 & ECX_SSE4_2) != 0;
-
-    flags->done = true;
-  }
-}
-
-#ifdef HACL_CAN_COMPILE_SIMD128
-static inline bool has_simd128(cpu_flags *flags) {
-  // For now this is Intel-only, could conceivably be #ifdef'd to something
-  // else.
-  return flags->sse && flags->sse2 && flags->sse3 && flags->sse41 && flags->sse42 && flags->cmov;
-}
-#endif
-
-#ifdef HACL_CAN_COMPILE_SIMD256
-static inline bool has_simd256(cpu_flags *flags) {
-  return flags->avx && flags->avx2;
-}
-#endif
-
 // Small mismatch between the variable names Python defines as part of configure
 // at the ones HACL* expects to be set in order to enable those headers.
 #define HACL_CAN_COMPILE_VEC128 HACL_CAN_COMPILE_SIMD128
@@ -154,9 +59,31 @@ PyDoc_STRVAR(blake2mod__doc__,
 typedef struct {
     PyTypeObject* blake2b_type;
     PyTypeObject* blake2s_type;
-    cpu_flags flags;
+
+    bool can_run_simd128;
+    bool can_run_simd256;
 } Blake2State;
 
+static void
+blake2_init_cpu_features(Blake2State *state)
+{
+    py_cpuid_features flags;
+    _Py_cpuid_detect_features(&flags);
+#if HACL_CAN_COMPILE_SIMD128
+    state->can_run_simd128 = flags.sse && flags.sse2 && flags.sse3
+                             && flags.sse41 && flags.sse42
+                             && flags.cmov;
+#else
+    state->can_run_simd128 = false;
+#endif
+
+#if HACL_CAN_COMPILE_SIMD256
+    state->can_run_simd256 = flags.avx && flags.avx2;
+#else
+    state->can_run_simd256 = false;
+#endif
+}
+
 static inline Blake2State*
 blake2_get_state(PyObject *module)
 {
@@ -224,10 +151,7 @@ static int
 blake2_exec(PyObject *m)
 {
     Blake2State* st = blake2_get_state(m);
-
-    // This is called at module initialization-time, and so appears to be as
-    // good a place as any to probe the CPU flags.
-    detect_cpu_features(&st->flags);
+    blake2_init_cpu_features(st);
 
     st->blake2b_type = (PyTypeObject *)PyType_FromModuleAndSpec(
         m, &blake2b_type_spec, NULL);
@@ -332,14 +256,14 @@ static inline blake2_impl type_to_impl(PyTypeObject *type) {
 #endif
     if (!strcmp(type->tp_name, blake2b_type_spec.name)) {
 #ifdef HACL_CAN_COMPILE_SIMD256
-      if (has_simd256(&st->flags))
+      if (st->can_run_simd256)
         return Blake2b_256;
       else
 #endif
         return Blake2b;
     } else if (!strcmp(type->tp_name, blake2s_type_spec.name)) {
 #ifdef HACL_CAN_COMPILE_SIMD128
-      if (has_simd128(&st->flags))
+      if (st->can_run_simd128)
         return Blake2s_128;
       else
 #endif

Modules/hmacmodule.c

@@ -17,6 +17,7 @@
 #endif
 
 #include "Python.h"
+#include "pycore_cpuinfo.h"             // py_cpuid_features
 #include "pycore_hashtable.h"
 #include "pycore_strhex.h"              // _Py_strhex()
 
@@ -1682,73 +1683,19 @@ hmacmodule_init_strings(hmacmodule_state *state)
 static void
 hmacmodule_init_cpu_features(hmacmodule_state *state)
 {
-    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];
-#endif
-    // fmt: off
-    (void)eax1; (void)ebx1; (void)ecx1; (void)edx1;
-    (void)eax7; (void)ebx7; (void)ecx7; (void)edx7;
-    // fmt: on
-
-#define EBX_AVX2 (1 << 5)
-#define ECX_SSE3 (1 << 0)
-#define ECX_SSSE3 (1 << 9)
-#define ECX_SSE4_1 (1 << 19)
-#define ECX_SSE4_2 (1 << 20)
-#define ECX_AVX (1 << 28)
-#define EDX_SSE (1 << 25)
-#define EDX_SSE2 (1 << 26)
-#define EDX_CMOV (1 << 15)
-
-    bool avx = (ecx1 & ECX_AVX) != 0;
-    bool avx2 = (ebx7 & EBX_AVX2) != 0;
-
-    bool sse = (edx1 & EDX_SSE) != 0;
-    bool sse2 = (edx1 & EDX_SSE2) != 0;
-    bool cmov = (edx1 & EDX_CMOV) != 0;
-
-    bool sse3 = (ecx1 & ECX_SSE3) != 0;
-    bool sse41 = (ecx1 & ECX_SSE4_1) != 0;
-    bool sse42 = (ecx1 & ECX_SSE4_2) != 0;
-
-#undef EDX_CMOV
-#undef EDX_SSE2
-#undef EDX_SSE
-#undef ECX_AVX
-#undef ECX_SSE4_2
-#undef ECX_SSE4_1
-#undef ECX_SSSE3
-#undef ECX_SSE3
-#undef EBX_AVX2
-
+    py_cpuid_features flags;
+    _Py_cpuid_detect_features(&flags);
 #if HACL_CAN_COMPILE_SIMD128
-    // TODO(picnixz): use py_cpuid_features (gh-125022) to improve detection
-    state->can_run_simd128 = sse && sse2 && sse3 && sse41 && sse42 && cmov;
+    state->can_run_simd128 = flags.sse && flags.sse2 && flags.sse3
+                             && flags.sse41 && flags.sse42
+                             && flags.cmov;
 #else
-    // fmt: off
-    (void)sse; (void)sse2; (void)sse3; (void)sse41; (void)sse42; (void)cmov;
-    // fmt: on
     state->can_run_simd128 = false;
 #endif
 
 #if HACL_CAN_COMPILE_SIMD256
-    // TODO(picnixz): use py_cpuid_features (gh-125022) to improve detection
-    state->can_run_simd256 = state->can_run_simd128 && avx && avx2;
+    state->can_run_simd256 = flags.avx && flags.avx2;
 #else
-    // fmt: off
-    (void)avx; (void)avx2;
-    // fmt: on
     state->can_run_simd256 = false;
 #endif
 }