Avoids the sharing of GL contexts between nested mediapipe graphs.

PiperOrigin-RevId: 730837827

Author: MediaPipe Team

mediapipe/framework/BUILD

@@ -376,6 +376,7 @@ cc_library(
         "//mediapipe/framework/tool:validate",
         "//mediapipe/framework/tool:validate_name",
         "//mediapipe/gpu:gpu_service",
+        "//mediapipe/gpu:gpu_shared_data_internal",
         "//mediapipe/gpu:graph_support",
         "//mediapipe/util:cpu_util",
         "@com_google_absl//absl/base:core_headers",

mediapipe/framework/calculator_graph.cc

@@ -80,6 +80,7 @@
 #include "mediapipe/framework/validated_graph_config.h"
 #include "mediapipe/framework/vlog_overrides.h"
 #include "mediapipe/gpu/gpu_service.h"
+#include "mediapipe/gpu/gpu_shared_data_internal.h"
 #include "mediapipe/gpu/graph_support.h"
 #include "mediapipe/util/cpu_util.h"
 
@@ -139,14 +140,34 @@ CalculatorGraph::CalculatorGraph() : CalculatorGraph(/*cc=*/nullptr) {}
 // Adopt all services from the CalculatorContext / parent graph.
 CalculatorGraph::CalculatorGraph(CalculatorContext* cc)
     : counter_factory_(std::make_unique<BasicCounterFactory>()),
-      service_manager_(cc != nullptr ? cc->GetGraphServiceManager() : nullptr),
       profiler_(std::make_shared<ProfilingContext>()),
       scheduler_(this) {
   if (cc != nullptr) {
     // Nested graphs should not create default initialized services to avoid
     // collisions between newly created and inherited graphs.
     // TODO b/368015341- Use factory method to avoid CHECK in constructor.
     ABSL_CHECK_OK(DisallowServiceDefaultInitialization());
+
+    // Adopt all services from the parent graph except for GpuResources.
+    const auto parent_service_manager = cc->GetGraphServiceManager();
+    const auto parent_service_packets =
+        parent_service_manager->ServicePackets();
+    GraphServiceManager::ServiceMap service_packets;
+    for (const auto& [key, packet] : parent_service_packets) {
+      if (key == kGpuService.key) {
+        // To avoid deadlocks when sharing the same GPU thread between
+        // multiple graphs, we create a new GpuResources instance for each
+        // sub-graph with a dedicated GL context / thread.
+        auto resources = mediapipe::GpuResources::Create(
+            *packet.Get<std::shared_ptr<mediapipe::GpuResources>>());
+        ABSL_CHECK_OK(resources);
+        service_packets[key] =
+            MakePacket<std::shared_ptr<mediapipe::GpuResources>>(*resources);
+      } else {
+        service_packets[key] = packet;
+      }
+    }
+    service_manager_.SetServicePackets(service_packets);
   }
   SetVLogOverrides();
 }

mediapipe/framework/graph_service_manager.h

@@ -4,6 +4,7 @@
 #include <map>
 #include <memory>
 #include <string>
+#include <utility>
 
 #include "absl/log/absl_check.h"
 #include "absl/status/status.h"
@@ -15,21 +16,10 @@ namespace mediapipe {
 
 class GraphServiceManager {
  public:
-  GraphServiceManager() = default;
-
-  explicit GraphServiceManager(
-      const GraphServiceManager* external_graph_manager) {
-    if (external_graph_manager != nullptr) {
-      // Nested graphs inherit all graph services from their parent graph and
-      // disable the registration of new services in the nested graph. This
-      // ensures that all services are created during the initialization of
-      // parent graph.
-      service_packets_ = external_graph_manager->ServicePackets();
-    }
-  }
-
   using ServiceMap = std::map<std::string, Packet>;
 
+  GraphServiceManager() = default;
+
   template <typename T>
   absl::Status SetServiceObject(const GraphService<T>& service,
                                 std::shared_ptr<T> object) {
@@ -41,6 +31,10 @@ class GraphServiceManager {
 
   absl::Status SetServicePacket(const GraphServiceBase& service, Packet p);
 
+  void SetServicePackets(const ServiceMap& service_packets) {
+    service_packets_ = service_packets;
+  }
+
   template <typename T>
   std::shared_ptr<T> GetServiceObject(const GraphService<T>& service) const {
     Packet p = GetServicePacket(service);

mediapipe/gpu/gpu_shared_data_internal.cc

@@ -108,6 +108,13 @@ GpuResources::StatusOrGpuResources GpuResources::Create(
   return gpu_resources;
 }
 
+GpuResources::StatusOrGpuResources GpuResources::Create(
+    const GpuResources& gpu_resources,
+    const MultiPoolOptions* gpu_buffer_pool_options) {
+  return Create(gpu_resources.gl_context()->native_context(),
+                gpu_buffer_pool_options);
+}
+
 GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context,
                            const MultiPoolOptions* gpu_buffer_pool_options)
     : gl_key_context_(new GlContextMapType(),
@@ -241,9 +248,10 @@ absl::Status GpuResources::PrepareGpuNode(CalculatorNode* node) {
 // TODO: expose and use an actual ID instead of using the
 // canonicalized name.
 const std::shared_ptr<GlContext>& GpuResources::gl_context(
-    CalculatorContext* cc) {
+    CalculatorContext* cc) const {
   if (cc) {
-    auto it = gl_key_context_->find(node_key_[cc->NodeName()]);
+    const auto node_key_it = node_key_.find(cc->NodeName());
+    const auto it = gl_key_context_->find(node_key_it->second);
     if (it != gl_key_context_->end()) {
       return it->second;
     }

mediapipe/gpu/gpu_shared_data_internal.h

@@ -57,6 +57,12 @@ class GpuResources {
       PlatformGlContext external_context,
       const MultiPoolOptions* gpu_buffer_pool_options = nullptr);
 
+  // Creates a GpuResources instance that is shared with the GL context provided
+  // by the gpu_resources argument.
+  static StatusOrGpuResources Create(
+      const GpuResources& gpu_resources,
+      const MultiPoolOptions* gpu_buffer_pool_options = nullptr);
+
   // The destructor must be defined in the implementation file so that on iOS
   // the correct ARC release calls are generated.
   ~GpuResources();
@@ -65,9 +71,11 @@ class GpuResources {
 
   // Shared GL context for calculators.
   // TODO: require passing a context or node identifier.
-  const std::shared_ptr<GlContext>& gl_context() { return gl_context(nullptr); }
+  const std::shared_ptr<GlContext>& gl_context() const {
+    return gl_context(nullptr);
+  }
 
-  const std::shared_ptr<GlContext>& gl_context(CalculatorContext* cc);
+  const std::shared_ptr<GlContext>& gl_context(CalculatorContext* cc) const;
 
   // Shared buffer pool.
   GpuBufferMultiPool& gpu_buffer_pool() { return gpu_buffer_pool_; }