Fix AARCH64 builds (#121700)

After #119992 was landed Pull Request resolved: #121700 Approved by: https://github.com/janeyx99, https://github.com/huydhn
pytorch · Mar 12, 2024 · 703e83e · 703e83e
1 parent 159f303
commit 703e83e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -560,6 +560,12 @@ if(CMAKE_COMPILER_IS_GNUCXX)
   set_source_files_properties(${TORCH_SRC_DIR}/csrc/jit/tensorexpr/llvm_codegen.cpp PROPERTIES COMPILE_FLAGS "-Wno-init-list-lifetime")
 endif()
 
+# Enable conditional FP16 arithmetic intrinsics
+if(CPU_AARCH64 AND LINUX)
+set_source_files_properties(${TORCH_ROOT}/aten/src/ATen/native/BlasKernel.cpp PROPERTIES COMPILE_FLAGS "-march=armv8.2-a+fp16")
+endif()
+
+
 if(NOT INTERN_DISABLE_MOBILE_INTERP)
   set(MOBILE_SRCS
      ${TORCH_SRC_DIR}/csrc/jit/mobile/function.cpp