diff --git a/WORKSPACE b/WORKSPACE index ea741c31c7f8d8..f520e8a5aff61d 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -36,6 +36,13 @@ load( bazel_toolchains_repositories() +http_archive( + name = "io_bazel_rules_docker", + sha256 = "aed1c249d4ec8f703edddf35cbe9dfaca0b5f5ea6e4cd9e83e99f3b0d1136c3d", + strip_prefix = "rules_docker-0.7.0", + urls = ["https://github.com/bazelbuild/rules_docker/archive/v0.7.0.tar.gz"], + ) + load( "@io_bazel_rules_docker//repositories:repositories.bzl", container_repositories = "repositories", diff --git a/tensorflow/compiler/jit/xla_lib/BUILD b/tensorflow/compiler/jit/xla_lib/BUILD new file mode 100644 index 00000000000000..b9d1c5573627f1 --- /dev/null +++ b/tensorflow/compiler/jit/xla_lib/BUILD @@ -0,0 +1,64 @@ +licenses(["notice"]) # Apache 2.0 + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", + "tf_cc_binary", + "tf_cc_shared_object", +) +load( + "//tensorflow/core/platform:build_config.bzl", + "tf_additional_all_protos", + "tf_proto_library", + "tf_proto_library_cc", + "tf_proto_library_py", +) +load("//tensorflow/compiler/xla:xla.bzl", "xla_py_proto_library") +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") + +cc_library( + name = "xla_runtime_util", + srcs = [ + "xla_runtime_util.cc", + "xla_runtime_util.h", + ], + copts = if_cuda(["-DGOOGLE_CUDA=1"]), + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/stream_executor:stream", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service/cpu:cpu_executable", + "//tensorflow/compiler/xla/service/gpu:gpu_executable", + ], +) + +tf_cc_shared_object( + name = "libxla_core.so", + linkopts = select({ + "//tensorflow:windows": [], + "//conditions:default": [ + "-z muldefs", + ], + }), + visibility = ["//visibility:public"], + deps = [ + ":xla_runtime_util", + "//tensorflow/compiler/jit:xla_tensor", + "//tensorflow/compiler/xla/client", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client:xla_computation", + "//tensorflow/compiler/xla/client:executable_build_options", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/client:client_library", + "//tensorflow/compiler/xla/service", + "//tensorflow/compiler/xla/service:cpu_plugin", +# "//tensorflow/compiler/jit:xla_cpu_jit", + "//tensorflow/compiler/tf2xla:tf2xla_util", + "//tensorflow/compiler/tf2xla/lib:util", + "//tensorflow/core:lib", + "@com_google_absl//absl/strings", + ] + if_cuda([ + "//tensorflow/compiler/xla/service:gpu_plugin", +# "//tensorflow/compiler/jit:xla_gpu_jit", + ]), +) diff --git a/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc new file mode 100644 index 00000000000000..23a37120d20e69 --- /dev/null +++ b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc @@ -0,0 +1,110 @@ +#if GOOGLE_CUDA +#include "tensorflow/stream_executor/gpu/gpu_stream.h" +#endif // GOOGLE_CUDA +#include "tensorflow/compiler/xla/shape_tree.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/service/hlo_value.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_executable.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_executable.h" + +#include "tensorflow/compiler/jit/xla_lib/xla_runtime_util.h" + +namespace xla { + +#if GOOGLE_CUDA +void SwapGpuStreamHandle(se::Stream *stream, void **gpu_stream) { + void **cuda_stream = se::gpu::AsGpuStream(stream)->GpuStreamMemberHack(); + void *tmp_stream = *cuda_stream; + *cuda_stream = *gpu_stream; + *gpu_stream = tmp_stream; +} +#endif // GOOGLE_CUDA + +inline size_t Align(int alignment, size_t size) { + return (size + alignment - 1) / alignment * alignment; +} + +gpu::GpuExecutable *AsGpuExecutable(Executable *executable) { + return dynamic_cast(executable); +} + +cpu::CpuExecutable *AsCpuExecutable(Executable *executable) { + return dynamic_cast(executable); +} + +const BufferAssignment *GetBufferAssignment(Executable *executable) { + const BufferAssignment *assignment = nullptr; + if (auto *gpu_executable = AsGpuExecutable(executable)) { + assignment = (gpu_executable->GetBufferAssignment()).get(); + } else if (auto *cpu_executable = AsCpuExecutable(executable)) { + assignment = &(cpu_executable->buffer_assignment()); + } else { + LOG(FATAL) << "Only support CPU or GPU executable."; + } + return assignment; +} + +size_t CalcWorkspaceByteSize(LocalExecutable *local_executable) { + const BufferAssignment *assignment = + GetBufferAssignment(local_executable->executable()); + CHECK_NOTNULL(assignment); + + size_t workspace_bytes = 0; + for (int i = 0; i < assignment->Allocations().size(); ++i) { + const BufferAllocation& allocation = assignment->GetAllocation(i); + // if (!allocation.is_entry_computation_parameter()) { + if (allocation.IsPreallocatedTempBuffer() || allocation.is_tuple()) { + workspace_bytes += Align(64/*alignment*/, allocation.size()); + } + } + return workspace_bytes; +} + +Status ResultAllocationIndices(LocalExecutable *local_executable, + std::vector *indices) { + const BufferAssignment *assignment = + GetBufferAssignment(local_executable->executable()); + CHECK_NOTNULL(assignment); + + std::unordered_map allocation_order; + for (int i = 0; i < assignment->Allocations().size(); ++i) { + const BufferAllocation& allocation = assignment->GetAllocation(i); + if ((allocation.maybe_live_out() || + allocation.IsPreallocatedTempBuffer()) && + !allocation.is_entry_computation_parameter()) { + allocation_order.emplace(i, allocation_order.size()); + } + } + + const HloModule &module = local_executable->executable()->module(); + const HloInstruction* root = module.entry_computation()->root_instruction(); + const InstructionValueSet &value_set = + assignment->dataflow_analysis().GetInstructionValueSet(root); + + CHECK(root->shape().IsTuple()); + int tuple_size = root->shape().tuple_shapes_size(); + std::vector allocation_indices(tuple_size); + for (int i = 0; i < tuple_size; ++i) { + const auto& sources = value_set.element({i}); + CHECK_EQ(1, sources.values().size()); + auto instruction = sources.values()[0]->instruction(); + + TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice slice, + assignment->GetUniqueSlice( + instruction, sources.values()[0]->index())); + if (slice.allocation()->is_entry_computation_parameter()) { + // Output aliased with input will not be allocated, so we assign the + // allocation index -1 + allocation_indices[i] = -1; + } else { + // Slice index is the allocation index + allocation_indices[i] = allocation_order.at(slice.index()); + } + } + + *indices = std::move(allocation_indices); + return Status::OK(); +} + +} // namespace xla diff --git a/tensorflow/compiler/jit/xla_lib/xla_runtime_util.h b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.h new file mode 100644 index 00000000000000..dd951f367dff2c --- /dev/null +++ b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.h @@ -0,0 +1,20 @@ +#ifndef TENSORFLOW_COMPILER_JIT_XLA_LIB_RUNTIME_WORKSPACE_BYTES_H_ +#define TENSORFLOW_COMPILER_JIT_XLA_LIB_RUNTIME_WORKSPACE_BYTES_H_ + +#include + +#include "tensorflow/stream_executor/stream.h" +#include "tensorflow/compiler/xla/client/local_client.h" + +namespace xla { + +void SwapGpuStreamHandle(se::Stream *stream, void **gpu_stream); + +size_t CalcWorkspaceByteSize(LocalExecutable *local_executable); + +Status ResultAllocationIndices(LocalExecutable *local_executable, + std::vector *indices); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_JIT_XLA_LIB_RUNTIME_WORKSPACE_BYTES_H_ diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index 24446d846cfdf9..0f3ab4bda0c5a4 100755 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -714,7 +714,7 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""): LLVM_SHA256 = "a21b752ee1866e195f3f72c7931c79f8c4ecc0f14861488284bdc2bdf14d6fe9" LLVM_URLS = [ "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT), - "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT), + "https://oneflow-static.oss-cn-beijing.aliyuncs.com/deps/llvm-project-7e825abd5704ce28b166f9463d4bd304348fd2a9.tar.gz".format(commit = LLVM_COMMIT), ] tf_http_archive( name = "llvm-project", diff --git a/third_party/repo.bzl b/third_party/repo.bzl index a4d2b899f801ff..1799b1f6aa0d16 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -12,6 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. + +def convert_url_to_oss_key1(url): + path = url[len("https://")::] + return "/".join(["third_party_mirror", "https", path]) + + +def should_be_mirrored(url): + return ( + url.endswith(("gz", "tar", "zip")) + and not "mirror.tensorflow.org" in url + and not "mirror.bazel.build" in url + and not "aliyuncs.com" in url + ) + + +def convert_url_to_oss_https_url(url): + if should_be_mirrored(url): + key = convert_url_to_oss_key1(url) + return "https://oneflow-static.oss-cn-beijing.aliyuncs.com/" + key + else: + return url + """Utilities for defining TensorFlow Bazel dependencies.""" _SINGLE_URL_WHITELIST = depset([ @@ -186,7 +208,7 @@ def _third_party_http_archive(ctx): else: ctx.download_and_extract( - ctx.attr.urls, + [convert_url_to_oss_https_url(u) for u in ctx.attr.urls], "", ctx.attr.sha256, ctx.attr.type,