AVX2 now includes FMA3

Raghuveer Devulapalli · Raghuveer Devulapalli · commit a46b082627ef · 2024-06-20T08:51:42.000-07:00
diff --git a/meson_cpu/x86/meson.build b/meson_cpu/x86/meson.build
@@ -59,13 +59,14 @@ FMA3 = mod_features.new(
   'FMA3', 24, implies: F16C, args: '-mfma',
   test_code: files(source_root + '/numpy/distutils/checks/cpu_fma3.c')[0]
 )
+# match this to HWY_AVX2
 AVX2 = mod_features.new(
-  'AVX2', 25, implies: F16C, args: '-march=skylake',
+  'AVX2', 25, implies: FMA3, args: ['-mavx2', '-maes', '-mpclmul', '-mbmi', '-mbmi2'],
   test_code: files(source_root + '/numpy/distutils/checks/cpu_avx2.c')[0]
 )
 # 25-40 left as margin for any extra features
 AVX512F = mod_features.new(
-  'AVX512F', 40, implies: [FMA3, AVX2],
+  'AVX512F', 40, implies: [AVX2],
   # Disables mmx because of stack corruption that may happen during mask
   # conversions.
   # TODO (seiko2plus): provide more clarification
@@ -93,7 +94,7 @@ AVX512_KNM = mod_features.new(
 )
 AVX512_SKX = mod_features.new(
   'AVX512_SKX', 50, implies: AVX512CD,
-  args: ['-march=skylake-avx512'],
+  args: ['-mavx512vl', '-mavx512bw', '-mavx512dq'],
   group: ['AVX512VL', 'AVX512BW', 'AVX512DQ'],
   test_code: files(source_root + '/numpy/distutils/checks/cpu_avx512_skx.c')[0],
   extra_tests: {
diff --git a/numpy/_core/src/common/npy_cpu_features.c b/numpy/_core/src/common/npy_cpu_features.c
@@ -474,6 +474,8 @@ npy__cpu_init_features(void)
     // third call to the cpuid to get extended AVX2 & AVX512 feature bits
     npy__cpu_cpuid(reg, 7);
     npy__cpu_have[NPY_CPU_FEATURE_AVX2]   = (reg[1] & (1 << 5))  != 0;
+    npy__cpu_have[NPY_CPU_FEATURE_AVX2]   = npy__cpu_have[NPY_CPU_FEATURE_AVX2] &&
+                                            npy__cpu_have[NPY_CPU_FEATURE_FMA3];
     if (!npy__cpu_have[NPY_CPU_FEATURE_AVX2])
         return;
     // detect AVX2 & FMA3
@@ -641,7 +643,7 @@ static void
 npy__cpu_init_features(void)
 {
     memset(npy__cpu_have, 0, sizeof(npy__cpu_have[0]) * NPY_CPU_FEATURE_MAX);
-    
+
     unsigned int hwcap = getauxval(AT_HWCAP);
     if ((hwcap & HWCAP_S390_VX) == 0) {
         return;
@@ -653,7 +655,7 @@ npy__cpu_init_features(void)
        npy__cpu_have[NPY_CPU_FEATURE_VXE2] = 1;
        return;
     }
-    
+
     npy__cpu_have[NPY_CPU_FEATURE_VXE] = (hwcap & HWCAP_S390_VXE) != 0;
 
     npy__cpu_have[NPY_CPU_FEATURE_VX]  = 1;
