xref: /aosp_15_r20/external/pytorch/c10/core/RefcountedDeleter.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <c10/core/Storage.h>
4 #include <c10/macros/Export.h>
5 #include <c10/util/UniqueVoidPtr.h>
6 
7 #include <atomic>
8 #include <memory>
9 
10 namespace c10 {
11 
12 // A RefcountedDeleterContext object is used as the `ctx` argument for DataPtr
13 // to implement a shared DataPtr. Normally, a DataPtr is unique, but we use
14 // this custom context and the `refcounted_deleter` function below to make the
15 // DataPtr act like a non-unique DataPtr. This context object holds onto an
16 // inner context and deleter function which handle the actual deletion of the
17 // data when the refcount reaches 0.
18 //
19 // This shared DataPtr feature is only used when storages are shared between
20 // multiple Python interpreters in MultiPy. Before storages had PyObject
21 // preservation, interpreters could just share the same StorageImpl instance.
22 // But now a StorageImpl can only be associated with one interpreter in order
23 // to properly manage a zombie PyObject. So we share storages across Python
24 // interpreters by creating a different StorageImpl instance for each one, but
25 // they all point to the same data.
26 struct C10_API RefcountedDeleterContext {
RefcountedDeleterContextRefcountedDeleterContext27   RefcountedDeleterContext(void* other_ctx, c10::DeleterFnPtr other_deleter)
28       : other_ctx(other_ctx, other_deleter), refcount(1) {}
29 
30   std::unique_ptr<void, c10::DeleterFnPtr> other_ctx;
31   std::atomic_int refcount;
32 };
33 
34 // `refcounted_deleter` is used as the `ctx_deleter` for DataPtr to implement
35 // a shared DataPtr.
36 //
37 // Warning: This should only be called on a pointer to
38 // a RefcountedDeleterContext that was allocated on the heap with `new`,
39 // because when the refcount reaches 0, the context is deleted with `delete`
40 C10_API void refcounted_deleter(void* ctx_);
41 
42 // If the storage's DataPtr does not use `refcounted_deleter`, replace it with
43 // a DataPtr that does, so it can be shared between multiple StorageImpls
44 C10_API void maybeApplyRefcountedDeleter(const c10::Storage& storage);
45 
46 // Create a new StorageImpl that points to the same data. If the original
47 // StorageImpl's DataPtr does not use `refcounted_deleter`, it will be replaced
48 // with one that does
49 C10_API c10::Storage newStorageImplFromRefcountedDataPtr(
50     const c10::Storage& storage);
51 
52 } // namespace c10
53