Add AVX512 support in ATen & remove AVX support

.circleci/config.yml

@@ -7170,14 +7170,14 @@ workflows:
             - pytorch_linux_bionic_py3_8_gcc9_coverage_build
           build_environment: "pytorch-linux-bionic-py3.8-gcc9-coverage-test1"
           docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-bionic-py3.8-gcc9"
-          resource_class: large
+          resource_class: xlarge
       - pytorch_linux_test:
           name: pytorch_linux_bionic_py3_8_gcc9_coverage_test2
           requires:
             - pytorch_linux_bionic_py3_8_gcc9_coverage_build
           build_environment: "pytorch-linux-bionic-py3.8-gcc9-coverage-test2"
           docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-bionic-py3.8-gcc9"
-          resource_class: large
+          resource_class: xlarge
       - pytorch_linux_build:
           name: pytorch_linux_bionic_rocm3_9_py3_6_build
           requires:

.jenkins/pytorch/test.sh

@@ -128,10 +128,10 @@ if [[ "$BUILD_ENVIRONMENT" == *asan* ]]; then
     (cd test && ! get_exit_code python -c "import torch; torch._C._crash_if_aten_asan(3)")
 fi
 
-if [[ "${BUILD_ENVIRONMENT}" == *-NO_AVX-* ]]; then
+if [[ "${BUILD_ENVIRONMENT}" == *-NO_AVX-* ]] || [[ "${BUILD_ENVIRONMENT}" == *-NO_AVX2-* ]]; then
   export ATEN_CPU_CAPABILITY=default
-elif [[ "${BUILD_ENVIRONMENT}" == *-NO_AVX2-* ]]; then
-  export ATEN_CPU_CAPABILITY=avx
+elif [[ "${BUILD_ENVIRONMENT}" == *-NO_AVX512-* ]]; then
+  export ATEN_CPU_CAPABILITY=avx2
 fi
 
 if [ -n "$IN_PULL_REQUEST" ] && [[ "$BUILD_ENVIRONMENT" != *coverage* ]]; then

.jenkins/pytorch/win-test-helpers/test_python.bat

@@ -1,3 +1,3 @@
 call %SCRIPT_HELPERS_DIR%\setup_pytorch_env.bat
-cd test && python run_test.py --exclude-jit-executor --verbose --determine-from="%1" && cd ..
+cd test && python run_test.py --exclude-jit-executor --verbose --continue-through-error --determine-from="%1" && cd ..
 if ERRORLEVEL 1 exit /b 1

aten.bzl

@@ -1,9 +1,8 @@
 load("@rules_cc//cc:defs.bzl", "cc_library")
 
-CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX", "AVX2"]
+CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX2"]
 CAPABILITY_COMPILER_FLAGS = {
     "AVX2": ["-mavx2", "-mfma"],
-    "AVX": ["-mavx"],
     "DEFAULT": [],
 }
 

aten/src/ATen/CMakeLists.txt

@@ -41,7 +41,7 @@ if(NOT BUILD_LITE_INTERPRETER)
 endif()
 EXCLUDE(ATen_CORE_SRCS "${ATen_CORE_SRCS}" ${ATen_CORE_TEST_SRCS})
 
-file(GLOB base_h "*.h" "detail/*.h" "cpu/*.h" "cpu/vec/vec256/*.h" "cpu/vec/*.h" "quantized/*.h")
+file(GLOB base_h "*.h" "detail/*.h" "cpu/*.h" "cpu/vec/vec512/*.h" "cpu/vec/vec256/*.h" "cpu/vec/*.h" "quantized/*.h")
 file(GLOB base_cpp "*.cpp" "detail/*.cpp" "cpu/*.cpp")
 file(GLOB cuda_h "cuda/*.h" "cuda/detail/*.h" "cuda/*.cuh" "cuda/detail/*.cuh")
 file(GLOB cuda_cpp "cuda/*.cpp" "cuda/detail/*.cpp")

aten/src/ATen/Version.cpp

@@ -108,12 +108,12 @@ std::string used_cpu_capability() {
     case native::CPUCapability::DEFAULT:
       ss << "NO AVX";
       break;
-    case native::CPUCapability::AVX:
-      ss << "AVX";
-      break;
     case native::CPUCapability::AVX2:
       ss << "AVX2";
       break;
+    case native::CPUCapability::AVX512:
+      ss << "AVX512";
+      break;
 #endif
     default:
       break;

aten/src/ATen/cpu/FlushDenormal.cpp

@@ -1,6 +1,11 @@
 #include <ATen/cpu/FlushDenormal.h>
 
+#if defined(CPU_CAPABILITY_AVX512)
+#include <ATen/cpu/vec/vec512/intrinsics.h>
+#else
 #include <ATen/cpu/vec/vec256/intrinsics.h>
+#endif
+
 #include <cpuinfo.h>
 
 namespace at { namespace cpu {