Revert "Add function to materialize COW storages (#113396)"

This reverts commit e2f0900.

Reverted #113396 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](#113396 (comment)))

aten/src/ATen/EmptyTensor.cpp

@@ -311,7 +311,6 @@ struct MetaAllocator final : public at::Allocator {
   DeleterFnPtr raw_deleter() const override {
     return deleter;
   }
-  void copy_data(void* dest, const void* src, std::size_t count) const final {}
 };
 
 static MetaAllocator g_meta_alloc;

aten/src/ATen/cuda/CachingHostAllocator.cpp

@@ -301,10 +301,6 @@ class CUDAHostAllocator {
     }
   }
 
-  void copy_data(void* dest, const void* src, std::size_t count) const {
-    TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for CUDAHostAllocator");
-  }
-
  private:
   void process_events() {
     while (true) {
@@ -500,10 +496,6 @@ struct CUDAHostAllocatorWrapper final : public at::Allocator {
         &CUDAHostAllocatorDeleter,
         at::DeviceType::CPU};
   }
-
-  void copy_data(void* dest, const void* src, std::size_t count) const final {
-    getCUDAHostAllocator().copy_data(dest, src, count);
-  }
 };
 
 static CUDAHostAllocatorWrapper cuda_host_allocator;

aten/src/ATen/hip/impl/HIPAllocatorMasqueradingAsCUDA.h

@@ -23,9 +23,6 @@ class HIPAllocatorMasqueradingAsCUDA final : public Allocator {
   DeleterFnPtr raw_deleter() const override {
     return allocator_->raw_deleter();
   }
-  void copy_data(void* dest, const void* src, std::size_t count) const final {
-    allocator_->copy_data(dest, src, count);
-  }
 };
 
 }} // namespace c10::hip

aten/src/ATen/mps/MPSAllocator.mm

@@ -819,10 +819,6 @@ bool waitForEvents(c10::ArrayRef<const void*> buffers) const override {
     return _getAllocImpl().format_size(size);
   }
 
-  void copy_data(void* dest, const void* src, std::size_t count) const final {
-    default_copy_data(dest, src, count);
-  }
-
  private:
   bool m_has_unified_memory;
   uint32_t m_usage;

aten/src/ATen/native/Resize.cpp

@@ -94,14 +94,13 @@ void resize_bytes_cpu(StorageImpl* storage, size_t size_bytes) {
   if (size_bytes != 0) {
     new_data = storage->allocator()->allocate(size_bytes);
   }
-  const at::DataPtr& old_data = storage->data_ptr();
+  at::DataPtr old_data = storage->set_data_ptr(std::move(new_data));
   const auto old_capacity = storage->nbytes();
+  storage->set_nbytes(size_bytes);
   const auto copy_capacity = std::min(size_bytes, old_capacity);
   if (old_data != nullptr && copy_capacity > 0) {
-    memcpy(new_data.get(), old_data.get(), copy_capacity);
+    memcpy(storage->mutable_data(), old_data.get(), copy_capacity);
   }
-  storage->set_data_ptr_noswap(std::move(new_data));
-  storage->set_nbytes(size_bytes);
 }
 
 // Call the sparse implementation in SparseTensor.cpp directly.

aten/src/ATen/native/TensorFactories.h

@@ -129,7 +129,6 @@ struct ZeroTensorAllocator final : public at::Allocator {
   DeleterFnPtr raw_deleter() const override {
     return deleter;
   }
-  void copy_data(void* dest, const void* src, std::size_t count) const final {}
   at::Device device_;
 };
 

aten/src/ATen/test/CMakeLists.txt

