1 /** 2 * Copyright 2020-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_BACKEND_KERNEL_COMPILER_AKG_AKG_KERNEL_JSON_GENERATOR_H_ 18 #define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_AKG_KERNEL_JSON_GENERATOR_H_ 19 #include <map> 20 #include <memory> 21 #include <string> 22 #include <unordered_map> 23 #include <utility> 24 #include <vector> 25 #include "backend/kernel_compiler/oplib/oplib.h" 26 #include "nlohmann/json.hpp" 27 28 namespace mindspore { 29 namespace kernel { 30 // json key 31 constexpr auto kJsonKeyOpDesc = "op_desc"; 32 constexpr auto kJsonKeyAttr = "attr"; 33 constexpr auto kJsonKeyInputDesc = "input_desc"; 34 constexpr auto kJsonKeyFormat = "format"; 35 constexpr auto kJsonKeyInferDataType = "infer_data_type"; 36 constexpr auto kJsonKeyInferShape = "infer_shape"; 37 constexpr auto kJsonKeyShape = "shape"; 38 constexpr auto kJsonKeyDataType = "data_type"; 39 constexpr auto kJsonKeyOutputDesc = "output_desc"; 40 constexpr auto kJsonKeyName = "name"; 41 constexpr auto kJsonKeyTensorName = "tensor_name"; 42 constexpr auto kJsonKeyValue = "value"; 43 constexpr auto kJsonKeyImplPath = "impl_path"; 44 constexpr auto kJsonKeyProcess = "process"; 45 constexpr auto kJsonKeyComposite = "composite"; 46 constexpr auto kJsonKeyId = "id"; 47 constexpr auto kJsonKeyOp = "op"; 48 constexpr auto kJsonKeyPtrAddress = "ptr_address"; 49 constexpr auto kJsonKeyCompositeGraph = "composite_graph"; 50 constexpr auto kJsonKeyPlatform = "platform"; 51 constexpr auto kJsonKeyOpFullName = "op_full_name"; 52 constexpr auto kJsonKeyParallelFusion = "parallel_fusion"; 53 constexpr auto kJsonKeyFusionType = "fusion_type"; 54 constexpr auto kJsonKeySubGraph = "sub_graph"; 55 constexpr auto kJsonKeyCoreNum = "core_num"; 56 constexpr auto kJsonKeyTypeInfo = "type_info"; 57 constexpr auto kJsonKeyRecomputeOps = "recompute_ops"; 58 constexpr auto kJsonKeyBufferStitch = "buffer_stitch"; 59 constexpr auto kJsonKeyStitchOp = "stitch_op"; 60 constexpr auto kJsonKeyStitchAtomicOp = "stitch_atomic_op"; 61 constexpr auto kJsonKeyComputeCapability = "compute_capability"; 62 63 constexpr auto kAttrInputNames = "input_names"; 64 65 // dump option 66 struct DumpOption { 67 bool is_before_select_kernel = false; 68 bool save_ptr_address = false; 69 bool extract_opinfo_from_anfnode = false; 70 bool get_compute_capability = false; 71 }; 72 73 class ComputeCapability { 74 public: Get()75 static const std::string &Get() { 76 static std::unique_ptr<ComputeCapability> instance = nullptr; 77 if (instance == nullptr) { 78 instance = std::make_unique<ComputeCapability>(); 79 instance->GetComputeCapability(); 80 } 81 return instance->compute_capability_; 82 } 83 84 private: 85 void GetComputeCapability(); 86 std::string compute_capability_; 87 }; 88 89 class AkgKernelJsonGenerator { 90 public: AkgKernelJsonGenerator()91 AkgKernelJsonGenerator() { Clear(); } AkgKernelJsonGenerator(DumpOption dump_option)92 explicit AkgKernelJsonGenerator(DumpOption dump_option) : dump_option_(dump_option) { Clear(); } 93 ~AkgKernelJsonGenerator() = default; 94 95 bool CollectJson(const AnfNodePtr &anf_node, nlohmann::json *kernel_json); 96 bool CollectFusedJson(const std::vector<AnfNodePtr> &anf_nodes, const std::vector<AnfNodePtr> &input_list, 97 const std::vector<AnfNodePtr> &output_list, nlohmann::json *kernel_json); 98 bool CollectJson(const AnfNodePtr &anf_node); 99 bool CollectFusedJson(const std::vector<AnfNodePtr> &anf_nodes, const std::vector<AnfNodePtr> &input_list, 100 const std::vector<AnfNodePtr> &output_list); 101 bool GenerateSingleKernelJson(const AnfNodePtr &anf_node, nlohmann::json *node_json); kernel_name()102 std::string kernel_name() const { return kernel_name_; } kernel_json()103 nlohmann::json kernel_json() const { return kernel_json_; } kernel_json_str()104 std::string kernel_json_str() const { return kernel_json_.dump(); } input_size_list()105 const std::vector<size_t> &input_size_list() const { return input_size_list_; } output_size_list()106 const std::vector<size_t> &output_size_list() const { return output_size_list_; } Clear()107 void Clear() { 108 input_tensor_idx_.clear(); 109 address_node_map_.clear(); 110 output_tensor_idx_ = 0; 111 } set_dump_option(DumpOption dump_option)112 void set_dump_option(DumpOption dump_option) { dump_option_ = dump_option; } address_node_map()113 std::map<std::string, AnfNodePtr> address_node_map() { return address_node_map_; } 114 115 private: 116 bool CreateInputDescJson(const AnfNodePtr &anf_node, const OpInfoPtr &op_info, nlohmann::json *inputs_json); 117 bool CreateOutputDescJson(const AnfNodePtr &anf_node, const OpInfoPtr &op_info, nlohmann::json *outputs_json); 118 void GetAttrJson(const AnfNodePtr &anf_node, const std::vector<int> &dyn_input_sizes, const OpAttrPtr &op_attr, 119 nlohmann::json *attr_json, const ValuePtr &attr_value); 120 bool CreateAttrDescJson(const AnfNodePtr &anf_node, const OpInfoPtr &op_info, nlohmann::json *attrs_json); 121 void GenStitchJson(const std::vector<AnfNodePtr> &anf_nodes, std::map<AnfNodePtr, nlohmann::json> *node_json_map, 122 nlohmann::json *kernel_json); 123 bool GetIOSize(const nlohmann::json &node_json, std::vector<size_t> *input_size, 124 std::vector<size_t> *output_size) const; 125 bool GenSingleJsons(const std::vector<AnfNodePtr> &anf_nodes, std::map<AnfNodePtr, nlohmann::json> *node_json_map); 126 void UpdateTensorName(const std::vector<AnfNodePtr> &anf_nodes, std::map<AnfNodePtr, nlohmann::json> *node_json_map); 127 nlohmann::json CreateInputsJson(const std::vector<AnfNodePtr> &anf_nodes, const std::vector<AnfNodePtr> &input_list, 128 const std::map<AnfNodePtr, nlohmann::json> &node_json_map); 129 nlohmann::json CreateOutputsJson(const std::vector<AnfNodePtr> &anf_nodes, const std::vector<AnfNodePtr> &input_list, 130 const std::vector<AnfNodePtr> &output_list, const nlohmann::json &inputs_json, 131 const std::map<AnfNodePtr, nlohmann::json> &node_json_map); 132 133 int GetOpCntInc(); 134 size_t GetInputTensorIdxInc(const AnfNodePtr &anf_node, size_t input_idx); 135 size_t GetOutputTensorIdxInc(); 136 void SetTensorName(const std::string &tag, const std::string &new_name, const std::pair<size_t, size_t> &position, 137 nlohmann::json *node_json) const; 138 std::string GetTensorName(const nlohmann::json &node_json, const std::string &tag, 139 const std::pair<size_t, size_t> &position) const; 140 TypeId GetInputDataType(const AnfNodePtr &anf_node, size_t real_index) const; 141 std::vector<size_t> GetInputShape(const AnfNodePtr &anf_node, size_t real_index) const; 142 std::string GetInputFormat(const AnfNodePtr &anf_node, size_t real_index) const; 143 TypeId GetOutputDataType(const AnfNodePtr &anf_node, size_t index) const; 144 std::vector<size_t> GetOutputShape(const AnfNodePtr &anf_node, size_t index) const; 145 std::string GetOutputFormat(const AnfNodePtr &anf_node, size_t index) const; 146 void SaveNodeAddress(const AnfNodePtr &anf_node, nlohmann::json *node_json); 147 OpInfoPtr ExtractOpInfo(const AnfNodePtr &anf_node) const; 148 void CollectParallelDimInfo(const AnfNodePtr &anf_node); 149 void GenParallelJson(const std::vector<AnfNodePtr> &anf_nodes, const std::vector<AnfNodePtr> &input_list, 150 const std::vector<AnfNodePtr> &output_list, 151 const std::map<AnfNodePtr, nlohmann::json> &node_json_map, nlohmann::json *kernel_json); 152 153 DumpOption dump_option_; 154 static int op_cnt_; 155 // lock for variable fusionOpCnt in singleton mode 156 static std::mutex op_cnt_mtx_; 157 std::string kernel_name_; 158 std::unordered_map<AnfNodePtr, size_t> input_tensor_idx_; 159 size_t output_tensor_idx_; 160 nlohmann::json kernel_json_; 161 std::vector<size_t> input_size_list_; 162 std::vector<size_t> output_size_list_; 163 std::map<std::string, AnfNodePtr> address_node_map_; 164 bool is_basic_op_{false}; 165 }; 166 } // namespace kernel 167 } // namespace mindspore 168 #endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_AKG_KERNEL_JSON_GENERATOR_H_ 169