1 #ifndef TH_GENERIC_FILE 2 #define TH_GENERIC_FILE "torch/csrc/generic/serialization.cpp" 5 #ifdef THC_GENERIC_FILE 6 #include <c10/cuda/CUDAGuard.h> 10 void THPStorage_(writeFileRaw)(THWStorage *
self, io fd)
12 #ifdef THC_GENERIC_FILE 17 int64_t size = THWStorage_(size)(LIBRARY_STATE
self);
18 #ifndef THC_GENERIC_FILE 19 data = THWStorage_(data)(LIBRARY_STATE
self);
21 std::unique_ptr<char[]> cpu_data(
new char[size *
sizeof(scalar_t)]);
22 data = (scalar_t*)cpu_data.get();
23 THCudaCheck(cudaMemcpy(data, THWStorage_(data)(LIBRARY_STATE
self), size *
sizeof(scalar_t), cudaMemcpyDeviceToHost));
25 doWrite(fd, &size,
sizeof(int64_t));
27 if (
sizeof(scalar_t) == 1 || THP_nativeByteOrder() == THPByteOrder::THP_LITTLE_ENDIAN) {
28 doWrite(fd, data,
sizeof(scalar_t) * size);
30 int64_t buffer_size = std::min(size, (int64_t)5000);
31 std::unique_ptr<uint8_t[]> le_buffer(
new uint8_t[buffer_size *
sizeof(scalar_t)]);
32 for (int64_t i = 0; i < size; i += buffer_size) {
33 size_t to_convert = std::min(size - i, buffer_size);
34 if (
sizeof(scalar_t) == 2) {
35 THP_encodeInt16Buffer((uint8_t*)le_buffer.get(),
36 (
const int16_t*)data + i,
37 THPByteOrder::THP_LITTLE_ENDIAN,
39 }
else if (
sizeof(scalar_t) == 4) {
40 THP_encodeInt32Buffer((uint8_t*)le_buffer.get(),
41 (
const int32_t*)data + i,
42 THPByteOrder::THP_LITTLE_ENDIAN,
44 }
else if (
sizeof(scalar_t) == 8) {
45 THP_encodeInt64Buffer((uint8_t*)le_buffer.get(),
46 (
const int64_t*)data + i,
47 THPByteOrder::THP_LITTLE_ENDIAN,
50 doWrite(fd, le_buffer.get(), to_convert *
sizeof(scalar_t));
55 template void THPStorage_(writeFileRaw<int>)(THWStorage *
self,
int fd);
56 template void THPStorage_(writeFileRaw<PyObject*>)(THWStorage *
self, PyObject* fd);
59 THWStorage * THPStorage_(readFileRaw)(io file, THWStorage *_storage)
61 #ifdef THC_GENERIC_FILE 63 if (_storage !=
nullptr) {
70 doRead(file, &size,
sizeof(int64_t));
71 THWStoragePtr storage;
72 if (_storage ==
nullptr) {
73 storage = THWStorage_(newWithSize)(LIBRARY_STATE size);
75 THPUtils_assert(THWStorage_(size)(LIBRARY_STATE _storage) == size,
76 "storage has wrong size: expected %ld got %ld",
77 size, THWStorage_(size)(LIBRARY_STATE _storage));
81 #ifndef THC_GENERIC_FILE 82 data = THWStorage_(data)(LIBRARY_STATE storage);
84 std::unique_ptr<char[]> cpu_data(
new char[size *
sizeof(scalar_t)]);
85 data = (scalar_t*)cpu_data.get();
89 if (
sizeof(scalar_t) == 1 || THP_nativeByteOrder() == THPByteOrder::THP_LITTLE_ENDIAN) {
90 doRead(file, data,
sizeof(scalar_t) * THWStorage_(size)(LIBRARY_STATE storage));
92 int64_t buffer_size = std::min(size, (int64_t)5000);
93 std::unique_ptr<uint8_t[]> le_buffer(
new uint8_t[buffer_size *
sizeof(scalar_t)]);
96 for (int64_t i = 0; i < size; i += buffer_size) {
97 size_t to_convert = std::min(size - i, buffer_size);
98 doRead(file, le_buffer.get(),
sizeof(scalar_t) * to_convert);
100 if (
sizeof(scalar_t) == 2) {
101 THP_decodeInt16Buffer((int16_t*)data + i,
103 THPByteOrder::THP_LITTLE_ENDIAN,
105 }
else if (
sizeof(scalar_t) == 4) {
106 THP_decodeInt32Buffer((int32_t*)data + i,
108 THPByteOrder::THP_LITTLE_ENDIAN,
110 }
else if (
sizeof(scalar_t) == 8) {
111 THP_decodeInt64Buffer((int64_t*)data + i,
113 THPByteOrder::THP_LITTLE_ENDIAN,
119 #ifdef THC_GENERIC_FILE 120 THCudaCheck(cudaMemcpy(THWStorage_(data)(LIBRARY_STATE storage), data, size *
sizeof(scalar_t), cudaMemcpyHostToDevice));
122 return storage.release();
125 template THWStorage* THPStorage_(readFileRaw<int>)(
int fd, THWStorage* storage);
126 template THWStorage* THPStorage_(readFileRaw<PyObject*>)(PyObject* fd, THWStorage* storage);
void set_device(Device device)
Sets the CUDA device to the given device, initializing the guard if it is not already initialized...
A variant of OptionalDeviceGuard that is specialized for CUDA.
A variant of DeviceGuard that is specialized for CUDA.