[XLA:CPU] Wire up support for Aarch64 on linux

I don't have an environment to test this, but I believe it should be enough to get the JIT to work. This came up in google/jax#5679 PiperOrigin-RevId: 356578451 Change-Id: I99a2aa0e87739b9edce81074fce1ca5c0bd25115
NVIDIA · May 17, 2022 · 5a7e197 · 5a7e197
1 parent 92162f5
commit 5a7e197
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 deletions.
diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD
@@ -191,6 +191,9 @@ cc_library(
         "@llvm-project//llvm:target",
         "@llvm-project//llvm:x86_code_gen",  # fixdeps: keep
     ] + select({
+        "//tensorflow:linux_aarch64": [
+            "@llvm-project//llvm:aarch64_code_gen",  # fixdeps: keep
+        ],
         "//tensorflow:linux_ppc64le": [
             "@llvm-project//llvm:powerpc_code_gen",  # fixdeps: keep
         ],

diff --git a/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h
@@ -15,18 +15,19 @@ limitations under the License.
 #ifndef TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_
 #define TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_
 
-#if (defined(__powerpc__) || \
-     defined(__ppc__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__))
+#if defined(__aarch64__)
+static const char kTargetCpuForHost[] = "aarch64";
+static const char kTargetTripleForHost[] = "aarch64-unknown-linux-gnu",
+#elif (defined(__powerpc__) || \
+       defined(__ppc__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__))
 static const char kTargetCpuForHost[] = "ppc";
 static const char kTargetTripleForHost[] = "ppc64le-ibm-linux-gnu";
-#else
-#if defined(__s390x__)
+#elif defined(__s390x__)
 static const char kTargetCpuForHost[] = "s390x";
 static const char kTargetTripleForHost[] = "systemz-none-linux-gnu";
 #else
 static const char kTargetCpuForHost[] = "";
 static const char kTargetTripleForHost[] = "x86_64-pc-linux";
 #endif
-#endif
 
 #endif
diff --git a/third_party/llvm/llvm.bzl b/third_party/llvm/llvm.bzl
@@ -311,6 +311,14 @@ llvm_all_cmake_vars = select({
             darwin_cmake_vars,
         ),
     ),
+    "@org_tensorflow//tensorflow:linux_aarch64": cmake_var_string(
+        _dict_add(
+            cmake_vars,
+            llvm_target_cmake_vars("AArch64", "aarch64-unknown-linux_gnu"),
+            posix_cmake_vars,
+            linux_cmake_vars,
+        ),
+    ),
     "@org_tensorflow//tensorflow:linux_ppc64le": cmake_var_string(
         _dict_add(
             cmake_vars,