1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
#include <torch/csrc/utils/tensor_memoryformats.h>
#include <c10/core/MemoryFormat.h>
#include <torch/csrc/DynamicTypes.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/MemoryFormat.h>
#include <torch/csrc/python_headers.h>
#include <torch/csrc/utils/object_ptr.h>
namespace torch {
namespace utils {
namespace {
// Intentionally leaked
std::array<PyObject*, static_cast<int>(at::MemoryFormat::NumOptions)>
memory_format_registry = {};
} // anonymous namespace
PyObject* getTHPMemoryFormat(at::MemoryFormat memory_format) {
return py::reinterpret_borrow<py::object>(
memory_format_registry[static_cast<size_t>(memory_format)])
.release()
.ptr();
}
void initializeMemoryFormats() {
auto torch_module = THPObjectPtr(PyImport_ImportModule("torch"));
if (!torch_module) {
throw python_error();
}
auto add_memory_format = [&](at::MemoryFormat format, const char* name) {
std::string module_name = "torch.";
PyObject* memory_format = THPMemoryFormat_New(format, module_name + name);
Py_INCREF(memory_format);
if (PyModule_AddObject(torch_module, name, memory_format) != 0) {
Py_DECREF(memory_format);
throw python_error();
}
Py_INCREF(memory_format);
memory_format_registry[static_cast<size_t>(format)] = memory_format;
};
add_memory_format(at::MemoryFormat::Preserve, "preserve_format");
add_memory_format(at::MemoryFormat::Contiguous, "contiguous_format");
add_memory_format(at::MemoryFormat::ChannelsLast, "channels_last");
add_memory_format(at::MemoryFormat::ChannelsLast3d, "channels_last_3d");
}
} // namespace utils
} // namespace torch
|
