diff --git a/inference-engine/thirdparty/clDNN/api/cldnn/graph/network.hpp b/inference-engine/thirdparty/clDNN/api/cldnn/graph/network.hpp index ea3af7000ad..6f8062683de 100644 --- a/inference-engine/thirdparty/clDNN/api/cldnn/graph/network.hpp +++ b/inference-engine/thirdparty/clDNN/api/cldnn/graph/network.hpp @@ -168,6 +168,7 @@ public: void execute_primitive(const std::shared_ptr& primitive, const std::vector& events); void allocate_primitives(); + void configure_primitives_second_output(); void build_insts_deps(); uint32_t get_id() const { return net_id; } stream& get_stream() const { return *_stream; } diff --git a/inference-engine/thirdparty/clDNN/src/network.cpp b/inference-engine/thirdparty/clDNN/src/network.cpp index 59703aed498..c6aa1b78fbc 100644 --- a/inference-engine/thirdparty/clDNN/src/network.cpp +++ b/inference-engine/thirdparty/clDNN/src/network.cpp @@ -214,6 +214,7 @@ network::network(program::ptr program, stream::ptr stream, bool is_internal, boo } allocate_primitives(); + configure_primitives_second_output(); check_names(); build_insts_deps(); build_exec_order(); @@ -552,6 +553,41 @@ void network::allocate_primitives() { } } +void network::configure_primitives_second_output() { + std::map> mutable_datas_ptrs; + for (auto& inst : _primitives) { + auto& node = inst.second->get_node(); + + if (!node.is_type()) + continue; + + mutable_datas_ptrs[node.as().get_attached_memory_ptr()].push_back(&node); + } + + for (auto item : mutable_datas_ptrs) { + if (item.second.size() != 2) + continue; + + auto is_first_node_input_md = [&](const cldnn::program_node* first, + const cldnn::program_node* second) { + for (auto user : first->get_users()) { + for (auto next_user : user->get_users()) { + if (next_user == second) + return true; + } + } + return false; + }; + + auto is_first_node_input = is_first_node_input_md(item.second[0], item.second[1]); + + auto input_md_inst = is_first_node_input ? _primitives[item.second[0]->id()] : _primitives[item.second[1]->id()]; + auto output_md_inst = is_first_node_input ? _primitives[item.second[1]->id()] : _primitives[item.second[0]->id()]; + + output_md_inst->set_output_memory(input_md_inst->output_memory_ptr(), false); + } +} + void network::build_insts_deps() { for (auto& inst : _primitives) { inst.second->build_deps(); diff --git a/inference-engine/thirdparty/clDNN/tests/test_cases/set_output_memory_gpu_test.cpp b/inference-engine/thirdparty/clDNN/tests/test_cases/set_output_memory_gpu_test.cpp index 6bfb4a5db12..c4901b7b660 100644 --- a/inference-engine/thirdparty/clDNN/tests/test_cases/set_output_memory_gpu_test.cpp +++ b/inference-engine/thirdparty/clDNN/tests/test_cases/set_output_memory_gpu_test.cpp @@ -327,7 +327,7 @@ TEST(set_output_memory_gpu, basic_opt) { network.set_input_data("input1", input1); network.set_input_data("input2", input2); network.set_output_memory(outputID, output_mem); - + auto outputs = network.execute(); auto output = outputs.at(outputID).get_memory(); // check for correct output memory setting @@ -349,14 +349,14 @@ TEST(set_output_memory_gpu, mutable_output_data) { auto& engine = get_test_engine(); const int top_k = 2; auto input = engine.allocate_memory({ data_types::f32, format::bfyx,{ batch_num, feature_num, x_size , y_size } }); - auto second_input = engine.allocate_memory({ data_types::f32, format::bfyx, { top_k, feature_num, x_size , y_size } }); - auto final_output = engine.allocate_memory({ data_types::f32, format::bfyx,{ 1, 1, 1 , 1 } }); + auto final_output = engine.allocate_memory({ data_types::f32, format::bfyx, { top_k, feature_num, x_size , y_size } }); + auto second_input = engine.allocate_memory({ data_types::f32, format::bfyx,{ 1, 1, 1 , 1 } }); topology topology; topology.add(input_layout("Add_1396", input->get_layout())); topology.add(cldnn::mutable_data("second_input", second_input)); topology.add(cldnn::mutable_data("12220_md_write", final_output)); - topology.add(arg_max_min("arg_max", { "Add_1396", "12220_md_write", "second_input" }, arg_max_min::min, top_k, arg_max_min::batch)); + topology.add(arg_max_min("arg_max", { "Add_1396", "second_input", "12220_md_write" }, arg_max_min::min, top_k, arg_max_min::batch)); topology.add(cldnn::mutable_data("pred/sink_port_0", {"arg_max"},final_output) ); std::vector input_vec = {