Add llvm support for s390x, fix xla and aot issue

tensorflow/compiler/aot/BUILD

@@ -71,6 +71,7 @@ cc_library(
         "@llvm-project//llvm:powerpc_code_gen",  # fixdeps: keep
         "@llvm-project//llvm:target",
         "@llvm-project//llvm:x86_code_gen",  # fixdeps: keep
+        "@llvm-project//llvm:system_z_code_gen",  # fixdeps: keep
         "//tensorflow/core:regexp_internal",
     ] + if_llvm_aarch64_available([
         "//third_party/llvm/llvm-project/llvm:aarch64_target",  # fixdeps: keep
@@ -113,6 +114,7 @@ cc_library(
         "@llvm-project//llvm:powerpc_code_gen",  # fixdeps: keep
         "@llvm-project//llvm:target",
         "@llvm-project//llvm:x86_code_gen",  # fixdeps: keep
+        "@llvm-project//llvm:system_z_code_gen",  # fixdeps: keep
     ] + if_llvm_aarch64_available([
         "//third_party/llvm/llvm-project/llvm:aarch64_target",  # fixdeps: keep
     ]),

tensorflow/compiler/aot/compile.cc

@@ -176,6 +176,10 @@ static void InitializeTargets() {
   LLVMInitializePowerPCTargetInfo();
   LLVMInitializePowerPCTargetMC();
   LLVMInitializePowerPCAsmPrinter();
+  LLVMInitializeSystemZTarget();
+  LLVMInitializeSystemZTargetInfo();
+  LLVMInitializeSystemZTargetMC();
+  LLVMInitializeSystemZAsmPrinter();
   LLVMInitializeX86Target();
   LLVMInitializeX86TargetInfo();
   LLVMInitializeX86TargetMC();

tensorflow/compiler/aot/tfcompile.bzl

@@ -434,5 +434,6 @@ def target_llvm_triple():
         "//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu",
         "//tensorflow:macos": "x86_64-none-darwin",
         "//tensorflow:windows": "x86_64-none-windows",
+        "//tensorflow:linux_s390x": "systemz-none-linux-gnu",
         "//conditions:default": "x86_64-pc-linux",
     })

tensorflow/compiler/xla/service/cpu/BUILD

@@ -191,6 +191,9 @@ cc_library(
         "//tensorflow:linux_ppc64le": [
             "@llvm-project//llvm:powerpc_code_gen",  # fixdeps: keep
         ],
+        "//tensorflow:linux_s390x": [
+            "@llvm-project//llvm:system_z_code_gen",  # fixdeps: keep
+        ],
         "//conditions:default": [
         ],
     }),

third_party/llvm/llvm.autogenerated.BUILD

@@ -33,6 +33,7 @@ llvm_targets = [
     "NVPTX",
     "PowerPC",
     "X86",
+    "SystemZ",
 ]
 
 llvm_target_asm_parsers = llvm_targets
@@ -538,6 +539,22 @@ llvm_target_list = [
             ("-gen-disassembler", "lib/Target/PowerPC/PPCGenDisassemblerTables.inc"),
         ],
     },
+    {
+        "name": "SystemZ",
+        "lower_name": "system_z",
+        "short_name": "SystemZ",
+        "tbl_outs": [
+            ("-gen-asm-writer", "lib/Target/SystemZ/SystemZGenAsmWriter.inc"),
+            ("-gen-asm-matcher", "lib/Target/SystemZ/SystemZGenAsmMatcher.inc"),
+            ("-gen-emitter", "lib/Target/SystemZ/SystemZGenMCCodeEmitter.inc"),
+            ("-gen-register-info", "lib/Target/SystemZ/SystemZGenRegisterInfo.inc"),
+            ("-gen-instr-info", "lib/Target/SystemZ/SystemZGenInstrInfo.inc"),
+            ("-gen-dag-isel", "lib/Target/SystemZ/SystemZGenDAGISel.inc"),
+            ("-gen-callingconv", "lib/Target/SystemZ/SystemZGenCallingConv.inc"),
+            ("-gen-subtarget", "lib/Target/SystemZ/SystemZGenSubtargetInfo.inc"),
+            ("-gen-disassembler", "lib/Target/SystemZ/SystemZGenDisassemblerTables.inc"),
+        ],
+    },
     {
         "name": "X86",
         "lower_name": "x86",
@@ -655,6 +672,7 @@ cc_library(
         ":nvptx_code_gen",
         ":powerpc_code_gen",
         ":x86_code_gen",
+        ":system_z_code_gen",
     ],
 )
 
@@ -3902,6 +3920,7 @@ cc_library(
         ":mc",
         ":support",
         ":system_z_info",
+        ":system_z_target_gen",
     ],
 )
 
@@ -3935,6 +3954,7 @@ cc_library(
         "lib/Target/SystemZ/TargetInfo/*.c",
         "lib/Target/SystemZ/TargetInfo/*.cpp",
         "lib/Target/SystemZ/TargetInfo/*.inc",
+        "lib/Target/SystemZ/MCTargetDesc/*.h",
     ]),
     hdrs = glob([
         "include/llvm/Target/SystemZ/TargetInfo/*.h",
@@ -3946,6 +3966,7 @@ cc_library(
     deps = [
         ":config",
         ":support",
+        ":system_z_target_gen",
     ],
 )
 

third_party/llvm/llvm.bzl

@@ -333,6 +333,14 @@ llvm_all_cmake_vars = select({
             posix_cmake_vars,
         ),
     ),
+    "@org_tensorflow//tensorflow:linux_s390x": cmake_var_string(
+        _dict_add(
+            cmake_vars,
+            llvm_target_cmake_vars("SystemZ", "systemz-unknown-linux_gnu"),
+            posix_cmake_vars,
+       linux_cmake_vars,
+        ),
+    ),
     "//conditions:default": cmake_var_string(
         _dict_add(
             cmake_vars,