Make TensorFlow build on Apple Silicon

.bazelrc

@@ -17,6 +17,9 @@
 #     ios_x86_64:
 #     ios_fat:
 #
+# Macosx options
+#     darwin_arm64:
+#
 # Compiler options:
 #     cuda_clang:             Use clang when building CUDA code.
 #     c++17:                  Build with C++17 options (links with libc++)
@@ -134,6 +137,11 @@ build:ios_x86_64 --cpu=ios_x86_64
 build:ios_fat --config=ios
 build:ios_fat --ios_multi_cpus=armv7,arm64,i386,x86_64
 
+# Enables all the macos config options for macos_arm64
+build:macos_arm64 --config=macos
+build:macos_arm64 --apple_platform_type=macos 
+build:macos_arm64 --cpu=darwin_arm64
+
 # Config to use a mostly-static build and disable modular op registration
 # support (this will revert to loading TensorFlow with RTLD_GLOBAL in Python).
 # By default, TensorFlow will build with a dependence on
@@ -272,6 +280,7 @@ build:c++1z_gcc --config=c++17_gcc
 build --enable_platform_specific_config
 build:android --noenable_platform_specific_config
 build:ios --noenable_platform_specific_config
+build:macos_arm64 --noenable_platform_specific_config
 
 # Suppress C++ compiler warnings, otherwise build logs become 10s of MBs.
 build:android --copt=-w

tensorflow/BUILD

