mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
b4a35632f9
Summary: From Kurt Mohler, see https://github.com/pytorch/pytorch/pull/113396 (manually imported due to ghimport problems) Test Plan: sandcastle, OSS CI Differential Revision: D52610522 Pull Request resolved: https://github.com/pytorch/pytorch/pull/117053 Approved by: https://github.com/malfet, https://github.com/kurtamohler
67 lines
2.0 KiB
C++
67 lines
2.0 KiB
C++
#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
|