fixes for AVX detection (#17915)

Summary:
Our AVX2 routines use functions such as _mm256_extract_epi64
that do not exist on 32 bit systems even when they have AVX2.
This disables AVX2 when _mm256_extract_epi64 does not exist.

This fixes the "local" part of #17901 (except disabling FBGEMM),
but there also is sleef to be updated and NNPACK to be fixed,
see the bug report for further discussion.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17915

Differential Revision: D14437338

Pulled By: soumith

fbshipit-source-id: d4ef7e0801b5d1222a855a38ec207dd88b4680da

diff --git a/caffe2/quantization/server/CMakeLists.txt b/caffe2/quantization/server/CMakeLists.txt
index 35bde8c..e0aed79 100644 (file)
--- a/caffe2/quantization/server/CMakeLists.txt
+++ b/caffe2/quantization/server/CMakeLists.txt
@@ -59,7 +59,7 @@ list(APPEND Caffe2_CPU_SRCS
   #"${CMAKE_CURRENT_SOURCE_DIR}/sigmoid_test.cc")
   #"${CMAKE_CURRENT_SOURCE_DIR}/tanh_test.cc")
 
-if (NOT MSVC)
+if (NOT MSVC AND CAFFE2_COMPILER_SUPPORTS_AVX2_EXTENSIONS)
   add_library(caffe2_dnnlowp_avx2_ops OBJECT ${caffe2_dnnlowp_avx2_ops_SRCS})
   add_dependencies(caffe2_dnnlowp_avx2_ops fbgemm Caffe2_PROTO c10)
   target_include_directories(caffe2_dnnlowp_avx2_ops BEFORE

diff --git a/cmake/MiscCheck.cmake b/cmake/MiscCheck.cmake
index 0d2e61c..06ff17b 100644 (file)
--- a/cmake/MiscCheck.cmake
+++ b/cmake/MiscCheck.cmake
@@ -167,6 +167,8 @@ CHECK_CXX_SOURCE_COMPILES(
        a = _mm256_set1_epi8 (1);
        b = a;
        _mm256_add_epi8 (a,a);
+       __m256 x;
+       _mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code
        return 0;
      }" CAFFE2_COMPILER_SUPPORTS_AVX2_EXTENSIONS)
 if (CAFFE2_COMPILER_SUPPORTS_AVX2_EXTENSIONS)

diff --git a/cmake/Modules/FindAVX.cmake b/cmake/Modules/FindAVX.cmake
index 7afffd1..5b7272f 100644 (file)
--- a/cmake/Modules/FindAVX.cmake
+++ b/cmake/Modules/FindAVX.cmake
@@ -19,6 +19,8 @@ SET(AVX2_CODE "
   {
     __m256i a = {0};
     a = _mm256_abs_epi16(a);
+    __m256 x;
+    _mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code
     return 0;
   }
 ")