feat: bump xla version

.github/workflows/release.yml

@@ -86,7 +86,7 @@ jobs:
   macos_arm:
     name: "aarch64-darwin-cpu (cross-compiled)"
     needs: [create_draft_release]
-    runs-on: macos-13-xlarge
+    runs-on: macos-latest-large
     permissions:
       contents: write
     steps:

Makefile

@@ -11,7 +11,7 @@ BUILD_MODE ?= opt # can also be dbg
 BUILD_CACHE ?= $(TEMP)/xla_extension
 OPENXLA_GIT_REPO ?= https://github.com/openxla/xla.git
 
-OPENXLA_GIT_REV ?= 771e38178340cbaaef8ff20f44da5407c15092cb
+OPENXLA_GIT_REV ?= 0b476d54327ff20efddd50cb2bf0a0422fbb16e1
 
 # Private configuration
 BAZEL_FLAGS = --define "framework_shared_object=false" -c $(BUILD_MODE)
@@ -26,6 +26,7 @@ $(BUILD_ARCHIVE): $(OPENXLA_DIR) extension/BUILD
 	rm -f $(OPENXLA_XLA_EXTENSION_DIR) && \
 		ln -s "$(ROOT_DIR)/extension" $(OPENXLA_XLA_EXTENSION_DIR) && \
 		cd $(OPENXLA_DIR) && \
+        cat $(ROOT_DIR)/WORKSPACE >> WORKSPACE && \
 		bazel build $(BAZEL_FLAGS) $(BUILD_FLAGS) $(BUILD_INTERNAL_FLAGS) //$(OPENXLA_XLA_EXTENSION_NS):xla_extension && \
 		mkdir -p $(dir $(BUILD_ARCHIVE)) && \
 		cp -f $(OPENXLA_XLA_BUILD_ARCHIVE) $(BUILD_ARCHIVE)

WORKSPACE

@@ -0,0 +1,13 @@
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive", "http_file")
+http_archive(
+    name = "build_bazel_rules_apple",
+    sha256 = "34c41bfb59cdaea29ac2df5a2fa79e5add609c71bb303b2ebb10985f93fa20e7",
+    url = "https://github.com/bazelbuild/rules_apple/releases/download/3.1.1/rules_apple.3.1.1.tar.gz",
+)
+
+load(
+    "@build_bazel_rules_apple//apple:repositories.bzl",
+    "apple_rules_dependencies",
+)
+
+apple_rules_dependencies()

extension/BUILD

@@ -69,7 +69,6 @@ cc_static_library(
     "@llvm-project//mlir:SparseTensorDialect",
     "@tsl//tsl/platform:errors",
     "@tsl//tsl/platform:fingerprint",
-    "@tsl//tsl/platform:float8",
     "@tsl//tsl/platform:statusor",
     "@tsl//tsl/platform:env_impl",
     "@tsl//tsl/platform:tensor_float_32_utils",
@@ -80,7 +79,6 @@ cc_static_library(
     "@tsl//tsl/protobuf:dnn_proto_cc_impl",
     "@tsl//tsl/framework:allocator",
     "@tsl//tsl/framework:allocator_registry_impl",
-    "@tsl//tsl/util:determinism",
   ]
   # GRPC Dependencies (needed for PjRt distributed)
   + tsl_grpc_cc_dependencies()
@@ -100,13 +98,19 @@ cc_static_library(
 # of unused headers as well
 transitive_hdrs(
   name = "xla_extension_dep_headers",
-  deps = [
+  deps =  [
     "//xla:xla_proto_cc_impl",
     "//xla:xla_data_proto_cc_impl",
     "//xla/service:hlo_proto_cc_impl",
-    "//xla/service:memory_space_assignment_proto_cc_impl",
+    "//xla/service/memory_space_assignment:memory_space_assignment_proto_cc_impl",
+    "//xla/service:buffer_assignment_proto_cc_impl",
     "//xla/service/gpu:backend_configs_cc_impl",
-    "//xla/stream_executor:dnn_proto_cc_impl",
+    "//xla/service/gpu/model:hlo_op_profile_proto_cc_impl",
+    "//xla/stream_executor:device_description_proto_cc_impl",
+    "//xla:autotune_results_proto_cc_impl",
+    "//xla/stream_executor:stream_executor_impl",
+    "//xla/stream_executor/gpu:gpu_init_impl",
+    "//xla/stream_executor/host:host_platform",
     "//xla:literal",
     "//xla:shape_util",
     "//xla:status",
@@ -131,13 +135,11 @@ transitive_hdrs(
     "//xla/pjrt:pjrt_compiler",
     "//xla/pjrt:tfrt_cpu_pjrt_client",
     "//xla/pjrt:pjrt_c_api_client",
-    "//xla/pjrt:tpu_client",
-    "//xla/pjrt:pjrt_plugin_device_client",
-    "//xla/pjrt:pjrt_plugin_device_client_headers",
     "//xla/pjrt/distributed",
     "//xla/pjrt/gpu:se_gpu_pjrt_client",
     "//xla/pjrt/distributed:client",
     "//xla/pjrt/distributed:service",
+    "//xla:autotuning_proto_cc_impl",
     "@com_google_absl//absl/types:span",
     "@com_google_absl//absl/types:optional",
     "@com_google_absl//absl/base:log_severity",
@@ -149,24 +151,21 @@ transitive_hdrs(
     "@llvm-project//mlir:Pass",
     "@llvm-project//mlir:ReconcileUnrealizedCasts",
     "@llvm-project//mlir:SparseTensorDialect",
-    "@tf_runtime//:core_runtime",
-    "@tf_runtime//:hostcontext",
     "@tsl//tsl/platform:errors",
     "@tsl//tsl/platform:fingerprint",
-    "@tsl//tsl/platform:float8",
     "@tsl//tsl/platform:statusor",
     "@tsl//tsl/platform:env_impl",
     "@tsl//tsl/platform:tensor_float_32_utils",
     "@tsl//tsl/profiler/utils:time_utils_impl",
     "@tsl//tsl/profiler/backends/cpu:annotation_stack_impl",
     "@tsl//tsl/profiler/backends/cpu:traceme_recorder_impl",
-    "@tsl//tsl/protobuf:autotuning_proto_cc_impl",
     "@tsl//tsl/protobuf:protos_all_cc_impl",
     "@tsl//tsl/protobuf:dnn_proto_cc_impl",
     "@tsl//tsl/framework:allocator",
     "@tsl//tsl/framework:allocator_registry_impl",
-    "@tsl//tsl/util:determinism",
   ]
+  # GRPC Dependencies (needed for PjRt distributed)
+  + tsl_grpc_cc_dependencies()
 )
 
 # This is the genrule used by TF install headers to correctly
@@ -268,3 +267,4 @@ genrule(
     tar czf "$@" xla_extension
   """
 )
+