diff --git a/Makefile b/Makefile index d8e5a12..be420eb 100644 --- a/Makefile +++ b/Makefile @@ -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) diff --git a/extension/BUILD b/extension/BUILD index e205930..6f4a164 100644 --- a/extension/BUILD +++ b/extension/BUILD @@ -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"]) @@ -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, ) @@ -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 = """