1f dep v2.3.0

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",

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",
+    ]),
+)

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<gpu::GpuExecutable *>(executable);
+}
+
+cpu::CpuExecutable *AsCpuExecutable(Executable *executable) {
+  return dynamic_cast<cpu::CpuExecutable *>(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<int64_t> *indices) {
+  const BufferAssignment *assignment =
+      GetBufferAssignment(local_executable->executable());
+  CHECK_NOTNULL(assignment);
+
+  std::unordered_map<int, int> 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<int64_t> 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

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 <vector>
+
+#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<int64_t> *indices);
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_JIT_XLA_LIB_RUNTIME_WORKSPACE_BYTES_H_

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",

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,