pytorch

1
#include <torch/csrc/utils/tensor_memoryformats.h>
2

3
#include <c10/core/MemoryFormat.h>
4
#include <torch/csrc/DynamicTypes.h>
5
#include <torch/csrc/Exceptions.h>
6
#include <torch/csrc/MemoryFormat.h>
7

8
#include <torch/csrc/python_headers.h>
9
#include <torch/csrc/utils/object_ptr.h>
10

11
namespace torch {
12
namespace utils {
13

14
namespace {
15
// Intentionally leaked
16
std::array<PyObject*, static_cast<int>(at::MemoryFormat::NumOptions)>
17
    memory_format_registry = {};
18
} // anonymous namespace
19

20
PyObject* getTHPMemoryFormat(at::MemoryFormat memory_format) {
21
  return py::reinterpret_borrow<py::object>(
22
             memory_format_registry[static_cast<size_t>(memory_format)])
23
      .release()
24
      .ptr();
25
}
26

27
void initializeMemoryFormats() {
28
  auto torch_module = THPObjectPtr(PyImport_ImportModule("torch"));
29
  if (!torch_module) {
30
    throw python_error();
31
  }
32

33
  auto add_memory_format = [&](at::MemoryFormat format, const char* name) {
34
    std::string module_name = "torch.";
35
    PyObject* memory_format = THPMemoryFormat_New(format, module_name + name);
36
    Py_INCREF(memory_format);
37
    if (PyModule_AddObject(torch_module, name, memory_format) != 0) {
38
      Py_DECREF(memory_format);
39
      throw python_error();
40
    }
41
    Py_INCREF(memory_format);
42
    memory_format_registry[static_cast<size_t>(format)] = memory_format;
43
  };
44

45
  add_memory_format(at::MemoryFormat::Preserve, "preserve_format");
46
  add_memory_format(at::MemoryFormat::Contiguous, "contiguous_format");
47
  add_memory_format(at::MemoryFormat::ChannelsLast, "channels_last");
48
  add_memory_format(at::MemoryFormat::ChannelsLast3d, "channels_last_3d");
49
}
50

51
} // namespace utils
52
} // namespace torch
53