1 /** 2 * Copyright 2019 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_KERNEL_INFO_H_ 18 #define MINDSPORE_DEVICE_KERNEL_INFO_H_ 19 20 #include <vector> 21 #include <memory> 22 #include <utility> 23 #include "ir/kernel_info_dev.h" 24 #include "kernel/kernel_build_info.h" 25 #include "kernel/kernel.h" 26 #include "include/backend/device_address.h" 27 #include "include/backend/visible.h" 28 29 namespace mindspore { 30 const uint32_t kInvalidGraphId = UINT32_MAX; 31 const uint32_t kInvalidDistincLabel = UINT32_MAX; 32 namespace device { 33 using kernel::KernelTensorPtr; 34 class BACKEND_EXPORT KernelInfo : public KernelInfoDevice { 35 public: KernelInfo()36 KernelInfo() { 37 kernel_mod_ = nullptr; 38 is_feature_map_ = false; 39 select_kernel_build_info_ = nullptr; 40 output_address_list_ = {}; 41 workspace_address_list_ = {}; 42 stream_id_ = 0; 43 stream_distinction_label_ = kInvalidDistincLabel; 44 graph_id_ = kInvalidGraphId; 45 } 46 virtual ~KernelInfo() = default; 47 has_build_info()48 bool has_build_info() const override { return select_kernel_build_info_ != nullptr; } 49 const kernel::KernelBuildInfo *select_kernel_build_info() const; 50 kernel::KernelBuildInfoPtr GetMutableSelectKernelBuildInfo() const; set_select_kernel_build_info(const kernel::KernelBuildInfoPtr & select_kernel_build_info)51 void set_select_kernel_build_info(const kernel::KernelBuildInfoPtr &select_kernel_build_info) { 52 select_kernel_build_info_ = select_kernel_build_info; 53 } set_feature_map_flag(bool flag)54 void set_feature_map_flag(bool flag) { is_feature_map_ = flag; } 55 const KernelTensorPtr &GetOutputKernelTensor(size_t index) const; 56 bool SetOutputKernelTensor(const KernelTensorPtr &kernel_tensor, size_t index); 57 bool OutputKernelTensorExist(size_t index) const; 58 const DeviceAddress *GetOutputAddr(size_t index) const; 59 DeviceAddressPtr GetMutableOutputAddr(size_t index) const; 60 bool OutputAddrExist(size_t index) const; 61 bool SetOutputAddr(const DeviceAddressPtr &output_address, size_t index); 62 DeviceAddress *GetWorkspaceAddr(size_t index) const; 63 DeviceAddressPtr GetMutableWorkspaceAddr(size_t index) const; 64 bool WorkspaceAddrExist(size_t index) const; 65 bool SetWorkspaceAddr(const DeviceAddressPtr &output_address, size_t index); 66 // The number of workspace may change after kernel Resize. set_workspace_address_list(const std::vector<DeviceAddressPtr> & workspace_address_list)67 void set_workspace_address_list(const std::vector<DeviceAddressPtr> &workspace_address_list) { 68 workspace_address_list_ = workspace_address_list; 69 } 70 void set_kernel_mod(const kernel::KernelModPtr &kernel_mod); 71 kernel::KernelMod *MutableKernelMod() const; 72 kernel::KernelModPtr GetKernelMod() const; 73 const kernel::KernelMod *kernel_mod() const; stream_id()74 uint32_t stream_id() const { return stream_id_; } set_stream_id(uint32_t stream_id)75 void set_stream_id(uint32_t stream_id) { stream_id_ = stream_id; } stream_distinction_label()76 uint32_t stream_distinction_label() const { return stream_distinction_label_; } set_stream_distinction_label(uint32_t stream_distinction_label)77 void set_stream_distinction_label(uint32_t stream_distinction_label) { 78 stream_distinction_label_ = stream_distinction_label; 79 } set_graph_id(uint32_t graph_id)80 void set_graph_id(uint32_t graph_id) { graph_id_ = graph_id; } graph_id()81 uint32_t graph_id() const { return graph_id_; } 82 bool operator==(const KernelInfo &other) const; is_feature_map()83 bool is_feature_map() const { return is_feature_map_; } 84 output_address_list()85 const std::vector<std::shared_ptr<DeviceAddress>> &output_address_list() const { return output_address_list_; } workspace_address_list()86 const std::vector<std::shared_ptr<DeviceAddress>> &workspace_address_list() const { return workspace_address_list_; } 87 88 // Set output and input reference. If all_ref is set true, each output is a reference to the input with the same 89 // index. 90 void set_ref_map(const bool &all_ref, const OutputInputRefMap &ref_map); out_in_ref_map()91 const OutputInputRefMap &out_in_ref_map() const { return out_in_ref_map_; } AddRefMap(size_t out,size_t in)92 void AddRefMap(size_t out, size_t in) { out_in_ref_map_[out] = in; } 93 94 // The interface of somas. 95 bool SetSomasResult(std::vector<std::pair<size_t, size_t>> &&output_somas_result, 96 std::vector<std::pair<size_t, size_t>> &&workspace_somas_result); 97 size_t GetTensorSomasOffset(const std::vector<std::pair<size_t, size_t>> &somas_result, size_t tensor_index) const; 98 size_t GetTensorSomasAlignedSize(const std::vector<std::pair<size_t, size_t>> &somas_result, 99 size_t tensor_index) const; 100 bool IsTensorEnableSomas(const std::vector<std::pair<size_t, size_t>> &somas_result, size_t tensor_index) const; somas_output_result()101 const std::vector<std::pair<size_t, size_t>> &somas_output_result() const { return somas_output_result_; } somas_workspace_result()102 const std::vector<std::pair<size_t, size_t>> &somas_workspace_result() const { return somas_workspace_result_; } 103 104 private: 105 bool is_feature_map_; 106 kernel::KernelBuildInfoPtr select_kernel_build_info_; 107 std::vector<KernelTensorPtr> output_kernel_tensor_list_; 108 std::vector<std::shared_ptr<DeviceAddress>> output_address_list_; 109 std::vector<std::shared_ptr<DeviceAddress>> workspace_address_list_; 110 // pair<size_t, size_t> : (offset, aligned_size) 111 // aligned_size of 0 means no memory allocation 112 std::vector<std::pair<size_t, size_t>> somas_output_result_; 113 // pair<size_t, size_t> : (offset, aligned_size) 114 // aligned_size of 0 means no memory allocation 115 std::vector<std::pair<size_t, size_t>> somas_workspace_result_; 116 kernel::KernelModPtr kernel_mod_; 117 // stream_id_ is the index of stream object vector 118 uint32_t stream_id_; 119 // stream_distinction_label_ is used mark different op in different stream 120 uint32_t stream_distinction_label_; 121 // record which graph the node belong to 122 uint32_t graph_id_; 123 // The map between kernel's output and input ref relationship. 124 OutputInputRefMap out_in_ref_map_; 125 }; 126 } // namespace device 127 } // namespace mindspore 128 #endif // MINDSPORE_DEVICE_KERNEL_INFO_H_ 129