Revert "Add lazy_clone_storage to create COW storages (#110192)"

This reverts commit 33f151348684bd74fbc9939f00c39408ef92074d.

Reverted https://github.com/pytorch/pytorch/pull/110192 on behalf of https://github.com/kit1980 due to revert to work around some importing issues ([comment](https://github.com/pytorch/pytorch/pull/110192#issuecomment-1762430374))
This commit is contained in:
PyTorch MergeBot
2023-10-14 00:48:45 +00:00
parent 4f4e2c1c08
commit 482782406a
5 changed files with 2 additions and 248 deletions

View File

@ -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,22 +92,6 @@ 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(

View File

@ -1,111 +0,0 @@
#include <c10/core/impl/cow/COW.h>
#include <c10/core/Allocator.h>
#include <c10/core/CPUAllocator.h>
#include <c10/core/StorageImpl.h>
#include <c10/core/alignment.h>
#include <c10/core/impl/cow/COWDeleter.h>
#include <c10/util/Exception.h>
#include <c10/util/UniqueVoidPtr.h>
#include <memory>
#include <optional>
namespace c10::impl::cow {
namespace {
// Wraps a DataPtr with a copy-on-write DataPtr.
at::DataPtr make_data_ptr(
at::DataPtr const& data_ptr,
cow::COWDeleterContext& ctx) {
return at::DataPtr(data_ptr.get(), &ctx, cow::cow_deleter, data_ptr.device());
}
/// Copies a copy-on-write DataPtr.
at::DataPtr copy_data_ptr(at::DataPtr const& data_ptr) {
auto* ctx = data_ptr.cast_context<cow::COWDeleterContext>(cow::cow_deleter);
TORCH_INTERNAL_ASSERT(ctx != nullptr);
ctx->increment_refcount();
return make_data_ptr(data_ptr, *ctx);
}
} // namespace
bool has_simple_data_ptr(const c10::StorageImpl& storage) {
const c10::DataPtr& data_ptr = storage.data_ptr();
if (storage.allocator() == c10::GetDefaultMobileCPUAllocator()) {
return reinterpret_cast<size_t>(data_ptr.get()) ==
reinterpret_cast<size_t>(data_ptr.get_context()) + c10::gAlignment;
} else {
return data_ptr.get() == data_ptr.get_context();
}
}
bool is_cow_data_ptr(const c10::DataPtr& data_ptr) {
return (void*)data_ptr.get_deleter() == (void*)&cow::cow_deleter;
}
c10::intrusive_ptr<StorageImpl> lazy_clone_storage(StorageImpl& storage) {
const at::DataPtr& data_ptr = storage.data_ptr();
// There are three possible circumstances:
//
// 1) The storage has a normal data pointer with no out of the ordinary
// context. In this case we know that there are no blind aliases to the
// storage impl: they all will be public aliases and the user is expected
// to synchronize manually.
//
// No locking is required in this case.
//
// 2) The storage already has a copy on write context. There
// is a potential race condition with a blind alias (i.e. an
// alias that the user is not required to synchronize
// with). Because our input storage is bound to a live reference
// to the data, we know that it isn't going away. A blind alias
// could be copying from it right now, but we will grab the
// context's mutex to protect us.
//
// We do not need to lock in this case either, because we're just
// wrapping a context that we know isn't going away.
//
// 3) The storage has a context that is not the copy on write
// context. This is not supported, so we just return null.
//
// No locking is required in this case.
std::optional<DataPtr> new_data_ptr; // must be set below
if (has_simple_data_ptr(storage)) {
// Case 1) We have a simple data pointer: wrap it.
std::unique_ptr<void, DeleterFnPtr> original_ctx =
storage.mutable_data_ptr().move_context();
TORCH_INTERNAL_ASSERT(original_ctx.get() == data_ptr.get());
// Save this for the result.
new_data_ptr = make_data_ptr(
data_ptr, *new cow::COWDeleterContext(std::move(original_ctx)));
// Update this storage to the new copy on write context.
storage.set_data_ptr_noswap(copy_data_ptr(*new_data_ptr));
} else if (is_cow_data_ptr(data_ptr)) {
// Case 2): there is already a copy on write context. Just return a
// new storage impl.
new_data_ptr = copy_data_ptr(data_ptr);
} else {
// Case 3) There is a context and it's not copy-on-write. Nothing
// we can do here.
return nullptr;
}
TORCH_INTERNAL_ASSERT(new_data_ptr.has_value());
return make_intrusive<StorageImpl>(
StorageImpl::use_byte_size_t(),
storage.sym_nbytes(),
*std::move(new_data_ptr),
storage.allocator(),
storage.resizable());
}
} // namespace c10::impl::cow

View File

@ -1,29 +0,0 @@
#pragma once
#include <c10/macros/Macros.h>
#include <c10/util/intrusive_ptr.h>
namespace c10 {
struct StorageImpl;
class DataPtr;
}; // namespace c10
namespace c10::impl::cow {
// Creates a Copy-on-write (COW) clone of the given storage. This will also
// convert the given storage into a COW storage if it is not COW already.
//
// Converting the storage into a COW storage will not be successful if the
// storage's DataPtr has some context (`DataPtr::get_context()`) which is not
// equal to the data pointer (`DataPtr::get()`). In this case, a nullptr is
// returned.
C10_API c10::intrusive_ptr<StorageImpl> lazy_clone_storage(
StorageImpl& storage);
// Check if a storage has a simple DataPtr with no abnormal context
C10_API bool has_simple_data_ptr(const c10::StorageImpl& storage);
// Check if a DataPtr is COW
C10_API bool is_cow_data_ptr(const c10::DataPtr& data_ptr);
} // namespace c10::impl::cow

View File

@ -21,7 +21,6 @@ def define_targets(rules):
"//c10/core:base",
"//c10/util:base",
"//c10/core:CPUAllocator",
"//c10/core:impl_cow",
"@com_google_googletest//:gtest_main",
],
)

