Automated rollback of commit 185aa89912376d4088c22615908696cd30f9951b

PiperOrigin-RevId: 213394522

4 files changed, 10 insertions, 57 deletions

diff --git a/tensorflow/core/framework/allocator.cc b/tensorflow/core/framework/allocator.cc
index 84cee5569c..2a7ee16a16 100644
--- a/tensorflow/core/framework/allocator.cc
+++ b/tensorflow/core/framework/allocator.cc
@@ -196,7 +196,7 @@ class CPUAllocatorFactory : public AllocatorFactory {
   class CPUSubAllocator : public SubAllocator {
    public:
     explicit CPUSubAllocator(CPUAllocator* cpu_allocator)
-        : SubAllocator({}, {}), cpu_allocator_(cpu_allocator) {}
+        : cpu_allocator_(cpu_allocator) {}
 
     void* Alloc(size_t alignment, size_t num_bytes) override {
       return cpu_allocator_->AllocateRaw(alignment, num_bytes);
@@ -222,22 +222,4 @@ Allocator* cpu_allocator() {
   }
   return cpu_alloc;
 }
-
-SubAllocator::SubAllocator(const std::vector<Visitor>& alloc_visitors,
-                           const std::vector<Visitor>& free_visitors)
-    : alloc_visitors_(alloc_visitors), free_visitors_(free_visitors) {}
-
-void SubAllocator::VisitAlloc(void* ptr, int index, size_t num_bytes) {
-  for (const auto& v : alloc_visitors_) {
-    v(ptr, index, num_bytes);
-  }
-}
-
-void SubAllocator::VisitFree(void* ptr, int index, size_t num_bytes) {
-  // Although we don't guarantee any order of visitor application, strive
-  // to apply free visitors in reverse order of alloc visitors.
-  for (int i = free_visitors_.size() - 1; i >= 0; --i) {
-    free_visitors_[i](ptr, index, num_bytes);
-  }
-}
 }  // namespace tensorflow
diff --git a/tensorflow/core/framework/allocator.h b/tensorflow/core/framework/allocator.h
index 8c23604625..ded120b704 100644
--- a/tensorflow/core/framework/allocator.h
+++ b/tensorflow/core/framework/allocator.h
@@ -24,7 +24,6 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_handle.h"
 #include "tensorflow/core/framework/type_traits.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
@@ -388,36 +387,13 @@ void EnableCPUAllocatorStats(bool enable);
 // full statistics. By default, it's disabled.
 void EnableCPUAllocatorFullStats(bool enable);
 
-// An object that does the underlying suballoc/free of memory for a higher-level
-// allocator.  The expectation is that the higher-level allocator is doing some
-// kind of cache or pool management so that it will call SubAllocator::Alloc and
-// Free relatively infrequently, compared to the number of times its own
-// AllocateRaw and Free methods are called.
+// Abstract interface of an object that does the underlying suballoc/free of
+// memory for a higher-level allocator.
 class SubAllocator {
  public:
-  // Visitor gets called with a pointer to a memory area and its
-  // size in bytes.  The index value will be numa_node for a CPU
-  // allocator and GPU id for a GPU allocator.
-  typedef std::function<void(void*, int index, size_t)> Visitor;
-
-  SubAllocator(const std::vector<Visitor>& alloc_visitors,
-               const std::vector<Visitor>& free_visitors);
-
   virtual ~SubAllocator() {}
   virtual void* Alloc(size_t alignment, size_t num_bytes) = 0;
   virtual void Free(void* ptr, size_t num_bytes) = 0;
-
- protected:
-  // Implementation of Alloc() method must call this on newly allocated
-  // value.
-  void VisitAlloc(void* ptr, int index, size_t num_bytes);
-
-  // Implementation of Free() method must call this on value to be
-  // freed immediately before deallocation.
-  void VisitFree(void* ptr, int index, size_t num_bytes);
-
-  const std::vector<Visitor> alloc_visitors_;
-  const std::vector<Visitor> free_visitors_;
 };
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/framework/device_base.h b/tensorflow/core/framework/device_base.h
index 53ac639b4c..794250a2c1 100644
--- a/tensorflow/core/framework/device_base.h
+++ b/tensorflow/core/framework/device_base.h
@@ -214,12 +214,10 @@ class DeviceBase {
 
   // This is overridden by GPU devices to reinitialize the derived
   // type returned by MakeGpuDevice.
-  virtual Status ReinitializeGpuDevice(OpKernelContext* /*context*/,
-                                       PerOpGpuDevice* /*device*/,
-                                       DeviceContext* /*dc*/,
-                                       Allocator* /*allocator*/) {
-    return Status::OK();
-  }
+  virtual void ReinitializeGpuDevice(OpKernelContext* /*context*/,
+                                     PerOpGpuDevice* /*device*/,
+                                     DeviceContext* /*dc*/,
+                                     Allocator* /*allocator*/) {}
 
   // Unimplemented by default
   virtual const DeviceAttributes& attributes() const;
diff --git a/tensorflow/core/framework/op_kernel.cc b/tensorflow/core/framework/op_kernel.cc
index 3e34bf0418..80f2b12987 100644
--- a/tensorflow/core/framework/op_kernel.cc
+++ b/tensorflow/core/framework/op_kernel.cc
@@ -265,12 +265,9 @@ OpKernelContext::OpKernelContext(Params* params, int num_outputs)
   params_->ensure_eigen_gpu_device();
   if (params_->eigen_gpu_device != nullptr) {
     Allocator* eigen_gpu_allocator = get_allocator(AllocatorAttributes());
-    Status s = params_->device->ReinitializeGpuDevice(
-        this, params_->eigen_gpu_device, params_->op_device_context,
-        eigen_gpu_allocator);
-    if (!s.ok()) {
-      SetStatus(s);
-    }
+    params_->device->ReinitializeGpuDevice(this, params_->eigen_gpu_device,
+                                           params_->op_device_context,
+                                           eigen_gpu_allocator);
   }
   if (params_->record_tensor_accesses) {
     referenced_tensors_.Init();