Avoids the creation of two "default" GpuExecutor instances

PiperOrigin-RevId: 703430623
google-ai-edge · Dec 6, 2024 · 5fbb13d · 5fbb13d
1 parent 29e87ad
commit 5fbb13d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/mediapipe/gpu/gpu_shared_data_internal.cc b/mediapipe/gpu/gpu_shared_data_internal.cc
@@ -42,6 +42,17 @@
 
 namespace mediapipe {
 
+namespace {
+
+inline constexpr char kGpuExecutorName[] = "__gpu";
+
+// Returns the executor name from a context key .
+std::string GetExecutorNameFromContextKey(const std::string& context_key) {
+  return absl::StrCat(kGpuExecutorName, "_", context_key);
+}
+
+}  // namespace
+
 #if __APPLE__
 static constexpr bool kGlContextUseDedicatedThread = false;
 #elif defined(__EMSCRIPTEN__)
@@ -130,7 +141,9 @@ GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context,
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 {
   gl_key_context_->insert({SharedContextKey(), gl_context});
-  named_executors_[kGpuExecutorName] =
+  const std::string executor_name =
+      GetExecutorNameFromContextKey(SharedContextKey());
+  named_executors_[executor_name] =
       std::make_shared<GlContextExecutor>(gl_context.get());
 #if __APPLE__
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
@@ -144,7 +157,9 @@ GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context,
 
 absl::StatusOr<std::shared_ptr<Executor>> GpuResources::GetDefaultGpuExecutor()
     const {
-  const auto it = named_executors_.find(kGpuExecutorName);
+  const std::string executor_name =
+      GetExecutorNameFromContextKey(SharedContextKey());
+  const auto it = named_executors_.find(executor_name);
   RET_CHECK(it != named_executors_.end()) << "Can't find default gpu executor.";
   return it->second;
 }
@@ -209,8 +224,8 @@ absl::Status GpuResources::PrepareGpuNode(CalculatorNode* node) {
                       GetOrCreateGlContext(context_key));
 
   if (kGlContextUseDedicatedThread) {
-    std::string executor_name =
-        absl::StrCat(kGpuExecutorName, "_", context_key);
+    const std::string executor_name =
+        GetExecutorNameFromContextKey(context_key);
     node->SetExecutor(executor_name);
     if (!ContainsKey(named_executors_, executor_name)) {
       named_executors_.emplace(

diff --git a/mediapipe/gpu/graph_support.h b/mediapipe/gpu/graph_support.h
@@ -20,7 +20,6 @@ namespace mediapipe {
 
 inline constexpr char kGpuSharedTagName[] = "GPU_SHARED";
 inline constexpr char kGpuSharedSidePacketName[] = "gpu_shared";
-inline constexpr char kGpuExecutorName[] = "__gpu";
 
 }  // namespace mediapipe