Caffe2 - C++ API
A deep learning, cross platform ML framework
THCStorageCopy.cpp
1 #ifndef THC_GENERIC_FILE
2 #define THC_GENERIC_FILE "THC/generic/THCStorageCopy.cpp"
3 #else
4 
5 void THCStorage_(copyCPU)(THCState *state, THCStorage *self, struct THStorage *src)
6 {
7  THArgCheck(self->numel() == src->numel(), 2, "size does not match");
8  cudaStream_t stream = THCState_getCurrentStream(state);
9  THCudaCheck(cudaMemcpyAsync(THCStorage_(data)(state, self),
10  THStorage_(data)(src),
11  self->numel() * sizeof(scalar_t),
12  cudaMemcpyHostToDevice,
13  stream));
14  THCudaCheck(cudaStreamSynchronize(stream));
15 }
16 
17 #define TH_CUDA_STORAGE_IMPLEMENT_COPY(TYPEC) \
18 void THCStorage_(copy##TYPEC)(THCState *state, THCStorage *self, struct TH##TYPEC##Storage *src) \
19 { \
20  THCTensor* selfTensor = \
21  THCTensor_(newWithStorage1d)(state, self, 0, self->numel(), 1); \
22  struct TH##TYPEC##Tensor* srcTensor = \
23  TH##TYPEC##Tensor_newWithStorage1d(src, 0, src->numel(), 1); \
24  THCTensor_(copy)(state, selfTensor, srcTensor); \
25  TH##TYPEC##Tensor_free(srcTensor); \
26  THCTensor_(free)(state, selfTensor); \
27 }
28 TH_CUDA_STORAGE_IMPLEMENT_COPY(Byte)
29 TH_CUDA_STORAGE_IMPLEMENT_COPY(Char)
30 TH_CUDA_STORAGE_IMPLEMENT_COPY(Short)
31 TH_CUDA_STORAGE_IMPLEMENT_COPY(Int)
32 TH_CUDA_STORAGE_IMPLEMENT_COPY(Long)
33 TH_CUDA_STORAGE_IMPLEMENT_COPY(Float)
34 TH_CUDA_STORAGE_IMPLEMENT_COPY(Half)
35 TH_CUDA_STORAGE_IMPLEMENT_COPY(Double)
36 TH_CUDA_STORAGE_IMPLEMENT_COPY(Bool)
37 
38 void THStorage_(copyCuda)(THCState *state, THStorage *self, struct THCStorage *src)
39 {
40  THArgCheck(self->numel() == src->numel(), 2, "size does not match");
41  cudaStream_t stream = THCState_getCurrentStream(state);
42  THCudaCheck(cudaMemcpyAsync(THStorage_(data)(self),
43  THCStorage_(data)(state, src),
44  self->numel() * sizeof(scalar_t),
45  cudaMemcpyDeviceToHost,
46  stream));
47  THCudaCheck(cudaStreamSynchronize(stream));
48 }
49 
50 #define TH_CUDA_STORAGE_IMPLEMENT_COPYTO(TYPEC) \
51 void TH_CONCAT_4(TH,TYPEC,Storage_copyCuda,Real)(THCState *state, TH##TYPEC##Storage *self, struct THCStorage *src) \
52 { \
53  TH##TYPEC##Tensor* selfTensor = \
54  TH##TYPEC##Tensor_newWithStorage1d(self, 0, self->numel(), 1); \
55  struct THCTensor* srcTensor = \
56  THCTensor_(newWithStorage1d)(state, src, 0, src->numel(), 1); \
57  THCTensor_(copy)(state, selfTensor, srcTensor); \
58  THCTensor_(free)(state, srcTensor); \
59  TH##TYPEC##Tensor_free(selfTensor); \
60 }
61 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Byte)
62 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Char)
63 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Short)
64 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Int)
65 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Long)
66 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Float)
67 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Half)
68 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Double)
69 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Bool)
70 
71 #undef TH_CUDA_STORAGE_IMPLEMENT_COPY
72 #undef TH_CUDA_STORAGE_IMPLEMENT_COPYTO
73 
74 #endif