1 /** 2 * Copyright 2019-2021 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_DEVICE_TENSOR_H 18 #define MINDSPORE_DEVICE_TENSOR_H 19 20 #include <string> 21 #include <vector> 22 #include <memory> 23 #include <map> 24 #include <utility> 25 #include "ir/dtype.h" 26 #include "ir/device_sync.h" 27 #include "utils/shape_utils.h" 28 29 namespace mindspore { 30 namespace device { 31 class Bucket; 32 namespace cpu { 33 class CPUSimpleMemPlan; 34 class CPUMemoryManager; 35 class CPUKernelRuntime; 36 class CPUDeviceContext; 37 } // namespace cpu 38 namespace ascend { 39 class AscendKernelRuntime; 40 class AscendMemoryManager; 41 #ifndef ENABLE_SECURITY 42 class DataDumper; 43 #endif 44 namespace tasksink { 45 class TaskGenerator; 46 } // namespace tasksink 47 } // namespace ascend 48 namespace gpu { 49 class GPUKernelRuntime; 50 class GPUMemoryManager; 51 class GPUDeviceContext; 52 } // namespace gpu 53 } // namespace device 54 } // namespace mindspore 55 56 namespace mindspore { 57 namespace device { 58 using KernelWithIndex = std::pair<AnfNodePtr, size_t>; 59 enum class DeviceAddressStatus { kInDevice, kInHost, kInDeviceToHost, kInHostToDevice }; 60 enum class DeviceAddressType { kUnknown, kAscend, kCPU, kGPU }; 61 static const std::map<DeviceAddressType, std::string> kDeviceTypeToName = {{DeviceAddressType::kUnknown, "Unknown"}, 62 {DeviceAddressType::kAscend, "Ascend"}, 63 {DeviceAddressType::kCPU, "CPU"}, 64 {DeviceAddressType::kGPU, "GPU"}}; 65 66 class DeviceAddress : public mindspore::DeviceSync { 67 public: DeviceAddress(void * ptr,size_t size)68 explicit DeviceAddress(void *ptr, size_t size) : ptr_(ptr), size_(size) {} DeviceAddress(void * ptr,size_t size,const string & format,TypeId type_id)69 explicit DeviceAddress(void *ptr, size_t size, const string &format, TypeId type_id) 70 : ptr_(ptr), size_(size), format_(format), type_id_(type_id) {} DeviceAddress(void * ptr,size_t size,const std::string & format,TypeId type_id,const KernelWithIndex & node_index)71 explicit DeviceAddress(void *ptr, size_t size, const std::string &format, TypeId type_id, 72 const KernelWithIndex &node_index) 73 : ptr_(ptr), size_(size), format_(format), type_id_(type_id), node_index_(node_index) {} ~DeviceAddress()74 virtual ~DeviceAddress() { ptr_ = nullptr; } GetPtr()75 const void *GetPtr() const { return ptr_; } GetSize()76 size_t GetSize() const { return size_; } SetSize(size_t size)77 void SetSize(size_t size) { size_ = size; } format()78 std::string format() const { return format_; } type_id()79 TypeId type_id() const { return type_id_; } from_mem_pool()80 bool from_mem_pool() const { return from_mem_pool_; } set_host_shape(const ShapeVector & shape)81 void set_host_shape(const ShapeVector &shape) { host_shape_ = shape; } set_status(DeviceAddressStatus status)82 virtual void set_status(DeviceAddressStatus status) {} status()83 virtual DeviceAddressStatus status() const { return DeviceAddressStatus::kInDevice; } DeviceType()84 virtual DeviceAddressType DeviceType() const { return DeviceAddressType::kUnknown; } GetMutablePtr()85 void *GetMutablePtr() const override { return ptr_; } SetNodeIndex(const AnfNodePtr & node,size_t out_index)86 virtual void SetNodeIndex(const AnfNodePtr &node, size_t out_index) { node_index_ = {node, out_index}; } DumpMemToFile(const std::string & filepath,const std::string & host_fmt,const ShapeVector & host_shape,TypeId host_type,bool trans_flag)87 virtual bool DumpMemToFile(const std::string &filepath, const std::string &host_fmt, const ShapeVector &host_shape, 88 TypeId host_type, bool trans_flag) const { 89 return true; 90 } 91 #ifdef ENABLE_DEBUGGER LoadMemToHost(const std::string & tensor_name,int execution_order,const std::string & host_fmt,const ShapeVector & host_shape,TypeId host_type,size_t slot,bool keep_prev)92 virtual bool LoadMemToHost(const std::string &tensor_name, int execution_order, const std::string &host_fmt, 93 const ShapeVector &host_shape, TypeId host_type, size_t slot, bool keep_prev) const { 94 return true; 95 } 96 #endif 97 98 protected: ptr()99 const void *ptr() const { return ptr_; } size()100 size_t size() const { return size_; } set_ptr(void * ptr)101 void set_ptr(void *ptr) { ptr_ = ptr; } GetNodeIndex()102 KernelWithIndex GetNodeIndex() const { 103 return node_index_.first.expired() ? KernelWithIndex{nullptr, node_index_.second} 104 : KernelWithIndex{node_index_.first.lock(), node_index_.second}; 105 } 106 mutable void *ptr_{nullptr}; 107 size_t size_{0}; 108 string format_{"DefaultFormat"}; 109 TypeId type_id_{kNumberTypeFloat16}; 110 mutable bool from_mem_pool_{false}; 111 uint8_t *communication_ptr_{nullptr}; 112 ShapeVector host_shape_{}; 113 // {node, out_index} 114 std::pair<AnfNodeWeakPtr, size_t> node_index_{AnfNodePtr(nullptr), 0}; 115 friend class KernelRuntime; 116 friend class MemoryManager; 117 friend class mindspore::device::ascend::tasksink::TaskGenerator; 118 friend class mindspore::device::cpu::CPUSimpleMemPlan; 119 friend class mindspore::device::cpu::CPUMemoryManager; 120 friend class mindspore::device::cpu::CPUKernelRuntime; 121 friend class mindspore::device::cpu::CPUDeviceContext; 122 friend class mindspore::device::gpu::GPUKernelRuntime; 123 friend class mindspore::device::gpu::GPUMemoryManager; 124 friend class mindspore::device::gpu::GPUDeviceContext; 125 friend class mindspore::device::ascend::AscendKernelRuntime; 126 friend class mindspore::device::ascend::AscendMemoryManager; 127 #ifndef ENABLE_SECURITY 128 friend class mindspore::device::ascend::DataDumper; 129 #endif 130 friend class mindspore::device::Bucket; 131 }; 132 133 using DeviceAddressPtr = std::shared_ptr<DeviceAddress>; 134 using DeviceAddressPtrList = std::vector<DeviceAddressPtr>; 135 } // namespace device 136 } // namespace mindspore 137 #endif // MINDSPORE_DEVICE_TENSOR_H 138