From e2b346d5af6801d68e790d396e7a6b7575bcaa32 Mon Sep 17 00:00:00 2001
From: anthonyaje <anthony2@huawei.com>
Date: Wed, 19 Aug 2020 16:48:31 -0400
Subject: [PATCH] Add Size() and Capacity() in gpu queue.

---
 .../runtime/device/gpu/blocking_queue.cc      |  2 ++
 .../ccsrc/runtime/device/gpu/blocking_queue.h |  9 +++--
 .../runtime/device/gpu/gpu_buffer_mgr.cc      | 34 +++++++++++++++++++
 .../ccsrc/runtime/device/gpu/gpu_buffer_mgr.h |  8 +++++
 4 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc b/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc
index 5cb97ac1e7..89028a7631 100644
--- a/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc
+++ b/mindspore/ccsrc/runtime/device/gpu/blocking_queue.cc
@@ -52,6 +52,7 @@ BlockQueueStatus_T GpuQueue::Push(const std::vector<DataItemGpu> &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 +70,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<void(void *)> &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<DataItemGpu> &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<size_t> shape_;
   size_t len_;
+  size_t size_;
   size_t capacity_;
   cudaStream_t stream_;
   std::unique_ptr<NodeInfo[]> 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;