diff --git a/ge/graph/build/task_generator.cc b/ge/graph/build/task_generator.cc index 4eda4020..13fc2601 100755 --- a/ge/graph/build/task_generator.cc +++ b/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 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); diff --git a/tests/ut/ge/CMakeLists.txt b/tests/ut/ge/CMakeLists.txt index 80636a20..453269da 100755 --- a/tests/ut/ge/CMakeLists.txt +++ b/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" diff --git a/tests/ut/ge/graph/build/task_generator_unittest.cc b/tests/ut/ge/graph/build/task_generator_unittest.cc new file mode 100644 index 00000000..95e75eb7 --- /dev/null +++ b/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 +#include + +#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); +}