[CpuInductor] Enable NEON ISA detection on Linux ARM (#129075)

malfet · pytorchmergebot · commit b2a9b8d48534 · 2024-06-20T16:22:57.000Z
Also, cleanup code a bit to use `x in [y, z]` instead of `x == y or x == z` And do not redefine `at_align`, but instead use `alignas(64)` as was suggested in https://github.com/pytorch/pytorch/pull/128686/files#r1639365978 Test plan: `python3 -c "import torch._inductor.codecache as cc; isa = cc.valid_vec_isa_list()[0];print(str(isa), bool(isa))"` Pull Request resolved: #129075 Approved by: https://github.com/jansel
diff --git a/torch/_inductor/codecache.py b/torch/_inductor/codecache.py
@@ -1306,18 +1306,7 @@ class VecISA:
 #include <ATen/cpu/vec/vec.h>
 #endif
 
-#ifdef __APPLE__
-// Fix Mac OS UT failed.
-__attribute__((aligned(64))) float in_out_ptr0[16] = {0.0};
-#else
-#if defined(_WIN32)
-#define __at_align__ __declspec(align(64))
-#else
-#define __at_align__ __attribute__((aligned(64)))
-#endif
-
-__at_align__ float in_out_ptr0[16] = {0.0};
-#endif
+alignas(64) float in_out_ptr0[16] = {0.0};
 
 extern "C" void __avx_chk_kernel() {
     auto tmp0 = at::vec::Vectorized<float>(1);
@@ -1510,12 +1499,11 @@ def _check_and_append_supported_isa(
 # we only cache some key isa information.
 @functools.lru_cache(None)
 def valid_vec_isa_list() -> List[VecISA]:
+    isa_list: List[VecISA] = []
     if sys.platform == "darwin" and platform.processor() == "arm":
-        return [VecNEON()]
+        isa_list.append(VecNEON())
 
-    isa_list: List[VecISA] = []
-    cur_os = sys.platform
-    if cur_os != "linux" and cur_os != "win32":
+    if sys.platform not in ["linux", "win32"]:
         return isa_list
 
     arch = platform.machine()
@@ -1532,17 +1520,16 @@ def valid_vec_isa_list() -> List[VecISA]:
                         if re.search(r"[\^ ]+vxe[\$ ]+", group):
                             isa_list.append(VecZVECTOR())
                             break
-        return isa_list
-
-    if arch == "x86_64" or arch == "AMD64":
+    elif arch == "aarch64":
+        isa_list.append(VecNEON())
+    elif arch in ["x86_64", "AMD64"]:
         """
         arch value is x86_64 on Linux, and the value is AMD64 on Windows.
         """
         _cpu_supported_x86_isa = x86_isa_checker()
         for isa in supported_vec_isa_list:
             if str(isa) in _cpu_supported_x86_isa and isa:
                 isa_list.append(isa)
-        return isa_list
 
     return isa_list
 