@@ -178,7 +178,7 @@ config_setting(
 )
 
 config_setting(
-    name = "macos",
+    name = "macos_x86_64",
     values = {
         "apple_platform_type": "macos",
         "cpu": "darwin",
@@ -186,6 +186,33 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+config_setting(
+    name = "macos_arm64e",
+    values = {
+        "apple_platform_type": "macos",
+        "cpu": "darwin_arm64e",
+    },
+    visibility = ["//visibility:public"],
+)
+
+config_setting(
+    name = "macos_arm64",
+    values = {
+        "apple_platform_type": "macos",
+        "cpu": "darwin_arm64",
+    },
+    visibility = ["//visibility:public"],
+)
+
+selects.config_setting_group(
+    name = "macos",
+    match_any = [
+        ":macos_x86_64",
+        ":macos_arm64",
+        ":macos_arm64e",
+    ],
+)
+
 config_setting(
     name = "ios",
     values = {"apple_platform_type": "ios"},
@@ -391,7 +418,6 @@ config_setting(
     },
     values = {
         "apple_platform_type": "macos",
-        "cpu": "darwin",
     },
     visibility = ["//visibility:public"],
 )
@@ -719,7 +745,9 @@ tf_cc_shared_object(
     name = "tensorflow_framework",
     framework_so = [],
     linkopts = select({
-        "//tensorflow:macos": [],
+        "//tensorflow:macos": [
+            "-Wl,-rename_section,__TEXT,text_env,__TEXT,__text"
+        ],
         "//tensorflow:windows": [],
         "//tensorflow:freebsd": [
             "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)",

tensorflow/compiler/aot/tfcompile.bzl

@@ -432,7 +432,8 @@ def target_llvm_triple():
         "//tensorflow:ios": "arm64-none-ios",
         "//tensorflow:ios_x86_64": "x86_64-apple-ios",
         "//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu",
-        "//tensorflow:macos": "x86_64-none-darwin",
+        "//tensorflow:macos_x86_64": "x86_64-none-darwin",
+        "//tensorflow:macos_arm64": "aarch64-none-darwin",
         "//tensorflow:windows": "x86_64-none-windows",
         "//tensorflow:linux_s390x": "systemz-none-linux-gnu",
         "//conditions:default": "x86_64-pc-linux",

tensorflow/core/platform/default/build_config_root.bzl

@@ -56,6 +56,7 @@ def if_static(extra_deps, otherwise = [], macos = []):
     }
     if macos:
         ret[str(Label("//tensorflow:macos_with_framework_shared_object"))] = macos
+    
     return select(ret)
 
 def if_static_and_not_mobile(extra_deps, otherwise = []):

tensorflow/lite/kernels/internal/BUILD

@@ -186,6 +186,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "darwin_arm64",
+    values = {
+        "cpu": "darwin_arm64",
+    },
+)
+
 config_setting(
     name = "freebsd",
     values = {
@@ -806,6 +813,7 @@ cc_library(
             ":ios_armv7",
             ":ios_arm64",
             ":ios_arm64e",
+            ":darwin_arm64",
             ":raspberry_pi_with_neon",
         ): [":neon_tensor_utils"],
         (

tensorflow/tensorflow.bzl

@@ -446,6 +446,7 @@ def _rpath_linkopts(name):
     return select({
         clean_dep("//tensorflow:macos"): [
             "-Wl,%s" % (_make_search_paths("@loader_path", levels_to_root),),
+            "-Wl,-rename_section,__TEXT,text_env,__TEXT,__text"
         ],
         clean_dep("//tensorflow:windows"): [],
         "//conditions:default": [
@@ -1920,22 +1921,24 @@ def pywrap_tensorflow_macro(
 
     if not version_script:
         version_script = select({
-            "@local_config_cuda//cuda:darwin": clean_dep("//tensorflow:tf_exported_symbols.lds"),
+            "//tensorflow:macos": clean_dep("//tensorflow:tf_exported_symbols.lds"),
             "//conditions:default": clean_dep("//tensorflow:tf_version_script.lds"),
         })
     vscriptname = name + "_versionscript"
     _append_init_to_versionscript(
         name = vscriptname,
         is_version_script = select({
-            "@local_config_cuda//cuda:darwin": False,
+            "//tensorflow:macos": False,
             "//conditions:default": True,
         }),
         module_name = module_name,
         template_file = version_script,
     )
     extra_linkopts = select({
-        "@local_config_cuda//cuda:darwin": [
+        clean_dep("//tensorflow:macos"): [
-            "-Wl,-exported_symbols_list,$(location %s.lds)" % vscriptname,
+            # TODO: the -w suppresses a wall of harmless warnings about hidden typeinfo symbols 
+            # not being exported.  There should be a better way to deal with this.
+            "-Wl,-w", "-Wl,-exported_symbols_list,$(location %s.lds)" % vscriptname,
         ],
         clean_dep("//tensorflow:windows"): [],
         "//conditions:default": [
@@ -1944,9 +1947,6 @@ def pywrap_tensorflow_macro(
         ],
     })
     extra_deps += select({
-        "@local_config_cuda//cuda:darwin": [
-            "%s.lds" % vscriptname,
-        ],
         clean_dep("//tensorflow:windows"): [],
         "//conditions:default": [
             "%s.lds" % vscriptname,
@@ -2578,8 +2578,10 @@ def pybind_extension(
             ],
         }),
         linkopts = linkopts + _rpath_linkopts(name) + select({
-            "@local_config_cuda//cuda:darwin": [
+            clean_dep("//tensorflow:macos"): [
-                "-Wl,-exported_symbols_list,$(location %s)" % exported_symbols_file,
+                # TODO: the -w suppresses a wall of harmless warnings about hidden typeinfo symbols 
+                # not being exported.  There should be a better way to deal with this.
+                "-Wl,-w", "-Wl,-exported_symbols_list,$(location %s)" % exported_symbols_file,
             ],
             clean_dep("//tensorflow:windows"): [],
             "//conditions:default": [

tensorflow/tools/pip_package/build_pip_package.sh

@@ -216,10 +216,10 @@ function prepare_src() {
   if [ -d "${TMPDIR}/tensorflow/_api/v1/" ]
   then
     cp -r ${TMPDIR}/tensorflow/python/keras/api/_v1/keras/ ${TMPDIR}/tensorflow/keras/
-    sed -i'.original' -e 's/.python.keras.api._v1/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py
+    find . -type f -maxdepth 0 -exec sed -i'.original' -e 's/.python.keras.api._v1/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py {} +
   else
     cp -r ${TMPDIR}/tensorflow/python/keras/api/_v2/keras/ ${TMPDIR}/tensorflow/keras/
-    sed -i'.original' -e 's/.python.keras.api._v2/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py
+    find . -type f -maxdepth 0 -exec sed -i'.original' -e 's/.python.keras.api._v2/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py {} +
   fi
 }
 

third_party/cpuinfo/BUILD.bazel

@@ -105,6 +105,7 @@ cc_library(
         ":linux_mips64": COMMON_SRCS + LINUX_SRCS,
         ":linux_s390x": COMMON_SRCS + LINUX_SRCS,
         ":macos_x86_64": COMMON_SRCS + X86_SRCS + MACH_SRCS + MACH_X86_SRCS,
+        ":darwin_arm64": COMMON_SRCS + MACH_SRCS + MACH_ARM_SRCS,
         ":windows_x86_64": COMMON_SRCS + X86_SRCS + WINDOWS_X86_SRCS,
         ":android_armv7": COMMON_SRCS + ARM_SRCS + LINUX_SRCS + LINUX_ARM32_SRCS + ANDROID_ARM_SRCS,
         ":android_arm64": COMMON_SRCS + ARM_SRCS + LINUX_SRCS + LINUX_ARM64_SRCS + ANDROID_ARM_SRCS,
@@ -293,6 +294,23 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "darwin_arm64",
+    values = {
+        "apple_platform_type": "macos",
+        "cpu": "darwin_arm64",
+    },
+)
+
+config_setting(
+    name = "macos_arm64",
+    values = {
+        "apple_platform_type": "macos",
+        "cpu": "darwin_arm64",
+    },
+)
+
+
 config_setting(
     name = "ios_x86",
     values = {

third_party/llvm/llvm.bzl

@@ -299,7 +299,7 @@ win32_cmake_vars = {
 # TODO(phawkins): use a better method to select the right host triple, rather
 # than hardcoding x86_64.
 llvm_all_cmake_vars = select({
-    "@org_tensorflow//tensorflow:macos": cmake_var_string(
+    "@org_tensorflow//tensorflow:macos_x86_64": cmake_var_string(
         _dict_add(
             cmake_vars,
             llvm_target_cmake_vars("X86", "x86_64-apple-darwin"),
@@ -307,6 +307,14 @@ llvm_all_cmake_vars = select({
             darwin_cmake_vars,
         ),
     ),
+    "@org_tensorflow//tensorflow:macos_arm64": cmake_var_string(
+        _dict_add(
+            cmake_vars,
+            llvm_target_cmake_vars("AArch64", "arm64-apple-darwin"),
+            posix_cmake_vars,
+            darwin_cmake_vars,
+        ),
+    ),
     "@org_tensorflow//tensorflow:linux_ppc64le": cmake_var_string(
         _dict_add(
             cmake_vars,

third_party/nccl/archive.BUILD

@@ -89,7 +89,10 @@ cc_library(
     ],
     hdrs = ["src/nccl.h"],
     include_prefix = "third_party/nccl",
-    linkopts = ["-lrt"],
+    linkopts = select({
+        "//tensorflow:macos": [],
+        "//conditions:default": ["-lrt"],
+    }),
     strip_include_prefix = "src",
     visibility = ["//visibility:public"],
     deps = [