From 31674c47fc5be37cf1d2fa345d87688604c7e99f Mon Sep 17 00:00:00 2001
From: zhaodezan <zhaodezan@huawei.com>
Date: Sat, 11 Sep 2021 10:03:07 +0800
Subject: [PATCH] fix tf model parser

---
 .../lite/tools/converter/parser/tf/tf_conv_base_parser.cc     | 3 +++
 mindspore/lite/tools/converter/parser/tf/tf_model_parser.cc   | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/mindspore/lite/tools/converter/parser/tf/tf_conv_base_parser.cc b/mindspore/lite/tools/converter/parser/tf/tf_conv_base_parser.cc
index 42cb2e56f95..d2963b7f808 100644
--- a/mindspore/lite/tools/converter/parser/tf/tf_conv_base_parser.cc
+++ b/mindspore/lite/tools/converter/parser/tf/tf_conv_base_parser.cc
@@ -32,6 +32,7 @@ constexpr size_t NCHWTopPadPos = 4;
 
 STATUS TFConvBaseParser::ParseKernels(const tensorflow::NodeDef &node_def, const mindspore::Format &format,
                                       std::vector<int64_t> *kernel) {
+  MS_ASSERT(kernel != nullptr);
   tensorflow::AttrValue attr_value;
   if (!TensorFlowUtils::FindAttrValue(node_def, "value", &attr_value)) {
     MS_LOG(ERROR) << "The kernels should be specified";
@@ -51,6 +52,7 @@ STATUS TFConvBaseParser::ParseKernels(const tensorflow::NodeDef &node_def, const
 
 STATUS TFConvBaseParser::ParseStrides(const tensorflow::NodeDef &node_def, const mindspore::Format &format,
                                       std::vector<int64_t> *strides) {
+  MS_ASSERT(strides != nullptr);
   tensorflow::AttrValue attr_value;
   if (!TensorFlowUtils::FindAttrValue(node_def, "strides", &attr_value)) {
     strides->at(0) = 1;
@@ -97,6 +99,7 @@ STATUS TFConvBaseParser::ParseExplicitPaddings(const tensorflow::NodeDef &node_d
 
 STATUS TFConvBaseParser::ParseDilations(const tensorflow::NodeDef &node_def, const mindspore::Format &format,
                                         std::vector<int64_t> *dilations) {
+  MS_ASSERT(dilations != nullptr);
   tensorflow::AttrValue attr_value;
   if (!TensorFlowUtils::FindAttrValue(node_def, "dilations", &attr_value)) {
     dilations->at(0) = 1;
diff --git a/mindspore/lite/tools/converter/parser/tf/tf_model_parser.cc b/mindspore/lite/tools/converter/parser/tf/tf_model_parser.cc
index dc323de2275..4f969ce391d 100644
--- a/mindspore/lite/tools/converter/parser/tf/tf_model_parser.cc
+++ b/mindspore/lite/tools/converter/parser/tf/tf_model_parser.cc
@@ -97,6 +97,7 @@ int GetShapeSize(const tensorflow::TensorProto &tensor_proto) {
   auto &tensor_shape = tensor_proto.tensor_shape();
   int shape_size = 1;
   for (int i = 0; i < tensor_shape.dim_size(); i++) {
+    MS_CHECK_INT_MUL_NOT_OVERFLOW(shape_size, tensor_shape.dim(i).size(), 0);
     shape_size *= tensor_shape.dim(i).size();
   }
   return shape_size;
@@ -582,6 +583,7 @@ FuncGraphPtr TFModelParser::Parse(const converter::ConverterParameters &flag) {
     return nullptr;
   }
   auto unify_format = std::make_shared<UnifyFormatToNHWC>(kFmkTypeTf, false);
+  MS_CHECK_TRUE_RET(unify_format != nullptr, nullptr);
   if (!unify_format->Run(res_graph_)) {
     MS_LOG(ERROR) << "Run insert transpose failed.";
     return nullptr;
@@ -733,6 +735,7 @@ STATUS TFModelParser::ConvertSubgraph() {
     }
 
     FuncGraphPtr sub_func_graph = std::make_shared<FuncGraph>();
+    MS_CHECK_TRUE_RET(sub_func_graph != nullptr, RET_ERROR);
     sub_func_graph->set_attr("graph_name", MakeValue(sub_graph_name));
     sub_func_graph->set_attr("fmk", MakeValue(static_cast<int>(converter::kFmkTypeTf)));
     std::unordered_map<std::string, AnfNodePtr> anf_sub_node_map;
@@ -1139,6 +1142,7 @@ STATUS TFModelParser::MakeAnfGraphOutputs(std::vector<AnfNodePtr> *output_nodes,
 int TFModelParser::TF2AnfAdjust(const std::set<FuncGraphPtr> &all_func_graphs) {
   for (const auto &func_graph : all_func_graphs) {
     auto functionalize_control_op_pass = std::make_shared<opt::FunctionalizeControlOpPass>();
+    MS_CHECK_TRUE_RET(functionalize_control_op_pass != nullptr, RET_ERROR);
     if (!functionalize_control_op_pass->Run(func_graph)) {
       MS_LOG(ERROR) << "functionalize control op pass failed.";
       ReturnCode::GetSingleReturnCode()->UpdateReturnCode(RET_ERROR);