Use sleef on macOS Apple silicon by default (pytorch#126509)

Use sleef ~~for aarch64~~ on macOS Apple silicon by default.

Pull Request resolved: pytorch#126509
Approved by: https://github.com/digantdesai, https://github.com/malfet

Author: helloguo

CMakeLists.txt

@@ -892,6 +892,14 @@ endif()
 
 if(USE_SLEEF_FOR_ARM_VEC256)
   string(APPEND CMAKE_CXX_FLAGS " -DAT_BUILD_ARM_VEC256_WITH_SLEEF")
+  add_definitions(-DAT_BUILD_ARM_VEC256_WITH_SLEEF)
+endif()
+
+# Enable sleef on macOS with Apple silicon by default
+if((${CMAKE_SYSTEM_NAME} STREQUAL "Darwin") AND (${CMAKE_SYSTEM_PROCESSOR} STREQUAL "arm64"))
+  message(STATUS "Running on macOS with Apple silicon")
+  string(APPEND CMAKE_CXX_FLAGS " -DAT_BUILD_ARM_VEC256_WITH_SLEEF")
+  add_definitions(-DAT_BUILD_ARM_VEC256_WITH_SLEEF)
 endif()
 
 if(USE_XNNPACK)

torch/_inductor/codecache.py

@@ -1327,6 +1327,8 @@ def __bool__(self) -> bool:
 class VecNEON(VecISA):
     _bit_width = 256  # This is required to leverage the compute implemented in aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h
     _macro = "-DCPU_CAPABILITY_NEON"
+    if sys.platform == "darwin" and platform.processor() == "arm":
+        _macro += " -DAT_BUILD_ARM_VEC256_WITH_SLEEF"
     _arch_flags = ""  # Unused
     _dtype_nelements = {torch.float: 8, torch.bfloat16: 16, torch.float16: 16}