[GPU][PoC] shape_info common fix

openvinotoolkit · Dec 20, 2024 · 498b793 · 498b793
1 parent 940b0cb
commit 498b793
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/plugins/intel_gpu/include/intel_gpu/graph/network.hpp b/src/plugins/intel_gpu/include/intel_gpu/graph/network.hpp
@@ -216,6 +216,7 @@ struct network {
     bool _is_dynamic = false;
     bool _enable_profiling = false;
     bool _reset_arguments;
+    memory::ptr _ptr;
 
     std::unordered_map<primitive_id, std::shared_ptr<primitive_inst>> _primitives;
     std::vector<shared_mem_type> _in_out_shared_mem_types;

diff --git a/src/plugins/intel_gpu/src/graph/network.cpp b/src/plugins/intel_gpu/src/graph/network.cpp
@@ -289,19 +289,18 @@ void network::preallocate_shape_info_buffers() {
     }
 
     std::cout << "Sum of shape elements " << sum << std::endl;
+
     if(sum) {
         auto& eng = get_engine();
-        auto ptr = eng.allocate_memory(layout{{sum * 128}, data_types::i32, format::bfyx}, false);
-        //auto ptr = eng.allocate_memory(layout{{sum}, data_types::i32, format::bfyx}, false);
+        _ptr = eng.allocate_memory(layout{{sum}, data_types::i32, format::bfyx}, false);
         int new_sum = 0;
         for (auto const& prim : _exec_order) {
             auto& node = prim->get_node();
             int64_t shape_elements = node.get_total_shape_info_size();
             if(shape_elements == 0) continue;
-	    auto new_ptr = ptr;
-	    auto new_mem = eng.reinterpret_subbuffer(*ptr, layout{{shape_elements}, data_types::i32, format::bfyx}, new_sum);
+	    auto new_mem = eng.reinterpret_subbuffer(*_ptr, layout{{shape_elements}, data_types::i32, format::bfyx}, new_sum);
             prim->set_shape_info_memory_ptr(new_mem);
-            new_sum += 4096;
+            new_sum += shape_elements * 4;
         }
     }
 }