include avx512vl to avx512 code path (#14733)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14733

We often also want to use AVX512VL instruction sets.
We already included AVX512F, AVX512DQ.
Skylake also has AVX512BW, AVX512CD we may want to later.

Reviewed By: duc0

Differential Revision: D13317282

fbshipit-source-id: 82c8e401d82d5c3a5452fb4ccb6e5cb88d242bda

diff --git a/caffe2/perfkernels/CMakeLists.txt b/caffe2/perfkernels/CMakeLists.txt
index 18ae10d..f2e2f86 100644 (file)
--- a/caffe2/perfkernels/CMakeLists.txt
+++ b/caffe2/perfkernels/CMakeLists.txt
@@ -41,7 +41,7 @@ if (NOT MSVC AND CAFFE2_COMPILER_SUPPORTS_AVX2_EXTENSIONS)
       add_library(Caffe2_perfkernels_avx512 OBJECT ${avx512_srcs})
       add_dependencies(Caffe2_perfkernels_avx512 Caffe2_PROTO c10)
       set_target_properties(
-          Caffe2_perfkernels_avx512 PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512dq -mavx2 -mfma -mavx -mf16c")
+          Caffe2_perfkernels_avx512 PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512dq -mavx512vl -mavx2 -mfma -mavx -mf16c")
       set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS}
           $<TARGET_OBJECTS:Caffe2_perfkernels_avx512>)
   endif()

diff --git a/caffe2/perfkernels/common.h b/caffe2/perfkernels/common.h
index 7fdf64f..91e9e86 100644 (file)
--- a/caffe2/perfkernels/common.h
+++ b/caffe2/perfkernels/common.h
@@ -43,10 +43,11 @@ In foo.cc, do:
 // During build time:
 //    The build system should provide flags CAFFE2_PERF_WITH_AVX512,
 //    CAFFE2_PERF_WITH_AVX2, and CAFFE2_PERF_WITH_AVX that corresponds to the
-//    __AVX512F__, __AVX512DQ__, __AVX__, and __AVX2__ flags the compiler
-//    provides. Note that we do not use the compiler flags but rely on the build
-//    system flags, because the common files (like foo.cc above) will always be
-//    built without __AVX512F__, __AVX512DQ__, __AVX__ and __AVX2__.
+//    __AVX512F__, __AVX512DQ__, __AVX512VL__, __AVX__, and __AVX2__ flags the
+//    compiler provides. Note that we do not use the compiler flags but rely on
+//    the build system flags, because the common files (like foo.cc above) will
+//    always be built without __AVX512F__, __AVX512DQ__, __AVX512VL__, __AVX__
+//    and __AVX2__.
 // During run time:
 //    we use cpuid to identify cpu support and run the proper functions.
 
@@ -62,7 +63,8 @@ In foo.cc, do:
 #ifdef CAFFE2_PERF_WITH_AVX512
 #define AVX512_DO(funcname, ...)                       \
   decltype(funcname##__base) funcname##__avx512;       \
-  if (GetCpuId().avx512f() && GetCpuId().avx512dq()) { \
+  if (GetCpuId().avx512f() && GetCpuId().avx512dq() && \
+      GetCpuId().avx512vl()) {                         \
     return funcname##__avx512(__VA_ARGS__);            \
   }
 #else // CAFFE2_PERF_WITH_AVX512

diff --git a/caffe2/perfkernels/common_avx512.cc b/caffe2/perfkernels/common_avx512.cc
index 055f95d..273e146 100644 (file)
--- a/caffe2/perfkernels/common_avx512.cc
+++ b/caffe2/perfkernels/common_avx512.cc
@@ -1,23 +1,23 @@
 // This file is here merely to check that the flags are not mixed up: for
-// example, if your compiler did not specify -mavx512f and -mavx512dq,
-// you should not provide the CAFFE2_PERF_WITH_AVX512 macro.
+// example, if your compiler did not specify -mavx512f, -mavx512dq, and
+// -mavx512vl you should not provide the CAFFE2_PERF_WITH_AVX512 macro.
 
 #include "caffe2/core/common.h"
 
 #ifdef CAFFE2_PERF_WITH_AVX512
-#if !defined(__AVX512F__) || !defined(__AVX512DQ__)
+#if !defined(__AVX512F__) || !defined(__AVX512DQ__) || !defined(__AVX512VL__)
 #error( \
     "You found a build system error: CAFFE2_PERF_WITH_AVX512 is defined" \
-    "but __AVX512F__ or __AVX512DQ__ is not defined" \
-    "(via e.g. -mavx512f and -mavx512dq).");
+    "but __AVX512F__, __AVX512DQ__, or __AVX512VL is not defined" \
+    "(via e.g. -mavx512f, -mavx512dq, and -mavx512vl).");
 #endif
 #endif // CAFFE2_PERF_WITH_AVX512
 
-#if defined(__AVX512F__) && defined(__AVX512DQ__)
+#if defined(__AVX512F__) && defined(__AVX512DQ__) && defined(__AVX512VL__)
 #ifndef CAFFE2_PERF_WITH_AVX512
 #error( \
-    "You found a build system error: __AVX512F__ and __AVX512DQ__ is defined" \
-    "(via e.g. -mavx512f and -mavx512dq) " \
+    "You found a build system error: __AVX512F__, __AVX512DQ__, __AVX512VL__ " \
+    "is defined (via e.g. -mavx512f, -mavx512dq, and -mavx512vl) " \
     "but CAFFE2_PERF_WITH_AVX512 is not defined.");
 #endif // CAFFE2_PERF_WITH_AVX512
 #endif