!32416 RL add TensorsQueue

Merge pull request !32416 from VectorSL/grads_queue_for_merge

mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.cc

@@ -23,11 +23,11 @@
 namespace mindspore {
 namespace device {
 namespace cpu {
-void *CPUTensorArray::CreateMemory(const size_t size) { return CPUMemoryPool::GetInstance().AllocTensorMem(size); }
+void *CPUTensorArray::AllocateMemory(const size_t size) { return CPUMemoryPool::GetInstance().AllocTensorMem(size); }
 
 void CPUTensorArray::ClearMemory(void *addr, const size_t size) { (void)memset_s(addr, size, 0, size); }
 
-void CPUTensorArray::ReleaseMemory(const DeviceMemPtr addr) { CPUMemoryPool::GetInstance().FreeTensorMem(addr); }
+void CPUTensorArray::FreeMemory(const DeviceMemPtr addr) { CPUMemoryPool::GetInstance().FreeTensorMem(addr); }
 }  // namespace cpu
 }  // namespace device
 }  // namespace mindspore

mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.h

@@ -30,8 +30,8 @@ class CPUTensorArray : public TensorArray {
   CPUTensorArray(const string &name, const TypePtr &dtype, const std::vector<size_t> &shapes)
       : TensorArray(name, dtype, shapes) {}
   ~CPUTensorArray() override = default;
-  void ReleaseMemory(const DeviceMemPtr addr) override;
-  void *CreateMemory(const size_t size) override;
+  void FreeMemory(const DeviceMemPtr addr) override;
+  void *AllocateMemory(const size_t size) override;
   void ClearMemory(void *addr, const size_t size) override;
 };
 using CPUTensorArrayPtr = std::shared_ptr<CPUTensorArray>;

mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.cc

