From d7fdc962bed7b5a006db16d0f4683c54977bf69c Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Fri, 8 Mar 2024 19:26:09 -0500
Subject: [PATCH] Update xla

---
 WORKSPACE                        | 10 +++++-----
 patches/xla.patch                | 10 ----------
 src/enzyme_ad/jax/enzyme_call.cc |  3 ---
 3 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/WORKSPACE b/WORKSPACE
index e4bad745e..233c579ab 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -13,8 +13,8 @@ load("@rules_cc//cc:repositories.bzl", "rules_cc_dependencies")
 
 rules_cc_dependencies()
 
-LLVM_COMMIT = "5932f3f861f84305bd01050d0af8e0dcb232a8b3"
-LLVM_SHA256 = "ffbb065b6b9c2aef72949e84484ce5db3a86f682e7f4910a79eb5236856d259a"
+LLVM_COMMIT = "e946b5a87b2db307da076093d0a9a72ecb4ec089"
+LLVM_SHA256 = "7528e725bd118fec655f378efea2abaec032309fecc04a7ea052fd63e68c4dc8"
 LLVM_TARGETS = ["X86", "AArch64", "AMDGPU", "NVPTX"]
 
 http_archive(
@@ -30,8 +30,8 @@ http_archive(
 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
 llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)
 
-XLA_COMMIT = "6ee7005b0dbe29ba0cd077a690db1555ec6de346"
-XLA_SHA256 = "76f36ca2eecb246eb7931f2d77c5e9a32859c42aeb5b235f5c17ef6fc6fa71ae"
+XLA_COMMIT = "541962e88f52237bc6050e4c8d7270e7c7e12b4e"
+XLA_SHA256 = "ca67f68edad0d898241b65cbd85869de2560e2fdba700d964c707b427158583d"
 
 http_archive(
     name = "xla",
@@ -70,7 +70,7 @@ http_archive(
     urls = ["https://github.com/EnzymeAD/Enzyme/archive/{commit}.tar.gz".format(commit = ENZYME_COMMIT)],
 )
 
-JAX_COMMIT = "9a098e922aff62a3b49bd673b9518d97ee599248"
+JAX_COMMIT = "5e039f7af54539eed3268610b737b9f38621feb1"
 JAX_SHA256 = ""
 
 http_archive(
diff --git a/patches/xla.patch b/patches/xla.patch
index 175553b26..6cf78a0d1 100644
--- a/patches/xla.patch
+++ b/patches/xla.patch
@@ -41,13 +41,3 @@
          # copybara:uncomment_begin(google-only)
          # # TODO(ezhulenev): Clean up dependencies that are leforvers from Autofusion project.
 
---- a/xla/mlir/math/BUILD
-+++ b/xla/mlir/math/BUILD
-@@ -1,6 +1,7 @@
- package_group(
-     name = "friends",
-     packages = [
-+        "public",
-         "//xla/mlir/...",
-         # copybara:uncomment_begin(google-only)
-         # # TODO(ezhulenev): Clean up dependencies that are leforvers from Autofusion project.
diff --git a/src/enzyme_ad/jax/enzyme_call.cc b/src/enzyme_ad/jax/enzyme_call.cc
index 8ae759d57..5cfcce119 100644
--- a/src/enzyme_ad/jax/enzyme_call.cc
+++ b/src/enzyme_ad/jax/enzyme_call.cc
@@ -40,11 +40,9 @@
 
 #include "mlir/InitAllPasses.h"
 #include "xla/mlir/backends/cpu/transforms/passes.h"
-#include "xla/mlir/math/transforms/passes.h"
 #include "xla/mlir/memref/transforms/passes.h"
 #include "xla/mlir/runtime/transforms/passes.h"
 #include "xla/mlir_hlo/deallocation/transforms/passes.h"
-#include "xla/mlir_hlo/lhlo/IR/lhlo_ops.h"
 #include "xla/mlir_hlo/lhlo/transforms/passes.h"
 #include "xla/mlir_hlo/lhlo_gpu/IR/lhlo_gpu_ops.h"
 #include "xla/mlir_hlo/mhlo/transforms/passes.h"
@@ -1015,7 +1013,6 @@ PYBIND11_MODULE(enzyme_call, m) {
   xla::cpu::registerCpuTransformsPasses();
   mlir::hlo::registerLMHLOTransformsPasses();
   xla::runtime::registerRuntimeTransformsPasses();
-  xla::registerMathTransformsPasses();
   xla::registerMemrefTransformsPasses();
 
   mlir::registerShapePasses();