No public description

PiperOrigin-RevId: 584349220

mediapipe/gpu/BUILD

@@ -669,6 +669,8 @@ cc_library(
         "//mediapipe/framework/port:ret_check",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/status",
     ] + select({
         "//conditions:default": [],
         "//mediapipe:apple": [

mediapipe/gpu/gpu_shared_data_internal.cc

@@ -14,10 +14,14 @@
 
 #include "mediapipe/gpu/gpu_shared_data_internal.h"
 
+#include <memory>
+#include <utility>
+
 #include "absl/base/attributes.h"
 #include "absl/log/absl_check.h"
+#include "absl/log/absl_log.h"
+#include "absl/status/status.h"
 #include "mediapipe/framework/deps/no_destructor.h"
-#include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/gpu/gl_context.h"
 #include "mediapipe/gpu/gl_context_options.pb.h"
 #include "mediapipe/gpu/graph_support.h"
@@ -83,8 +87,25 @@ GpuResources::StatusOrGpuResources GpuResources::Create(
 }
 
 GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context)
+    : gl_key_context_(new GlContextMapType(),
+                      [](auto* map) {
+                        // This flushes all pending jobs in all GL contexts,
+                        // ensuring that all GL contexts not referenced
+                        // elsewhere are destroyed as part of this destructor.
+                        // Failure to do this may cause GL threads to outlast
+                        // this destructor and execute jobs after the
+                        // GpuResources object is destroyed.
+                        for (auto& [key, context] : *map) {
+                          const auto status = std::move(context)->Run(
+                              []() { return absl::OkStatus(); });
+                          ABSL_LOG_IF(ERROR, !status.ok())
+                              << "Failed to flush GlContext jobs: " << status;
+                        }
+                        delete map;
+                      })
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-    : texture_caches_(std::make_shared<CvTextureCacheManager>()),
+      ,
+      texture_caches_(std::make_shared<CvTextureCacheManager>()),
       gpu_buffer_pool_(
           [tc = texture_caches_](const internal::GpuBufferSpec& spec,
                                  const MultiPoolOptions& options) {
@@ -92,7 +113,7 @@ GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context)
           })
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 {
-  gl_key_context_[SharedContextKey()] = gl_context;
+  gl_key_context_->insert({SharedContextKey(), gl_context});
   named_executors_[kGpuExecutorName] =
       std::make_shared<GlContextExecutor>(gl_context.get());
 #if __APPLE__
@@ -104,6 +125,15 @@ GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context)
 }
 
 GpuResources::~GpuResources() {
+  // This flushes all pending jobs in all GL contexts,
+  // ensuring that all existing jobs, which may refer GpuResource and kept their
+  // gpu resources (e.g. GpuResources::gpu_buffer_pool_) through a raw pointer,
+  // have finished before kept gpu resources get deleted.
+  for (auto& [key, context] : *gl_key_context_) {
+    const auto status = context->Run([]() { return absl::OkStatus(); });
+    ABSL_LOG_IF(ERROR, !status.ok())
+        << "Failed to flush GlContext jobs: " << status;
+  }
 #if __APPLE__
   // Note: on Apple platforms, this object contains Objective-C objects.
   // The destructor will release them, but ARC must be on.
@@ -111,7 +141,7 @@ GpuResources::~GpuResources() {
 #error This file must be built with ARC.
 #endif
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-  for (auto& kv : gl_key_context_) {
+  for (auto& kv : *gl_key_context_) {
     texture_caches_->UnregisterTextureCache(kv.second->cv_texture_cache());
   }
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
@@ -173,23 +203,24 @@ absl::Status GpuResources::PrepareGpuNode(CalculatorNode* node) {
 const std::shared_ptr<GlContext>& GpuResources::gl_context(
     CalculatorContext* cc) {
   if (cc) {
-    auto it = gl_key_context_.find(node_key_[cc->NodeName()]);
+    auto it = gl_key_context_->find(node_key_[cc->NodeName()]);
-    if (it != gl_key_context_.end()) {
+    if (it != gl_key_context_->end()) {
       return it->second;
     }
   }
 
-  return gl_key_context_[SharedContextKey()];
+  return gl_key_context_->at(SharedContextKey());
 }
 
 GlContext::StatusOrGlContext GpuResources::GetOrCreateGlContext(
     const std::string& key) {
-  auto it = gl_key_context_.find(key);
+  auto it = gl_key_context_->find(key);
-  if (it == gl_key_context_.end()) {
+  if (it == gl_key_context_->end()) {
-    MP_ASSIGN_OR_RETURN(std::shared_ptr<GlContext> new_context,
+    MP_ASSIGN_OR_RETURN(
-                        GlContext::Create(*gl_key_context_[SharedContextKey()],
+        std::shared_ptr<GlContext> new_context,
-                                          kGlContextUseDedicatedThread));
+        GlContext::Create(*gl_key_context_->at(SharedContextKey()),
-    it = gl_key_context_.emplace(key, new_context).first;
+                          kGlContextUseDedicatedThread));
+    it = gl_key_context_->emplace(key, new_context).first;
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
     texture_caches_->RegisterTextureCache(it->second->cv_texture_cache());
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER

mediapipe/gpu/gpu_shared_data_internal.h

@@ -21,6 +21,8 @@
 #ifndef MEDIAPIPE_GPU_GPU_SHARED_DATA_INTERNAL_H_
 #define MEDIAPIPE_GPU_GPU_SHARED_DATA_INTERNAL_H_
 
+#include <memory>
+
 #include "mediapipe/framework/calculator_context.h"
 #include "mediapipe/framework/calculator_node.h"
 #include "mediapipe/framework/executor.h"
@@ -82,7 +84,10 @@ class GpuResources {
   const std::string& ContextKey(const std::string& canonical_node_name);
 
   std::map<std::string, std::string> node_key_;
-  std::map<std::string, std::shared_ptr<GlContext>> gl_key_context_;
+
+  using GlContextMapType = std::map<std::string, std::shared_ptr<GlContext>>;
+  std::unique_ptr<GlContextMapType, void (*)(GlContextMapType*)>
+      gl_key_context_;
 
 #ifdef MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   std::shared_ptr<CvTextureCacheManager> texture_caches_;