1 #include <torch/csrc/python_headers.h>
2 #ifdef _MSC_VER
3 #include <c10/util/win32-headers.h>
4 #endif
5 #include <structmember.h>
6
7 #include <c10/core/CPUAllocator.h>
8 #include <libshm.h>
9 #include <torch/csrc/CudaIPCTypes.h>
10 #include <torch/csrc/Device.h>
11 #include <torch/csrc/DynamicTypes.h>
12 #include <torch/csrc/THP.h>
13 #include <torch/csrc/autograd/utils/wrap_outputs.h>
14 #include <torch/csrc/copy_utils.h>
15
16 #include <c10/util/intrusive_ptr.h>
17 #include <fmt/format.h>
18
19 #include <torch/csrc/Storage.h>
20 #include <torch/csrc/StorageMethods.h>
21
22 #include <ATen/ATen.h>
23 #include <ATen/MapAllocator.h>
24 #include <ATen/StorageUtils.h>
25 #include <torch/csrc/utils/pycfunction_helpers.h>
26 #include <torch/csrc/utils/python_arg_parser.h>
27 #include <torch/csrc/utils/python_numbers.h>
28
29 #ifdef USE_CUDA
30 #include <ATen/native/cuda/Resize.h>
31 #include <cuda_runtime.h>
32 #endif
33
34 #include <ATen/detail/PrivateUse1HooksInterface.h>
35 #include <ATen/native/Resize.h>
36
37 #ifdef _MSC_VER
38 #define LSEEK _lseeki64
39 #else
40 #define LSEEK lseek
41 #endif
42
THPStorage_nbytes(PyObject * self,PyObject * noargs)43 static PyObject* THPStorage_nbytes(PyObject* self, PyObject* noargs) {
44 HANDLE_TH_ERRORS
45 THPStorage_assertNotNull(self);
46 return py::cast(THPStorage_Unpack(self).sym_nbytes()).release().ptr();
47 END_HANDLE_TH_ERRORS
48 }
49
THPStorage_dataPtr(PyObject * self,PyObject * noargs)50 static PyObject* THPStorage_dataPtr(PyObject* self, PyObject* noargs) {
51 HANDLE_TH_ERRORS
52 // PyLong_FromVoidPtr should not need to mutate the pointer in order
53 // to extract a new long object from it.
54
55 auto self_ = THPStorage_Unpack(self);
56 // See Note [Invalid Python Storages]
57 auto invalid = self_.data() == nullptr &&
58 self_.device_type() != c10::DeviceType::Meta && self_.sym_nbytes() != 0;
59 TORCH_CHECK(
60 !invalid,
61 "Attempted to access the data pointer on an invalid python storage.")
62 return PyLong_FromVoidPtr(self_.mutable_data());
63 END_HANDLE_TH_ERRORS
64 }
65
THPStorage_resizable(PyObject * self,PyObject * noargs)66 static PyObject* THPStorage_resizable(PyObject* self, PyObject* noargs) {
67 HANDLE_TH_ERRORS
68 THPStorage_assertNotNull(self);
69 return PyBool_FromLong(THPStorage_Unpack(self).resizable());
70 END_HANDLE_TH_ERRORS
71 }
72
THPStorage_copy_(PyObject * self,PyObject * args,PyObject * kwargs)73 static PyObject* THPStorage_copy_(
74 PyObject* self,
75 PyObject* args,
76 PyObject* kwargs) {
77 HANDLE_TH_ERRORS
78 THPStorage_assertNotNull(self);
79
80 at::Storage self_ = torch::createStorage(self);
81
82 static torch::PythonArgParser parser({
83 "copy_(Storage src, bool? non_blocking=None)",
84 });
85 torch::ParsedArgs<2> parsed_args;
86 auto r = parser.parse(args, kwargs, parsed_args);
87
88 at::Storage src = r.storage(0);
89 bool non_blocking = r.toBoolOptional(1).value_or(false);
90
91 // See Note [Invalid Python Storages]
92 auto invalid = src.data() == nullptr &&
93 src.device_type() != c10::DeviceType::Meta && src.sym_nbytes() != 0;
94 TORCH_CHECK(
95 !invalid, "Attempted to call copy_() on an invalid python storage.")
96
97 TORCH_CHECK(
98 self_.nbytes() == src.nbytes(),
99 "size does not match, self was ",
100 self_.nbytes(),
101 " bytes but src was ",
102 src.nbytes(),
103 " bytes");
104
105 at::storage_copy(self_, src, non_blocking);
106
107 Py_INCREF(self);
108 return self;
109
110 END_HANDLE_TH_ERRORS
111 }
112
THPStorage_elementSize(PyObject * _self,PyObject * noargs)113 static PyObject* THPStorage_elementSize(PyObject* _self, PyObject* noargs) {
114 HANDLE_TH_ERRORS
115 THPStorage_assertNotNull(_self);
116 return THPUtils_packInt64(sizeof(uint8_t));
117 END_HANDLE_TH_ERRORS
118 }
119
THPStorage_new(PyObject * self,PyObject * noargs)120 static PyObject* THPStorage_new(PyObject* self, PyObject* noargs) {
121 HANDLE_TH_ERRORS
122 THPStorage_assertNotNull(self);
123 c10::Allocator* allocator = THPStorage_Unpack(self).allocator();
124 auto new_storage = c10::make_intrusive<at::StorageImpl>(
125 c10::StorageImpl::use_byte_size_t(),
126 0,
127 allocator,
128 /*resizable=*/true);
129
130 return THPStorage_Wrap(std::move(new_storage));
131 END_HANDLE_TH_ERRORS
132 }
133
THPStorage_resize_(PyObject * self,PyObject * number_arg)134 static PyObject* THPStorage_resize_(PyObject* self, PyObject* number_arg) {
135 HANDLE_TH_ERRORS
136 THPStorage_assertNotNull(self);
137 const auto& storage = THPStorage_Unpack(self);
138 // See Note [Invalid Python Storages]
139 auto invalid = storage.data() == nullptr &&
140 storage.device_type() != c10::DeviceType::Meta &&
141 storage.sym_nbytes() != 0;
142 TORCH_CHECK(
143 !invalid, "Attempted to call resize_() on an invalid python storage.")
144 TORCH_CHECK(
145 THPUtils_checkLong(number_arg),
146 "resize_ expects an int, "
147 "but got ",
148 THPUtils_typename(number_arg));
149 int64_t newsize = THPUtils_unpackLong(number_arg);
150 c10::DeviceType device_type = storage.device_type();
151 if (device_type == at::kCUDA) {
152 #ifdef USE_CUDA
153 ptrdiff_t size_bytes_i = newsize;
154 TORCH_CHECK(
155 !c10::overflows<size_t>(size_bytes_i),
156 "Requested storage size (",
157 size_bytes_i,
158 ") cannot be represented as a size_t");
159 const auto size_bytes = static_cast<size_t>(size_bytes_i);
160 at::native::resize_bytes_cuda(storage.unsafeGetStorageImpl(), size_bytes);
161 #else
162 TORCH_CHECK(false, "built without USE_CUDA");
163 #endif
164 } else {
165 at::native::resize_bytes_nocuda(storage, newsize);
166 }
167 Py_INCREF(self);
168 return self;
169 END_HANDLE_TH_ERRORS
170 }
171
THPStorage_fill_(PyObject * self,PyObject * number_arg)172 static PyObject* THPStorage_fill_(PyObject* self, PyObject* number_arg) {
173 HANDLE_TH_ERRORS
174 THPStorage_assertNotNull(self);
175 const auto& storage = THPStorage_Unpack(self);
176 // See Note [Invalid Python Storages]
177 auto invalid = storage.data() == nullptr &&
178 storage.device_type() != c10::DeviceType::Meta &&
179 storage.sym_nbytes() != 0;
180 TORCH_CHECK(
181 !invalid, "Attempted to call fill_() on an invalid python storage.")
182 TORCH_CHECK(
183 THPByteUtils_checkReal(number_arg),
184 "fill_ expects int, "
185 "but got ",
186 THPUtils_typename(number_arg));
187 storage_fill(storage, THPByteUtils_unpackReal(number_arg));
188 Py_INCREF(self);
189 return self;
190 END_HANDLE_TH_ERRORS
191 }
192
THPStorage_fromBuffer(PyObject * _unused,PyObject * args,PyObject * keywds)193 static PyObject* THPStorage_fromBuffer(
194 PyObject* _unused,
195 PyObject* args,
196 PyObject* keywds) {
197 HANDLE_TH_ERRORS
198 PyObject* obj = nullptr;
199 const char* byte_order_str = nullptr;
200 Py_ssize_t count = -1, offset = 0;
201 PyObject* dtype_obj = nullptr;
202 c10::ScalarType scalar_type = at::kByte;
203 Py_buffer buffer = {};
204 // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
205 constexpr const char* kwlist[] = {
206 "buffer", "byte_order", "count", "offset", "dtype", nullptr};
207 constexpr const char* argtypes = "O|snnO";
208
209 if (!PyArg_ParseTupleAndKeywords(
210 args,
211 keywds,
212 argtypes,
213 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
214 const_cast<char**>(kwlist),
215 &obj,
216 &byte_order_str,
217 &count,
218 &offset,
219 &dtype_obj)) {
220 return nullptr;
221 }
222 TORCH_CHECK(dtype_obj != nullptr, "argument 'dtype' cannot be None");
223 TORCH_CHECK(
224 THPDtype_Check(dtype_obj),
225 "argument 'dtype' must be of type torch.dtype");
226 auto dtype = reinterpret_cast<THPDtype*>(dtype_obj);
227 scalar_type = dtype->scalar_type;
228
229 const bool is_endian_independent = (scalar_type == at::kByte) ||
230 (scalar_type == at::kChar) || (scalar_type == at::kFloat8_e5m2) ||
231 (scalar_type == at::kFloat8_e5m2fnuz) ||
232 (scalar_type == at::kFloat8_e4m3fn) ||
233 (scalar_type == at::kFloat8_e4m3fnuz);
234
235 TORCH_CHECK(
236 is_endian_independent || (byte_order_str != nullptr),
237 "function missing required argument 'byte_order' (pos 2)");
238 size_t element_size = c10::elementSize(scalar_type);
239
240 bool do_byte_swap = false;
241 if (!is_endian_independent) {
242 if (strcmp(byte_order_str, "native") == 0) {
243 do_byte_swap = false;
244 } else if (strcmp(byte_order_str, "big") == 0) {
245 do_byte_swap =
246 (torch::utils::THP_LITTLE_ENDIAN ==
247 torch::utils::THP_nativeByteOrder());
248 } else if (strcmp(byte_order_str, "little") == 0) {
249 do_byte_swap =
250 (torch::utils::THP_BIG_ENDIAN == torch::utils::THP_nativeByteOrder());
251 } else {
252 PyErr_Format(
253 PyExc_ValueError,
254 "invalid byte_order '%s' (expected 'big', 'little', or 'native')",
255 byte_order_str);
256 return nullptr;
257 }
258 }
259
260 if (PyObject_GetBuffer(obj, &buffer, PyBUF_SIMPLE) < 0)
261 return nullptr;
262
263 if (offset < 0 || offset > buffer.len) {
264 PyErr_SetString(
265 PyExc_ValueError,
266 fmt::format(
267 "offset must be non-negative and no greater than buffer length ({}) , but got {}",
268 offset,
269 buffer.len));
270 PyBuffer_Release(&buffer);
271 return nullptr;
272 }
273
274 size_t size_bytes = 0;
275 if (count < 0) {
276 if ((buffer.len - offset) % element_size != 0) {
277 PyErr_SetString(
278 PyExc_ValueError,
279 fmt::format(
280 "buffer size ({}) must be a multiple of element size ({})",
281 buffer.len,
282 element_size));
283 PyBuffer_Release(&buffer);
284 return nullptr;
285 }
286 size_bytes = buffer.len - offset;
287 count = static_cast<Py_ssize_t>(size_bytes / element_size);
288 } else {
289 size_bytes = count * element_size;
290 }
291
292 if (offset + (count * (Py_ssize_t)element_size) > buffer.len) {
293 PyErr_SetString(
294 PyExc_ValueError,
295 fmt::format(
296 "buffer has only {} elements after offset {}, but specified a size of {}",
297 buffer.len - offset,
298 offset,
299 count));
300 PyBuffer_Release(&buffer);
301 return nullptr;
302 }
303
304 uint8_t* src = (uint8_t*)buffer.buf;
305 auto storage = c10::make_intrusive<at::StorageImpl>(
306 c10::StorageImpl::use_byte_size_t(),
307 size_bytes,
308 c10::GetDefaultCPUAllocator(),
309 /*resizable=*/true);
310
311 if (is_endian_independent) {
312 memcpy(storage->mutable_data(), src + offset, count);
313 } else if (scalar_type == at::kBool) {
314 // Because of ASAN checks, that are failing whenever
315 // we are trying to get a value which is not 0 or 1, we have to manually
316 // convert original values to boolean ones.
317 torch::utils::THP_decodeBoolBuffer(
318 static_cast<bool*>(storage->mutable_data()), src + offset, count);
319 } else if (scalar_type == at::kShort) {
320 torch::utils::THP_decodeInt16Buffer(
321 static_cast<int16_t*>(storage->mutable_data()),
322 src + offset,
323 do_byte_swap,
324 count);
325 } else if (scalar_type == at::kInt) {
326 torch::utils::THP_decodeInt32Buffer(
327 static_cast<int32_t*>(storage->mutable_data()),
328 src + offset,
329 do_byte_swap,
330 count);
331 } else if (scalar_type == at::kLong) {
332 torch::utils::THP_decodeInt64Buffer(
333 static_cast<int64_t*>(storage->mutable_data()),
334 src + offset,
335 do_byte_swap,
336 count);
337 } else if (scalar_type == at::kHalf) {
338 torch::utils::THP_decodeHalfBuffer(
339 static_cast<c10::Half*>(storage->mutable_data()),
340 src + offset,
341 do_byte_swap,
342 count);
343 } else if (scalar_type == at::kBFloat16) {
344 torch::utils::THP_decodeBFloat16Buffer(
345 static_cast<c10::BFloat16*>(storage->mutable_data()),
346 src + offset,
347 do_byte_swap,
348 count);
349 } else if (scalar_type == at::kFloat) {
350 torch::utils::THP_decodeFloatBuffer(
351 static_cast<float*>(storage->mutable_data()),
352 src + offset,
353 do_byte_swap,
354 count);
355 } else if (scalar_type == at::kDouble) {
356 torch::utils::THP_decodeDoubleBuffer(
357 static_cast<double*>(storage->mutable_data()),
358 src + offset,
359 do_byte_swap,
360 count);
361 } else if (scalar_type == at::kComplexFloat) {
362 torch::utils::THP_decodeComplexFloatBuffer(
363 static_cast<c10::complex<float>*>(storage->mutable_data()),
364 src + offset,
365 do_byte_swap,
366 count);
367 } else if (scalar_type == at::kComplexDouble) {
368 torch::utils::THP_decodeComplexDoubleBuffer(
369 static_cast<c10::complex<double>*>(storage->mutable_data()),
370 src + offset,
371 do_byte_swap,
372 count);
373 } else {
374 TORCH_CHECK(false, "Unknown type: ", scalar_type);
375 }
376
377 PyBuffer_Release(&buffer);
378 return THPStorage_Wrap(storage);
379 END_HANDLE_TH_ERRORS
380 }
381
THPStorage_fromFile(PyObject * _unused,PyObject * args,PyObject * keywds)382 static PyObject* THPStorage_fromFile(
383 PyObject* _unused,
384 PyObject* args,
385 PyObject* keywds) {
386 HANDLE_TH_ERRORS
387 const char* filename = nullptr;
388 Py_ssize_t nbytes = 0;
389 int shared = 0;
390 // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
391 constexpr const char* kwlist[] = {"filename", "shared", "nbytes", nullptr};
392 if (!PyArg_ParseTupleAndKeywords(
393 args,
394 keywds,
395 "s|in",
396 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
397 const_cast<char**>(kwlist),
398 &filename,
399 &shared,
400 &nbytes)) {
401 return nullptr;
402 }
403 if (shared)
404 shared = at::ALLOCATOR_MAPPED_SHARED;
405
406 size_t actual_nbytes = -1;
407 auto storage = c10::make_intrusive<at::StorageImpl>(
408 c10::StorageImpl::use_byte_size_t(),
409 nbytes,
410 at::MapAllocator::makeDataPtr(filename, shared, nbytes, &actual_nbytes),
411 /*allocator=*/nullptr,
412 /*resizable=*/false);
413
414 if (nbytes <= 0) {
415 storage->set_nbytes(actual_nbytes);
416 }
417
418 return THPStorage_NewWithStorage(
419 THPStorageClass,
420 std::move(storage),
421 c10::impl::PyInterpreterStatus::TAGGED_BY_US);
422 END_HANDLE_TH_ERRORS
423 }
424
THPStorage_writeFile(PyObject * self,PyObject * args)425 PyObject* THPStorage_writeFile(PyObject* self, PyObject* args) {
426 HANDLE_TH_ERRORS
427 THPStorage_assertNotNull(self);
428 const auto& storage = THPStorage_Unpack(self);
429 // See Note [Invalid Python Storages]
430 auto invalid = storage.data() == nullptr &&
431 storage.device_type() != c10::DeviceType::Meta &&
432 storage.sym_nbytes() != 0;
433 TORCH_CHECK(
434 !invalid, "Attempted to call _write_file() on an invalid python storage.")
435 PyObject* file = PyTuple_GetItem(args, 0);
436 bool is_real_file = PyTuple_GetItem(args, 1) == Py_True;
437 bool save_size = PyTuple_GetItem(args, 2) == Py_True;
438 PyObject* element_size_obj = PyTuple_GET_ITEM(args, 3);
439
440 TORCH_CHECK(
441 element_size_obj != Py_None, "_write_file: need to specify element size");
442 uint64_t element_size = THPUtils_unpackUInt64(element_size_obj);
443
444 if (!is_real_file) {
445 THPStorage_writeFileRaw<PyObject*>(
446 storage.unsafeGetStorageImpl(), file, save_size, element_size);
447 Py_RETURN_NONE;
448 }
449
450 int fd = PyObject_AsFileDescriptor(file);
451 TORCH_CHECK(
452 fd != -1,
453 "_write_file couldn't retrieve a file descriptor "
454 "from given object");
455 THPStorage_writeFileRaw(
456 storage.unsafeGetStorageImpl(), fd, save_size, element_size);
457 Py_RETURN_NONE;
458 END_HANDLE_TH_ERRORS
459 }
460
THPStorage_newWithFile(PyObject * _unused,PyObject * args)461 PyObject* THPStorage_newWithFile(PyObject* _unused, PyObject* args) {
462 HANDLE_TH_ERRORS
463 TORCH_CHECK(
464 PyTuple_Size(args) == 2, "_new_with_file takes exactly two arguments");
465 int fd = PyObject_AsFileDescriptor(PyTuple_GetItem(args, 0));
466 TORCH_CHECK(
467 fd != -1,
468 "_new_with_file couldn't retrieve a file "
469 "descriptor from given object");
470 PyObject* element_size_obj = PyTuple_GET_ITEM(args, 1);
471 TORCH_CHECK(
472 element_size_obj != Py_None,
473 "_new_with_file: need to specify element size");
474 uint64_t element_size = THPUtils_unpackUInt64(element_size_obj);
475
476 auto storage = THPStorage_readFileRaw<int>(fd, {}, element_size);
477 if (!storage.defined())
478 return nullptr;
479 return THPStorage_Wrap(std::move(storage));
480 END_HANDLE_TH_ERRORS
481 }
482
THPStorage_setFromFile(PyObject * self,PyObject * args)483 static PyObject* THPStorage_setFromFile(PyObject* self, PyObject* args) {
484 HANDLE_TH_ERRORS
485 THPStorage_assertNotNull(self);
486 const auto& storage = THPStorage_Unpack(self);
487 PyObject* file = PyTuple_GET_ITEM(args, 0);
488 PyObject* offset = PyTuple_GET_ITEM(args, 1);
489 bool is_real_file = PyTuple_GET_ITEM(args, 2) == Py_True;
490
491 PyObject* element_size_obj = PyTuple_GET_ITEM(args, 3);
492
493 TORCH_CHECK(
494 element_size_obj != Py_None,
495 "_set_from_file: need to specify element size");
496 uint64_t element_size = THPUtils_unpackUInt64(element_size_obj);
497
498 if (!is_real_file) {
499 // offset can be implemented with a call to the Python object's seek()
500 // but it is currently unnecessary to support this.
501 TORCH_CHECK(
502 offset == Py_None,
503 "_set_from_file: offset is NYI for filelike objects");
504
505 auto self_storage_impl = c10::intrusive_ptr<c10::StorageImpl>::reclaim_copy(
506 storage.unsafeGetStorageImpl());
507 auto storage_impl = THPStorage_readFileRaw<PyObject*>(
508 file, std::move(self_storage_impl), element_size);
509 if (!storage_impl.defined()) {
510 return nullptr;
511 }
512 Py_INCREF(self);
513 return (PyObject*)self;
514 }
515
516 // file is backed by a fd
517 const int fd = PyObject_AsFileDescriptor(file);
518 const auto fd_original_pos = LSEEK(fd, 0, SEEK_CUR);
519 if (offset != Py_None) {
520 LSEEK(fd, THPUtils_unpackLong(offset), SEEK_SET);
521 }
522 TORCH_CHECK(
523 fd != -1,
524 "_set_from_file couldn't retrieve a file "
525 "descriptor from given object");
526 auto self_storage_impl = c10::intrusive_ptr<c10::StorageImpl>::reclaim_copy(
527 storage.unsafeGetStorageImpl());
528 auto storage_impl =
529 THPStorage_readFileRaw<int>(fd, self_storage_impl, element_size);
530 if (!storage_impl.defined())
531 return nullptr;
532 Py_INCREF(self);
533
534 // the file descriptor is returned to original position and
535 // the file handle at python call-site needs updating to the
536 // advanced position
537 const auto fd_current_pos = LSEEK(fd, 0, SEEK_CUR);
538 LSEEK(fd, fd_original_pos, SEEK_SET);
539 const auto seek_return =
540 PyObject_CallMethod(file, "seek", "Li", (long long)fd_current_pos, 0);
541 if (seek_return == nullptr) {
542 return nullptr;
543 }
544 Py_DECREF(seek_return);
545
546 return self;
547 END_HANDLE_TH_ERRORS
548 }
549
THPStorage__setCdata(PyObject * _self,PyObject * new_cdata)550 PyObject* THPStorage__setCdata(PyObject* _self, PyObject* new_cdata) {
551 HANDLE_TH_ERRORS
552 auto self = (THPStorage*)_self;
553 TORCH_CHECK(
554 THPUtils_checkLong(new_cdata),
555 "given an invalid argument to "
556 "_set_cdata - expected an int or long, but got ",
557 THPUtils_typename(new_cdata));
558 c10::StorageImpl* ptr = (c10::StorageImpl*)PyLong_AsVoidPtr(new_cdata);
559 self->cdata.~MaybeOwned<c10::Storage>();
560 self->cdata = c10::MaybeOwned<c10::Storage>::owned(
561 c10::Storage(c10::intrusive_ptr<c10::StorageImpl>::reclaim_copy(ptr)));
562 Py_INCREF(self);
563 return (PyObject*)self;
564 END_HANDLE_TH_ERRORS
565 }
566
THPStorage_byteswap(PyObject * self,PyObject * args)567 PyObject* THPStorage_byteswap(PyObject* self, PyObject* args) {
568 HANDLE_TH_ERRORS
569 TORCH_CHECK(PyTuple_GET_SIZE(args) == 1, "tuple of 1 item expected");
570 PyObject* _elem_size = PyTuple_GET_ITEM(args, 0);
571 TORCH_CHECK(
572 THPUtils_checkLong(_elem_size), "_byteswap(): arg must be an 'int'");
573 auto elem_size = THPUtils_unpackLong(_elem_size);
574 TORCH_CHECK(
575 elem_size == 1 || elem_size == 2 || elem_size == 4 || elem_size == 8,
576 "elem_size must be 1, 2, 4, or 8");
577
578 const auto& storage = THPStorage_Unpack(self);
579 const auto nbytes = static_cast<uint64_t>(storage.nbytes());
580 const uint64_t count = nbytes / elem_size;
581
582 if (elem_size == 1) {
583 Py_RETURN_NONE;
584 }
585 TORCH_CHECK(
586 nbytes % elem_size == 0,
587 "the length of data is not a multiple of ",
588 elem_size);
589
590 if (elem_size == 2) {
591 auto buffer = static_cast<uint16_t*>(storage.mutable_data());
592 for (uint64_t i = 0; i < count; i++, buffer++) {
593 *buffer = thp_bswap16(*buffer);
594 }
595 } else if (elem_size == 4) {
596 auto buffer = static_cast<uint32_t*>(storage.mutable_data());
597 for (uint64_t i = 0; i < count; i++, buffer++) {
598 *buffer = thp_bswap32(*buffer);
599 }
600 } else if (elem_size == 8) {
601 auto buffer = static_cast<uint64_t*>(storage.mutable_data());
602 for (uint64_t i = 0; i < count; i++, buffer++) {
603 *buffer = thp_bswap64(*buffer);
604 }
605 }
606
607 Py_RETURN_NONE;
608 END_HANDLE_TH_ERRORS
609 }
610
THPStorage_fix_weakref(PyObject * self,PyObject * noargs)611 static PyObject* THPStorage_fix_weakref(PyObject* self, PyObject* noargs) {
612 const auto& storage = THPStorage_Unpack(self);
613 Py_DECREF(THPStorage_Wrap(storage));
614 Py_RETURN_NONE;
615 }
616
THPStorage__get_filename(PyObject * self,PyObject * noargs)617 static PyObject* THPStorage__get_filename(PyObject* self, PyObject* noargs) {
618 HANDLE_TH_ERRORS
619
620 const auto& self_ = THPStorage_Unpack(self);
621 const c10::DataPtr& data_ptr = self_.data_ptr();
622 at::MapAllocator* map_allocator = at::MapAllocator::fromDataPtr(data_ptr);
623
624 if (map_allocator == nullptr ||
625 !(map_allocator->flags() & at::ALLOCATOR_MAPPED_SHARED)) {
626 Py_RETURN_NONE;
627 }
628 std::string filename = map_allocator->filename();
629
630 return THPUtils_packString(filename);
631 END_HANDLE_TH_ERRORS
632 }
633
634 // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
635 static PyMethodDef THPStorage_methods[] = {
636 {"copy_",
637 castPyCFunctionWithKeywords(THPStorage_copy_),
638 METH_VARARGS | METH_KEYWORDS,
639 nullptr},
640 {"element_size", THPStorage_elementSize, METH_NOARGS, nullptr},
641 {"fill_", THPStorage_fill_, METH_O, nullptr},
642 {"new", THPStorage_new, METH_NOARGS, nullptr},
643 {"resize_", THPStorage_resize_, METH_O, nullptr},
644 {"nbytes", THPStorage_nbytes, METH_NOARGS, nullptr},
645 {"data_ptr", THPStorage_dataPtr, METH_NOARGS, nullptr},
646 {"resizable", THPStorage_resizable, METH_NOARGS, nullptr},
647 {"_write_file", THPStorage_writeFile, METH_VARARGS, nullptr},
648 {"_new_with_file",
649 THPStorage_newWithFile,
650 METH_VARARGS | METH_STATIC,
651 nullptr},
652 {"_set_from_file", THPStorage_setFromFile, METH_VARARGS, nullptr},
653 {"from_buffer",
654 castPyCFunctionWithKeywords(THPStorage_fromBuffer),
655 METH_VARARGS | METH_KEYWORDS | METH_STATIC,
656 nullptr},
657 {"from_file",
658 castPyCFunctionWithKeywords(THPStorage_fromFile),
659 METH_VARARGS | METH_KEYWORDS | METH_STATIC,
660 nullptr},
661 {"_set_cdata", THPStorage__setCdata, METH_O, nullptr},
662 {"_byteswap", THPStorage_byteswap, METH_VARARGS, nullptr},
663 {"_fix_weakref", THPStorage_fix_weakref, METH_NOARGS, nullptr},
664 {"_get_filename", THPStorage__get_filename, METH_NOARGS, nullptr},
665 {nullptr}};
666
THPStorage_getMethods()667 PyMethodDef* THPStorage_getMethods() {
668 return THPStorage_methods;
669 }
670