static PyObject * THPStorage_(newSharedFd)(PyObject *_unused, PyObject *args)
{
HANDLE_TH_ERRORS
THPUtils_assert(PyTuple_GET_SIZE(args) == 2, "tuple of 2 items expected");
PyObject *_tmp_fd = PyTuple_GET_ITEM(args, 0);
PyObject *_size = PyTuple_GET_ITEM(args, 1);
if (!THPUtils_checkLong(_tmp_fd) || !THPUtils_checkLong(_size)) {
THPUtils_invalidArguments(args, NULL, "_new_shared in file descriptor mode",
1, "a file descriptor (int) and storage size (int)");
return NULL;
}
int fd;
int tmp_fd = (int) THPUtils_unpackLong(_tmp_fd);
int64_t size = THPUtils_unpackLong(_size);
if ((fd = dup(tmp_fd)) == -1) {
THPUtils_setError("could not duplicate a shared memory file descriptor");
return NULL;
}
int flags = TH_ALLOCATOR_MAPPED_SHAREDMEM |
TH_ALLOCATOR_MAPPED_NOCREATE |
TH_ALLOCATOR_MAPPED_KEEPFD |
TH_ALLOCATOR_MAPPED_FROMFD;
THMapAllocatorContext *ctx = THMapAllocatorContext_newWithFd(NULL, fd, flags);
return THPStorage_(New)(THStorage_(newWithAllocator)(size, &THMapAllocator,
(void*)ctx));
END_HANDLE_TH_ERRORS
}
static bool THPModule_loadClasses(PyObject *self)
{
#define ASSERT_NOT_NULL(ptr) if (!(ptr)) { THPUtils_setError("couldn't load classes"); return false; }
PyObject *torch_module = PyImport_ImportModule("torch");
if (!torch_module) {
THPUtils_setError("class loader couldn't access torch module");
return false;
}
ASSERT_NOT_NULL(tensor_classes = PyObject_GetAttrString(torch_module, "_tensor_classes"));
if (!THPDoubleTensor_postInit(torch_module)) return false;
if (!THPFloatTensor_postInit(torch_module)) return false;
if (!THPHalfTensor_postInit(torch_module)) return false;
if (!THPLongTensor_postInit(torch_module)) return false;
if (!THPIntTensor_postInit(torch_module)) return false;
if (!THPShortTensor_postInit(torch_module)) return false;
if (!THPCharTensor_postInit(torch_module)) return false;
if (!THPByteTensor_postInit(torch_module)) return false;
ASSERT_NOT_NULL(THPDoubleStorageClass = PyObject_GetAttrString(torch_module,(char*)"DoubleStorage"));
ASSERT_NOT_NULL(THPFloatStorageClass = PyObject_GetAttrString(torch_module,(char*)"FloatStorage"));
ASSERT_NOT_NULL(THPHalfStorageClass = PyObject_GetAttrString(torch_module,(char*)"HalfStorage"));
ASSERT_NOT_NULL(THPLongStorageClass = PyObject_GetAttrString(torch_module,(char*)"LongStorage"));
ASSERT_NOT_NULL(THPIntStorageClass = PyObject_GetAttrString(torch_module,(char*)"IntStorage"));
ASSERT_NOT_NULL(THPShortStorageClass = PyObject_GetAttrString(torch_module,(char*)"ShortStorage"));
ASSERT_NOT_NULL(THPCharStorageClass = PyObject_GetAttrString(torch_module,(char*)"CharStorage"));
ASSERT_NOT_NULL(THPByteStorageClass = PyObject_GetAttrString(torch_module,(char*)"ByteStorage"));
return true;
#undef ASSERT_NOT_NULL
}
static bool THDPModule_loadClasses(PyObject *self)
{
#ifdef WITH_DISTRIBUTED_MW
#define ASSERT_NOT_NULL(ptr) if (!(ptr)) { THPUtils_setError("couldn't load classes"); return false; }
PyObject *torch_module = PyImport_ImportModule("torch.distributed");
if (!torch_module) {
THPUtils_setError("class loader couldn't access torch.distributed module");
return false;
}
if (!THDPDoubleTensor_postInit(torch_module)) return false;
if (!THDPFloatTensor_postInit(torch_module)) return false;
if (!THDPHalfTensor_postInit(torch_module)) return false;
if (!THDPLongTensor_postInit(torch_module)) return false;
if (!THDPIntTensor_postInit(torch_module)) return false;
if (!THDPShortTensor_postInit(torch_module)) return false;
if (!THDPCharTensor_postInit(torch_module)) return false;
if (!THDPByteTensor_postInit(torch_module)) return false;
ASSERT_NOT_NULL(THDPDoubleStorageClass = PyObject_GetAttrString(torch_module,(char*)"DoubleStorage"));
ASSERT_NOT_NULL(THDPFloatStorageClass = PyObject_GetAttrString(torch_module,(char*)"FloatStorage"));
ASSERT_NOT_NULL(THDPHalfStorageClass = PyObject_GetAttrString(torch_module,(char*)"HalfStorage"));
ASSERT_NOT_NULL(THDPLongStorageClass = PyObject_GetAttrString(torch_module,(char*)"LongStorage"));
ASSERT_NOT_NULL(THDPIntStorageClass = PyObject_GetAttrString(torch_module,(char*)"IntStorage"));
ASSERT_NOT_NULL(THDPShortStorageClass = PyObject_GetAttrString(torch_module,(char*)"ShortStorage"));
ASSERT_NOT_NULL(THDPCharStorageClass = PyObject_GetAttrString(torch_module,(char*)"CharStorage"));
ASSERT_NOT_NULL(THDPByteStorageClass = PyObject_GetAttrString(torch_module,(char*)"ByteStorage"));
#undef ASSERT_NOT_NULL
#endif
return true;
}
static THStorage* THPStorage_(newWithAllocator)(int64_t size, THAllocator* allocator)
{
#if defined(THC_GENERIC_FILE) || defined(THD_GENERIC_FILE)
THPUtils_setError(THPStorageStr " does not support custom allocators");
return NULL;
#else
return THStorage_(newWithAllocator)(LIBRARY_STATE size, allocator, NULL);
#endif
}
static PyObject * THPStorage_(pynew)(PyTypeObject *type, PyObject *args, PyObject *kwargs)
{
HANDLE_TH_ERRORS
Py_ssize_t num_args = args ? PyTuple_Size(args) : 0;
THPStoragePtr self((THPStorage *)type->tp_alloc(type, 0));
THPUtils_assert(self, "failed to allocate a " THPStorageStr " object");
THAllocator* allocator = NULL;
// Internally we allow constructing with a keywoard only argument cdata
if (kwargs != NULL) {
PyObject *allocator_ptr = PyDict_GetItemString(kwargs, "allocator");
if (allocator_ptr) {
THPUtils_assert(THPUtils_checkLong(allocator_ptr), "invalid allocator");
allocator = (THAllocator*) PyLong_AsVoidPtr(allocator_ptr);
PyDict_DelItemString(kwargs, "allocator");
}
Py_ssize_t num_kwargs = PyDict_Size(kwargs);
if (num_args == 0) {
PyObject *cdata_ptr = PyDict_GetItemString(kwargs, "cdata");
if (num_kwargs == 1 && cdata_ptr && THPUtils_checkLong(cdata_ptr)) {
THStorage *ptr = (THStorage*)PyLong_AsVoidPtr(cdata_ptr);
self->cdata = ptr;
return (PyObject*)self.release();
}
}
THPUtils_assert(num_kwargs == 0, THPStorageStr "(): invalid keyword arguments");
}
// torch.Storage()
if (num_args == 0) {
if (allocator) {
self->cdata = THPStorage_(newWithAllocator)(0, allocator);
} else {
self->cdata = THStorage_(new)(LIBRARY_STATE_NOARGS);
}
return (PyObject*)self.release();
}
PyObject *first_arg = PyTuple_GET_ITEM(args, 0);
// torch.Storage(size)
if (num_args == 1 && THPUtils_checkLong(first_arg)) {
int64_t size = THPUtils_unpackLong(first_arg);
if (allocator) {
self->cdata = THPStorage_(newWithAllocator)(size, allocator);
} else {
self->cdata = THStorage_(newWithSize)(LIBRARY_STATE size);
}
return (PyObject*)self.release();
}
// torch.Storage(view_source, [offset, [size]])
if (num_args < 4 && THPStorage_(Check)(first_arg)) {
#ifdef THD_GENERIC_FILE
THPUtils_setError("distributed storages don't support storage views");
return NULL;
#else
THPStorage *storage_arg = (THPStorage *)first_arg;
int64_t numel = storage_arg->cdata->size;
int64_t offset = 0;
if (num_args >= 2) {
PyObject *second_arg = PyTuple_GET_ITEM(args, 1);
if (!THPUtils_checkLong(second_arg))
goto invalid_arguments;
offset = THPUtils_unpackLong(second_arg);
}
int64_t size = numel - offset;
if (num_args >= 3) {
PyObject *third_arg = PyTuple_GET_ITEM(args, 2);
if (!THPUtils_checkLong(third_arg))
goto invalid_arguments;
size = THPUtils_unpackLong(third_arg);
}
THPUtils_assert(offset >= 0 && offset <= numel, "specified an offset of "
"%" PRId64 ", but the viewed storage has only %" PRId64 " element(s)", offset, numel);
THPUtils_assert(size >= 1 && size <= numel - offset, "specified a size of "
"%" PRId64 ", but the viewed storage has only %" PRId64 " element(s) after offset %" PRId64,
size, numel - offset, offset);
real *data_ptr = THStorage_(data)(LIBRARY_STATE storage_arg->cdata) + offset;
THStoragePtr storage(THStorage_(newWithData)(LIBRARY_STATE data_ptr, size));
storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_VIEW;
storage->view = storage_arg->cdata;
THStorage_(retain)(LIBRARY_STATE storage_arg->cdata);
self->cdata = storage.release();
return (PyObject*)self.release();
#endif
}
// torch.Storage(sequence)
if (num_args == 1 && PySequence_Check(first_arg)) {
#ifdef THD_GENERIC_FILE
THPUtils_setError("distributed storages don't support construction from a sequence");
#else
Py_ssize_t length = PySequence_Length(first_arg);
THPUtils_assert(length >= 0, "couldn't obtain the length of %s",
THPUtils_typename(first_arg));
self->cdata = THStorage_(newWithSize)(LIBRARY_STATE length);
THPObjectPtr item;
try {
for (Py_ssize_t i = 0; i < length; i++) {
item = PySequence_GetItem(first_arg, i);
real value = THPUtils_(unpackReal)(item.get());
#if !defined(THC_GENERIC_FILE)
self->cdata->unsafe_data<real>()[i] = value;
#else
// TODO: this might be slow - consider batched updates?
THCStorage_(set)(LIBRARY_STATE self->cdata, i, value);
#endif
}
} catch (std::runtime_error &e) {
THPUtils_setError("tried to construct a storage from a sequence (%s), "
"but one of the items was of type %s instead of %s",
THPUtils_typename(first_arg),
THPUtils_typename(item.get()),
THPUtils_typeTraits<real>::python_type_str);
return NULL;
}
return (PyObject*)self.release();
#endif
}
#ifndef THD_GENERIC_FILE
invalid_arguments:
#endif
THPUtils_invalidArguments(args, kwargs, THPStorageStr " constructor", 6,
"no arguments",
"(int size)",
"(Sequence data)",
"(" THPStorageStr " view_source)",
"(" THPStorageStr " view_source, int offset)",
"(" THPStorageStr " view_source, int offset, int size)");
return NULL;
END_HANDLE_TH_ERRORS
}
