Support multi-target builds for Apple M1

CMakeLists.txt

@@ -6,20 +6,6 @@ cmake_minimum_required(VERSION 3.17)
 
 project(taichi)
 
-# Taichi does not set target architecture explicitly,
-# but rather rely on CMake to detect the host arch.
-#
-# However on Mac m1, there are two available architectures namely x86_64 and arm64.
-# On some combination of "OSX version" and "CMake version", CMake will use x86_64 as default architecture even if it's on m1 chip.
-# This causes conflicts with the precompiled LLVM/Clang binaries downloaded from Taichi's repo (pre-built for arm64)
-#
-# Therefore we force CMake to choose arm64 architecture on arm64 chips.
-if (APPLE)
-  if( "${CMAKE_HOST_SYSTEM_PROCESSOR}" STREQUAL "arm64" )
-    set(CMAKE_OSX_ARCHITECTURES ${CMAKE_HOST_SYSTEM_PROCESSOR})
-  endif()
-endif()
-
 if (NOT DEFINED TI_VERSION_MAJOR)
     message(WARNING "It seems that you are running cmake manually, which may cause issues. Please use setup.py to build taichi from source, see https://docs.taichi-lang.org/docs/dev_install for more details.")
     set(TI_VERSION_MAJOR 0)

cmake/TaichiCXXFlags.cmake

@@ -81,8 +81,15 @@ if ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64" OR "${CMAKE_SYSTEM_PROCESSOR}"
     if (MSVC)
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /D \"TI_ARCH_x64\"")
     else()
-        message("Setting -march=nehalem for x86_64 processors")
-        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=nehalem -DTI_ARCH_x64")
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_ARCH_x64")
+        if ("arm64" IN_LIST CMAKE_OSX_ARCHITECTURES)
+            # TODO: (penguinliong) Will probably need this in a future version
+            # of Clang. Clang11 doesn't recognize this.
+            #set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mcpu=apple-m1")
+        else()
+            message("Setting -march=nehalem for x86_64 processors")
+            set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=nehalem")
+        endif()
     endif()
     set(ARCH "x64")
 elseif ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "aarch64" OR "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "arm64")

setup.py

@@ -8,6 +8,7 @@
 import glob
 import multiprocessing
 import os
+import platform
 import shutil
 import sys
 from distutils.command.clean import clean
@@ -131,6 +132,11 @@ def get_cmake_args():
 
     if sys.platform != 'win32':
         os.environ['SKBUILD_BUILD_OPTIONS'] = f'-j{num_threads}'
+    if sys.platform == "darwin":
+        if platform.architecture() == "arm64":
+            cmake_args += ["-DCMAKE_OSX_ARCHITECTURES=arm64"]
+        else:
+            cmake_args += ["-DCMAKE_OSX_ARCHITECTURES=x86_64"]
     return cmake_args
 
 

taichi/program/program.cpp

@@ -40,7 +40,7 @@
 #include "taichi/rhi/dx/dx_api.h"
 #endif
 
-#if defined(TI_ARCH_x64)
+#if defined(TI_ARCH_x64) && !(defined(__arm64__) || defined(__aarch64__))
 // For _MM_SET_FLUSH_ZERO_MODE
 #include <xmmintrin.h>
 #endif
@@ -55,7 +55,7 @@ Program::Program(Arch desired_arch) : snode_rw_accessors_bank_(this) {
   // For performance considerations and correctness of QuantFloatType
   // operations, we force floating-point operations to flush to zero on all
   // backends (including CPUs).
-#if defined(TI_ARCH_x64)
+#if defined(TI_ARCH_x64) && !(defined(__arm64__) || defined(__aarch64__))
   _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON);
 #else
   // Enforce flush to zero on arm64 CPUs

taichi/system/timer.cpp

@@ -220,7 +220,7 @@ uint64 Time::get_cycles() {
 #else
 
 uint64 Time::get_cycles() {
-#if defined(TI_ARCH_x64)
+#if defined(TI_ARCH_x64) && !(defined(__arm64__) || defined(__aarch64__))
   unsigned int lo, hi;
   __asm__ __volatile__("rdtsc" : "=a"(lo), "=d"(hi));
   return ((uint64)hi << 32) | lo;