diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/data/dataset_init_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/data/dataset_init_kernel.h index f8cc9b19ea..82b6463e88 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/data/dataset_init_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/data/dataset_init_kernel.h @@ -49,7 +49,7 @@ class DatasetInitKernel : public GpuKernel { std::vector workspace_size_list_; // The capacity of buffer Q. - size_t buffer_q_capacity_{2}; + size_t buffer_q_capacity_{1}; }; MS_REG_GPU_KERNEL(InitDataSetQueue, DatasetInitKernel) diff --git a/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc b/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc index 5cb97ac1e7..8aae3eeda0 100644 --- a/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc +++ b/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc @@ -22,7 +22,15 @@ namespace mindspore { namespace device { GpuQueue::GpuQueue(void *addr, const std::vector &shape, const size_t &capacity) - : buffer_(addr), head_(0), tail_(0), shape_(shape), len_(0), capacity_(capacity), stream_(0), node_info_(nullptr) { + : buffer_(addr), + head_(0), + tail_(0), + shape_(shape), + len_(0), + size_(0), + capacity_(capacity), + stream_(0), + node_info_(nullptr) { CHECK_CUDA_RET_WITH_ERROR(cudaStreamCreate(&stream_), "Cuda Create Stream Failed"); node_info_ = std::make_unique(capacity); for (auto item : shape) { @@ -52,6 +60,7 @@ BlockQueueStatus_T GpuQueue::Push(const std::vector &data) { CHECK_CUDA_RET_WITH_ERROR(cudaEventCreate(&(*(node_info_[tail_].event_))), "Cuda Create Event Failed"); node_info_[tail_].data_ = data; tail_ = (tail_ + 1) % (capacity_); + ++size_; return SUCCESS; } @@ -69,6 +78,7 @@ BlockQueueStatus_T GpuQueue::Front(void **addr, size_t *len) const { BlockQueueStatus_T GpuQueue::Pop() { head_ = (head_ + 1) % (capacity_); + --size_; return SUCCESS; } diff --git a/mindspore/ccsrc/runtime/device/gpu/blocking_queue.h b/mindspore/ccsrc/runtime/device/gpu/blocking_queue.h index f5a33d36ca..94dee538bf 100644 --- a/mindspore/ccsrc/runtime/device/gpu/blocking_queue.h +++ b/mindspore/ccsrc/runtime/device/gpu/blocking_queue.h @@ -44,13 +44,15 @@ class GpuQueue { void RegisterRelease(const std::function &func) { host_release_ = func; } - inline bool IsEmpty() const { return head_ == tail_; } - inline bool IsFull() const { return head_ == ((tail_ + 1) % (capacity_)); } + inline bool IsEmpty() const { return size_ == 0; } + inline bool IsFull() const { return size_ == capacity_; } BlockQueueStatus_T Push(const std::vector &data); BlockQueueStatus_T Front(void **ptr, size_t *len) const; BlockQueueStatus_T Pop(); bool Destroy(); + size_t Size() { return size_; } + size_t Capacity() { return capacity_; } private: struct NodeInfo { @@ -63,6 +65,7 @@ class GpuQueue { size_t tail_; std::vector shape_; size_t len_; + size_t size_; size_t capacity_; cudaStream_t stream_; std::unique_ptr node_info_; @@ -83,6 +86,8 @@ class BlockingQueue { BlockQueueStatus_T Front(void **ptr, size_t *len); BlockQueueStatus_T Pop(); bool Destroy(); + size_t Size() { return queue_->Size(); } + size_t Capacity() { return queue_->Capacity(); } private: std::mutex mutex_; diff --git a/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.cc b/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.cc index 7ca7878d56..ba3578a0b0 100644 --- a/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.cc +++ b/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.cc @@ -187,5 +187,39 @@ bool GpuBufferMgr::CloseNotify() { } void GpuBufferMgr::CloseConfirm() { sema.Signal(); } + +size_t GpuBufferMgr::Size(unsigned int handle) { + if (handle == HandleMgr::INVALID_HANDLE) { + MS_LOG(ERROR) << "handle is invalid"; + return 0; + } + return handle_queue_map_.at(handle)->Size(); +} + +size_t GpuBufferMgr::Size(unsigned int device_id, const std::string &channel_name) { + std::string name = std::to_string(device_id) + std::string("_") + channel_name; + if (!name_queue_map_.count(name)) { + MS_LOG(ERROR) << "Queue not exist " << name; + return 0; + } + return name_queue_map_.at(name)->Size(); +} + +size_t GpuBufferMgr::Capacity(unsigned int handle) { + if (handle == HandleMgr::INVALID_HANDLE) { + MS_LOG(ERROR) << "handle is invalid"; + return 0; + } + return handle_queue_map_.at(handle)->Capacity(); +} + +size_t GpuBufferMgr::Capacity(unsigned int device_id, const std::string &channel_name) { + std::string name = std::to_string(device_id) + std::string("_") + channel_name; + if (!name_queue_map_.count(name)) { + MS_LOG(ERROR) << "Queue not exist " << name; + return 0; + } + return name_queue_map_.at(name)->Capacity(); +} } // namespace device } // namespace mindspore diff --git a/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.h b/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.h index 610836bcf6..de25b948c6 100644 --- a/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.h +++ b/mindspore/ccsrc/runtime/device/gpu/gpu_buffer_mgr.h @@ -111,6 +111,14 @@ class GpuBufferMgr { // call for dataset send thread EXPORT void CloseConfirm(); + EXPORT size_t Size(unsigned int handle); + + EXPORT size_t Size(unsigned int device_id, const std::string &channel_name); + + EXPORT size_t Capacity(unsigned int handle); + + EXPORT size_t Capacity(unsigned int device_id, const std::string &channel_name); + private: void set_device() const;