add dump_om_model flag

microsoft · Mar 26, 2024 · ca66d6e · ca66d6e
1 parent 7d976cf
commit ca66d6e
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 1 deletion.
diff --git a/include/onnxruntime/core/providers/cann/cann_provider_options.h b/include/onnxruntime/core/providers/cann/cann_provider_options.h
@@ -16,6 +16,7 @@ struct OrtCANNProviderOptions {
   int enable_cann_graph;                                   // Flag indicating if prioritizing the use of
                                                            // CANN's graph-running capabilities
   int dump_graphs;                                         // Flag indicating if dumping graphs
+  int dump_om_model;                                       // Flag indicating if dumping om model
   std::string precision_mode;                              // Operator Precision Mode
   std::string op_select_impl_mode;                         // Operator-level model compilation options:
                                                            // Mode selection

diff --git a/onnxruntime/core/providers/cann/cann_execution_provider_info.cc b/onnxruntime/core/providers/cann/cann_execution_provider_info.cc
@@ -21,6 +21,7 @@ constexpr const char* kMemLimit = "npu_mem_limit";
 constexpr const char* kArenaExtendStrategy = "arena_extend_strategy";
 constexpr const char* kEnableCannGraph = "enable_cann_graph";
 constexpr const char* kDumpGraphs = "dump_graphs";
+constexpr const char* kDumpOmModel = "dump_om_model";
 constexpr const char* kPrecisionMode = "precision_mode";
 constexpr const char* kOpSelectImplMode = "op_select_impl_mode";
 constexpr const char* kOpTypeListForImplMode = "optypelist_for_implmode";
@@ -58,6 +59,7 @@ CANNExecutionProviderInfo CANNExecutionProviderInfo::FromProviderOptions(const P
               arena_extend_strategy_mapping, info.arena_extend_strategy)
           .AddAssignmentToReference(cann::provider_option_names::kEnableCannGraph, info.enable_cann_graph)
           .AddAssignmentToReference(cann::provider_option_names::kDumpGraphs, info.dump_graphs)
+          .AddAssignmentToReference(cann::provider_option_names::kDumpOmModel, info.dump_om_model)
           .AddAssignmentToReference(cann::provider_option_names::kPrecisionMode, info.precision_mode)
           .AddAssignmentToReference(cann::provider_option_names::kOpSelectImplMode, info.op_select_impl_mode)
           .AddAssignmentToReference(cann::provider_option_names::kOpTypeListForImplMode, info.optypelist_for_implmode)
@@ -73,6 +75,7 @@ ProviderOptions CANNExecutionProviderInfo::ToProviderOptions(const CANNExecution
        EnumToName(arena_extend_strategy_mapping, info.arena_extend_strategy)},
       {cann::provider_option_names::kEnableCannGraph, MakeStringWithClassicLocale(info.enable_cann_graph)},
       {cann::provider_option_names::kDumpGraphs, MakeStringWithClassicLocale(info.dump_graphs)},
+      {cann::provider_option_names::kDumpOmModel, MakeStringWithClassicLocale(info.dump_om_model)},
       {cann::provider_option_names::kPrecisionMode, MakeStringWithClassicLocale(info.precision_mode)},
       {cann::provider_option_names::kOpSelectImplMode, MakeStringWithClassicLocale(info.op_select_impl_mode)},
       {cann::provider_option_names::kOpTypeListForImplMode, MakeStringWithClassicLocale(info.optypelist_for_implmode)}};
@@ -87,6 +90,7 @@ ProviderOptions CANNExecutionProviderInfo::ToProviderOptions(const OrtCANNProvid
        EnumToName(arena_extend_strategy_mapping, ArenaExtendStrategy(info.arena_extend_strategy))},
       {cann::provider_option_names::kEnableCannGraph, MakeStringWithClassicLocale(info.enable_cann_graph)},
       {cann::provider_option_names::kDumpGraphs, MakeStringWithClassicLocale(info.dump_graphs)},
+      {cann::provider_option_names::kDumpOmModel, MakeStringWithClassicLocale(info.dump_om_model)},
       {cann::provider_option_names::kPrecisionMode, MakeStringWithClassicLocale(info.precision_mode)},
       {cann::provider_option_names::kOpSelectImplMode, MakeStringWithClassicLocale(info.op_select_impl_mode)},
       {cann::provider_option_names::kOpTypeListForImplMode, MakeStringWithClassicLocale(info.optypelist_for_implmode)}};

diff --git a/onnxruntime/core/providers/cann/cann_execution_provider_info.h b/onnxruntime/core/providers/cann/cann_execution_provider_info.h
@@ -19,6 +19,7 @@ struct CANNExecutionProviderInfo {
   ArenaExtendStrategy arena_extend_strategy{ArenaExtendStrategy::kNextPowerOfTwo};
   bool enable_cann_graph{true};
   bool dump_graphs{false};
+  bool dump_om_model{true};
   std::string precision_mode;
   std::string op_select_impl_mode;
   std::string optypelist_for_implmode;

diff --git a/onnxruntime/core/providers/cann/cann_graph.cc b/onnxruntime/core/providers/cann/cann_graph.cc
@@ -115,7 +115,9 @@ Status BuildONNXModel(ge::Graph& graph, std::string input_shape, const char* soc
   options.emplace(ge::ir_option::INPUT_SHAPE, input_shape.c_str());
   CANN_GRAPH_RETURN_IF_ERROR(ge::aclgrphBuildModel(graph, options, model));
 
-  CANN_GRAPH_RETURN_IF_ERROR(ge::aclgrphSaveModel(file_name.c_str(), model));
+  if (info.dump_om_model) {
+    CANN_GRAPH_RETURN_IF_ERROR(ge::aclgrphSaveModel(file_name.c_str(), model));
+  }
 
   return Status::OK();
 }

diff --git a/onnxruntime/core/providers/cann/cann_provider_factory.cc b/onnxruntime/core/providers/cann/cann_provider_factory.cc
@@ -77,6 +77,7 @@ struct CANN_Provider : Provider {
     info.arena_extend_strategy = params->arena_extend_strategy;
     info.enable_cann_graph = params->enable_cann_graph != 0;
     info.dump_graphs = params->dump_graphs != 0;
+    info.dump_om_model = params->dump_om_model != 0;
     info.precision_mode = params->precision_mode;
     info.op_select_impl_mode = params->op_select_impl_mode;
     info.optypelist_for_implmode = params->optypelist_for_implmode;
@@ -94,6 +95,7 @@ struct CANN_Provider : Provider {
     cann_options.arena_extend_strategy = internal_options.arena_extend_strategy;
     cann_options.enable_cann_graph = internal_options.enable_cann_graph;
     cann_options.dump_graphs = internal_options.dump_graphs;
+    cann_options.dump_om_model = internal_options.dump_om_model;
     cann_options.precision_mode = internal_options.precision_mode;
     cann_options.op_select_impl_mode = internal_options.op_select_impl_mode;
     cann_options.optypelist_for_implmode = internal_options.optypelist_for_implmode;

diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc
@@ -2403,6 +2403,7 @@ ORT_API_STATUS_IMPL(OrtApis::CreateCANNProviderOptions, _Outptr_ OrtCANNProvider
   options->arena_extend_strategy = static_cast<onnxruntime::ArenaExtendStrategy>(0);
   options->enable_cann_graph = 1;
   options->dump_graphs = 0;
+  options->dump_om_model = 1;
   options->default_memory_arena_cfg = nullptr;
   *out = options.release();
   return nullptr;