[GPU] Update loop inst ids instead of getting from prim

Author: kelvinchoi-intel

src/plugins/intel_gpu/src/graph/impls/common/loop.cpp

@@ -130,17 +130,17 @@ struct loop_impl : typed_primitive_impl<loop> {
         }
 
         body_network->set_shape_predictor(outer_network.get_shape_predictor());
-        OPENVINO_ASSERT(!primitive->num_iteration_id.empty(), "loop operation should have num_iteration_id");
+        OPENVINO_ASSERT(!instance.get_num_iterations_id().empty(), "loop operation should have num_iteration_id");
 
         // shortcut of execution_condition memory in body network
         memory::ptr body_execution_condition_mem = nullptr;
-        if (!primitive->body_execution_condition_id.empty()) {
-            body_execution_condition_mem = body_network->get_primitive(primitive->body_execution_condition_id)->output_memory_ptr();
+        if (!instance.get_condition_id().empty()) {
+            body_execution_condition_mem = body_network->get_primitive(instance.get_condition_id())->output_memory_ptr();
         }
 
         // shortcut of current_iteration memory in body network
-        if (!primitive->body_current_iteration_id.empty()) {
-            memory::ptr body_current_iteration_mem = body_network->get_primitive(primitive->body_current_iteration_id)->output_memory_ptr();
+        if (!instance.get_current_iteration_id().empty()) {
+            memory::ptr body_current_iteration_mem = body_network->get_primitive(instance.get_current_iteration_id())->output_memory_ptr();
             write_scalar_value(body_current_iteration_mem, body_network->get_stream(), 0);
         }
 
@@ -149,11 +149,11 @@ struct loop_impl : typed_primitive_impl<loop> {
 
         // read trip_count from outer network
         int64_t trip_count = -1;
-        if (!primitive->trip_count_id.empty()) {
-            memory::ptr trip_count_mem = outer_network.get_primitive(primitive->trip_count_id)->output_memory_ptr();
+        if (!instance.get_trip_count_id().empty()) {
+            memory::ptr trip_count_mem = outer_network.get_primitive(instance.get_trip_count_id())->output_memory_ptr();
             trip_count = read_scalar_value(std::move(trip_count_mem), stream);
         } else {
-            OPENVINO_ASSERT(!primitive->body_execution_condition_id.empty()
+            OPENVINO_ASSERT(!instance.get_condition_id().empty()
                             || num_iterations > 0 || primitive->max_num_iterations > 0,
                             "num_iterations should be positive when trip_count_id is not existed");
             // If trip_count_id is not existed, the original ngraph operation is TensorIterator.
@@ -166,19 +166,19 @@ struct loop_impl : typed_primitive_impl<loop> {
 
         // read initial execution condition from outer network
         int64_t execution_condition = 1;
-        if (!primitive->first_execution_condition_id.empty()) {
+        if (!instance.get_initial_execution_id().empty()) {
             // Wait for completion of the execution_condition of outer_network
-            if (outer_network.has_event(primitive->first_execution_condition_id))
-                outer_network.get_primitive_event(primitive->first_execution_condition_id)->wait();
-            memory::ptr first_execution_condition_mem = outer_network.get_primitive(primitive->first_execution_condition_id)->output_memory_ptr();
+            if (outer_network.has_event(instance.get_initial_execution_id()))
+                outer_network.get_primitive_event(instance.get_initial_execution_id())->wait();
+            memory::ptr first_execution_condition_mem = outer_network.get_primitive(instance.get_initial_execution_id())->output_memory_ptr();
             execution_condition = read_scalar_value(first_execution_condition_mem, stream);
         }
         GPU_DEBUG_LOG << "execution_condition: " << execution_condition << std::endl;
 
         // When execution_condition is false or trip_count is zero, return execute_impl without any body_network execution.
         if (!execution_condition || trip_count == 0) {
             // Update num_iterations (actual number of iterations)
-            memory::ptr num_actual_iterations_mem = outer_network.get_primitive(primitive->num_iteration_id)->output_memory_ptr();
+            memory::ptr num_actual_iterations_mem = outer_network.get_primitive(instance.get_num_iterations_id())->output_memory_ptr();
             write_scalar_value(num_actual_iterations_mem, stream, current_iteration_idx);
 
             instance.update_output_layout();
@@ -255,7 +255,7 @@ struct loop_impl : typed_primitive_impl<loop> {
 
             // execution condition is the result of body network execution
             if (body_execution_condition_mem != nullptr) {
-                auto execution_id = primitive->body_execution_condition_id;
+                auto execution_id = instance.get_condition_id();
                 if (body_network->has_event(execution_id)) {
                     auto ev = body_network->get_primitive_event(execution_id);
                     if (ev) ev->wait();
@@ -275,9 +275,9 @@ struct loop_impl : typed_primitive_impl<loop> {
 
         // Update actual num iteration
         // update num_iterations (actual number of iterations)
-        memory::ptr num_actual_iterations_mem = outer_network.get_primitive(primitive->num_iteration_id)->output_memory_ptr();
+        memory::ptr num_actual_iterations_mem = outer_network.get_primitive(instance.get_num_iterations_id())->output_memory_ptr();
         write_scalar_value(num_actual_iterations_mem, stream, current_iteration_idx);
-        GPU_DEBUG_LOG << "current_iteration_idx(" << primitive->num_iteration_id << ", "
+        GPU_DEBUG_LOG << "current_iteration_idx(" << instance.get_num_iterations_id() << ", "
                         << num_actual_iterations_mem << ")  : " << current_iteration_idx << std::endl;
 
         if (is_dynamic)

src/plugins/intel_gpu/src/graph/include/loop_inst.h

@@ -22,6 +22,12 @@ struct typed_program_node<loop> : public typed_program_node_base<loop> {
 private:
     using parent = typed_program_node_base<loop>;
 
+    primitive_id trip_count_id;
+    primitive_id initial_execution_id;
+    primitive_id current_iteration_id;
+    primitive_id execution_condition_id;
+    primitive_id num_iterations_id;
+
     std::vector<loop::io_primitive_map>& input_primitive_maps;
     std::vector<loop::io_primitive_map>& output_primitive_maps;
     std::vector<loop::backedge_mapping>& back_edges;
@@ -31,22 +37,34 @@ struct typed_program_node<loop> : public typed_program_node_base<loop> {
         parent(prim, prog),
         input_primitive_maps(prim->input_primitive_maps),
         output_primitive_maps(prim->output_primitive_maps),
-        back_edges(prim->back_edges) {}
+        back_edges(prim->back_edges) {
+        set_primitive_ids(prim);
+    }
 
     program::ptr get_body_program() const { return get_primitive()->body_program; }
 
-    const primitive_id& get_trip_count_id() const { return get_primitive()->trip_count_id; }
-    const primitive_id& get_initial_execution_id() const { return get_primitive()->first_execution_condition_id; }
-    const primitive_id& get_current_iteration_id() const { return get_primitive()->body_current_iteration_id; }
-    const primitive_id& get_execution_condition_id() const { return get_primitive()->body_execution_condition_id; }
-    const primitive_id& get_num_iterations_id() const { return get_primitive()->num_iteration_id; }
+    const primitive_id& get_trip_count_id() const { return trip_count_id; }
+    const primitive_id& get_initial_execution_id() const { return initial_execution_id; }
+    const primitive_id& get_current_iteration_id() const { return current_iteration_id; }
+    const primitive_id& get_execution_condition_id() const { return execution_condition_id; }
+    const primitive_id& get_num_iterations_id() const { return num_iterations_id; }
+
     const int32_t get_max_num_iteration() const { return get_primitive()->max_num_iterations; }
 
     const std::vector<loop::io_primitive_map>& get_input_primitive_maps() const { return input_primitive_maps; }
     const std::vector<loop::io_primitive_map>& get_output_primitive_maps() const { return output_primitive_maps; }
     const std::vector<loop::backedge_mapping>& get_back_edges() const { return back_edges;}
 
+    void set_primitive_ids(std::shared_ptr<loop> prim) {
+        trip_count_id           = prim->trip_count_id;
+        initial_execution_id    = prim->first_execution_condition_id;
+        current_iteration_id    = prim->body_current_iteration_id;
+        execution_condition_id  = prim->body_execution_condition_id;
+        num_iterations_id       = prim->num_iteration_id;
+    }
+
     void update_primitive_map(const primitive_id& prevID, const primitive_id& newID, bool external_id = true) {
+        std::cout << "======= update_primitive_map prevID:" << prevID << ", newID:" << newID << ", external_id:" << external_id << std::endl;
         if (external_id) {
             for (auto& pm : input_primitive_maps) {
                 if (pm.external_id.pid == prevID) {
@@ -78,6 +96,18 @@ struct typed_program_node<loop> : public typed_program_node_base<loop> {
                 }
             }
         }
+
+        // Update ids
+        if (get_trip_count_id() == prevID)
+            trip_count_id = newID;
+        if (get_initial_execution_id() == prevID)
+            initial_execution_id = newID;
+        if (get_current_iteration_id() == prevID)
+            current_iteration_id = newID;
+        if (get_execution_condition_id() == prevID)
+            execution_condition_id = newID;
+        if (get_num_iterations_id() == prevID)
+            num_iterations_id = newID;
     }
 
     // current_iteration is necessary to calculate output layout in dynamic shape
@@ -329,6 +359,12 @@ class typed_primitive_inst<loop> : public typed_primitive_inst_base<loop> {
     std::vector<event::ptr> preprocess_memory_for_body_network(int64_t current_iteration_idx);
     std::vector<event::ptr> postprocess_memory_for_body_network(int64_t current_iteration_idx);
 
+    primitive_id get_trip_count_id()        { return _trip_count_id; };
+    primitive_id get_initial_execution_id() { return _initial_execution_id; };
+    primitive_id get_current_iteration_id() { return _current_iteration_id; };
+    primitive_id get_condition_id()         { return _condition_id; };
+    primitive_id get_num_iterations_id()    { return _num_iterations_id; };
+
 private:
     network::ptr body_network;
     memory::ptr get_external_memory(const primitive_id& external_id, size_t mem_idx = 0) const;

src/plugins/intel_gpu/tests/unit/test_cases/loop_gpu_test.cpp

@@ -787,3 +787,182 @@ TEST(loop_gpu, support_loop_w_dynamic_input_w_various_shapes) {
         std::vector<float>(),
         2, 3);
 }
+
+static void test_loop_gpu_wo_trip_count_update_primitive_id(ov::PartialShape body_input_layout,
+                                        std::vector<ov::PartialShape> whole_layouts,
+                                        std::vector<std::vector<float>> input_data_list,
+                                        std::vector<float> expected_output_data,
+                                        size_t axis,
+                                        size_t exit_value,
+                                        bool is_caching_test = false) {
+    auto& engine = get_test_engine();
+
+    auto b_input_layout = cldnn::layout{ body_input_layout, data_types::f32, format::bfyx };
+
+    ov::PartialShape sliced_input_shape = body_input_layout;
+    sliced_input_shape[axis] = 1;
+    auto sliced_input_layout = cldnn::layout{ sliced_input_shape, data_types::f32, format::bfyx };
+
+    auto const_layout = cldnn::layout{ {}, data_types::i64, format::bfyx };
+
+    auto e_initial_condition_mem = engine.allocate_memory(const_layout);
+    auto e_num_iteration_mem = engine.allocate_memory(const_layout);
+    auto b_exit_value_mem = engine.allocate_memory(const_layout);
+    auto b_index_inc_mem = engine.allocate_memory(const_layout);
+    auto init_mean = engine.allocate_memory(const_layout);
+
+    // initialize input buffers
+    set_values(e_initial_condition_mem, {1});
+    set_values(b_exit_value_mem, {exit_value});
+    set_values(b_index_inc_mem, {1});
+    set_values(e_num_iteration_mem, {0});
+    set_values(init_mean, {1});
+
+    primitive_id body_current_iteration_id = "b_index";
+    primitive_id body_execution_condition_id = "b_cond_exit_value";
+
+    cldnn::topology body(
+        input_layout(body_current_iteration_id, const_layout),
+        input_layout("b_add_data", sliced_input_layout),
+        input_layout("b_mul_data", sliced_input_layout),
+        data("b_exit_value", b_exit_value_mem),
+        data("b_index_inc", b_index_inc_mem),
+        eltwise("b_index_update", input_info(body_current_iteration_id), input_info("b_index_inc"), eltwise_mode::sum),
+        reorder("b_index_cast", input_info("b_index_update"),
+                    cldnn::format::any, data_types::f32, {}, cldnn::reorder_mean_mode::subtract, cldnn::padding(), true),
+        eltwise(body_execution_condition_id, input_info("b_index"), input_info("b_exit_value"), eltwise_mode::lt),
+        eltwise("b_add", input_info("b_add_data"), input_info("b_index_cast"), eltwise_mode::sum),
+        eltwise("b_mul", input_info("b_mul_data"), input_info("b_index_cast"), eltwise_mode::prod));
+
+    primitive_id trip_count_id = "";
+    primitive_id actual_iteration_count_id = "actual_iteration_count";
+    primitive_id initial_mean = "initial_mean";
+
+    primitive_id initial_condition_id = "initial_condition";
+    primitive_id initial_condition_id_elt = "initial_condition_elt";
+    primitive_id initial_condition_id_reorder = "initial_condition_reorder";
+    primitive_id initial_condition_id_reorder2 = "initial_condition_reorder2";
+    int64_t num_iterations = -1;
+
+    std::vector<loop::io_primitive_map> input_primitive_maps {
+        loop::io_primitive_map("input", "b_add_data", axis),
+        loop::io_primitive_map("input", "b_mul_data", axis),
+        loop::io_primitive_map(actual_iteration_count_id, body_current_iteration_id) };
+    std::vector<loop::io_primitive_map> output_primitive_maps {
+        loop::io_primitive_map(cldnn::input_info("loop", 0), cldnn::input_info("b_add", 0), axis),
+        loop::io_primitive_map(cldnn::input_info("loop", 1), cldnn::input_info("b_mul", 0), axis) };
+    std::vector<loop::backedge_mapping> back_edges {
+        loop::backedge_mapping("b_index_update", body_current_iteration_id) };
+
+    auto body_program = build_program(engine, body, body_execution_condition_id, output_primitive_maps, back_edges, true);
+
+    auto const_shape = engine.allocate_memory({ov::PartialShape{4}, data_types::i32, format::bfyx});
+    
+
+    std::vector<int32_t> body_input_layouts;
+    for (size_t i = 0; i < body_input_layout.size(); i++) {
+        if (body_input_layout[i].is_dynamic())
+            body_input_layouts.push_back(-1);
+        else
+            body_input_layouts.push_back(body_input_layout[i].get_length());
+    }
+    set_values<int32_t>(const_shape, body_input_layouts);
+    
+
+    cldnn::topology topology(
+        input_layout("input_origin", b_input_layout),
+        input_layout(initial_condition_id, e_initial_condition_mem->get_layout()),
+        mutable_data(actual_iteration_count_id, e_num_iteration_mem),
+        data(initial_mean, init_mean),
+
+        reorder(initial_condition_id_reorder, input_info(initial_condition_id), cldnn::format::any, data_types::f32/*, initial_mean*/),
+        reorder(initial_condition_id_reorder2, input_info(initial_condition_id_reorder), cldnn::format::any, data_types::i32),  // should be fused to test updating input id of loop
+
+        shape_of("shape_of_input", input_info("input_origin"), data_types::i32),
+        reduce("reduced_shape", input_info("shape_of_input"), reduce_mode::prod, {0}, true),
+        reshape("reshape1", input_info("input_origin"), input_info("reduced_shape"), false, ov::PartialShape::dynamic(1)),
+        data("const", const_shape),
+        reshape("input", input_info("reshape1"), input_info("const"), false, ov::PartialShape::dynamic(4)),
+
+        loop("loop", { input_info(actual_iteration_count_id), input_info(initial_condition_id_reorder2), input_info("input") }, body_program,
+             trip_count_id, initial_condition_id_reorder2, actual_iteration_count_id,
+             input_primitive_maps, output_primitive_maps, back_edges,
+             num_iterations, body_current_iteration_id, body_execution_condition_id, 2),
+        eltwise("out_sum", input_info("loop", 0), input_info("loop", 1), eltwise_mode::sum));
+
+    ExecutionConfig config = get_test_default_config(engine);
+    config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
+
+    cldnn::network::ptr network = get_network(engine, topology, config, get_test_stream_ptr(), is_caching_test);
+
+    for (size_t i = 0 ; i < whole_layouts.size(); i++) {
+        auto whole_layout = whole_layouts[i];
+        auto input_data = input_data_list[i];
+
+        // initialize input buffers
+        set_values(e_initial_condition_mem, {1});
+        set_values(b_exit_value_mem, {exit_value});
+        set_values(b_index_inc_mem, {1});
+        set_values(e_num_iteration_mem, {0});
+
+        auto e_input_layout = cldnn::layout{ whole_layout, data_types::f32, format::bfyx };
+        auto e_input_mem = engine.allocate_memory(e_input_layout); // b,f,x,y
+        auto expected_output_layout = whole_layout;
+        set_values(e_input_mem, input_data);
+        network->set_input_data("input_origin", e_input_mem);
+
+        network->set_input_data(initial_condition_id, e_initial_condition_mem);
+
+        auto outputs = network->execute();
+        ASSERT_EQ(outputs.size(), 2);
+
+        auto expected_num_iterations = (exit_value + 1);
+        expected_output_layout[axis] = expected_num_iterations;
+        auto e_output_layout = cldnn::layout{ expected_output_layout, data_types::f32, format::bfyx };
+
+        auto num_iter_mem = network->get_output_memory(actual_iteration_count_id);
+        if (num_iter_mem != nullptr) {
+            mem_lock<int64_t> num_iter_ptr{ num_iter_mem, get_test_stream() };
+            ASSERT_EQ(num_iter_ptr.data()[0], expected_num_iterations);
+        }
+
+        std::vector<float> expected(input_data.size());
+        if (expected_output_data.size() == 0) {
+            size_t unit = 1;
+            for (size_t k = axis; k < whole_layout.size(); k++) {
+                unit *= whole_layout[k].get_length();
+            }
+
+            for (size_t j = 0; j < input_data.size(); j++) {
+                auto val = static_cast<size_t>((j % unit) / 4) + 1;
+                expected[j] = static_cast<float>(input_data[j] + val) + static_cast<float>(input_data[j] * val);
+            }
+        } else {
+            expected = expected_output_data;
+        }
+
+        auto output_mem = outputs.begin()->second.get_memory();
+        auto output_layout = output_mem->get_layout();
+        ASSERT_EQ(output_layout.batch(), e_output_layout.batch());
+        ASSERT_EQ(output_layout.feature(), e_output_layout.feature());
+        ASSERT_EQ(output_layout.spatial(0), e_output_layout.spatial(0));
+        ASSERT_EQ(output_layout.spatial(1), e_output_layout.spatial(1));
+        // value check
+        {
+            mem_lock<float> output_ptr{ output_mem, get_test_stream() };
+            for (size_t i = 0, iend = output_layout.count(); i < iend; ++i) {
+                ASSERT_FLOAT_EQ(output_ptr[i], expected.at(i));
+            }
+        }
+    }
+}
+
+
+TEST(loop_gpu, support_loop_w_dynamic_input_update_primitive_id) {
+    test_loop_gpu_wo_trip_count_update_primitive_id(
+        { 1, -1, 4, 4 },
+        {{ 1, 1, 4, 4 }},   // axis value should be iter_num = (exit_value + 1)
+        {input_data_4_4, input_data_2_4_4},
+        std::vector<float>(),
+        2, 3);
+}