[GPU] Fix premature termination of async compilation task in tests (#19035)

src/plugins/intel_gpu/src/graph/compilation_context.cpp

@@ -7,6 +7,7 @@
 #include <mutex>
 #include <atomic>
 #include <unordered_set>
+#include <future>
 #include "intel_gpu/runtime/utils.hpp"
 
 namespace cldnn {
@@ -20,11 +21,18 @@ public:
         if (_stop_compilation)
             return;
 
+        auto promise = std::make_shared<std::promise<void>>();
+        futures.emplace_back(promise->get_future());
+
         std::lock_guard<std::mutex> lock(_mutex);
         if (_task_keys.find(key) == _task_keys.end()) {
             _task_keys.insert(key);
-            if (_task_executor != nullptr)
-                _task_executor->run(task);
+            if (_task_executor != nullptr) {
+                _task_executor->run([task, promise] {
+                    task();
+                    promise->set_value();
+                });
+            }
         }
     }
 
@@ -60,12 +68,19 @@ public:
         }
     }
 
+    void wait_all() override {
+        for (auto&& future : futures) {
+            future.wait();
+        }
+    }
+
 private:
     ov::threading::IStreamsExecutor::Config _task_executor_config;
     std::shared_ptr<ov::threading::IStreamsExecutor> _task_executor;
     std::mutex _mutex;
     std::unordered_set<kernel_impl_params, kernel_impl_params::Hasher> _task_keys;
     std::atomic_bool _stop_compilation{false};
+    std::vector<std::future<void>> futures;
 };
 
 std::unique_ptr<ICompilationContext> ICompilationContext::create(ov::threading::IStreamsExecutor::Config task_executor_config) {

src/plugins/intel_gpu/src/graph/graph_optimizer/post_optimize_weights.cpp

@@ -44,7 +44,7 @@ void post_optimize_weights::optimize_weights(T& node, program& p) {
         }
         // TODO: To relax current limitation w.r.t the future optimization of weight reorder process
         // In dynamic shape, selected weight format can change in runtime. However reordering blocked format to blocked format is not fully verified yet.
-        // So we need to enable other primiives such as convolution with verifying reorder b/w the possible layouts
+        // So we need to enable other primitives such as convolution with verifying reorder b/w the possible layouts
         // Also we skip weight reorder for onednn impl because onednn fully connected layer is using simple format, therefore
         // reordering to cldnn shape_agnostic_kernel's preferred blocked format at build time does not helpful for the performance.
         // This situation might be changed once onednn shape agnostic kernel is used in the future.

src/plugins/intel_gpu/src/graph/include/compilation_context.hpp

@@ -19,6 +19,7 @@ public:
     virtual ~ICompilationContext() = default;
     virtual bool is_stopped() = 0;
     virtual void cancel() = 0;
+    virtual void wait_all() = 0;
 
     static std::unique_ptr<ICompilationContext> create(ov::threading::IStreamsExecutor::Config task_executor_config);
 };

src/plugins/intel_gpu/src/graph/primitive_inst.cpp

@@ -588,7 +588,8 @@ bool primitive_inst::update_impl() {
             o.data_padding.set_dynamic_pad(tensor(0));
         }
 
-        auto& cache = get_network().get_program()->get_implementations_cache();
+        const auto& prog = get_network().get_program();
+        auto& cache = prog->get_implementations_cache();
         std::shared_ptr<primitive_impl> cached_impl = nullptr;
         {
             cached_impl = cache.get(updated_params_no_dyn_pad);
@@ -604,7 +605,7 @@ bool primitive_inst::update_impl() {
         if (!cached_impl) {
             if (_dynamic_impl) {
                 if (use_async_compilation()) {
-                    auto& compilation_context = get_network().get_program()->get_compilation_context();
+                    auto& compilation_context = prog->get_compilation_context();
                     compilation_context.push_task(updated_params_no_dyn_pad, [this, &compilation_context, updated_params_no_dyn_pad]() {
                         if (compilation_context.is_stopped())
                             return;
@@ -637,7 +638,7 @@ bool primitive_inst::update_impl() {
                 _impl = _node->type()->choose_impl(*_node, updated_params_no_dyn_pad);
                 _impl->set_node_params(*_node);
                 if (!can_be_optimized()) {
-                    auto& kernels_cache = get_network().get_program()->get_kernels_cache();
+                    auto& kernels_cache = prog->get_kernels_cache();
                     auto kernels = kernels_cache.compile(updated_params_no_dyn_pad, _impl->get_kernels_source());
                     _impl->set_kernels(std::move(kernels));
                     cache.add(updated_params_no_dyn_pad, _impl->clone());

src/plugins/intel_gpu/tests/unit/module_tests/kernel_impl_params_relevance_test.cpp

@@ -44,7 +44,7 @@ TEST(kernel_impl_params_relevance, weights_layout) {
     network.set_input_data("input", actual_input_data);
 
     // 2. Force reference `fully_connected_gpu_bfyx_ref` kernel impl before execution,
-    //    so during _node->type()->choose_impl(*_node, updated_params); call for static kernel vesrion reference
+    //    so during _node->type()->choose_impl(*_node, updated_params); call for static kernel version reference
     //    impl will be used. Call execute() to trigger desired kernel compilation
     auto fc_ref_impl = ov::intel_gpu::ImplementationDesc(format::bfyx, "fully_connected_gpu_bfyx_ref", impl_types::ocl);
     auto force_impl_prop = ov::intel_gpu::force_implementations(ov::intel_gpu::ImplForcingMap{ {"fc", fc_ref_impl} });
@@ -52,8 +52,8 @@ TEST(kernel_impl_params_relevance, weights_layout) {
 
     network.execute();
 
-    // 3. WA: Call cancel() to wait for all queued kernels compilation finish (including above `fully_connected_gpu_bfyx_ref`)
-    network.get_program()->get_compilation_context().cancel();
+    // 3. WA: Call wait_all() to wait for all queued kernels compilation finish (including above `fully_connected_gpu_bfyx_ref`)
+    network.get_program()->get_compilation_context().wait_all();
 
     // 4. Call execute() second time with same input shape to use pre-compiled `fully_connected_gpu_bfyx_ref` kernel
     network.execute();

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

@@ -1738,8 +1738,8 @@ TEST(fully_connected_onednn, impl_replacement_with_cldnn) {
     ASSERT_EQ(-2.25f, output_ptr[2]);
     ASSERT_EQ(3.0f, output_ptr[3]);
 
-    // WA: Call cancel() to wait for all queued kernels compilation finish
-    network.get_program()->get_compilation_context().cancel();
+    // WA: Call wait_all() to wait for all queued kernels compilation finish
+    network.get_program()->get_compilation_context().wait_all();
 
     // Check if OneDNN's impl is used for the next execute() call
     network.execute();

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

@@ -1559,8 +1559,8 @@ TEST(gemm_onednn, impl_replacement_with_cldnn) {
         ASSERT_FLOAT_EQ(output_ptr[i], out_data[i]);
     }
 
-    // WA: Call cancel() to wait for all queued kernels compilation finish
-    network.get_program()->get_compilation_context().cancel();
+    // WA: Call wait_all() to wait for all queued kernels compilation finish
+    network.get_program()->get_compilation_context().wait_all();
 
     // Check if OneDNN's impl is used for the next execute() call
     network.execute();