Rebase to later llvm/xla

EnzymeAD · Dec 21, 2023 · a5438c7 · a5438c7
1 parent 24eedb4
commit a5438c7
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -13,8 +13,8 @@ load("@rules_cc//cc:repositories.bzl", "rules_cc_dependencies")
 
 rules_cc_dependencies()
 
-LLVM_COMMIT = "668865789620f390fbad4d7093ed8ca6eb932c31"
-LLVM_SHA256 = "8d7cbbe492a17656c09af1e79b802303f11cb47d64768760b70d52f11ed4d9da"
+LLVM_COMMIT = "d3ef86708241a3bee902615c190dead1638c4e09"
+LLVM_SHA256 = ""
 LLVM_TARGETS = ["X86", "AArch64", "AMDGPU", "NVPTX"]
 
 http_archive(
@@ -30,14 +30,15 @@ http_archive(
 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
 llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)
 
-XLA_COMMIT = "a6e6c1f6a53d4a23451c649110519c7ba8581bf9"
-XLA_SHA256 = "5fe6dfa30621bd50b022a6cab026d6f4cde9883a3e150ce1b6fd52822a57c59a"
+XLA_COMMIT = "9f84844fa54020a6d99add1645a5d386603b6462"
+XLA_COMMIT = "9e3cb9fa5fbb6ecc6e631d581cbb851b018d0728"
+XLA_SHA256 = ""
 
 http_archive(
     name = "xla",
     sha256 = XLA_SHA256,
     strip_prefix = "xla-" + XLA_COMMIT,
-    urls = ["https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
+    urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
     patch_args = ["-p1"],
     patches = ["//:patches/xla.patch"],
 )
@@ -60,8 +61,8 @@ load("@rules_python//python/pip_install:repositories.bzl", "pip_install_dependen
 
 pip_install_dependencies()
 
-ENZYME_COMMIT = "cbb970161fd41ce55da028f0960a441382b07112"
-ENZYME_SHA256 = "ec0450fdbc7f18cab46492acd3288b8347fa222317f9ff475768f5f10c45478c"
+ENZYME_COMMIT = "bad7df07dd7657c0c6884667d62ad1d9bcfd1d16"
+ENZYME_SHA256 = ""
 
 http_archive(
     name = "enzyme",

diff --git a/src/enzyme_ad/jax/clang_compile.cc b/src/enzyme_ad/jax/clang_compile.cc
@@ -64,7 +64,7 @@
 #include "llvm/Linker/Linker.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/Host.h"
+#include "llvm/TargetParser/Host.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/raw_ostream.h"
 

diff --git a/src/enzyme_ad/jax/compile_with_xla.cc b/src/enzyme_ad/jax/compile_with_xla.cc
@@ -89,6 +89,7 @@ compile_mhlo_to_llvm_with_xla(llvm::StringRef mhlo_text, std::string &output) {
 
   xla::ExecutableBuildOptions build_options;
   build_options.mutable_debug_options()->set_xla_embed_ir_in_executable(true);
+  // build_options.mutable_debug_options()->set_xla_cpu_use_xla_runtime(true);
 
   if (build_options.device_ordinal() == -1) {
     build_options.set_device_ordinal(local_client->default_device_ordinal());