1 /** 2 * Copyright 2020-2023 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 #ifndef MINDSPORE_CCSRC_INCLUDE_COMMON_DEBUG_DUMP_PROTO_H_ 17 #define MINDSPORE_CCSRC_INCLUDE_COMMON_DEBUG_DUMP_PROTO_H_ 18 19 #include <map> 20 #include <memory> 21 #include <utility> 22 #include <algorithm> 23 #include <functional> 24 #include <fstream> 25 #include <string> 26 #include <set> 27 #include <list> 28 #include <unordered_map> 29 #include <vector> 30 31 #include "utils/hash_map.h" 32 #include "ir/tensor.h" 33 #include "ir/param_info.h" 34 #include "ir/func_graph.h" 35 #include "ir/quantization_param.h" 36 #include "proto/mind_ir.pb.h" 37 #include "utils/check_convert_utils.h" 38 #include "include/common/debug/common.h" 39 #include "include/common/visible.h" 40 #include "utils/ms_utils.h" 41 #include "include/common/utils/utils.h" 42 #ifndef MINDIR_EXPORT_TENSOR_LAYOUT_CLIP 43 #include "frontend/parallel/tensor_layout/tensor_layout.h" 44 #endif 45 #include "abstract/abstract_function.h" 46 #include "mindspore/core/utils/file_utils.h" 47 #include "mindspore/core/utils/system/env.h" 48 #include "ir/functor.h" 49 50 namespace mindspore { 51 using FloatPtr = std::shared_ptr<Float>; 52 using BFloatPtr = std::shared_ptr<BFloat>; 53 using IntPtr = std::shared_ptr<Int>; 54 using UIntPtr = std::shared_ptr<UInt>; 55 using ComplexPtr = std::shared_ptr<Complex>; 56 using ModelProtoPtr = std::shared_ptr<mind_ir::ModelProto>; 57 constexpr const size_t TOTAL_SAVE = 1024 * 1024 * 1024; 58 constexpr const int64_t OFFSET = 64; 59 constexpr const size_t PARA_ROUND = 1024; 60 61 // anf type to mindir type map 62 static mindspore::HashMap<int, mind_ir::TensorProto_DataType> g_data_type_map = { 63 {kNumberTypeBool, mind_ir::TensorProto_DataType_BOOL}, 64 {kNumberTypeInt4, mind_ir::TensorProto_DataType_QINT4X2}, 65 {kNumberTypeInt8, mind_ir::TensorProto_DataType_INT8}, 66 {kNumberTypeInt16, mind_ir::TensorProto_DataType_INT16}, 67 {kNumberTypeInt, mind_ir::TensorProto_DataType_INT32}, 68 {kNumberTypeInt32, mind_ir::TensorProto_DataType_INT32}, 69 {kNumberTypeInt64, mind_ir::TensorProto_DataType_INT64}, 70 {kNumberTypeUInt8, mind_ir::TensorProto_DataType_UINT8}, 71 {kNumberTypeUInt16, mind_ir::TensorProto_DataType_UINT16}, 72 {kNumberTypeUInt32, mind_ir::TensorProto_DataType_UINT32}, 73 {kNumberTypeUInt64, mind_ir::TensorProto_DataType_UINT64}, 74 {kNumberTypeFloat16, mind_ir::TensorProto_DataType_FLOAT16}, 75 {kNumberTypeFloat, mind_ir::TensorProto_DataType_FLOAT}, 76 {kNumberTypeFloat32, mind_ir::TensorProto_DataType_FLOAT}, 77 {kNumberTypeFloat64, mind_ir::TensorProto_DataType_DOUBLE}, 78 {kNumberTypeBFloat16, mind_ir::TensorProto_DataType_BFLOAT16}, 79 {kObjectTypeString, mind_ir::TensorProto_DataType_STRING}, 80 {kNumberTypeComplex64, mind_ir::TensorProto_DataType_COMPLEX64}, 81 {kNumberTypeComplex128, mind_ir::TensorProto_DataType_COMPLEX128}}; 82 83 static mindspore::HashMap<int, mind_ir::TensorProto_DataType> g_data_bits_int_map = { 84 {4, mind_ir::TensorProto_DataType_QINT4X2}, {8, mind_ir::TensorProto_DataType_INT8}, 85 {16, mind_ir::TensorProto_DataType_INT16}, {32, mind_ir::TensorProto_DataType_INT32}, 86 {64, mind_ir::TensorProto_DataType_INT64}, 87 }; 88 89 static mindspore::HashMap<int, mind_ir::TensorProto_DataType> g_data_bits_uint_map = { 90 {8, mind_ir::TensorProto_DataType_UINT8}, 91 {16, mind_ir::TensorProto_DataType_UINT16}, 92 {32, mind_ir::TensorProto_DataType_UINT32}, 93 {64, mind_ir::TensorProto_DataType_UINT64}, 94 }; 95 96 static mindspore::HashMap<int, mind_ir::TensorProto_DataType> g_data_bits_float_map = { 97 {16, mind_ir::TensorProto_DataType_FLOAT16}, 98 {32, mind_ir::TensorProto_DataType_FLOAT}, 99 {64, mind_ir::TensorProto_DataType_FLOAT64}, 100 }; 101 102 static mindspore::HashMap<int, mind_ir::TensorProto_DataType> g_data_bits_bfloat_map = { 103 {16, mind_ir::TensorProto_DataType_BFLOAT16}, 104 }; 105 106 static mindspore::HashMap<int, mind_ir::TensorProto_DataType> g_data_bits_complex_map = { 107 {64, mind_ir::TensorProto_DataType_COMPLEX64}, 108 {128, mind_ir::TensorProto_DataType_COMPLEX128}, 109 }; 110 111 static std::set<std::string> g_export_attr_blacklist = {kAttrDump}; 112 113 // Can build different builder according to format 114 class IrExportBuilder; 115 using IrExportBuilderPtr = std::shared_ptr<IrExportBuilder>; 116 117 class IrExporter { 118 public: IrExporter(IrExportBuilderPtr builder)119 explicit IrExporter(IrExportBuilderPtr builder) : builder_(std::move(builder)) {} 120 virtual ~IrExporter() = default; 121 std::string GetDumpString(const FuncGraphPtr &func_graph); 122 ModelProtoPtr GetDumpProto(const FuncGraphPtr &func_graph); 123 ModelProtoPtr GetDumpProto(const FuncGraphPtr &root_graph, const std::vector<FuncGraphPtr> &child_graphs, 124 const std::vector<AnfNodePtr> &isolated_nodes); 125 126 private: 127 IrExportBuilderPtr builder_; 128 }; 129 using IrExporterPtr = std::shared_ptr<IrExporter>; 130 131 class IrExportBuilder { 132 public: 133 explicit IrExportBuilder(const bool &incremental = false) model_(std::make_shared<mind_ir::ModelProto> ())134 : model_(std::make_shared<mind_ir::ModelProto>()), incremental_(incremental) {} 135 ~IrExportBuilder() = default; 136 std::string GetProtoString() const; 137 void BuildModelInfo(); 138 bool BuildModel(const FuncGraphPtr &func_graph); 139 bool BuildModel(const FuncGraphPtr &root_graph, const std::vector<FuncGraphPtr> &child_graphs, 140 const std::vector<AnfNodePtr> &isolated_nodes); Model()141 ModelProtoPtr Model() { return model_; } 142 bool BuildFuncGraph(const FuncGraphPtr &func_graph, mind_ir::GraphProto *const graph_proto); 143 bool BuildFuncGraphAttrs(const FuncGraphPtr &func_graph, mind_ir::GraphProto *const graph_proto); 144 bool BuildParameters(const FuncGraphPtr &func_graph, mind_ir::GraphProto *const graph_proto); 145 bool BuildNodes(const FuncGraphPtr &func_graph, mind_ir::GraphProto *const graph_proto); 146 bool BuildIsolatedNodes(const std::vector<AnfNodePtr> &isolated_nodes); 147 bool BuildIsolatedCNode(const AnfNodePtr &node, std::set<AnfNodePtr> *visited); 148 bool BuildOutput(const CNodePtr &node, mind_ir::GraphProto *const graph_proto); 149 bool BuildCNode(const CNodePtr &node, mind_ir::GraphProto *const graph_proto); 150 bool BuildValueNode(const ValueNodePtr &node, const std::string &node_name, mind_ir::GraphProto *const graph_proto); 151 std::string BuildInputNode(const AnfNodePtr &node, mind_ir::GraphProto *const graph_proto); 152 bool BuildCNodeAttr(const CNodePtr &node, mind_ir::NodeProto *const node_proto); 153 bool SetValueInfoProto(const AnfNodePtr &node, mind_ir::ValueInfoProto *const value_proto); 154 bool SetParamToTensorProto(const ParameterPtr ¶m, mind_ir::TensorProto *const tensor_proto); 155 bool ConvertMapParameterToMapTensorProto(const ParameterPtr &map_parameter, 156 mind_ir::MapTensorProto *const map_tensor_proto); 157 bool ConvertAbstractMapTensorToAttrProto(const AbstractBasePtr &abstract, mind_ir::AttributeProto *const attr_proto); 158 bool SetTensorProto(const AbstractBasePtr &abstract, mind_ir::TensorProto *const tensor_proto); 159 bool SetCSRTensorToProto(const AbstractBasePtr &abstract, mind_ir::AttributeProto *const attr_proto); 160 bool SetCOOTensorToProto(const AbstractBasePtr &abstract, mind_ir::AttributeProto *const attr_proto); 161 bool SetAttributeProto(const AnfNodePtr &node, mind_ir::NodeProto *const node_proto); 162 bool ExportSequence(const abstract::AbstractSequencePtr &abs, mind_ir::AttributeProto *const attr_proto); 163 bool SetAbstractToNodeProto(const CNodePtr &node, mind_ir::NodeProto *const node_proto); 164 bool SetAbstractToNodeProto(const abstract::AbstractBasePtr &abstract, mind_ir::AttributeProto *const attr_proto); 165 bool SetValueToAttributeProto(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto); 166 bool SetNamedValueToAttributeProto(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto) const; 167 bool SetTypeToAttributeProto(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto); 168 bool SetScalarToAttributeProto_ir(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto) const; 169 bool SetScalarToAttributeProtoForInt_ir(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto) const; 170 bool SetScalarToAttributeProto_irs(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto) const; 171 bool SetScalarToAttributeProtoForInt_irs(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto) const; 172 bool SetTypeToAttributeProto_irs(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto); 173 bool SetTensorToAttributeProto(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto); 174 bool SetSequenceToAttributeProto(const ValueSequencePtr &value, mind_ir::AttributeProto *const attr_proto); 175 bool SetDictToAttributeProto(const ValueDictionaryPtr &value, mind_ir::AttributeProto *const attr_proto); 176 bool SetSeqElemToAttributeProto(const ValuePtr &value, mind_ir::AttributeProto *const attr_proto); 177 bool SetQuantizationParamToAttrProto(const std::shared_ptr<QuantizationParam> &quantization_param, 178 mind_ir::TensorProto_QuantParamProto *const quant_param_proto); 179 bool SetFunctorToAttrProto(const FunctorPtr &value, mind_ir::AttributeProto *const attr_proto); 180 bool SetTensorTypeToAttributeProto(const ValuePtr &value, mind_ir::TensorProto *tensor_proto); 181 182 mind_ir::TensorProto_DataType GetMindirDataType(TypeId type_id) const; 183 mind_ir::TensorProto_DataType GetMindirDataBitsIntType(int bits) const; 184 mind_ir::TensorProto_DataType GetMindirDataBitsFloatType(int bits) const; 185 mind_ir::TensorProto_DataType GetMindirDataBitsBFloatType(int bits) const; 186 mind_ir::TensorProto_DataType GetMindirDataBitsUIntType(int bits) const; 187 mind_ir::TensorProto_DataType GetMindirDataBitsComplexType(int bits) const; 188 std::string GetNodeName(const AnfNodePtr &node) const; 189 std::string GetUniqueNodeName(const AnfNodePtr &node); 190 std::string GetOpTypeName(const AnfNodePtr &node); GetUniqueID()191 size_t GetUniqueID() { return ++unique_id_; } 192 193 private: 194 bool SetAbstractFuncToAttributeProto(const abstract::AbstractBasePtr &abstract, 195 mind_ir::AttributeProto *const attr_proto); 196 bool ExportWeight(const ParameterPtr ¶m, const std::string ¶m_name, mind_ir::GraphProto *const graph_proto); 197 std::string GetPrimitiveUniqueName(const PrimitivePtr &primitive_ptr); 198 bool BuildPrimitives(); 199 200 ModelProtoPtr model_; 201 mind_ir::NodeProto *last_node_{nullptr}; 202 std::list<FuncGraphPtr> todo_; 203 std::map<AnfNodePtr, std::string> node_name_map_; 204 std::map<PrimitivePtr, std::string> primitive_name_map_; 205 std::set<std::string> nodeName_; 206 size_t unique_id_{0}; 207 bool top_graph{true}; 208 std::map<FuncGraphPtr, mind_ir::GraphProto *> graph_protos_; 209 bool incremental_{false}; 210 bool is_kernel_graph_{false}; 211 }; 212 213 COMMON_EXPORT std::string GetFuncGraphProtoString(const FuncGraphPtr &func_graph); 214 215 std::string GetOnnxProtoString(const FuncGraphPtr &func_graph); 216 217 COMMON_EXPORT std::string GetBinaryProtoString(const FuncGraphPtr &func_graph, const bool &incremental = false); 218 219 COMMON_EXPORT ModelProtoPtr GenBinaryProto(const FuncGraphPtr &func_graph); 220 221 COMMON_EXPORT bool DumpBinaryProto(const FuncGraphPtr &func_graph, const std::string &file_path); 222 COMMON_EXPORT bool DumpBinaryProto(const FuncGraphPtr &root_graph, const std::vector<FuncGraphPtr> &child_graphs, 223 const std::vector<AnfNodePtr> &isolated_nodes, const std::string &file_path); 224 COMMON_EXPORT void DumpIRProto(const FuncGraphPtr &func_graph, const std::string &suffix); 225 226 COMMON_EXPORT std::string GetFuncGraphProtoJsonString(const FuncGraphPtr &func_graph); 227 228 class COMMON_EXPORT MindIRExporter { 229 public: MindIRExporter()230 MindIRExporter() {} MindIRExporter(bool is_export_model)231 explicit MindIRExporter(bool is_export_model) { is_export_model_ = is_export_model; } ~MindIRExporter()232 virtual ~MindIRExporter() { 233 if (data_fs_ != nullptr) { 234 data_fs_->close(); 235 delete data_fs_; 236 data_fs_ = nullptr; 237 } 238 } 239 240 bool ExportProto(const FuncGraphPtr &func_graph, const std::string &file_path, 241 const FuncGraphPtr ¶m_layout_fg = nullptr); 242 bool IsSystemLittleEndidan() const; 243 bool PreProcSaveTogether(const FuncGraphPtr &func_graph); 244 bool SaveProtoToFile(mind_ir::ModelProto *model_proto, const std::string &output_file); 245 246 private: 247 bool ParserPath(const std::string &output_path); 248 bool IfSaveTogether(bool *save_together); 249 bool SaveMindIRTogether(); 250 bool SplitSave(); 251 bool UpdateParamCount(const FuncGraphPtr &func_graph); 252 253 private: 254 bool ParamDict(const FuncGraphPtr &func_graph); 255 bool CreateParameterDir(); 256 std::shared_ptr<Parameter> GetFgParaAccordingToProtoName(const std::string &proto_name); 257 bool ChangeParaDataFile(const std::string &file); 258 std::string CreateExternalPath(const std::string &external_file); 259 260 private: 261 std::string model_name_; 262 std::string save_path_; 263 std::string save_model_path_; 264 std::string dir_name_; 265 std::string dir_path_; 266 bool save_together_ = true; 267 mind_ir::ModelProto model_proto_; 268 std::unordered_map<std::string, ParameterPtr> param_dict_{}; 269 std::unordered_map<tensor::TensorPtr, mind_ir::TensorProto *> para_proto_dict_{}; 270 std::fstream *data_fs_ = nullptr; 271 std::shared_ptr<system::FileSystem> fs_{}; 272 bool is_export_model_ = true; 273 }; 274 } // namespace mindspore 275 276 #endif // MINDSPORE_CCSRC_INCLUDE_COMMON_DEBUG_DUMP_PROTO_H_ 277