1 /** 2 * Copyright 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_CCSRC_RUNTIME_HARDWARE_DEVICE_CONTEXT_H_ 18 #define MINDSPORE_CCSRC_RUNTIME_HARDWARE_DEVICE_CONTEXT_H_ 19 20 #include <string> 21 #include <vector> 22 #include <memory> 23 #include "runtime/device/device_address.h" 24 #include "runtime/device/bucket.h" 25 #include "backend/session/kernel_graph.h" 26 #include "backend/session/anf_runtime_algorithm.h" 27 28 namespace mindspore { 29 namespace device { 30 using mindspore::kernel::AddressPtr; 31 using mindspore::kernel::KernelMod; 32 33 const size_t kDeviceContextsNumOne = 1; 34 const size_t kDeviceContextsNumTwo = 2; 35 36 struct DeviceContextKey { 37 // device type name, such as 'GPU' 'Ascend' 'CPU'. 38 std::string device_name_; 39 uint32_t device_id_{0}; 40 41 // Use the result of ToString() as key to look up DeviceContext 42 // in cache map which maintains created DeviceContext objects. ToStringDeviceContextKey43 std::string ToString() const { return device_name_ + "_" + std::to_string(device_id_); } 44 }; 45 46 // DeviceContext is unified interface of interaction with device. 47 class DeviceContext { 48 public: DeviceContext(const DeviceContextKey & device_context_key)49 explicit DeviceContext(const DeviceContextKey &device_context_key) : device_context_key_(device_context_key) {} 50 virtual ~DeviceContext() = default; 51 52 // Initialize the device context. 53 virtual void Initialize() = 0; 54 55 // Destroy device context and release device resource. Destroy()56 virtual void Destroy() {} 57 58 // Relevant function to allocate and free device memory. 59 virtual bool AllocateMemory(DeviceAddress *const &address, size_t size) const = 0; 60 virtual void FreeMemory(DeviceAddress *const &address) const = 0; 61 62 // Allocate continuous device memory end to end into 'addr_list'. 63 // Communication operators may need continuous memory for input and output 64 // to optimize the communication performance. AllocateContinuousMemory(const std::vector<DeviceAddressPtr> & addr_list,size_t total_size,const std::vector<size_t> & size_list)65 virtual bool AllocateContinuousMemory(const std::vector<DeviceAddressPtr> &addr_list, size_t total_size, 66 const std::vector<size_t> &size_list) const { 67 return true; 68 } 69 70 // Create concrete device address according different device type. 71 virtual DeviceAddressPtr CreateDeviceAddress(void *const device_ptr, size_t device_size, const string &format, 72 TypeId type_id) const = 0; 73 74 // Get device address type according different device type, such GPU, Ascend. 75 virtual DeviceAddressType GetDeviceAddressType() const = 0; 76 77 // Optimize the kernel graph for graph mode. OptimizeGraph(const KernelGraphPtr & graph)78 virtual void OptimizeGraph(const KernelGraphPtr &graph) const {} 79 80 // Optimize the single operator graph for PyNative mode. OptimizeSingleOpGraph(const KernelGraphPtr & graph)81 virtual void OptimizeSingleOpGraph(const KernelGraphPtr &graph) const {} 82 83 // Select the matching backend kernels according to the data type and format of input and output for all 84 // execution operators, and set final device data type and format information for backend kernels, device 85 // data type and format which replace original data type and format will use for executing kernels. 86 virtual void SetOperatorInfo(const std::vector<CNodePtr> &nodes) const = 0; 87 88 // Generate 'KernelMod' for all kernels and set 'KernelMod' into kernel, 89 // 'KernelMod' is real executive object of kernel. 90 virtual void CreateKernel(const std::vector<CNodePtr> &nodes) const = 0; 91 92 // Adjust kernel graph before run graph, used in Graph Mode. PreprocessBeforeRunGraph(const KernelGraphPtr & graph)93 virtual void PreprocessBeforeRunGraph(const KernelGraphPtr &graph) const {} 94 // Adjust single op kernel graph before run graph, used in PyNative Mode. PreprocessBeforeRunSingleOpGraph(const KernelGraphPtr & graph)95 virtual void PreprocessBeforeRunSingleOpGraph(const KernelGraphPtr &graph) const {} 96 97 // Infer kernel shape and update abstract info for dynamic shape kernel. UpdateDynamicShape(const CNodePtr & kernel)98 virtual void UpdateDynamicShape(const CNodePtr &kernel) const { AnfAlgo::InferShape(kernel); } 99 100 // Launch a kernel via 'KernelMod' of the kernel. 101 virtual bool LaunchKernel(const CNodePtr &kernel, const std::vector<AddressPtr> &inputs, 102 const std::vector<AddressPtr> &workspace, const std::vector<AddressPtr> &outputs, 103 bool is_dynamic_shape = false) const = 0; 104 105 // Synchronize stream, device such as GPU and Ascend need stream to launch kernel asynchronously, 106 // using 'SyncStream' to block thread and wait for completing all tasks in stream. 107 // Devices that do not need stream could ignore the implementation of this function. 108 virtual bool SyncStream(size_t stream_id = 0) const { return true; } 109 110 // Get device_context_key_ to obtain device name and device id. device_context_key()111 const DeviceContextKey &device_context_key() const { return device_context_key_; } 112 113 // Get rank id for distributed training. GetRankID()114 virtual uint32_t GetRankID() const { return 0; } 115 116 // Create and initialize bucket for every allreduce operator. Bucket is used in PyNative distributed training mode, 117 // one bucket handles all resource to launch and sync allreduce operator. CreateBucket(uint32_t bucket_id,uint32_t bucket_size)118 virtual std::shared_ptr<Bucket> CreateBucket(uint32_t bucket_id, uint32_t bucket_size) const { return nullptr; } 119 120 protected: 121 DeviceContextKey device_context_key_; 122 }; 123 using DeviceContextPtr = std::shared_ptr<DeviceContext>; 124 } // namespace device 125 } // namespace mindspore 126 127 #endif // MINDSPORE_CCSRC_RUNTIME_HARDWARE_DEVICE_CONTEXT_H_ 128