xref: /aosp_15_r20/external/pytorch/torch/csrc/Storage.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #ifndef THP_STORAGE_INC
2 #define THP_STORAGE_INC
3 
4 #include <Python.h>
5 #include <c10/core/Storage.h>
6 #include <torch/csrc/Exceptions.h>
7 #include <torch/csrc/Export.h>
8 #include <torch/csrc/Types.h>
9 
10 #define THPStorageStr "torch.UntypedStorage"
11 
12 struct THPStorage {
13   PyObject_HEAD;
14   c10::MaybeOwned<c10::Storage> cdata;
15   bool is_hermetic;
16 };
17 
18 TORCH_PYTHON_API PyObject* THPStorage_Wrap(c10::Storage storage);
19 TORCH_PYTHON_API PyObject* THPStorage_NewWithStorage(
20     PyTypeObject* type,
21     c10::Storage _storage,
22     c10::impl::PyInterpreterStatus status,
23     bool allow_preexisting_pyobj = false);
24 extern PyTypeObject* THPStorageClass;
25 
THPStorage_CheckTypeExact(PyTypeObject * tp)26 inline bool THPStorage_CheckTypeExact(PyTypeObject* tp) {
27   return tp == THPStorageClass;
28 }
29 
THPStorage_CheckExact(PyObject * obj)30 inline bool THPStorage_CheckExact(PyObject* obj) {
31   return THPStorage_CheckTypeExact(Py_TYPE(obj));
32 }
33 
THPStorage_Check(PyObject * obj)34 inline bool THPStorage_Check(PyObject* obj) {
35   if (!THPStorageClass)
36     return false;
37 
38   const auto result = PyObject_IsInstance(obj, (PyObject*)THPStorageClass);
39   if (result == -1)
40     throw python_error();
41   return result;
42 }
43 
44 bool THPStorage_init(PyObject* module);
45 void THPStorage_postInit(PyObject* module);
46 
47 void THPStorage_assertNotNull(THPStorage* storage);
48 void THPStorage_assertNotNull(PyObject* obj);
49 
50 extern PyTypeObject THPStorageType;
51 
THPStorage_Unpack(THPStorage * storage)52 inline const c10::Storage& THPStorage_Unpack(THPStorage* storage) {
53   return *storage->cdata;
54 }
55 
THPStorage_Unpack(PyObject * obj)56 inline const c10::Storage& THPStorage_Unpack(PyObject* obj) {
57   return THPStorage_Unpack(reinterpret_cast<THPStorage*>(obj));
58 }
59 
60 #endif
61