Merge pull request tensorflow#8216 from dandelionmane/branch_149567254

Branch 149567254

Author: Dandelion Mané

BUILD

@@ -1,28 +0,0 @@
-# Description:
-# TensorFlow is an open source software library for numerical computation using
-# data flow graphs.
-
-package(
-    default_visibility = [
-        "//tensorflow:internal",
-        "//tensorflow_models:__subpackages__",
-    ],
-)
-
-licenses(["notice"])  # Apache 2.0
-
-exports_files(["LICENSE"])
-
-load(
-    "//tensorflow:tensorflow.bzl",
-    "cc_header_only_library",
-)
-
-cc_header_only_library(
-    name = "protobuf_headers",
-    includes = ["external/protobuf/src"],
-    visibility = ["//visibility:public"],
-    deps = [
-        "@protobuf//:protobuf",
-    ],
-)

tensorflow/BUILD

@@ -221,6 +221,11 @@ filegroup(
         "//tensorflow/core/debug:all_files",
         "//tensorflow/core/distributed_runtime:all_files",
         "//tensorflow/core/distributed_runtime/rpc:all_files",
+        "//tensorflow/core/grappler:all_files",
+        "//tensorflow/core/grappler/clusters:all_files",
+        "//tensorflow/core/grappler/costs:all_files",
+        "//tensorflow/core/grappler/inputs:all_files",
+        "//tensorflow/core/grappler/optimizers:all_files",
         "//tensorflow/core/kernels:all_files",
         "//tensorflow/core/kernels/cloud:all_files",
         "//tensorflow/core/kernels/hexagon:all_files",

tensorflow/compiler/jit/kernels/xla_local_launch_op.cc

@@ -39,6 +39,8 @@ namespace gpu = perftools::gputools;
 namespace tensorflow {
 
 // Adapter class that wraps a Tensorflow allocator as an XLA allocator.
+// Assumes that the Tensorflow allocator permits asynchronous deallocation:
+// see comment on `AllowsAsynchronousDeallocation()`.
 class XlaAllocator : public xla::DeviceMemoryAllocator {
  public:
   XlaAllocator(const perftools::gputools::Platform* platform,
@@ -54,6 +56,15 @@ class XlaAllocator : public xla::DeviceMemoryAllocator {
   Status MakeTensorFromBuffer(gpu::DeviceMemoryBase buffer, DataType dtype,
                               const TensorShape& shape, Tensor* tensor) const;
 
+  // The Tensorflow BFC allocator used on GPU allows host-side deallocation
+  // before GPU execution takes place. Tensorflow uses the ordering of the main
+  // compute stream to enforce a happens-before relationship between a memory
+  // allocation and code that reuses the same memory. If Tensorflow adds
+  // support for multiple GPU streams or allocators with different ordering
+  // requirements, this code may need to change.
+  // (This attribute has no effect on CPU.)
+  bool AllowsAsynchronousDeallocation() const override { return true; }
+
  private:
   OpKernelContext* const op_context_;
 

tensorflow/compiler/jit/kernels/xla_local_launch_op.h

@@ -31,8 +31,9 @@ namespace tensorflow {
 // Once all inputs are present, and their shapes are known, the op can
 // use a 'XlaCompilationCache' to compile and execute code which is specific
 // to the shapes of input Tensors.
-// XlaLocalLaunchOp uses xla::LocalClient::ExecuteLocally and passes
-// arguments into/out of XLA in device memory.
+// XlaLocalLaunchOp uses xla::LocalClient::Compile() and
+// xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device
+// memory.
 class XlaLocalLaunchOp : public OpKernel {
  public:
   explicit XlaLocalLaunchOp(OpKernelConstruction* ctx);

tensorflow/compiler/tests/BUILD

@@ -195,6 +195,21 @@ tf_xla_py_test(
     ],
 )
 
+tf_xla_py_test(
+    name = "pooling_ops_3d_test",
+    size = "medium",
+    srcs = ["pooling_ops_3d_test.py"],
+    shard_count = 10,
+    deps = [
+        ":xla_test",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:nn_ops",
+        "//tensorflow/python:nn_ops_gen",
+        "//tensorflow/python:platform_test",
+    ],
+)
+
 tf_xla_py_test(
     name = "random_ops_test",
     size = "small",
@@ -224,6 +239,17 @@ tf_xla_py_test(
     ],
 )
 
+tf_xla_py_test(
+    name = "reverse_ops_test",
+    size = "small",
+    srcs = ["reverse_ops_test.py"],
+    deps = [
+        ":xla_test",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework_for_generated_wrappers",
+    ],
+)
+
 tf_xla_py_test(
     name = "ternary_ops_test",
     size = "small",