@@ -25,8 +25,8 @@
 namespace mindspore {
 namespace device {
 namespace gpu {
-// ReleaseMemory() used in Free() in TensorArray.
-void GPUTensorArray::ReleaseMemory(const DeviceMemPtr addr) {
+// FreeMemory() used in Free() in TensorArray.
+void GPUTensorArray::FreeMemory(const DeviceMemPtr addr) {
   device::gpu::GPUMemoryAllocator::GetInstance().FreeTensorMem(addr);
 }
 
@@ -34,7 +34,30 @@ void GPUTensorArray::ClearMemory(void *addr, const size_t size) {
   CHECK_CUDA_RET_WITH_EXCEPT_NOTRACE(cudaMemsetAsync(addr, 0, size), "failed to set cuda memory with zeros.");
 }
 
-void *GPUTensorArray::CreateMemory(const size_t size) {
+void *GPUTensorArray::AllocateMemory(const size_t size) {
+  return device::gpu::GPUMemoryAllocator::GetInstance().AllocTensorMem(size);
+}
+
+void GPUTensorsQueue::CopyTensor(const mindspore::kernel::AddressPtr &dst, const mindspore::kernel::AddressPtr &src,
+                                 void *stream) {
+  if (dst->size != src->size) {
+    MS_LOG(EXCEPTION) << "For TensorsQueue Put/Get function, each tensor in element should have the same size, but get "
+                      << src->size << ", not equal to dst " << dst->size;
+  }
+  CHECK_CUDA_RET_WITH_EXCEPT_NOTRACE(
+    cudaMemcpyAsync(dst->addr, src->addr, dst->size, cudaMemcpyDeviceToDevice, reinterpret_cast<cudaStream_t>(stream)),
+    "Copy tensor failed");
+}
+// FreeMemory() used in Free() in TensorArray.
+void GPUTensorsQueue::FreeMemory(const DeviceMemPtr addr) {
+  device::gpu::GPUMemoryAllocator::GetInstance().FreeTensorMem(addr);
+}
+
+void GPUTensorsQueue::ClearMemory(void *addr, const size_t size) {
+  CHECK_CUDA_RET_WITH_EXCEPT_NOTRACE(cudaMemsetAsync(addr, 0, size), "failed to set cuda memory with zeros.");
+}
+
+void *GPUTensorsQueue::AllocateMemory(const size_t size) {
   return device::gpu::GPUMemoryAllocator::GetInstance().AllocTensorMem(size);
 }
 }  // namespace gpu

mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.h

@@ -22,6 +22,7 @@
 #include <memory>
 #include "plugin/device/gpu/hal/device/gpu_memory_allocator.h"
 #include "runtime/device/tensor_array.h"
+#include "runtime/device/tensors_queue.h"
 
 namespace mindspore {
 namespace device {
@@ -31,12 +32,25 @@ class GPUTensorArray : public TensorArray {
   GPUTensorArray(const string &name, const TypePtr &dtype, const std::vector<size_t> &shapes)
       : TensorArray(name, dtype, shapes) {}
   ~GPUTensorArray() override = default;
-  void ReleaseMemory(const DeviceMemPtr addr) override;
-  void *CreateMemory(const size_t size) override;
+  void FreeMemory(const DeviceMemPtr addr) override;
+  void *AllocateMemory(const size_t size) override;
   void ClearMemory(void *addr, const size_t size) override;
 };
-using GPUTensorArray = GPUTensorArray;
+
+class GPUTensorsQueue : public TensorsQueue {
+ public:
+  GPUTensorsQueue(const string &name, const TypePtr &dtype, const int64_t size, const int64_t elements_num,
+                  const std::vector<std::vector<int64_t>> &shapes)
+      : TensorsQueue(name, dtype, size, elements_num, shapes) {}
+  ~GPUTensorsQueue() override = default;
+  void FreeMemory(const DeviceMemPtr addr) override;
+  void *AllocateMemory(const size_t size) override;
+  void ClearMemory(void *addr, const size_t size) override;
+  void CopyTensor(const mindspore::kernel::AddressPtr &dst, const mindspore::kernel::AddressPtr &src,
+                  void *stream) override;
+};
 using GPUTensorArrayPtr = std::shared_ptr<GPUTensorArray>;
+using GPUTensorsQueuePtr = std::shared_ptr<GPUTensorsQueue>;
 }  // namespace gpu
 }  // namespace device
 }  // namespace mindspore

mindspore/ccsrc/plugin/device/gpu/kernel/rl/tensor_array_write_kernel.cc

@@ -43,7 +43,7 @@ bool TensorArrayWriteKernelMod::Init(const CNodePtr &kernel_node) {
 void TensorArrayWriteKernelMod::InitSizeLists() {
   input_size_list_.push_back(sizeof(int64_t));
   input_size_list_.push_back(sizeof(int64_t));
-  input_size_list_.push_back(sizeof(value_size_));
+  input_size_list_.push_back(value_size_);
   output_size_list_.push_back(sizeof(int64_t));
 }
 

mindspore/ccsrc/runtime/device/CMakeLists.txt

@@ -2,7 +2,7 @@ file(GLOB_RECURSE DEVICE_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "common/*
     "kernel_info.cc" "executor/dynamic_kernel.cc" "executor/executor_callback.cc" "kernel_runtime.cc"
     "memory_manager.cc" "kernel_runtime_manager.cc" "convert_tensor_utils.cc" "memory_scheduler.cc"
     "memory_offload_strategy.cc" "bucket.cc" "launch_kernel.cc" "launch_mul.cc" "tensor_array.cc"
-    "ms_device_shape_transfer.cc" "context_extends.cc" "stream_synchronizer.cc"
+    "ms_device_shape_transfer.cc" "context_extends.cc" "stream_synchronizer.cc" "tensors_queue.cc"
 )
 
 if("${ENABLE_HIDDEN}" STREQUAL "OFF")

mindspore/ccsrc/runtime/device/tensor_array.cc

@@ -69,7 +69,7 @@ bool TensorArray::Write(const int64_t index, const mindspore::kernel::AddressPtr
     size_t create_size = (LongToSize(index) > tensors_.size()) ? (LongToSize(index) - tensors_.size()) : 0;
     for (size_t i = 0; i < create_size; i++) {
       kernel::AddressPtr create_dev = std::make_shared<kernel::Address>();
-      create_dev->addr = CreateMemory(dev_value->size);
+      create_dev->addr = AllocateMemory(dev_value->size);
       create_dev->size = dev_value->size;
       tensors_.push_back(create_dev);
     }
@@ -99,7 +99,7 @@ void TensorArray::Free() {
   MS_LOG(DEBUG) << "Free device memory for " << name_;
   for (const auto &addr : tensors_) {
     if (addr != nullptr) {
-      ReleaseMemory(static_cast<DeviceMemPtr>(addr->addr));
+      FreeMemory(static_cast<DeviceMemPtr>(addr->addr));
     }
   }
 }

mindspore/ccsrc/runtime/device/tensor_array.h

@@ -51,8 +51,8 @@ class TensorArray {
   // These three func should by implied for different device due to the difference in memory usage.
   // Create/Release Memory is used for malloc/free a device memory, used in function Write().
   // ClearMemory is used to reset the input addr with zeros, used in function Free().
-  virtual void ReleaseMemory(const DeviceMemPtr addr) = 0;
-  virtual void *CreateMemory(const size_t size) = 0;
+  virtual void FreeMemory(const DeviceMemPtr addr) = 0;
+  virtual void *AllocateMemory(const size_t size) = 0;
   virtual void ClearMemory(void *addr, const size_t size) = 0;
 
   // Clear() will only set the valid size of TensorArray to zero. The memory in TensorArray is still

mindspore/ccsrc/runtime/device/tensor_array_manager.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2021 Huawei Technologies Co., Ltd
+ * Copyright 2021-2022 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,6 @@
 
 #ifndef MINDSPORE_CCSRC_RUNTIME_DEVICE_TENSOR_ARRAY_MANAGER_H_
 #define MINDSPORE_CCSRC_RUNTIME_DEVICE_TENSOR_ARRAY_MANAGER_H_
-
 #include <vector>
 #include <string>
 #include <atomic>
@@ -26,6 +25,7 @@
 #include "backend/common/session/anf_runtime_algorithm.h"
 #include "include/common/utils/anfalgo.h"
 #include "runtime/device/tensor_array.h"
+#include "runtime/device/tensors_queue.h"
 
 namespace mindspore {
 namespace device {
@@ -78,6 +78,56 @@ class TensorArrayMgr {
   // Used as an unique handle number for each TensorArray.
   std::atomic<int64_t> tensor_array_handle_count{0};
 };
+
+class TensorsQueueMgr {
+ public:
+  // TensorsQueueMgr is used to manage the TensorsQueues.
+  TensorsQueueMgr() {}
+  ~TensorsQueueMgr() = default;
+
+  static TensorsQueueMgr &GetInstance() noexcept {
+    static TensorsQueueMgr instance;
+    return instance;
+  }
+
+  TensorsQueueMgr(const TensorsQueueMgr &) = delete;
+  TensorsQueueMgr(const TensorsQueueMgr &&) = delete;
+
+  void AddTensorsQueue(const int64_t handle, const TensorsQueuePtr &tq) {
+    MS_LOG(DEBUG) << "Add a TensorsQueue to map, handle is " << handle;
+    tensorsqueue_map_.emplace(std::make_pair(handle, tq));
+    // Increase handle count when added a TensorsQueue.
+    tensors_queue_handle_count += 1;
+  }
+
+  TensorsQueuePtr GetTensorsQueue(const int64_t handle) {
+    if (!tensorsqueue_map_.count(handle)) {
+      MS_LOG(EXCEPTION) << "Error handle [" << handle << "] to get TensorsQueue";
+    } else {
+      MS_LOG(DEBUG) << "Get TensorsQueue succeed, handle is " << handle;
+      return tensorsqueue_map_[handle];
+    }
+  }
+
+  bool EraseTensorsQueue(const int64_t handle) {
+    if (tensorsqueue_map_.count(handle)) {
+      MS_LOG(DEBUG) << "Erase TensorsQueue from map, handle number is " << handle;
+      tensorsqueue_map_.erase(handle);
+      return true;
+    } else {
+      MS_LOG(ERROR) << "Erase TensorsQueue failed, no such handle " << handle;
+      return false;
+    }
+  }
+
+  int64_t GetHandleCount() const { return tensors_queue_handle_count; }
+
+ private:
+  // Store the TensorsQueues in a map, as pair(handle, TensorsQueuePtr).
+  std::map<const int64_t, TensorsQueuePtr> tensorsqueue_map_;
+  // Used as an unique handle number for each TensorsQueue.
+  std::atomic<int64_t> tensors_queue_handle_count{0};
+};
 }  // namespace device
 }  // namespace mindspore
 

mindspore/ccsrc/runtime/device/tensors_queue.cc

@@ -0,0 +1,150 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <functional>
+#include "runtime/device/tensors_queue.h"
+
+namespace mindspore {
+namespace device {
+void TensorsQueue::CreateTensorsQueue() {
+  // Store one element tensors' size.
+  // The whole TensorsQueue is like: [[tensor1, tensor2], [tensor3, tensor4]].
+  // One element means [tensor1, tensor2].
+  std::vector<int64_t> element_size_list;
+  for (auto shape : shapes_) {
+    int64_t item_size = std::accumulate(shape.begin(), shape.end(), GetTypeByte(dtype_), std::multiplies<int64_t>());
+    element_size_list.push_back(item_size);
+  }
+  // Create the elements in TensorsQueue when construct.
+  for (int64_t i = 0; i < size_; i++) {
+    mindspore::kernel::AddressPtrList element_addrs;
+    for (auto element_size : element_size_list) {
+      kernel::AddressPtr create_dev = std::make_shared<kernel::Address>();
+      create_dev->addr = AllocateMemory(element_size);
+      create_dev->size = element_size;
+      element_addrs.push_back(create_dev);
+      MS_LOG(DEBUG) << "Create  " << element_size << "bytes for " << name_;
+    }
+    tensors_q.push_back(element_addrs);
+  }
+  MS_LOG(DEBUG) << "Create a TensorsQueue: " << name_ << ", Q size is " << size_ << ", elements num is "
+                << elements_num_;
+}
+
+void TensorsQueue::CopyTensor(const mindspore::kernel::AddressPtr &dst, const mindspore::kernel::AddressPtr &src) {
+  MS_LOG(EXCEPTION) << "This should be overridden by subclass !";
+}
+void TensorsQueue::CopyTensor(const mindspore::kernel::AddressPtr &dst, const mindspore::kernel::AddressPtr &src,
+                              void *stream) {
+  MS_LOG(EXCEPTION) << "This should be overridden by subclass !";
+}
+
+size_t TensorsQueue::AvailableSize() { return (rear_ > front_) ? (rear_ - front_) : (size_ - front_ + rear_); }
+bool TensorsQueue::IsFull() { return (rear_ + 1) % size_ == front_; }
+bool TensorsQueue::IsEmpty() { return front_ == rear_; }
+
+bool TensorsQueue::Put(const mindspore::kernel::AddressPtrList &dev_addr) {
+  // When the tensor_q is full, put will failed.
+  if (IsFull()) {
+    MS_LOG(WARNING) << "The " << name_ << " is full, total size is " << size_;
+    return false;
+  }
+  // Get the element in position rear_ and change the value by input, the we increase the rear_.
+  // We can get a effect like a circle queue and reuse the addrs.
+  mindspore::kernel::AddressPtrList element = tensors_q[rear_];
+  for (int64_t i = 0; i < elements_num_; i++) {
+    CopyTensor(element[i], dev_addr[i + 1]);
+  }
+  rear_ = (rear_ + 1) % size_;
+  MS_LOG(DEBUG) << "Put an element into  " << name_ << ", now the avliable q size is [" << AvailableSize() << "/"
+                << size_ << "]";
+  return true;
+}
+
+bool TensorsQueue::Put(const mindspore::kernel::AddressPtrList &dev_addr, void *stream) {
+  if (IsFull()) {
+    MS_LOG(WARNING) << "The " << name_ << " is full, total size is " << size_;
+    return false;
+  }
+  mindspore::kernel::AddressPtrList element = tensors_q[rear_];
+  for (int64_t i = 0; i < elements_num_; i++) {
+    CopyTensor(element[i], dev_addr[i + 1], stream);
+  }
+  rear_ = (rear_ + 1) % size_;
+  MS_LOG(DEBUG) << "Put an element into  " << name_ << ", now the avliable q size is [" << AvailableSize() << "/"
+                << size_ << "]";
+  return true;
+}
+
+bool TensorsQueue::Get(const mindspore::kernel::AddressPtrList &dev_addr, const bool &pop_after_get, void *stream) {
+  // Get a tensor addrs list from the queue.
+  // If pop_after_get is true, we will pop the addrs from tensors_q.
+  if (IsEmpty()) {
+    MS_LOG(WARNING) << "The TensorsQueue " << name_ << " is empty";
+    return false;
+  }
+  mindspore::kernel::AddressPtrList element = tensors_q[front_];
+  for (int64_t i = 0; i < elements_num_; i++) {
+    CopyTensor(dev_addr[i], element[i], stream);
+  }
+  if (pop_after_get) {
+    front_ = (front_ + 1) % size_;
+  }
+  MS_LOG(DEBUG) << "Get an element from  " << name_ << ", pop_after_get is " << pop_after_get
+                << ", now the avliable q size is[" << AvailableSize() << " / " << size_ << "] ";
+  return true;
+}
+
+bool TensorsQueue::Get(const mindspore::kernel::AddressPtrList &dev_addr, const bool &pop_after_get) {
+  if (IsEmpty()) {
+    MS_LOG(WARNING) << "The TensorsQueue " << name_ << " is empty";
+    return false;
+  }
+  mindspore::kernel::AddressPtrList element = tensors_q.front();
+  for (int64_t i = 0; i < elements_num_; i++) {
+    CopyTensor(dev_addr[i], element[i]);
+  }
+  if (pop_after_get) {
+    front_ = (front_ + 1) % size_;
+  }
+  MS_LOG(DEBUG) << "Get an element from  " << name_ << ", pop_after_get is " << pop_after_get
+                << ", now the avliable q size is[" << AvailableSize() << " / " << size_ << "] ";
+  return true;
+}
+
+void TensorsQueue::Clear() {
+  // Clear the tensors_q and return the element addr back to tensors_store.
+  if (IsEmpty()) {
+    MS_LOG(WARNING) << "The TensorsQueue " << name_ << " is already empty when execute Clear.";
+  }
+  rear_ = 0;
+  front_ = 0;
+  MS_LOG(DEBUG) << "Clear the elements for " << name_;
+}
+
+void TensorsQueue::Free() {
+  while (!IsEmpty()) {
+    auto element = tensors_q[front_];
+    for (const auto &addr : element) {
+      if (addr != nullptr) {
+        FreeMemory(static_cast<DeviceMemPtr>(addr->addr));
+      }
+    }
+    front_ = (front_ + 1) % size_;
+  }
+  MS_LOG(DEBUG) << "Free the TensorsQueue's memory for " << name_;
+}
+}  // namespace device
+}  // namespace mindspore

mindspore/ccsrc/runtime/device/tensors_queue.h

@@ -0,0 +1,80 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_CCSRC_RUNTIME_DEVICE_TENSORS_QUEUE_H_
+#define MINDSPORE_CCSRC_RUNTIME_DEVICE_TENSORS_QUEUE_H_
+
+#include <vector>
+#include <queue>
+#include <string>
+#include <memory>
+#include "backend/common/session/kernel_graph.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "kernel/kernel.h"
+
+namespace mindspore {
+namespace device {
+class TensorsQueue {
+ public:
+  // Base TensorsQueue. Constructed by name, dtype, size, elements_num and shapes.
+  TensorsQueue(const string &name, const TypePtr &dtype, const int64_t size, const int64_t elements_num,
+               const std::vector<std::vector<int64_t>> &shapes)
+      : name_(name), dtype_(dtype), shapes_(shapes), size_(size), elements_num_(elements_num) {}
+  virtual ~TensorsQueue() = default;
+  virtual void CreateTensorsQueue();
+
+  // These three function (FreeMemory, AllocateMemory and ClearMemory) are related with devices.
+  // These should be achieved with different devices.
+  virtual void FreeMemory(const DeviceMemPtr addr) = 0;
+  virtual void *AllocateMemory(const size_t size) = 0;
+  virtual void ClearMemory(void *addr, const size_t size) = 0;
+
+  // When memory operations are involved, we need to determine whether to use streams according to the device.
+  virtual bool Put(const mindspore::kernel::AddressPtrList &dev_value);
+  virtual bool Put(const mindspore::kernel::AddressPtrList &dev_value, void *stream);
+  virtual void CopyTensor(const mindspore::kernel::AddressPtr &dst, const mindspore::kernel::AddressPtr &src);
+  virtual void CopyTensor(const mindspore::kernel::AddressPtr &dst, const mindspore::kernel::AddressPtr &src,
+                          void *stream);
+  virtual bool Get(const mindspore::kernel::AddressPtrList &outputs, const bool &pop_after_get);
+  virtual bool Get(const mindspore::kernel::AddressPtrList &outputs, const bool &pop_after_get, void *stream);
+
+  // Common functions for TensorsQueue which are device independent.
+  virtual void Clear();
+  virtual void Free();
+  virtual size_t AvailableSize();
+  virtual bool IsFull();
+  virtual bool IsEmpty();
+
+ protected:
+  std::string name_;
+  TypePtr dtype_;
+  std::vector<std::vector<int64_t>> shapes_;
+  int64_t size_;
+  int64_t elements_num_;
+
+ private:
+  // Using a vector of address list to store the tensors.
+  // Using to cursors to simulate the behavior of circular queue.
+  std::vector<mindspore::kernel::AddressPtrList> tensors_q;
+  size_t front_ = 0;
+  size_t rear_ = 0;
+};
+using TensorsQueuePtr = std::shared_ptr<TensorsQueue>;
+}  // namespace device
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_RUNTIME_DEVICE_TENSORS_QUEUE_H_