1 /** 2 * Copyright 2019 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 #ifndef MINDSPORE_CCSRC_RUNTIME_DEVICE_GPU_GPU_BUFFER_MGR_H_ 18 #define MINDSPORE_CCSRC_RUNTIME_DEVICE_GPU_GPU_BUFFER_MGR_H_ 19 20 #include <unistd.h> 21 #include <cstring> 22 #include <iostream> 23 #include <functional> 24 #include <map> 25 #include <vector> 26 #include <string> 27 #include <memory> 28 #include "runtime/device/gpu/blocking_queue.h" 29 30 #define EXPORT __attribute__((visibility("default"))) 31 32 namespace mindspore { 33 namespace device { 34 static const unsigned int MAX_WAIT_TIME_IN_SEC = 60; 35 36 class Semaphore { 37 public: count_(count)38 explicit Semaphore(int count = 0) : count_(count) {} 39 Signal()40 inline void Signal() { 41 std::unique_lock<std::mutex> lock(mutex_); 42 ++count_; 43 cv_.notify_one(); 44 } 45 Wait()46 inline bool Wait() { 47 std::unique_lock<std::mutex> lock(mutex_); 48 while (count_ == 0) { 49 if (cv_.wait_for(lock, std::chrono::seconds(MAX_WAIT_TIME_IN_SEC)) == std::cv_status::timeout) { 50 return false; 51 } 52 } 53 --count_; 54 return true; 55 } 56 57 private: 58 std::mutex mutex_; 59 std::condition_variable cv_; 60 int count_; 61 }; 62 63 class HandleMgr { 64 public: 65 static const unsigned int MAX_HANDLE_NUM = 32; 66 static const unsigned int INVALID_HANDLE = 0xffffffffUL; 67 68 unsigned int AllocHandle(); 69 void FreeHandle(unsigned int); 70 71 private: 72 bool handle_list_[MAX_HANDLE_NUM]; 73 }; 74 75 class GpuBufferMgr { 76 public: GpuBufferMgr()77 EXPORT GpuBufferMgr() : cur_dev_id_(0), init_(false), closed_(false), open_by_dataset_(0) {} 78 79 EXPORT virtual ~GpuBufferMgr() = default; 80 81 EXPORT static GpuBufferMgr &GetInstance() noexcept; 82 83 EXPORT BlockQueueStatus_T Create(unsigned int device_id, const std::string &channel_name, void *addr, 84 const std::vector<size_t> &shape, const size_t &capacity); 85 86 // call for Push thread 87 EXPORT unsigned int Open(unsigned int device_id, const std::string &channel_name, const std::vector<size_t> &shape, 88 std::function<void(void *, int32_t)> func); 89 90 // call for Front/Pop thread 91 EXPORT unsigned int Open(unsigned int device_id, const std::string &channel_name, const std::vector<size_t> &shape); 92 93 EXPORT BlockQueueStatus_T Push(unsigned int handle, const std::vector<DataItemGpu> &data, 94 unsigned int timeout_in_sec); 95 EXPORT BlockQueueStatus_T Front(unsigned int handle, void **addr, size_t *len); 96 EXPORT BlockQueueStatus_T Pop(unsigned int handle); 97 98 EXPORT void set_device_id(int device_id); 99 100 EXPORT void Close(unsigned int handle) noexcept; 101 102 EXPORT bool IsInit() const; 103 104 EXPORT bool IsClosed() const; 105 106 EXPORT bool Destroy(); 107 108 // call for Release GPU Resources 109 EXPORT bool CloseNotify(); 110 111 // call for dataset send thread 112 EXPORT void CloseConfirm(); 113 114 EXPORT size_t Size(unsigned int handle); 115 116 EXPORT size_t Size(unsigned int device_id, const std::string &channel_name); 117 118 EXPORT size_t Capacity(unsigned int handle); 119 120 EXPORT size_t Capacity(unsigned int device_id, const std::string &channel_name); 121 122 private: 123 void set_device() const; 124 125 int cur_dev_id_; 126 bool init_; 127 bool closed_; 128 std::mutex mutex_; 129 std::mutex close_mutex_; 130 // how many queues opened by dataset 131 int open_by_dataset_; 132 Semaphore sema; 133 134 HandleMgr handle_mgr_; 135 136 std::map<unsigned int, std::shared_ptr<BlockingQueue>> handle_queue_map_; 137 std::map<std::string, std::shared_ptr<BlockingQueue>> name_queue_map_; 138 139 inline bool isCreated(unsigned int device_id, const std::string &channel_name); 140 141 GpuBufferMgr(const GpuBufferMgr &) = delete; 142 GpuBufferMgr &operator=(const GpuBufferMgr &) = delete; 143 }; 144 } // namespace device 145 } // namespace mindspore 146 147 #endif // MINDSPORE_CCSRC_RUNTIME_DEVICE_GPU_GPU_BUFFER_MGR_H_ 148