mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
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:
@ -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(
|
||||
|
||||
@ -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
|
||||
@ -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
|
||||
@ -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",
|
||||
],
|
||||
)
|
||||
|
||||
@ -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*)
|
||||
|
||||
Reference in New Issue
Block a user