torch/csrc/MemoryFormat.cpp

#include <torch/csrc/MemoryFormat.h>

#include <torch/csrc/Exceptions.h>
#include <torch/csrc/utils/object_ptr.h>
#include <torch/csrc/utils/python_strings.h>

#include <c10/core/MemoryFormat.h>

#include <structmember.h>
#include <cstring>
#include <string>

PyObject *THPMemoryFormat_New(at::MemoryFormat memory_format, const std::string& name)
{
  auto type = (PyTypeObject*)&THPMemoryFormatType;
  auto self = THPObjectPtr{type->tp_alloc(type, 0)};
  if (!self) throw python_error();
  auto self_ = reinterpret_cast<THPMemoryFormat*>(self.get());
  self_->memory_format = memory_format;
  std::strncpy (self_->name, name.c_str(), MEMORY_FORMAT_NAME_LEN);
  self_->name[MEMORY_FORMAT_NAME_LEN] = '\0';
  return self.release();
}

PyObject *THPMemoryFormat_repr(THPMemoryFormat *self)
{
  return THPUtils_packString(self->name);
}

PyTypeObject THPMemoryFormatType = {
  PyVarObject_HEAD_INIT(nullptr, 0)
  "torch.memory_format",                       /* tp_name */
  sizeof(THPMemoryFormat),                     /* tp_basicsize */
  0,                                           /* tp_itemsize */
  nullptr,                                     /* tp_dealloc */
  nullptr,                                     /* tp_print */
  nullptr,                                     /* tp_getattr */
  nullptr,                                     /* tp_setattr */
  nullptr,                                     /* tp_reserved */
  (reprfunc)THPMemoryFormat_repr,              /* tp_repr */
  nullptr,                                     /* tp_as_number */
  nullptr,                                     /* tp_as_sequence */
  nullptr,                                     /* tp_as_mapping */
  nullptr,                                     /* tp_hash  */
  nullptr,                                     /* tp_call */
  nullptr,                                     /* tp_str */
  nullptr,                                     /* tp_getattro */
  nullptr,                                     /* tp_setattro */
  nullptr,                                     /* tp_as_buffer */
  Py_TPFLAGS_DEFAULT,                          /* tp_flags */
  nullptr,                                     /* tp_doc */
  nullptr,                                     /* tp_traverse */
  nullptr,                                     /* tp_clear */
  nullptr,                                     /* tp_richcompare */
  0,                                           /* tp_weaklistoffset */
  nullptr,                                     /* tp_iter */
  nullptr,                                     /* tp_iternext */
  nullptr,                                     /* tp_methods */
  nullptr,                                     /* tp_members */
  nullptr,                                     /* tp_getset */
  nullptr,                                     /* tp_base */
  nullptr,                                     /* tp_dict */
  nullptr,                                     /* tp_descr_get */
  nullptr,                                     /* tp_descr_set */
  0,                                           /* tp_dictoffset */
  nullptr,                                     /* tp_init */
  nullptr,                                     /* tp_alloc */
  nullptr,                                     /* tp_new */
};

void THPMemoryFormat_init(PyObject *module)
{
  if (PyType_Ready(&THPMemoryFormatType) < 0) {
    throw python_error();
  }
  Py_INCREF(&THPMemoryFormatType);
  if (PyModule_AddObject(module, "memory_format", (PyObject *)&THPMemoryFormatType) != 0) {
    throw python_error();
  }
}