View File

@ -1,4 +1,3 @@
#include <c10/core/impl/cow/COW.h>
#include <c10/core/impl/cow/COWDeleter.h>
#include <c10/core/CPUAllocator.h>
@ -79,94 +78,6 @@ TEST_F(ContextTest, cow_deleter) {
ASSERT_THAT(delete_count(), testing::Eq(1));
}
MATCHER(is_copy_on_write, "") {
const c10::StorageImpl& storage = std::ref(arg);
return cow::is_cow_data_ptr(storage.data_ptr());
}
TEST(lazy_clone_storage_test, no_context) {
StorageImpl original_storage(
{}, /*size_bytes=*/7, GetDefaultCPUAllocator(), /*resizable=*/false);
ASSERT_THAT(original_storage, testing::Not(is_copy_on_write()));
ASSERT_TRUE(cow::has_simple_data_ptr(original_storage));
intrusive_ptr<StorageImpl> new_storage =
cow::lazy_clone_storage(original_storage);
ASSERT_THAT(new_storage.get(), testing::NotNull());
// The original storage was modified in-place to now hold a copy on
// write context.
ASSERT_THAT(original_storage, is_copy_on_write());
// The result is a different storage impl.
ASSERT_THAT(&*new_storage, testing::Ne(&original_storage));
// But it is also copy-on-write.
ASSERT_THAT(*new_storage, is_copy_on_write());
// But they share the same data!
ASSERT_THAT(new_storage->data(), testing::Eq(original_storage.data()));
}
struct MyDeleterContext {
MyDeleterContext(void* bytes) : bytes(bytes) {}
~MyDeleterContext() {
delete[] static_cast<std::byte*>(bytes);
}
void* bytes;
};
void my_deleter(void* ctx) {
delete static_cast<MyDeleterContext*>(ctx);
}
TEST(lazy_clone_storage_test, different_context) {
void* bytes = new std::byte[5];
StorageImpl storage(
{},
/*size_bytes=*/5,
at::DataPtr(
/*data=*/bytes,
/*ctx=*/new MyDeleterContext(bytes),
/*ctx_deleter=*/my_deleter,
/*device=*/Device(Device::Type::CPU)),
/*allocator=*/nullptr,
/*resizable=*/false);
// We can't handle an arbitrary context.
ASSERT_THAT(cow::lazy_clone_storage(storage), testing::IsNull());
}
TEST(lazy_clone_storage_test, already_copy_on_write) {
std::unique_ptr<void, DeleterFnPtr> data(
new std::byte[5],
+[](void* bytes) { delete[] static_cast<std::byte*>(bytes); });
void* data_ptr = data.get();
StorageImpl original_storage(
{},
/*size_bytes=*/5,
at::DataPtr(
/*data=*/data_ptr,
/*ctx=*/new cow::COWDeleterContext(std::move(data)),
cow::cow_deleter,
Device(Device::Type::CPU)),
/*allocator=*/nullptr,
/*resizable=*/false);
ASSERT_THAT(original_storage, is_copy_on_write());
intrusive_ptr<StorageImpl> new_storage =
cow::lazy_clone_storage(original_storage);
ASSERT_THAT(new_storage.get(), testing::NotNull());
// The result is a different storage.
ASSERT_THAT(&*new_storage, testing::Ne(&original_storage));
// But it is also copy-on-write.
ASSERT_THAT(*new_storage, is_copy_on_write());
// But they share the same data!
ASSERT_THAT(new_storage->data(), testing::Eq(original_storage.data()));
}
} // namespace
} // namespace c10::impl
// NOLINTEND(clang-analyzer-cplusplus*)