1 /**
2 * Copyright 2020 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "ps/ps_cache/gpu/gpu_ps_cache.h"
18 #include "ps/ps_cache/ps_cache_factory.h"
19 #include "backend/kernel_compiler/gpu/cuda_impl/hash_impl.cuh"
20 #include "runtime/device/gpu/gpu_common.h"
21 #include "runtime/device/gpu/cuda_driver.h"
22 #include "runtime/device/gpu/gpu_memory_allocator.h"
23 #include "utils/ms_context.h"
24
25 namespace mindspore {
26 namespace ps {
27 namespace gpu {
28 MS_REG_PS_CACHE(kGPUDevice, GPUPsCache);
InitDevice(uint32_t device_id,const void *)29 bool GPUPsCache::InitDevice(uint32_t device_id, const void *) {
30 bool ret = device::gpu::CudaDriver::SetDevice(UintToInt(device_id));
31 if (!ret) {
32 MS_LOG(ERROR) << "Failed to set device id:" << device_id;
33 return false;
34 }
35 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(cudaStreamCreate(reinterpret_cast<CUstream_st **>(&stream_)),
36 "Cuda create stream failed");
37 return true;
38 }
39
MallocMemory(size_t size)40 void *GPUPsCache::MallocMemory(size_t size) {
41 return device::gpu::GPUMemoryAllocator::GetInstance().AllocTensorMem(size);
42 }
43
RecordEvent()44 bool GPUPsCache::RecordEvent() {
45 event_.reset(new cudaEvent_t());
46 MS_ERROR_IF_NULL_W_RET_VAL(event_, false);
47 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(cudaEventCreate(&(*event_)), "Cuda create event failed");
48 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(cudaEventRecord(*event_, reinterpret_cast<cudaStream_t>(stream_)),
49 "Cuda record event failed");
50 return true;
51 }
52
SynchronizeEvent()53 bool GPUPsCache::SynchronizeEvent() {
54 MS_ERROR_IF_NULL_W_RET_VAL(event_, false);
55 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(cudaEventSynchronize(*event_), "Cuda sync event failed");
56 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(cudaEventDestroy(*event_), "Cuda destroy event failed");
57 return true;
58 }
59
SynchronizeStream()60 bool GPUPsCache::SynchronizeStream() {
61 MS_ERROR_IF_NULL_W_RET_VAL(stream_, false);
62 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream_)),
63 "Cuda sync stream failed");
64 return true;
65 }
66
CopyHostMemToDevice(void * dst,const void * src,size_t size)67 bool GPUPsCache::CopyHostMemToDevice(void *dst, const void *src, size_t size) {
68 MS_ERROR_IF_NULL(dst);
69 MS_ERROR_IF_NULL(src);
70 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(
71 cudaMemcpyAsync(dst, src, size, cudaMemcpyHostToDevice, reinterpret_cast<cudaStream_t>(stream_)),
72 "Cuda memcpy failed");
73 return true;
74 }
75
CopyDeviceMemToHost(void * dst,const void * src,size_t size)76 bool GPUPsCache::CopyDeviceMemToHost(void *dst, const void *src, size_t size) {
77 MS_ERROR_IF_NULL(dst);
78 MS_ERROR_IF_NULL(src);
79 CHECK_CUDA_RET_WITH_RETURN_ERROR_NOTRACE(
80 cudaMemcpyAsync(dst, src, size, cudaMemcpyDeviceToHost, reinterpret_cast<cudaStream_t>(stream_)),
81 "Cuda memcpy failed");
82 return true;
83 }
84
HashSwapOut(void * hash_table_addr,void * swap_out_value_addr,void * swap_out_index_addr,size_t,size_t embedding_size,size_t swap_out_size)85 bool GPUPsCache::HashSwapOut(void *hash_table_addr, void *swap_out_value_addr, void *swap_out_index_addr, size_t,
86 size_t embedding_size, size_t swap_out_size) {
87 MS_ERROR_IF_NULL(hash_table_addr);
88 MS_ERROR_IF_NULL(swap_out_value_addr);
89 MS_ERROR_IF_NULL(swap_out_index_addr);
90 DoHashSwapOut(reinterpret_cast<float *>(hash_table_addr), reinterpret_cast<float *>(swap_out_value_addr),
91 reinterpret_cast<int *>(swap_out_index_addr), swap_out_size, embedding_size,
92 reinterpret_cast<cudaStream_t>(stream_));
93 return true;
94 }
95
HashSwapIn(void * hash_table_addr,void * swap_in_value_addr,void * swap_in_index_addr,size_t,size_t embedding_size,size_t swap_in_size)96 bool GPUPsCache::HashSwapIn(void *hash_table_addr, void *swap_in_value_addr, void *swap_in_index_addr, size_t,
97 size_t embedding_size, size_t swap_in_size) {
98 MS_ERROR_IF_NULL(hash_table_addr);
99 MS_ERROR_IF_NULL(swap_in_value_addr);
100 MS_ERROR_IF_NULL(swap_in_index_addr);
101 DoHashSwapIn(reinterpret_cast<float *>(hash_table_addr), reinterpret_cast<float *>(swap_in_value_addr),
102 reinterpret_cast<int *>(swap_in_index_addr), swap_in_size, embedding_size,
103 reinterpret_cast<cudaStream_t>(stream_));
104 return true;
105 }
106 } // namespace gpu
107 } // namespace ps
108 } // namespace mindspore
109