From dea9488e5f05ffcaff7e729f33d475af3a7021ba Mon Sep 17 00:00:00 2001 From: Shenghang Tsai Date: Mon, 3 Aug 2020 11:13:54 +0800 Subject: [PATCH 01/14] apply diff --- WORKSPACE | 7 ++ tensorflow/compiler/jit/xla_lib/BUILD | 63 +++++++++++ .../compiler/jit/xla_lib/xla_runtime_util.cc | 106 ++++++++++++++++++ .../compiler/jit/xla_lib/xla_runtime_util.h | 20 ++++ 4 files changed, 196 insertions(+) create mode 100644 tensorflow/compiler/jit/xla_lib/BUILD create mode 100644 tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc create mode 100644 tensorflow/compiler/jit/xla_lib/xla_runtime_util.h diff --git a/WORKSPACE b/WORKSPACE index 74ea14d0fd79b6..b93c286a54c9e5 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -34,6 +34,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..cd729553bec777 --- /dev/null +++ b/tensorflow/compiler/jit/xla_lib/BUILD @@ -0,0 +1,63 @@ +licenses(["notice"]) # Apache 2.0 + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", + "tf_cc_binary", + "tf_cc_shared_object", +) +load( + "//tensorflow/core:platform/default/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_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", + ], + 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..21f2f2de0b3b25 --- /dev/null +++ b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc @@ -0,0 +1,106 @@ +#include "tensorflow/stream_executor/gpu/gpu_stream.h" +#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 { + +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; +} + +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_ From d9de45b18e459f4069e612179ae8d8fab5e67b5e Mon Sep 17 00:00:00 2001 From: Shenghang Tsai Date: Wed, 2 Sep 2020 18:48:42 +0800 Subject: [PATCH 02/14] fix bazel --- tensorflow/compiler/jit/xla_lib/BUILD | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/jit/xla_lib/BUILD b/tensorflow/compiler/jit/xla_lib/BUILD index cd729553bec777..0b2817e1fa50fa 100644 --- a/tensorflow/compiler/jit/xla_lib/BUILD +++ b/tensorflow/compiler/jit/xla_lib/BUILD @@ -7,7 +7,7 @@ load( "tf_cc_shared_object", ) load( - "//tensorflow/core:platform/default/build_config.bzl", + "//tensorflow/core:platform/default:build_config.bzl", "tf_additional_all_protos", "tf_proto_library", "tf_proto_library_cc", From 3e93f003745d2c0f6c9c801f39ecabb1a299fff9 Mon Sep 17 00:00:00 2001 From: Shenghang Tsai Date: Wed, 2 Sep 2020 18:51:58 +0800 Subject: [PATCH 03/14] fix --- tensorflow/compiler/jit/xla_lib/BUILD | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/jit/xla_lib/BUILD b/tensorflow/compiler/jit/xla_lib/BUILD index 0b2817e1fa50fa..16baa0b797196f 100644 --- a/tensorflow/compiler/jit/xla_lib/BUILD +++ b/tensorflow/compiler/jit/xla_lib/BUILD @@ -7,7 +7,7 @@ load( "tf_cc_shared_object", ) load( - "//tensorflow/core:platform/default:build_config.bzl", + "//tensorflow/core/platform:build_config.bzl", "tf_additional_all_protos", "tf_proto_library", "tf_proto_library_cc", From ee1a8ee737487d57b66f6ea83cdd82c0136521b5 Mon Sep 17 00:00:00 2001 From: tsai Date: Wed, 2 Sep 2020 14:36:26 +0000 Subject: [PATCH 04/14] fix xla proto --- tensorflow/compiler/jit/xla_lib/BUILD | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/jit/xla_lib/BUILD b/tensorflow/compiler/jit/xla_lib/BUILD index 16baa0b797196f..d9a824029935c9 100644 --- a/tensorflow/compiler/jit/xla_lib/BUILD +++ b/tensorflow/compiler/jit/xla_lib/BUILD @@ -13,7 +13,7 @@ load( "tf_proto_library_cc", "tf_proto_library_py", ) -load("//tensorflow/compiler/xla:xla.bzl", "xla_proto_library") +load("//tensorflow/compiler/xla:xla.bzl", "xla_py_proto_library") load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") cc_library( From fe260b50dca0ae9d165c4575c77d30982db184cd Mon Sep 17 00:00:00 2001 From: tsai Date: Fri, 5 Mar 2021 13:36:04 +0800 Subject: [PATCH 05/14] add --- tensorflow/workspace.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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", From 797818e58b807689ae443ada6e0675a5a3ee25d9 Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 21:31:33 +0800 Subject: [PATCH 06/14] add mirror --- third_party/repo.bzl | 40 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index a4d2b899f801ff..5991fe71c585ba 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -12,6 +12,44 @@ # See the License for the specific language governing permissions and # limitations under the License. +from urllib.parse import urlparse +import os + + +def convert_url_to_oss_key(url): + parsed = urlparse(url) + assert parsed.scheme == "https", url + assert not parsed.params + assert not parsed.query + assert not parsed.port + assert not parsed.fragment + assert parsed.path.startswith("/") + path = parsed.path[1::] + return os.path.join("third_party_mirror", parsed.scheme, parsed.netloc, path) + + +def should_be_mirrored(url: str): + parsed = urlparse(url) + return ( + not parsed.port + and not parsed.query + and not parsed.params + and url.endswith(("gz", "tar", "zip")) + and url + 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_key(url) + prefix = "https://oneflow-static.oss-cn-beijing.aliyuncs.com/" + return os.path.join(prefix, key) + else: + return url + """Utilities for defining TensorFlow Bazel dependencies.""" _SINGLE_URL_WHITELIST = depset([ @@ -186,7 +224,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, From 8bffb83d701c08d31fc2ac310ad4fe5db43853ef Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 22:08:00 +0800 Subject: [PATCH 07/14] fix --- third_party/repo.bzl | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index 5991fe71c585ba..de3d74248cfd2c 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -12,30 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -from urllib.parse import urlparse -import os - -def convert_url_to_oss_key(url): - parsed = urlparse(url) - assert parsed.scheme == "https", url - assert not parsed.params - assert not parsed.query - assert not parsed.port - assert not parsed.fragment - assert parsed.path.startswith("/") - path = parsed.path[1::] - return os.path.join("third_party_mirror", parsed.scheme, parsed.netloc, path) +def convert_url_to_oss_key1(url: str): + https = "https://" + assert url.startswith(https) + path = url[len(https) :] + return "/".join(["third_party_mirror", "https", path]) def should_be_mirrored(url: str): - parsed = urlparse(url) return ( - not parsed.port - and not parsed.query - and not parsed.params - and url.endswith(("gz", "tar", "zip")) - and url + 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 @@ -44,7 +31,7 @@ def should_be_mirrored(url: str): def convert_url_to_oss_https_url(url): if should_be_mirrored(url): - key = convert_url_to_oss_key(url) + key = convert_url_to_oss_key1(url) prefix = "https://oneflow-static.oss-cn-beijing.aliyuncs.com/" return os.path.join(prefix, key) else: From fc70c85c6072f5db61a14f768cdc0b833af089d1 Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 22:10:46 +0800 Subject: [PATCH 08/14] fix import --- third_party/repo.bzl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index de3d74248cfd2c..98bc2e44a8d24e 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -32,8 +32,7 @@ def should_be_mirrored(url: str): def convert_url_to_oss_https_url(url): if should_be_mirrored(url): key = convert_url_to_oss_key1(url) - prefix = "https://oneflow-static.oss-cn-beijing.aliyuncs.com/" - return os.path.join(prefix, key) + return "https://oneflow-static.oss-cn-beijing.aliyuncs.com/" + key else: return url From 25061eb095baaa8ccaf12052a6e4030b430244c9 Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 22:16:31 +0800 Subject: [PATCH 09/14] fix assert --- third_party/repo.bzl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index 98bc2e44a8d24e..2f3829e1755906 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -14,9 +14,7 @@ def convert_url_to_oss_key1(url: str): - https = "https://" - assert url.startswith(https) - path = url[len(https) :] + path = url[len("https://") :] return "/".join(["third_party_mirror", "https", path]) From 864baba20a112462b82fb583406df1104294dea4 Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 22:40:54 +0800 Subject: [PATCH 10/14] add --- third_party/repo.bzl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index 2f3829e1755906..2e4d96f062353a 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -14,7 +14,8 @@ def convert_url_to_oss_key1(url: str): - path = url[len("https://") :] + path = url[8::] + url.split("/") return "/".join(["third_party_mirror", "https", path]) From 1852a92a5af581b14f0283a181e4f615101b88c0 Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 22:45:11 +0800 Subject: [PATCH 11/14] fix --- third_party/repo.bzl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index 2e4d96f062353a..b451cd1beabd3a 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -13,13 +13,13 @@ # limitations under the License. -def convert_url_to_oss_key1(url: str): - path = url[8::] +def convert_url_to_oss_key1(url): + path = url[len("https://")::] url.split("/") return "/".join(["third_party_mirror", "https", path]) -def should_be_mirrored(url: str): +def should_be_mirrored(url): return ( url.endswith(("gz", "tar", "zip")) and not "mirror.tensorflow.org" in url From fc42cf2a17e4af9f494278ddee66b6d17e1e9eaf Mon Sep 17 00:00:00 2001 From: tsai Date: Sat, 6 Mar 2021 22:46:12 +0800 Subject: [PATCH 12/14] rm unused code --- third_party/repo.bzl | 1 - 1 file changed, 1 deletion(-) diff --git a/third_party/repo.bzl b/third_party/repo.bzl index b451cd1beabd3a..1799b1f6aa0d16 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -15,7 +15,6 @@ def convert_url_to_oss_key1(url): path = url[len("https://")::] - url.split("/") return "/".join(["third_party_mirror", "https", path]) From b5e87ed74a2a1f870e11b3656e00f321c7c81224 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 22 Jun 2021 06:40:00 +0000 Subject: [PATCH 13/14] add GOOGLE_CUDA --- tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc index 21f2f2de0b3b25..23a37120d20e69 100644 --- a/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc +++ b/tensorflow/compiler/jit/xla_lib/xla_runtime_util.cc @@ -1,4 +1,6 @@ +#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" @@ -10,12 +12,14 @@ 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; From 7016a22292a607edc4175d07dae263faad31cd04 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 22 Jun 2021 08:41:20 +0000 Subject: [PATCH 14/14] fix cuda fail --- tensorflow/compiler/jit/xla_lib/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/compiler/jit/xla_lib/BUILD b/tensorflow/compiler/jit/xla_lib/BUILD index d9a824029935c9..b9d1c5573627f1 100644 --- a/tensorflow/compiler/jit/xla_lib/BUILD +++ b/tensorflow/compiler/jit/xla_lib/BUILD @@ -22,6 +22,7 @@ cc_library( "xla_runtime_util.cc", "xla_runtime_util.h", ], + copts = if_cuda(["-DGOOGLE_CUDA=1"]), visibility = ["//visibility:public"], deps = [ "//tensorflow/stream_executor:stream",