1 | #include <torch/csrc/tensor/python_tensor.h> |
2 | |
3 | #include <pybind11/pybind11.h> |
4 | #include <structmember.h> |
5 | #include <torch/csrc/utils/pybind.h> |
6 | |
7 | #include <torch/csrc/Dtype.h> |
8 | #include <torch/csrc/DynamicTypes.h> |
9 | #include <torch/csrc/Exceptions.h> |
10 | #include <torch/csrc/Layout.h> |
11 | #include <torch/csrc/autograd/generated/VariableType.h> |
12 | #include <torch/csrc/autograd/python_variable.h> |
13 | #include <torch/csrc/autograd/utils/wrap_outputs.h> |
14 | #include <torch/csrc/autograd/variable.h> |
15 | #include <torch/csrc/utils/cuda_enabled.h> |
16 | #include <torch/csrc/utils/cuda_lazy_init.h> |
17 | #include <torch/csrc/utils/python_strings.h> |
18 | #include <torch/csrc/utils/tensor_new.h> |
19 | #include <torch/csrc/utils/tensor_types.h> |
20 | |
21 | #include <ATen/ATen.h> |
22 | |
23 | #include <sstream> |
24 | #include <string> |
25 | #include <type_traits> |
26 | #include <vector> |
27 | |
28 | namespace torch { |
29 | namespace tensors { |
30 | |
31 | using namespace at; |
32 | using namespace torch::autograd; |
33 | |
34 | struct PyTensorType { |
35 | PyTypeObject py_type; |
36 | THPDtype* dtype; |
37 | THPLayout* layout; |
38 | bool is_cuda; |
39 | // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-magic-numbers,modernize-avoid-c-arrays) |
40 | char name[64]; |
41 | int backend; |
42 | int scalar_type; |
43 | |
44 | Backend get_backend() const { |
45 | return static_cast<Backend>(backend); |
46 | } |
47 | |
48 | DispatchKey get_dispatch_key() const { |
49 | return backendToDispatchKey(static_cast<Backend>(backend)); |
50 | } |
51 | |
52 | ScalarType get_scalar_type() const { |
53 | return static_cast<ScalarType>(scalar_type); |
54 | } |
55 | }; |
56 | |
57 | static_assert( |
58 | std::is_standard_layout<PyTensorType>::value, |
59 | "PyTensorType must be standard layout" ); |
60 | |
61 | static Backend default_backend = Backend::CPU; |
62 | |
63 | static void py_bind_tensor_types( |
64 | const std::vector<PyTensorType*>& tensor_types); |
65 | |
66 | static TypeError unavailable_type(const PyTensorType& type) { |
67 | return TypeError( |
68 | "type %s not available. Torch not compiled with CUDA enabled." , |
69 | type.name); |
70 | } |
71 | |
72 | static PyObject* Tensor_new( |
73 | PyTypeObject* type, |
74 | PyObject* args, |
75 | PyObject* kwargs) { |
76 | HANDLE_TH_ERRORS |
77 | auto& tensor_type = *((PyTensorType*)type); |
78 | if (tensor_type.is_cuda && !torch::utils::cuda_enabled()) { |
79 | throw unavailable_type(tensor_type); |
80 | } |
81 | return THPVariable_Wrap(torch::utils::legacy_tensor_ctor( |
82 | tensor_type.get_dispatch_key(), |
83 | tensor_type.get_scalar_type(), |
84 | args, |
85 | kwargs)); |
86 | END_HANDLE_TH_ERRORS |
87 | } |
88 | |
89 | // TODO: Deprecate this instancecheck entirely. It's here to make |
90 | // instanceof(t, torch.FloatTensor) work, but we are not going to keep |
91 | // adding torch.QuantizedIntTensor classes for every new tensor type |
92 | // we add... |
93 | static PyObject* Tensor_instancecheck(PyObject* _self, PyObject* arg) { |
94 | HANDLE_TH_ERRORS |
95 | auto self = (PyTensorType*)_self; |
96 | if (THPVariable_Check(arg)) { |
97 | const auto& var = THPVariable_Unpack(arg); |
98 | // NB: This is a little unfortunate, in that if I do an isinstance check |
99 | // against torch.cuda.FloatTensor, this will immediately initialize CUDA. |
100 | // I originally thought that it would not be possible for aten_type_ to |
101 | // be nullptr if you had a tensor of some type, in which case you can |
102 | // skip initializing aten_type(), but TestAutograd.test_type_conversions |
103 | // seems to violate this property (for whatever reason.) |
104 | // |
105 | // TODO: Stop using legacyExtractDispatchKey here (probably need to build |
106 | // in instanceof checking to Tensor class itself) |
107 | if (legacyExtractDispatchKey(var.key_set()) == self->get_dispatch_key() && |
108 | var.scalar_type() == static_cast<ScalarType>(self->scalar_type)) { |
109 | Py_RETURN_TRUE; |
110 | } |
111 | } |
112 | Py_RETURN_FALSE; |
113 | END_HANDLE_TH_ERRORS |
114 | } |
115 | |
116 | PyObject* Tensor_dtype(PyTensorType* self, void* unused) { |
117 | return torch::autograd::utils::wrap(self->dtype); |
118 | } |
119 | |
120 | PyObject* Tensor_layout(PyTensorType* self, void* unused) { |
121 | return torch::autograd::utils::wrap(self->layout); |
122 | } |
123 | |
124 | PyObject* Tensor_is_cuda(PyTensorType* self, void* unused) { |
125 | if (self->is_cuda) { |
126 | Py_RETURN_TRUE; |
127 | } else { |
128 | Py_RETURN_FALSE; |
129 | } |
130 | } |
131 | |
132 | PyObject* Tensor_is_sparse(PyTensorType* self, void* unused) { |
133 | if (self->layout->layout == at::Layout::Strided) { |
134 | Py_RETURN_FALSE; |
135 | } else { |
136 | Py_RETURN_TRUE; |
137 | } |
138 | } |
139 | |
140 | PyObject* Tensor_is_sparse_csr(PyTensorType* self, void* unused) { |
141 | if (self->layout->layout == at::Layout::SparseCsr) { |
142 | Py_RETURN_TRUE; |
143 | } else { |
144 | Py_RETURN_FALSE; |
145 | } |
146 | } |
147 | |
148 | // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables,modernize-avoid-c-arrays) |
149 | static struct PyMethodDef metaclass_methods[] = { |
150 | {"__instancecheck__" , Tensor_instancecheck, METH_O, nullptr}, |
151 | {nullptr}}; |
152 | |
153 | typedef PyObject* (*getter)(PyObject*, void*); |
154 | |
155 | // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables,modernize-avoid-c-arrays) |
156 | static struct PyGetSetDef metaclass_properties[] = { |
157 | {"dtype" , (getter)Tensor_dtype, nullptr, nullptr, nullptr}, |
158 | {"layout" , (getter)Tensor_layout, nullptr, nullptr, nullptr}, |
159 | {"is_cuda" , (getter)Tensor_is_cuda, nullptr, nullptr, nullptr}, |
160 | {"is_sparse" , (getter)Tensor_is_sparse, nullptr, nullptr, nullptr}, |
161 | {"is_sparse_csr" , (getter)Tensor_is_sparse_csr, nullptr, nullptr, nullptr}, |
162 | {nullptr}}; |
163 | |
164 | static PyTypeObject metaclass = { |
165 | PyVarObject_HEAD_INIT(nullptr, 0) "torch.tensortype" , /* tp_name */ |
166 | sizeof(PyTypeObject) /* tp_basicsize */ |
167 | }; |
168 | |
169 | static void py_initialize_metaclass(PyTypeObject& metaclass) { |
170 | metaclass.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE; |
171 | metaclass.tp_methods = metaclass_methods; |
172 | metaclass.tp_getset = metaclass_properties; |
173 | metaclass.tp_base = &PyType_Type; |
174 | if (PyType_Ready(&metaclass) < 0) { |
175 | throw python_error(); |
176 | } |
177 | } |
178 | |
179 | static PyTypeObject tensor_type_prototype = { |
180 | PyVarObject_HEAD_INIT(&metaclass, 0) nullptr, /* tp_name */ |
181 | sizeof(PyTensorType) /* tp_basicsize */ |
182 | }; |
183 | |
184 | static void py_initialize_tensor_type( |
185 | PyTypeObject& type, |
186 | const char* name, |
187 | PyObject* tp_dict) { |
188 | // NOTE: we don't use the typical static declaration of PyTypeObject because |
189 | // we need to initialize as many types as there are VariableType instances. |
190 | // We copy the basic object fields from a prototype definition and initialize |
191 | // the remaining fields below. |
192 | memcpy(&type, &tensor_type_prototype, sizeof(PyTypeObject)); |
193 | // Subclassing from torch.<ScalarType>Tensor isn't supported. |
194 | // (Py_TPFLAGS_BASETYPE omitted). Subclassing torch.Tensor still allowed. |
195 | type.tp_flags = Py_TPFLAGS_DEFAULT; |
196 | type.tp_name = name; |
197 | type.tp_new = Tensor_new; |
198 | if (PyType_Ready(&type) < 0) { |
199 | throw python_error(); |
200 | } |
201 | if (PyDict_Merge(type.tp_dict, tp_dict, 0) < 0) { |
202 | throw python_error(); |
203 | } |
204 | } |
205 | |
206 | static const char* get_module(Backend backend) { |
207 | switch (backend) { |
208 | case Backend::CPU: |
209 | return "torch" ; |
210 | case Backend::CUDA: |
211 | return "torch.cuda" ; |
212 | case Backend::SparseCPU: |
213 | return "torch.sparse" ; |
214 | case Backend::SparseCUDA: |
215 | return "torch.cuda.sparse" ; |
216 | default: |
217 | AT_ERROR("invalid backend: " , toString(backend)); |
218 | } |
219 | } |
220 | |
221 | static std::string get_name(Backend backend, ScalarType scalarType) { |
222 | std::ostringstream ss; |
223 | ss << get_module(backend) << "." << toString(scalarType) << "Tensor" ; |
224 | return ss.str(); |
225 | } |
226 | |
227 | static THPObjectPtr get_storage_obj(Backend backend, ScalarType dtype) { |
228 | auto module_name = get_module(backend); |
229 | auto module_obj = THPObjectPtr(PyImport_ImportModule(module_name)); |
230 | if (!module_obj) |
231 | throw python_error(); |
232 | |
233 | auto storage_name = std::string(toString(dtype)) + "Storage" ; |
234 | THPObjectPtr storage( |
235 | PyObject_GetAttrString(module_obj.get(), storage_name.c_str())); |
236 | if (!storage.get()) { |
237 | throw TypeError("couldn't find storage object %s" , storage_name.c_str()); |
238 | } |
239 | return storage; |
240 | } |
241 | |
242 | static void set_type( |
243 | PyTensorType& type_obj, |
244 | Backend backend, |
245 | ScalarType scalarType) { |
246 | // This field is lazily initialized from backend and scalar_type |
247 | type_obj.backend = static_cast<int>(backend); |
248 | type_obj.scalar_type = static_cast<int>(scalarType); |
249 | type_obj.layout = torch::getTHPLayout(layout_from_backend(backend)); |
250 | type_obj.dtype = torch::getTHPDtype(scalarType); |
251 | type_obj.is_cuda = |
252 | (backend == at::Backend::CUDA || backend == at::Backend::SparseCUDA); |
253 | } |
254 | |
255 | static void set_name(PyTensorType& type_obj, const std::string& name) { |
256 | size_t n = sizeof(type_obj.name); |
257 | strncpy(type_obj.name, name.c_str(), n); |
258 | type_obj.name[n - 1] = '\0'; |
259 | } |
260 | |
261 | static THPObjectPtr get_tensor_dict() { |
262 | auto torch = THPObjectPtr(PyImport_ImportModule("torch" )); |
263 | if (!torch) |
264 | throw python_error(); |
265 | |
266 | auto tensor_class = THPObjectPtr(PyObject_GetAttrString(torch, "Tensor" )); |
267 | if (!tensor_class) |
268 | throw python_error(); |
269 | |
270 | auto tensor_type = (PyTypeObject*)tensor_class.get(); |
271 | TORCH_CHECK(tensor_type->tp_base, "missing base type for Tensor" ); |
272 | |
273 | auto res = THPObjectPtr(PyDict_New()); |
274 | if (!res) |
275 | throw python_error(); |
276 | |
277 | if (PyDict_Merge(res.get(), tensor_type->tp_dict, 0) < 0) { |
278 | throw python_error(); |
279 | } |
280 | if (PyDict_Merge(res.get(), tensor_type->tp_base->tp_dict, 0) < 0) { |
281 | throw python_error(); |
282 | } |
283 | |
284 | return res; |
285 | } |
286 | |
287 | // A note about the lifetime of the various PyTensorType: normally |
288 | // PyTypeObject instances are statically allocated, but we want to create them |
289 | // dynamically at init time, because their exact number depends on |
290 | // torch::utils::all_declared_types(). The memory for each PyTensorType is |
291 | // allocated by initialize_aten_types() and never freed: technically it's a |
292 | // leak, but it's not a problem since we want them to be alive for the whole |
293 | // time of the process anyway. |
294 | // |
295 | // An alternative is to use a std::vector<PyTensorType> instead, and let |
296 | // std::vector to manage the lifetime of its items. This is problematic |
297 | // though, because it means that the memory of PyTensorType is deallocated at |
298 | // some point during the exit: if by chance we have another global destructor |
299 | // and/or atexit() function which tries to access the PyTensorTypes, we risk |
300 | // an use-after-free error. This happens for example if we embed CPython and |
301 | // call Py_Finalize inside an atexit() function which was registered before |
302 | // importing torch. |
303 | static std::vector<PyTensorType*> tensor_types; |
304 | |
305 | void set_default_storage_type(Backend backend, ScalarType dtype) { |
306 | THPObjectPtr storage = get_storage_obj(backend, dtype); |
307 | |
308 | auto torch_module = THPObjectPtr(PyImport_ImportModule("torch" )); |
309 | if (!torch_module) |
310 | throw python_error(); |
311 | |
312 | if (PyObject_SetAttrString(torch_module.get(), "Storage" , storage) != 0) { |
313 | throw python_error(); |
314 | } |
315 | } |
316 | |
317 | void set_default_tensor_type( |
318 | c10::optional<Backend> backend, |
319 | c10::optional<ScalarType> dtype) { |
320 | if (backend.has_value()) { |
321 | TORCH_CHECK_TYPE( |
322 | *backend != Backend::Undefined, "default type cannot be undefined" ); |
323 | TORCH_CHECK_TYPE( |
324 | !isSparse(*backend), |
325 | "only dense types are supported as the default type" ); |
326 | } |
327 | if (dtype.has_value()) { |
328 | TORCH_CHECK_TYPE( |
329 | at::isFloatingType(*dtype), |
330 | "only floating-point types are supported as the default type" ); |
331 | } |
332 | |
333 | // Try setting default storage in python first as it's the only operation that |
334 | // can fail |
335 | set_default_storage_type( |
336 | backend.value_or(default_backend), |
337 | dtype.value_or(at::get_default_dtype_as_scalartype())); |
338 | |
339 | if (dtype.has_value()) { |
340 | at::set_default_dtype(scalarTypeToTypeMeta(*dtype)); |
341 | } |
342 | if (backend.has_value()) { |
343 | default_backend = *backend; |
344 | } |
345 | } |
346 | |
347 | static void initialize_aten_types(std::vector<PyTensorType*>& tensor_types) { |
348 | // includes CUDA types even when PyTorch is not built with CUDA |
349 | auto declared_types = torch::utils::all_declared_types(); |
350 | tensor_types.resize(declared_types.size()); |
351 | |
352 | for (size_t i = 0, end = declared_types.size(); i != end; i++) { |
353 | tensor_types[i] = new PyTensorType(); |
354 | auto& tensor_type = *tensor_types[i]; |
355 | Backend backend = declared_types[i].first; |
356 | ScalarType scalar_type = declared_types[i].second; |
357 | set_type(tensor_type, backend, scalar_type); |
358 | set_name(tensor_type, get_name(backend, scalar_type)); |
359 | } |
360 | |
361 | set_default_tensor_type(Backend::CPU, ScalarType::Float); |
362 | } |
363 | |
364 | void initialize_python_bindings() { |
365 | // Initialize the at::Type* pointers, name, and properties of the PyTensorType |
366 | // vector. After this call, the vector must not be resized. |
367 | initialize_aten_types(tensor_types); |
368 | |
369 | // Initialize the Python metaclass for the torch.FloatTensor, etc. types. |
370 | // The metaclass handles __instancecheck__ checks and binds the dtype property |
371 | // on the type objects. |
372 | py_initialize_metaclass(metaclass); |
373 | |
374 | // Get the tp_dict of the Variable class. We copy function definitions |
375 | // onto each Tensor type object so that they can be accessed via e.g. |
376 | // `torch.FloatTensor.add`. |
377 | auto tensor_dict = get_tensor_dict(); |
378 | |
379 | // Initialize each Python type object torch.FloatTensor, torch.DoubleTensor, |
380 | // etc. |
381 | for (auto& tensor_type : tensor_types) { |
382 | py_initialize_tensor_type( |
383 | tensor_type->py_type, tensor_type->name, tensor_dict.get()); |
384 | } |
385 | |
386 | // Add the type objects to their corresponding modules. e.g. torch.FloatTensor |
387 | // is added to the `torch` module as `FloatTensor`. Also add all the type |
388 | // objects to the set torch._tensor_classes. |
389 | py_bind_tensor_types(tensor_types); |
390 | } |
391 | |
392 | static void py_bind_tensor_types( |
393 | const std::vector<PyTensorType*>& tensor_types) { |
394 | auto torch_module = THPObjectPtr(PyImport_ImportModule("torch" )); |
395 | if (!torch_module) |
396 | throw python_error(); |
397 | |
398 | auto tensor_classes = THPObjectPtr( |
399 | PyObject_GetAttrString(torch_module.get(), "_tensor_classes" )); |
400 | if (!tensor_classes) |
401 | throw python_error(); |
402 | |
403 | for (auto& tensor_type : tensor_types) { |
404 | auto name = std::string(tensor_type->name); |
405 | auto idx = name.rfind('.'); |
406 | auto type_name = name.substr(idx + 1); |
407 | auto module_name = name.substr(0, idx); |
408 | |
409 | auto module_obj = THPObjectPtr(PyImport_ImportModule(module_name.c_str())); |
410 | if (!module_obj) |
411 | throw python_error(); |
412 | |
413 | PyObject* type_obj = (PyObject*)tensor_type; |
414 | Py_INCREF(type_obj); |
415 | if (PyModule_AddObject(module_obj.get(), type_name.c_str(), type_obj) < 0) { |
416 | throw python_error(); |
417 | } |
418 | if (PySet_Add(tensor_classes.get(), type_obj) < 0) { |
419 | throw python_error(); |
420 | } |
421 | } |
422 | } |
423 | |
424 | static bool PyTensorType_Check(PyObject* obj) { |
425 | auto it = std::find_if( |
426 | tensor_types.begin(), tensor_types.end(), [obj](PyTensorType* x) { |
427 | return (PyObject*)x == obj; |
428 | }); |
429 | return it != tensor_types.end(); |
430 | } |
431 | |
432 | void py_set_default_tensor_type(PyObject* obj) { |
433 | // NOLINTNEXTLINE(cppcoreguidelines-init-variables) |
434 | TORCH_CHECK_TYPE( |
435 | PyTensorType_Check(obj), |
436 | "invalid type object: only floating-point types are supported as the default type" ); |
437 | PyTensorType* type = (PyTensorType*)obj; |
438 | if (type->is_cuda && !torch::utils::cuda_enabled()) { |
439 | throw unavailable_type(*type); |
440 | } |
441 | set_default_tensor_type(type->get_backend(), type->get_scalar_type()); |
442 | } |
443 | |
444 | void py_set_default_dtype(PyObject* obj) { |
445 | TORCH_CHECK_TYPE( |
446 | THPDtype_Check(obj), |
447 | "invalid dtype object: only floating-point types are supported as the default type" ); |
448 | auto scalar_type = ((THPDtype*)obj)->scalar_type; |
449 | set_default_tensor_type(/*backend=*/c10::nullopt, scalar_type); |
450 | } |
451 | |
452 | c10::DispatchKey get_default_dispatch_key() { |
453 | return backendToDispatchKey(default_backend); |
454 | } |
455 | |
456 | at::Device get_default_device() { |
457 | return at::Device(c10::backendToDeviceType(default_backend)); |
458 | } |
459 | |
460 | ScalarType get_default_scalar_type() { |
461 | return get_default_dtype_as_scalartype(); |
462 | } |
463 | |
464 | } // namespace tensors |
465 | } // namespace torch |
466 | |