diff --git a/mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.cc b/mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.cc index 98e7f7eb1d7..389be4f90d8 100644 --- a/mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.cc +++ b/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 diff --git a/mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.h b/mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.h index 45a28eb22f6..0a24ef10149 100644 --- a/mindspore/ccsrc/plugin/device/cpu/hal/device/cpu_tensor_array.h +++ b/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 &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; diff --git a/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.cc b/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.cc index 0b9248de009..c588a38f269 100644 --- a/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.cc +++ b/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(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 diff --git a/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.h b/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.h index 0e679737c82..8cc68641625 100644 --- a/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.h +++ b/mindspore/ccsrc/plugin/device/gpu/hal/device/gpu_tensor_array.h @@ -22,6 +22,7 @@ #include #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 &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> &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; +using GPUTensorsQueuePtr = std::shared_ptr; } // namespace gpu } // namespace device } // namespace mindspore diff --git a/mindspore/ccsrc/plugin/device/gpu/kernel/rl/tensor_array_write_kernel.cc b/mindspore/ccsrc/plugin/device/gpu/kernel/rl/tensor_array_write_kernel.cc index bb537d5f3e0..2a91829a664 100644 --- a/mindspore/ccsrc/plugin/device/gpu/kernel/rl/tensor_array_write_kernel.cc +++ b/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)); } diff --git a/mindspore/ccsrc/runtime/device/CMakeLists.txt b/mindspore/ccsrc/runtime/device/CMakeLists.txt index 3c1428af134..c19e38e7892 100644 --- a/mindspore/ccsrc/runtime/device/CMakeLists.txt +++ b/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") diff --git a/mindspore/ccsrc/runtime/device/tensor_array.cc b/mindspore/ccsrc/runtime/device/tensor_array.cc index a2295e028ae..be3e98d7432 100644 --- a/mindspore/ccsrc/runtime/device/tensor_array.cc +++ b/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(); - 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(addr->addr)); + FreeMemory(static_cast(addr->addr)); } } } diff --git a/mindspore/ccsrc/runtime/device/tensor_array.h b/mindspore/ccsrc/runtime/device/tensor_array.h index 2939f6d1a95..e4ae9d45482 100644 --- a/mindspore/ccsrc/runtime/device/tensor_array.h +++ b/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 diff --git a/mindspore/ccsrc/runtime/device/tensor_array_manager.h b/mindspore/ccsrc/runtime/device/tensor_array_manager.h index 7e345aef1ee..a552b441ad5 100644 --- a/mindspore/ccsrc/runtime/device/tensor_array_manager.h +++ b/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 #include #include @@ -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 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 tensorsqueue_map_; + // Used as an unique handle number for each TensorsQueue. + std::atomic tensors_queue_handle_count{0}; +}; } // namespace device } // namespace mindspore diff --git a/mindspore/ccsrc/runtime/device/tensors_queue.cc b/mindspore/ccsrc/runtime/device/tensors_queue.cc new file mode 100644 index 00000000000..34147555d15 --- /dev/null +++ b/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 +#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 element_size_list; + for (auto shape : shapes_) { + int64_t item_size = std::accumulate(shape.begin(), shape.end(), GetTypeByte(dtype_), std::multiplies()); + 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(); + 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(addr->addr)); + } + } + front_ = (front_ + 1) % size_; + } + MS_LOG(DEBUG) << "Free the TensorsQueue's memory for " << name_; +} +} // namespace device +} // namespace mindspore diff --git a/mindspore/ccsrc/runtime/device/tensors_queue.h b/mindspore/ccsrc/runtime/device/tensors_queue.h new file mode 100644 index 00000000000..984d651cfb8 --- /dev/null +++ b/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 +#include +#include +#include +#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> &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> 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 tensors_q; + size_t front_ = 0; + size_t rear_ = 0; +}; +using TensorsQueuePtr = std::shared_ptr; +} // namespace device +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_RUNTIME_DEVICE_TENSORS_QUEUE_H_