frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reorganize and rename COW files and APIs (#110191)

Browse Source 
This PR does the following:
* Combine `cow/context.<h/cpp>` and `cow/deleter.<h/cpp>` into `cow/COWDeleter.<h/cpp>`
* Rename `Context` to `COWDeleterContext`
* Rename `delete_context` to `cow_deleter`
* Remove the separate `impl_cow_context` bazel library, combining it with the base c10 core library
* Rename `context_test.cpp` to `cow_test.cpp`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/110191
Approved by: https://github.com/ezyang
This commit is contained in:
Kurt Mohler
2023-09-28 00:41:57 +00:00
committed by PyTorch MergeBot
parent c62be12061
commit f2c360e3e5
8 changed files with 95 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
c10/core/build.bzl
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,
@ -81,7 +81,6 @@ def define_targets(rules):
visibility = ["//visibility:public"],
deps = [
":ScalarType",
":impl_cow_context",
"//third_party/cpuinfo",
"//c10/macros",
"//c10/util:TypeCast",
@ -93,23 +92,6 @@ def define_targets(rules):
alwayslink = True,
)
rules.cc_library(
name = "impl_cow_context",
srcs = [
"impl/cow/context.cpp",
"impl/cow/deleter.cpp",
],
hdrs = [
"impl/cow/context.h",
"impl/cow/deleter.h",
],
deps = [
"//c10/macros",
"//c10/util:base",
],
visibility = ["//c10/test:__pkg__"],
)
rules.filegroup(
name = "headers",
srcs = rules.glob(

23
c10/core/impl/cow/context.cpp → c10/core/impl/cow/COWDeleter.cpp
View File

@ -1,23 +1,26 @@
#include <c10/core/impl/cow/context.h>
#include <c10/core/impl/cow/deleter.h>
#include <c10/core/impl/cow/COWDeleter.h>
#include <c10/util/Exception.h>
#include <mutex>
namespace c10::impl {
cow::Context::Context(std::unique_ptr<void, DeleterFnPtr> data)
: data_(std::move(data)) {
// We never wrap a Context.
TORCH_INTERNAL_ASSERT(data_.get_deleter() != cow::delete_context);
void cow::cow_deleter(void* ctx) {
static_cast<cow::COWDeleterContext*>(ctx)->decrement_refcount();
}
auto cow::Context::increment_refcount() -> void {
cow::COWDeleterContext::COWDeleterContext(
std::unique_ptr<void, DeleterFnPtr> data)
: data_(std::move(data)) {
// We never wrap a COWDeleterContext.
TORCH_INTERNAL_ASSERT(data_.get_deleter() != cow::cow_deleter);
}
auto cow::COWDeleterContext::increment_refcount() -> void {
auto refcount = ++refcount_;
TORCH_INTERNAL_ASSERT(refcount > 1);
}
auto cow::Context::decrement_refcount()
auto cow::COWDeleterContext::decrement_refcount()
-> std::variant<NotLastReference, LastReference> {
auto refcount = --refcount_;
TORCH_INTERNAL_ASSERT(refcount >= 0, refcount);
@ -32,7 +35,7 @@ auto cow::Context::decrement_refcount()
return std::shared_lock(mutex_);
}
cow::Context::~Context() {
cow::COWDeleterContext::~COWDeleterContext() {
TORCH_INTERNAL_ASSERT(refcount_ == 0);
}

66
c10/core/impl/cow/COWDeleter.h Normal file
View File

@ -0,0 +1,66 @@
#pragma once
#include <c10/macros/Export.h>
#include <c10/util/UniqueVoidPtr.h>
#include <atomic>
#include <cstdint>
#include <memory>
#include <shared_mutex>
#include <variant>
namespace c10::impl::cow {
// A COWDeleterContext object is used as the `ctx` argument for DataPtr
// to implement a Copy-on-write (COW) DataPtr.
class C10_API COWDeleterContext {
public:
// Creates an instance, holding the pair of data and original
// deleter.
//
// Note that the deleter will only be called in our destructor if
// the last reference to this goes away without getting
// materialized.
explicit COWDeleterContext(std::unique_ptr<void, DeleterFnPtr> data);
// Increments the current refcount.
void increment_refcount();
// See README.md in this directory to understand the locking
// strategy.
// Represents a reference to the context.
//
// This is returned by decrement_refcount to allow the caller to
// copy the data under the shared lock.
using NotLastReference = std::shared_lock<std::shared_mutex>;
// Represents the last reference to the context.
//
// This will be returned by decrement_refcount when it is the last
// reference remaining and after any pending copies have completed.
using LastReference = std::unique_ptr<void, DeleterFnPtr>;
// Decrements the refcount, returning a handle indicating what to
// do with it.
std::variant<NotLastReference, LastReference> decrement_refcount();
private:
// The destructor is hidden, this should only ever be used within
// UniqueVoidPtr using cow::delete_context as the deleter.
~COWDeleterContext();
std::shared_mutex mutex_;
std::unique_ptr<void, DeleterFnPtr> data_;
std::atomic<std::int64_t> refcount_ = 1;
};
// `cow_deleter` is used as the `ctx_deleter` for DataPtr to implement a COW
// DataPtr.
//
// Warning: This should only be called on a pointer to a COWDeleterContext that
// was allocated on the heap with `new`, because when the refcount reaches 0,
// the context is deleted with `delete`.
C10_API void cow_deleter(void* ctx);
} // namespace c10::impl::cow

57
c10/core/impl/cow/context.h
View File

@ -1,57 +0,0 @@
#pragma once
#include <c10/macros/Export.h>
#include <c10/util/UniqueVoidPtr.h>
#include <atomic>
#include <cstdint>
#include <memory>
#include <shared_mutex>
#include <variant>
namespace c10::impl::cow {
/// The c10::DataPtr context for copy-on-write storage.
class C10_API Context {
public:
/// Creates an instance, holding the pair of data and original
/// deleter.
///
/// Note that the deleter will only be called in our destructor if
/// the last reference to this goes away without getting
/// materialized.
explicit Context(std::unique_ptr<void, DeleterFnPtr> data);
/// Increments the current refcount.
auto increment_refcount() -> void;
// See README.md in this directory to understand the locking
// strategy.
/// Represents a reference to the context.
///
/// This is returned by decrement_refcount to allow the caller to
/// copy the data under the shared lock.
using NotLastReference = std::shared_lock<std::shared_mutex>;
/// Represents the last reference to the context.
///
/// This will be returned by decrement_refcount when it is the last
/// reference remaining and after any pending copies have completed.
using LastReference = std::unique_ptr<void, DeleterFnPtr>;
/// Decrements the refcount, returning a handle indicating what to
/// do with it.
auto decrement_refcount() -> std::variant<NotLastReference, LastReference>;
private:
// The destructor is hidden, this should only ever be used within
// UniqueVoidPtr using cow::delete_context as the deleter.
~Context();
std::shared_mutex mutex_;
std::unique_ptr<void, DeleterFnPtr> data_;
std::atomic<std::int64_t> refcount_ = 1;
};
} // namespace c10::impl::cow

14
c10/core/impl/cow/deleter.cpp
View File

@ -1,14 +0,0 @@
#include <c10/core/impl/cow/deleter.h>
#include <c10/core/impl/cow/context.h>
namespace c10::impl {
/// Deletes a copy-on-write context.
///
/// Requires: ctx is cow::Context.
auto cow::delete_context(void* ctx) -> void {
static_cast<cow::Context*>(ctx)->decrement_refcount();
}
} // namespace c10::impl

21
c10/core/impl/cow/deleter.h
View File

@ -1,21 +0,0 @@
// This is its own header to minimize code visible in other public
// headers in the system. This is beneficial for compilation times as
// well as to avoid issues with internal Meta builds that aren't using
// C++17.
#pragma once
#include <c10/macros/Export.h>
namespace c10 {
namespace impl {
namespace cow {
/// Deletes a copy-on-write context.
///
/// Requires: ctx is cow::Context.
auto C10_API delete_context(void* ctx) -> void;
} // namespace cow
} // namespace impl
} // namespace c10

10
c10/test/build.bzl
View File

@ -9,15 +9,6 @@ def define_targets(rules):
visibility = ["//:__pkg__"],
)
rules.cc_test(
name = "core_impl_cow_context_test",
srcs = ["core/impl/cow/context_test.cpp"],
deps = [
"//c10/core:impl_cow_context",
"@com_google_googletest//:gtest_main",
],
)
rules.cc_test(
name = "core_tests",
size = "small",
@ -29,6 +20,7 @@ def define_targets(rules):
deps = [
"//c10/core:base",
"//c10/util:base",
"//c10/core:CPUAllocator",
"@com_google_googletest//:gtest_main",
],
)

21
c10/test/core/impl/cow/context_test.cpp → c10/test/core/impl/cow_test.cpp
View File

@ -1,10 +1,14 @@
#include <c10/core/impl/cow/context.h>
#include <c10/core/impl/cow/COWDeleter.h>
#include <c10/core/impl/cow/deleter.h>
#include <c10/core/CPUAllocator.h>
#include <c10/core/StorageImpl.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <cstddef>
#include <memory>
namespace c10::impl {
namespace {
@ -35,7 +39,7 @@ class ContextTest : public testing::Test {
};
TEST_F(ContextTest, Basic) {
auto& context = *new cow::Context(new_delete_tracker());
auto& context = *new cow::COWDeleterContext(new_delete_tracker());
ASSERT_THAT(delete_count(), testing::Eq(0));
context.increment_refcount();
@ -45,7 +49,8 @@ TEST_F(ContextTest, Basic) {
// is expected to give us a shared lock.
auto result = context.decrement_refcount();
ASSERT_THAT(
std::holds_alternative<cow::Context::NotLastReference>(result),
std::holds_alternative<cow::COWDeleterContext::NotLastReference>(
result),
testing::IsTrue());
ASSERT_THAT(delete_count(), testing::Eq(0));
}
@ -53,7 +58,7 @@ TEST_F(ContextTest, Basic) {
{
auto result = context.decrement_refcount();
ASSERT_THAT(
std::holds_alternative<cow::Context::LastReference>(result),
std::holds_alternative<cow::COWDeleterContext::LastReference>(result),
testing::IsTrue());
// Result holds the DeleteTracker.
ASSERT_THAT(delete_count(), testing::Eq(0));
@ -63,12 +68,12 @@ TEST_F(ContextTest, Basic) {
ASSERT_THAT(delete_count(), testing::Eq(1));
}
TEST_F(ContextTest, delete_context) {
TEST_F(ContextTest, cow_deleter) {
// This is effectively the same thing as decrement_refcount() above.
auto& context = *new cow::Context(new_delete_tracker());
auto& context = *new cow::COWDeleterContext(new_delete_tracker());
ASSERT_THAT(delete_count(), testing::Eq(0));
cow::delete_context(&context);
cow::cow_deleter(&context);
ASSERT_THAT(delete_count(), testing::Eq(1));
}