bugfix for auto find fp

ge/graph/build/task_generator.cc

@@ -49,6 +49,7 @@ const char *const kIsLastNode = "is_last_node";
 const char *const kIsInputVar = "INPUT_IS_VAR";
 const char *const kIsOutputVar = "OUTPUT_IS_VAR";
 const char *const kProfilingMode = "PROFILING_MODE";
+const char *const kIteratorV2 = "IteratorV2";
 const uint32_t kProfilingArStep = 2;
 const uint64_t kProfilingFpStartLogid = 1;
 const uint64_t kProfilingBpEndLogid = 2;
@@ -57,6 +58,7 @@ const uint64_t kProfilingArEndLogid = 4;
 const uint64_t kProfilingIterEndLogid = 65535;
 const int64_t kHashFactor = 100000;
 const int64_t kInvalidGroupId = -1;
+const std::set<std::string> kFpNodeTypes = {ge::DATA, ge::GETNEXT, kIteratorV2};
 }  // namespace
 namespace ge {
 TaskGenerator::TaskGenerator(uint8_t *var_mem_base, uint64_t var_mem_size) {
@@ -621,8 +623,10 @@ Status TaskGenerator::AutoFindFpOpIndex(const ComputeGraphPtr &graph, ProfilingP
     if (op_kernel_lib_name.empty()) {
       continue;
     }
-
-    if (op_desc->GetType() == GETNEXT || op_desc->GetType() == DATA) {
+    auto type = op_desc->GetType();
+    std::string original_type;
+    (void)AttrUtils::GetStr(op_desc, ATTR_NAME_FRAMEWORK_ORIGINAL_TYPE, original_type);
+    if (kFpNodeTypes.find(type) != kFpNodeTypes.end() || kFpNodeTypes.find(original_type) != kFpNodeTypes.end()) {
       auto out_anchor = node->GetOutDataAnchor(0);
       for (auto &peer_in_anchor : out_anchor->GetPeerInDataAnchors()) {
         GE_CHECK_NOTNULL(peer_in_anchor);

tests/ut/ge/CMakeLists.txt

@@ -731,6 +731,7 @@ set(KERNEL_TEST_FILES
     "graph/passes/folding_kernel/gather_v2_kernel_unittest.cc"
     "graph/passes/folding_kernel/slice_kernel_unittest.cc"
     "graph/passes/folding_kernel/dynamic_stitch_kernel_unittest.cc"
+    "graph/passes/atomic_addr_clean_pass_unittest.cc"
 )
 
 set(MULTI_PARTS_TEST_FILES
@@ -760,6 +761,7 @@ set(MULTI_PARTS_TEST_FILES
     "graph/variable_accelerate_ctrl_unittest.cc"
     "graph/build/logical_stream_allocator_unittest.cc"
     "graph/build/mem_assigner_unittest.cc"
+    "graph/build/task_generator_unittest.cc"
     "graph/preprocess/graph_preprocess_unittest.cc"
     "graph/manager/hcom_util_unittest.cc"
     "graph/manager/graph_caching_allocator_unittest.cc"

tests/ut/ge/graph/build/task_generator_unittest.cc

@@ -0,0 +1,68 @@
+/**
+ * Copyright 2019-2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "graph/anchor.h"
+#include "graph/attr_value.h"
+#include "graph/debug/ge_attr_define.h"
+#include "graph/utils/graph_utils.h"
+#include "graph/utils/node_utils.h"
+#include "graph/utils/op_desc_utils.h"
+#include "graph/utils/tensor_utils.h"
+#include "omg/omg_inner_types.h"
+#include "../passes/graph_builder_utils.h"
+
+#define protected public
+#define private public
+#include "graph/build/task_generator.h"
+#undef protected
+#undef private
+
+using namespace std;
+using namespace testing;
+using namespace ge;
+
+class UtestTaskGeneratorTest : public testing::Test {
+ public:
+  ge::ComputeGraphPtr BuildGraphFpProfiling() {
+    ge::ut::GraphBuilder builder("graph");
+    auto data = builder.AddNode("data", "phony", 1, 1);
+    auto addn1 = builder.AddNode("addn1", "AddN", 1, 1);
+    auto netoutput = builder.AddNode("netoutput", "NetOutput", 2, 0);
+    auto op_desc = data->GetOpDesc();
+    (void)AttrUtils::SetStr(op_desc, ATTR_NAME_FRAMEWORK_ORIGINAL_TYPE, "IteratorV2");
+    op_desc->SetOpKernelLibName("GE");
+    builder.AddDataEdge(data, 0, addn1, 0);
+    builder.AddDataEdge(addn1, 0, netoutput, 0);
+    return builder.GetGraph();
+  }
+
+ protected:
+  void SetUp() {}
+  void TearDown() {}
+};
+
+TEST_F(UtestTaskGeneratorTest, AutoFindFpOpIndex) {
+  auto graph = BuildGraphFpProfiling();
+  TaskGenerator task_generator(nullptr, 0);
+  ProfilingPoint profiling_point;
+  profiling_point.fp_index = -1;
+  EXPECT_EQ(task_generator.AutoFindFpOpIndex(graph, profiling_point), SUCCESS);
+  // addn1 is fp
+  EXPECT_EQ(profiling_point.fp_index, 2);
+}