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
28namespace torch {
29namespace tensors {
30
31using namespace at;
32using namespace torch::autograd;
33
34struct 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
57static_assert(
58 std::is_standard_layout<PyTensorType>::value,
59 "PyTensorType must be standard layout");
60
61static Backend default_backend = Backend::CPU;
62
63static void py_bind_tensor_types(
64 const std::vector<PyTensorType*>& tensor_types);
65
66static 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
72static 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...
93static 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
116PyObject* Tensor_dtype(PyTensorType* self, void* unused) {
117 return torch::autograd::utils::wrap(self->dtype);
118}
119
120PyObject* Tensor_layout(PyTensorType* self, void* unused) {
121 return torch::autograd::utils::wrap(self->layout);
122}
123
124PyObject* 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
132PyObject* 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
140PyObject* 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)
149static struct PyMethodDef metaclass_methods[] = {
150 {"__instancecheck__", Tensor_instancecheck, METH_O, nullptr},
151 {nullptr}};
152
153typedef PyObject* (*getter)(PyObject*, void*);
154
155// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables,modernize-avoid-c-arrays)
156static 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
164static PyTypeObject metaclass = {
165 PyVarObject_HEAD_INIT(nullptr, 0) "torch.tensortype", /* tp_name */
166 sizeof(PyTypeObject) /* tp_basicsize */
167};
168
169static 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
179static PyTypeObject tensor_type_prototype = {
180 PyVarObject_HEAD_INIT(&metaclass, 0) nullptr, /* tp_name */
181 sizeof(PyTensorType) /* tp_basicsize */
182};
183
184static 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
206static 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
221static 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
227static 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
242static 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
255static 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
261static 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.
303static std::vector<PyTensorType*> tensor_types;
304
305void 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
317void 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
347static 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
364void 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
392static 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
424static 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
432void 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
444void 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
452c10::DispatchKey get_default_dispatch_key() {
453 return backendToDispatchKey(default_backend);
454}
455
456at::Device get_default_device() {
457 return at::Device(c10::backendToDeviceType(default_backend));
458}
459
460ScalarType get_default_scalar_type() {
461 return get_default_dtype_as_scalartype();
462}
463
464} // namespace tensors
465} // namespace torch
466