[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
https://github.com/google/jax/issues/5679

PiperOrigin-RevId: 356578451
Change-Id: I99a2aa0e87739b9edce81074fce1ca5c0bd25115
This commit is contained in:
Benjamin Kramer 2021-02-09 13:30:51 -08:00 committed by TensorFlower Gardener
parent 3b0a3b7e8d
commit 942f315def
3 changed files with 17 additions and 5 deletions

View File

@ -229,6 +229,9 @@ cc_library(
"@llvm-project//llvm:Target",
"@llvm-project//llvm:X86CodeGen", # fixdeps: keep
] + select({
"//tensorflow:linux_aarch64": [
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
],
"//tensorflow:linux_ppc64le": [
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
],

View File

@ -16,18 +16,19 @@ limitations under the License.
#ifndef TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_
#define TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_
#if (defined(__powerpc__) || \
#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

View File

@ -315,6 +315,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,