Add LLVM patch to fix AVX512 codegeneration problem

Author: yashk2810

third_party/xla/llvm_fix.patch

@@ -0,0 +1,32 @@
+diff --git a/third_party/llvm/llvm_jax_fix.patch b/third_party/llvm/llvm_jax_fix.patch
+new file mode 100644
+index 0000000000..5a2a60205e
+--- /dev/null
++++ b/third_party/llvm/llvm_jax_fix.patch
+@@ -0,0 +1,14 @@
++diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
++index 96be91256915d..8bcd8670879a9 100644
++--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
++@@ -59383,7 +59383,8 @@ static SDValue combineConcatVectorOps(const SDLoc &DL, MVT VT,
++ 
++   // We can always convert per-lane vXf64 shuffles into VSHUFPD.
++   if (!IsSplat &&
++-      (VT == MVT::v4f64 || (VT == MVT::v8f64 && Subtarget.useAVX512Regs())) &&
+++      ((NumOps == 2 && VT == MVT::v4f64) ||
+++       (NumOps == 4 && VT == MVT::v8f64 && Subtarget.useAVX512Regs())) &&
++       all_of(Ops, [](SDValue Op) { return Op.hasOneUse(); })) {
++     // Collect the individual per-lane v2f64/v4f64 shuffles.
++     MVT OpVT = Ops[0].getSimpleValueType();
+diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
+index ae0c1b550f..ce408f554a 100644
+--- a/third_party/llvm/workspace.bzl
++++ b/third_party/llvm/workspace.bzl
+@@ -22,6 +22,7 @@ def repo(name):
+             "//third_party/llvm:mathextras.patch",
+             "//third_party/llvm:toolchains.patch",
+             "//third_party/llvm:zstd.patch",
++            "//third_party/llvm:llvm_jax_fix.patch",
+         ],
+         link_files = {"//third_party/llvm:run_lit.sh": "mlir/run_lit.sh"},
+     )

third_party/xla/workspace.bzl

@@ -30,6 +30,7 @@ def repo():
         sha256 = XLA_SHA256,
         strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT),
         urls = tf_mirror_urls("https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)),
+        patch_file = ["//third_party/xla:llvm_fix.patch"],
     )
 
     # For development, one often wants to make changes to the TF repository as well