mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
5b78a5eadb
Summary: #19975 was separated by 2 PRs. This one: Introduce MemoryFormat argument to the `x.is_contiguous(memory_format=torch.channels_last)` and to the `y = x.contiguous(memory_format=torch.channels_last)` functions. At this moment both functions just operate with strides and doesn't store any tensor state. (Original RFC #19092) ----- Expands functionality of two tensor functions `.is_contiguous` and `.contiguous` (both python and c++ api). Note: We had several complaints about `.to(memory_format)` function, and decided not to support it. 1. `.contiguous` now support optional keyword-only argument - `memory_format`, which can be either `torch.contiguous_format` or `torch.channels_last`. - Using `torch.contiguous_format` will preserve existing `.contiguous()` behavior. - Calling `x.contiguous(memory_format=torch.channels_last)` returns new tensor which maintain same semantical layout (NCHW), but have different memory allocation pattern. `x.contiguous(memory_format=torch.channels_last)` expects input tensor to be 3d, 4d or 5d; and fails otherwise. 2. `.is_contiguous` now support optional keyword-only argument - `memory_format`, which can be either `torch.contiguous_format` or `torch.channels_last`. - `x.is_contiguous(memory_format=torch.contiguous_format)` preserves same functionality as `x.is_contiguous()` and remains unchanged. - `x.is_contiguous(memory_format=torch.channels_last)` returns true if A) input tensor is contiguous in memory AND B) allocated in the memory in NWHC (or similar for 3d,5d) format. Note: By the end of the phase one `x.is_contiguous(memory_format=torch.channels_last)` will calculate state of the Tensor on every call. This functionality going to be updated later. Pull Request resolved: https://github.com/pytorch/pytorch/pull/20455 Differential Revision: D15341577 Pulled By: VitalyFedyunin fbshipit-source-id: bbb6b4159a8a49149110ad321109a3742383185d
81 lines
3.5 KiB
C++
81 lines
3.5 KiB
C++
#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();
|
|
}
|
|
}
|