Binaries without AVX512 kernels shouldn't report CPU Capability as AVX512 on machines with AVX512 support (#66703)

imaginary-person · facebook-github-bot · commit b696d64ef47d · 2021-10-21T16:17:28.000-07:00
Summary: ### BUG If a PyTorch binary is built with a compiler that doesn't support all the AVX512 intrinsics in the codebase, then it won't have ATen AVX512 kernels, but at runtime, CPU capability would still be incorrectly returned as AVX512 on a machine that supports AVX512. It seems that PyTorch Linux releases are done on CentOS with `gcc 7.3`, so this bug would manifest in the 1.10 release, unless a fix such as this one is added. gcc versions below 9.0 don't support all the AVX512 intrinsics in the codebase, such as `_mm512_set_epi16`. ### FIX CPU Capability would be returned as AVX512 at runtime only if the binary was built with a compiler that supports all the AVX512 intrinsics in the codebase, and if the hardware the binary is being run on supports all the required AVX512 instruction sets. Pull Request resolved: #66703 Reviewed By: gchanan Differential Revision: D31732625 Pulled By: malfet fbshipit-source-id: e52d06b87fbe2af9b303a2e9c264189c8512d5ec
diff --git a/aten/src/ATen/native/DispatchStub.cpp b/aten/src/ATen/native/DispatchStub.cpp
@@ -16,12 +16,16 @@ static CPUCapability compute_cpu_capability() {
       return CPUCapability::VSX;
     }
 #else
+#ifdef HAVE_AVX512_CPU_DEFINITION
     if (strcmp(envar, "avx512") == 0) {
       return CPUCapability::AVX512;
     }
+#endif
+#ifdef HAVE_AVX2_CPU_DEFINITION
     if (strcmp(envar, "avx2") == 0) {
       return CPUCapability::AVX2;
     }
+#endif
 #endif
     if (strcmp(envar, "default") == 0) {
       return CPUCapability::DEFAULT;
@@ -31,13 +35,21 @@ static CPUCapability compute_cpu_capability() {
 
 #if !defined(__powerpc__) && !defined(__s390x__)
   if (cpuinfo_initialize()) {
+#ifdef HAVE_AVX512_CPU_DEFINITION
+    // GCC supports some AVX512 intrinsics such as _mm512_set_epi16 only in
+    // versions 9 & beyond. So, we want to ensure that only releases built with
+    // supported compilers on supported hardware return CPU Capability AVX512,
+    // if it's supported on the hardware PyTorch is running on.
     if (cpuinfo_has_x86_avx512vl() && cpuinfo_has_x86_avx512bw() &&  \
         cpuinfo_has_x86_avx512dq() && cpuinfo_has_x86_fma3()) {
       return CPUCapability::AVX512;
     }
+#endif
+#ifdef HAVE_AVX2_CPU_DEFINITION
     if (cpuinfo_has_x86_avx2() && cpuinfo_has_x86_fma3()) {
       return CPUCapability::AVX2;
     }
+#endif
   }
 #endif
 #ifdef HAVE_VSX_CPU_DEFINITION