@@ -13,7 +13,6 @@ list(APPEND ATen_CPU_TEST_SRCS
   ${CMAKE_CURRENT_SOURCE_DIR}/atest.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/basic.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/broadcast_test.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/cpu_allocator_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/cpu_generator_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/cpu_profiling_allocator_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/cpu_rng_test.cpp
@@ -55,7 +54,6 @@ list(APPEND ATen_CPU_TEST_SRCS
   )
 
 list(APPEND ATen_CUDA_TEST_SRCS
-  ${CMAKE_CURRENT_SOURCE_DIR}/cuda_allocator_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/cuda_apply_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/cuda_atomic_ops_test.cu
   ${CMAKE_CURRENT_SOURCE_DIR}/cuda_caching_host_allocator_test.cpp

aten/src/ATen/test/xla_tensor_test.cpp

@@ -2,8 +2,6 @@
 
 #include <ATen/ATen.h>
 
-#include <ATen/test/allocator_clone_test.h>
-
 using namespace at;
 
 void XLAFree(void *ptr) {
@@ -24,9 +22,6 @@ struct XLAAllocator final : public at::Allocator {
   at::DeleterFnPtr raw_deleter() const override {
     return &XLAFree;
   }
-  void copy_data(void* dest, const void* src, std::size_t count) const final {
-    default_copy_data(dest, src, count);
-  }
 };
 
 TEST(XlaTensorTest, TestNoStorage) {
@@ -38,11 +33,3 @@ TEST(XlaTensorTest, TestNoStorage) {
   at::Tensor t(std::move(tensor_impl));
   ASSERT_TRUE(t.device() == at::Device(DeviceType::XLA, 0));
 }
-
-TEST(XlaTensorTest, test_allocator_clone) {
-  if (!at::hasXLA()) {
-    return;
-  }
-  XLAAllocator allocator;
-  test_allocator_clone(&allocator);
-}

c10/core/Allocator.cpp

@@ -4,23 +4,6 @@
 
 namespace c10 {
 
-DataPtr Allocator::clone(const void* data, std::size_t n) const {
-  DataPtr new_data = allocate(n);
-  copy_data(new_data.mutable_get(), data, n);
-  return new_data;
-}
-
-void Allocator::default_copy_data(
-    void* dest,
-    const void* src,
-    std::size_t count) const {
-  std::memcpy(dest, src, count);
-}
-
-bool Allocator::is_simple_data_ptr(const DataPtr& data_ptr) const {
-  return data_ptr.get() == data_ptr.get_context();
-}
-
 static void deleteInefficientStdFunctionContext(void* ptr) {
   delete static_cast<InefficientStdFunctionContext*>(ptr);
 }

c10/core/Allocator.h

@@ -156,21 +156,6 @@ struct C10_API Allocator {
 
   virtual DataPtr allocate(size_t n) const = 0;
 
-  // Clones an allocation that came from this allocator.
-  //
-  // To perform the copy, this function calls `copy_data`, which
-  // must be implemented by derived classes.
-  //
-  // Note that this explicitly ignores any context that may have been
-  // attached to the input data.
-  //
-  // Requires: input data was allocated by the same allocator.
-  DataPtr clone(const void* data, std::size_t n) const;
-
-  // Checks if DataPtr has a simple context, not wrapped with any out of the
-  // ordinary contexts.
-  virtual bool is_simple_data_ptr(const DataPtr& data_ptr) const;
-
   // If this returns a non nullptr, it means that allocate()
   // is guaranteed to return a unique_ptr with this deleter attached;
   // it means the rawAllocate and rawDeallocate APIs are safe to use.
@@ -188,22 +173,6 @@ struct C10_API Allocator {
     AT_ASSERT(d);
     d(ptr);
   }
-
-  // Copies data from one allocation to another.
-  // Pure virtual, so derived classes must define behavior.
-  // Derived class implementation can simply call `default_copy_data`
-  // to use `std::memcpy`.
-  //
-  // Requires: src and dest were allocated by this allocator
-  // Requires: src and dest both have length >= count
-  virtual void copy_data(void* dest, const void* src, std::size_t count)
-      const = 0;
-
- protected:
-  // Uses `std::memcpy` to copy data.
-  // Child classes can use this as `copy_data` when an alternative copy
-  // API is not needed.
-  void default_copy_data(void* dest, const void* src, std::size_t count) const;
 };
 
 // This context is used to generate DataPtr which have arbitrary

c10/core/CPUAllocator.cpp

@@ -40,10 +40,6 @@ struct C10_API DefaultCPUAllocator final : at::Allocator {
   at::DeleterFnPtr raw_deleter() const override {
     return &ReportAndDelete;
   }
-
-  void copy_data(void* dest, const void* src, std::size_t count) const final {
-    default_copy_data(dest, src, count);
-  }
 };
 
 ProfiledCPUMemoryReporter& profiledCPUMemoryReporter() {
@@ -146,16 +142,6 @@ class DefaultMobileCPUAllocator final : public at::Allocator {
   DeleterFnPtr raw_deleter() const override {
     return deleter;
   }
-
-  bool is_simple_data_ptr(const c10::DataPtr& data_ptr) const final {
-    return reinterpret_cast<const uint8_t*>(data_ptr.get()) ==
-        reinterpret_cast<const uint8_t*>(data_ptr.get_context()) +
-        PreGuardBytes;
-  }
-
-  void copy_data(void* dest, const void* src, std::size_t count) const final {
-    default_copy_data(dest, src, count);
-  }
 };
 
 void NoDelete(void*) {}

c10/core/StorageImpl.h

@@ -3,8 +3,6 @@
 #include <c10/core/Allocator.h>
 #include <c10/core/SymInt.h>
 #include <c10/core/impl/PyObjectSlot.h>
-#include <c10/core/impl/cow/COW.h>
-#include <c10/core/impl/cow/COWDeleter.h>
 
 #include <c10/util/intrusive_ptr.h>
 
@@ -113,7 +111,6 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
   }
 
   at::DataPtr& mutable_data_ptr() {
-    maybe_materialize_cow();
     return data_ptr_;
   }
 
@@ -123,10 +120,9 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
 
   // Returns the previous data_ptr
   at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) {
-    // We need to materialize the old COW DataPtr because it is
-    // being returned as mutable.
-    maybe_materialize_cow();
-    return set_data_ptr_no_materialize_cow(std::move(data_ptr));
+    at::DataPtr old_data_ptr(std::move(data_ptr_));
+    data_ptr_ = std::move(data_ptr);
+    return old_data_ptr;
   }
 
   void set_data_ptr_noswap(at::DataPtr&& data_ptr) {
@@ -138,7 +134,6 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
   }
 
   void* mutable_data() {
-    maybe_materialize_cow();
     return data_ptr_.mutable_get();
   }
 
@@ -216,26 +211,7 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
     return &pyobj_slot_;
   }
 
- protected:
-  // materialize_cow_storage needs to call set_data_ptr_no_materlize_cow
-  friend void c10::impl::cow::materialize_cow_storage(StorageImpl& storage);
-
-  // Returns the previous data_ptr. If the old data_ptr was COW,
-  // this avoids materializing it
-  at::DataPtr set_data_ptr_no_materialize_cow(at::DataPtr&& data_ptr) {
-    at::DataPtr old_data_ptr(std::move(data_ptr_));
-    data_ptr_ = std::move(data_ptr);
-    return old_data_ptr;
-  }
-
  private:
-  // Triggers a copy if this is a copy-on-write tensor.
-  void maybe_materialize_cow() {
-    if (data_ptr_.get_deleter() == impl::cow::cow_deleter) {
-      impl::cow::materialize_cow_storage(*this);
-    }
-  }
-
   DataPtr data_ptr_;
   SymInt size_bytes_;
   bool size_bytes_is_heap_allocated_;

c10/core/build.bzl

@@ -58,22 +58,22 @@ def define_targets(rules):
             [
                 "*.cpp",
                 "impl/*.cpp",
-                "impl/cow/*.cpp",
             ],
             exclude = [
                 "CPUAllocator.cpp",
                 "impl/alloc_cpu.cpp",
+                "impl/cow/*.cpp",
             ],
         ),
         hdrs = rules.glob(
             [
                 "*.h",
                 "impl/*.h",
-                "impl/cow/*.h",
             ],
             exclude = [
                 "CPUAllocator.h",
                 "impl/alloc_cpu.h",
+                "impl/cow/*.h",
             ],
         ),
         linkstatic = True,
@@ -92,6 +92,22 @@ def define_targets(rules):
         alwayslink = True,
     )
 
+    rules.cc_library(
+        name = "impl_cow",
+        srcs = rules.glob([
+            "impl/cow/*.cpp",
+        ]),
+        hdrs = rules.glob([
+            "impl/cow/*.h",
+        ]),
+        deps = [
+            ":base",
+            ":CPUAllocator",
+        ],
+        visibility = ["//c10/test:__pkg__"],
+
+    )
+
     rules.filegroup(
         name = "headers",
         srcs = rules.glob(