Skip to content

Commit

Permalink
Transition to OpenXLA (#45)
Browse files Browse the repository at this point in the history
  • Loading branch information
@seanmor5
seanmor5 authored Jul 28, 2023
1 parent 8e582a6 commit 7930e2d
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 65 deletions.
58 changes: 27 additions & 31 deletions Makefile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,47 +9,43 @@ TEMP ?= $(HOME)/.cache
# Public configuration
BUILD_MODE ?= opt # can also be dbg
BUILD_CACHE ?= $(TEMP)/xla_extension
TENSORFLOW_GIT_REPO ?= https://github.com/tensorflow/tensorflow.git
OPENXLA_GIT_REPO ?= https://github.com/openxla/xla.git

# Tensorflow 2.11.0
TENSORFLOW_GIT_REV ?= d5b57ca93e506df258271ea00fc29cf98383a374
OPENXLA_GIT_REV ?= c1e4a16e77a7ba2000003ccade3ffba3749ada35

# Private configuration
BAZEL_FLAGS = --define "framework_shared_object=false" -c $(BUILD_MODE)

TENSORFLOW_NS = tf-$(TENSORFLOW_GIT_REV)
TENSORFLOW_DIR = $(BUILD_CACHE)/$(TENSORFLOW_NS)
TENSORFLOW_XLA_EXTENSION_NS = tensorflow/compiler/xla/extension
TENSORFLOW_XLA_EXTENSION_DIR = $(TENSORFLOW_DIR)/$(TENSORFLOW_XLA_EXTENSION_NS)
TENSORFLOW_XLA_BUILD_ARCHIVE = $(TENSORFLOW_DIR)/bazel-bin/$(TENSORFLOW_XLA_EXTENSION_NS)/xla_extension.tar.gz

$(BUILD_ARCHIVE): $(TENSORFLOW_DIR) extension/BUILD
rm -f $(TENSORFLOW_XLA_EXTENSION_DIR) && \
ln -s "$(ROOT_DIR)/extension" $(TENSORFLOW_XLA_EXTENSION_DIR) && \
cd $(TENSORFLOW_DIR) && \
bazel build $(BAZEL_FLAGS) $(BUILD_FLAGS) $(BUILD_INTERNAL_FLAGS) //$(TENSORFLOW_XLA_EXTENSION_NS):xla_extension && \
OPENXLA_NS = xla-$(OPENXLA_GIT_REV)
OPENXLA_DIR = $(BUILD_CACHE)/$(OPENXLA_NS)
OPENXLA_XLA_EXTENSION_NS = xla/extension
OPENXLA_XLA_EXTENSION_DIR = $(OPENXLA_DIR)/$(OPENXLA_XLA_EXTENSION_NS)
OPENXLA_XLA_BUILD_ARCHIVE = $(OPENXLA_DIR)/bazel-bin/$(OPENXLA_XLA_EXTENSION_NS)/xla_extension.tar.gz

$(BUILD_ARCHIVE): $(OPENXLA_DIR) extension/BUILD
rm -f $(OPENXLA_XLA_EXTENSION_DIR) && \
ln -s "$(ROOT_DIR)/extension" $(OPENXLA_XLA_EXTENSION_DIR) && \
cd $(OPENXLA_DIR) && \
bazel build $(BAZEL_FLAGS) $(BUILD_FLAGS) $(BUILD_INTERNAL_FLAGS) //$(OPENXLA_XLA_EXTENSION_NS):xla_extension && \
mkdir -p $(dir $(BUILD_ARCHIVE)) && \
cp -f $(TENSORFLOW_XLA_BUILD_ARCHIVE) $(BUILD_ARCHIVE)

# Clones tensorflow
$(TENSORFLOW_DIR):
mkdir -p $(TENSORFLOW_DIR) && \
cp extension/patch_tensorflow.sh $(TENSORFLOW_DIR) && \
cp extension/tensorflow-alpine.patch $(TENSORFLOW_DIR) && \
cd $(TENSORFLOW_DIR) && \
cp -f $(OPENXLA_XLA_BUILD_ARCHIVE) $(BUILD_ARCHIVE)

# Clones OPENXLA
$(OPENXLA_DIR):
mkdir -p $(OPENXLA_DIR) && \
cd $(OPENXLA_DIR) && \
git init && \
git remote add origin $(TENSORFLOW_GIT_REPO) && \
git fetch --depth 1 origin $(TENSORFLOW_GIT_REV) && \
git remote add origin $(OPENXLA_GIT_REPO) && \
git fetch --depth 1 origin $(OPENXLA_GIT_REV) && \
git checkout FETCH_HEAD && \
sh patch_tensorflow.sh && \
rm $(TENSORFLOW_DIR)/.bazelversion
rm $(OPENXLA_DIR)/.bazelversion

# Print Tensorflow Dir
# Print OPENXLA Dir
PTD:
@ echo $(TENSORFLOW_DIR)
@ echo $(OPENXLA_DIR)

clean:
cd $(TENSORFLOW_DIR) && bazel clean --expunge
rm -f $(TENSORFLOW_XLA_EXTENSION_DIR)
rm -rf $(TENSORFLOW_DIR)
cd $(OPENXLA_DIR) && bazel clean --expunge
rm -f $(OPENXLA_XLA_EXTENSION_DIR)
rm -rf $(OPENXLA_DIR)
rm -rf $(TARGET_DIR)
106 changes: 72 additions & 34 deletions extension/BUILD
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
load("@org_tensorflow//tensorflow:tensorflow.bzl", "if_cuda_or_rocm",)
load("@org_tensorflow//tensorflow:tensorflow.bzl", "if_with_tpu_support",)
load("//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm",)
load("@tsl//tsl:tsl.bzl", "if_with_tpu_support",)
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda",)
load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm",)
load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependencies",)
load("@org_tensorflow//tensorflow:tensorflow.bzl", "transitive_hdrs",)
load("@tsl//tsl:tsl.bzl", "tsl_grpc_cc_dependencies",)
load("@tsl//tsl:tsl.bzl", "transitive_hdrs",)
load("@bazel_tools//tools/build_defs/pkg:pkg.bzl", "pkg_tar",)

package(default_visibility=["//visibility:private"])
Expand All @@ -13,53 +13,90 @@ package(default_visibility=["//visibility:private"])
cc_binary(
name = "libxla_extension.so",
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/gpu:backend_configs_cc_impl",
"//xla/stream_executor:dnn_proto_cc_impl",
"//xla:literal",
"//xla:shape_util",
"//xla:status",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla/client:xla_computation",
"//xla/mlir/utils:error_util",
"//xla/mlir_hlo",
"//xla/mlir_hlo:all_passes",
"//xla/pjrt:mlir_to_hlo",
"//xla/client/lib:lu_decomposition",
"//xla/client/lib:math",
"//xla/client/lib:qr",
"//xla/client/lib:svd",
"//xla/client/lib:self_adjoint_eig",
"//xla/client/lib:sorting",
"//xla/mlir_hlo:mhlo_passes",
"//xla/translate/hlo_to_mhlo:hlo_to_mlir_hlo",
"//xla/pjrt:interpreter_device",
"//xla/pjrt:pjrt_client",
"//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",
"@com_google_absl//absl/types:span",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/base:log_severity",
"@com_google_protobuf//:protobuf",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:ReconcileUnrealizedCasts",
"@llvm-project//mlir:SparseTensorDialect",
"@tf_runtime//:core_runtime",
"@tf_runtime//:hostcontext",
"@tf_runtime//third_party/llvm_derived:unique_any",
"@org_tensorflow//tensorflow/compiler/xla/client/lib:lu_decomposition",
"@org_tensorflow//tensorflow/compiler/xla/client/lib:math",
"@org_tensorflow//tensorflow/compiler/xla/client/lib:qr",
"@org_tensorflow//tensorflow/compiler/xla/client/lib:svd",
"@org_tensorflow//tensorflow/compiler/xla/client/lib:self_adjoint_eig",
"@org_tensorflow//tensorflow/compiler/xla/client/lib:sorting",
"@org_tensorflow//tensorflow/compiler/xla:literal",
"@org_tensorflow//tensorflow/compiler/xla:shape_util",
"@org_tensorflow//tensorflow/compiler/xla:status",
"@org_tensorflow//tensorflow/compiler/xla:statusor",
"@org_tensorflow//tensorflow/compiler/xla:types",
"@org_tensorflow//tensorflow/compiler/xla:util",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:mlir_to_hlo",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:interpreter_device",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:pjrt_client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:pjrt_compiler",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:tfrt_cpu_pjrt_client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:pjrt_c_api_client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:tpu_client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt/distributed",
"@org_tensorflow//tensorflow/compiler/xla/pjrt/gpu:se_gpu_pjrt_client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt/distributed:client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt/distributed:service",
"@org_tensorflow//tensorflow/core:lib",
"@org_tensorflow//tensorflow/core:lib_internal_impl",
"@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)
+ tf_grpc_cc_dependencies()
+ tsl_grpc_cc_dependencies()
+ if_with_tpu_support([
"@org_tensorflow//tensorflow/core/tpu:tpu_api_dlsym_initializer",
"@tsl//tsl/core/tpu:tpu_api_dlsym_initializer",
])
+ if_cuda_or_rocm([
"@org_tensorflow//tensorflow/compiler/xla/service:gpu_plugin",
"//xla/service:gpu_plugin",
])
+ if_cuda([
"@org_tensorflow//tensorflow/compiler/xla/stream_executor:cuda_platform"
"//xla/stream_executor:cuda_platform"
])
+ if_rocm([
"@org_tensorflow//tensorflow/compiler/xla/stream_executor:rocm_platform"
"//xla/stream_executor:rocm_platform"
]),
copts = ["-fvisibility=default"],
linkopts = ["-shared"],
features = ["-use_header_modules"],
linkshared = 1,
)

Expand All @@ -79,6 +116,7 @@ genrule(
name = "xla_extension_headers",
srcs = [
":xla_extension_dep_headers",
"//xla/mlir_hlo:hlo_ops_enums_inc_gen",
],
outs = ["include"],
cmd = """
Expand Down

0 comments on commit 7930e2d

Please sign in to comment.