Make GpuBuffer a shared_ptr to a storage collection

PiperOrigin-RevId: 493519590

mediapipe/gpu/BUILD

@@ -289,7 +289,9 @@ cc_library(
     deps = [
         ":gpu_buffer_format",
         ":gpu_buffer_storage",
+        "@com_google_absl//absl/functional:bind_front",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
         "//mediapipe/framework/formats:image_frame",
         "//mediapipe/framework/port:logging",
         ":gpu_buffer_storage_image_frame",

mediapipe/gpu/gpu_buffer.cc

@@ -3,6 +3,7 @@
 #include <memory>
 #include <utility>
 
+#include "absl/functional/bind_front.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "mediapipe/framework/port/logging.h"
@@ -25,19 +26,27 @@ struct StorageTypeFormatter {
 }  // namespace
 
 std::string GpuBuffer::DebugString() const {
-  return absl::StrCat("GpuBuffer[",
+  return holder_ ? absl::StrCat("GpuBuffer[", width(), "x", height(), " ",
-                      absl::StrJoin(storages_, ", ", StorageTypeFormatter()),
+                                format(), " as ", holder_->DebugString(), "]")
-                      "]");
+                 : "GpuBuffer[invalid]";
 }
 
-internal::GpuBufferStorage* GpuBuffer::GetStorageForView(
+std::string GpuBuffer::StorageHolder::DebugString() const {
-    TypeId view_provider_type, bool for_writing) const {
+  absl::MutexLock lock(&mutex_);
-  const std::shared_ptr<internal::GpuBufferStorage>* chosen_storage = nullptr;
+  return absl::StrJoin(storages_, ", ", StorageTypeFormatter());
+}
 
+internal::GpuBufferStorage* GpuBuffer::StorageHolder::GetStorageForView(
+    TypeId view_provider_type, bool for_writing) const {
+  std::shared_ptr<internal::GpuBufferStorage> chosen_storage;
+  std::function<std::shared_ptr<internal::GpuBufferStorage>()> conversion;
+
+  {
+    absl::MutexLock lock(&mutex_);
     // First see if any current storage supports the view.
     for (const auto& s : storages_) {
       if (s->can_down_cast_to(view_provider_type)) {
-      chosen_storage = &s;
+        chosen_storage = s;
         break;
       }
     }
@@ -49,33 +58,69 @@ internal::GpuBufferStorage* GpuBuffer::GetStorageForView(
         if (auto converter = internal::GpuBufferStorageRegistry::Get()
                                  .StorageConverterForViewProvider(
                                      view_provider_type, s->storage_type())) {
-        if (auto new_storage = converter(s)) {
+          conversion = absl::bind_front(converter, s);
-          storages_.push_back(new_storage);
-          chosen_storage = &storages_.back();
           break;
         }
       }
     }
   }
 
+  // Avoid invoking a converter or factory while holding the mutex.
+  // Two reasons:
+  // 1. Readers that don't need a conversion will not be blocked.
+  // 2. We use mutexes to make sure GL contexts are not used simultaneously on
+  //    different threads, and we also rely on Mutex's deadlock detection
+  //    heuristic, which enforces a consistent mutex acquisition order.
+  //    This function is likely to be called within a GL context, and the
+  //    conversion function may in turn use a GL context, and this may cause a
+  //    false positive in the deadlock detector.
+  //    TODO: we could use Mutex::ForgetDeadlockInfo instead.
+  if (conversion) {
+    auto new_storage = conversion();
+    absl::MutexLock lock(&mutex_);
+    // Another reader might have already completed and inserted the same
+    // conversion. TODO: prevent this?
+    for (const auto& s : storages_) {
+      if (s->can_down_cast_to(view_provider_type)) {
+        chosen_storage = s;
+        break;
+      }
+    }
+    if (!chosen_storage) {
+      storages_.push_back(std::move(new_storage));
+      chosen_storage = storages_.back();
+    }
+  }
+
   if (for_writing) {
+    // This will temporarily hold storages to be released, and do so while the
+    // lock is not held (see above).
+    decltype(storages_) old_storages;
+    using std::swap;
     if (chosen_storage) {
       // Discard all other storages.
-      storages_ = {*chosen_storage};
+      absl::MutexLock lock(&mutex_);
-      chosen_storage = &storages_.back();
+      swap(old_storages, storages_);
+      storages_ = {chosen_storage};
     } else {
       // Allocate a new storage supporting the requested view.
       if (auto factory =
               internal::GpuBufferStorageRegistry::Get()
                   .StorageFactoryForViewProvider(view_provider_type)) {
-        if (auto new_storage = factory(width(), height(), format())) {
+        if (auto new_storage = factory(width_, height_, format_)) {
+          absl::MutexLock lock(&mutex_);
+          swap(old_storages, storages_);
           storages_ = {std::move(new_storage)};
-          chosen_storage = &storages_.back();
+          chosen_storage = storages_.back();
         }
       }
     }
   }
-  return chosen_storage ? chosen_storage->get() : nullptr;
+
+  // It is ok to return a non-owning storage pointer here because this object
+  // ensures the storage's lifetime. Overwriting a GpuBuffer while readers are
+  // active would violate this, but it's not allowed in MediaPipe.
+  return chosen_storage ? chosen_storage.get() : nullptr;
 }
 
 internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
@@ -84,8 +129,7 @@ internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
       GpuBuffer::GetStorageForView(view_provider_type, for_writing);
   CHECK(chosen_storage) << "no view provider found for requested view "
                         << view_provider_type.name() << "; storages available: "
-                        << absl::StrJoin(storages_, ", ",
+                        << (holder_ ? holder_->DebugString() : "invalid");
-                                         StorageTypeFormatter());
   DCHECK(chosen_storage->can_down_cast_to(view_provider_type));
   return *chosen_storage;
 }

mediapipe/gpu/gpu_buffer.h

@@ -15,9 +15,12 @@
 #ifndef MEDIAPIPE_GPU_GPU_BUFFER_H_
 #define MEDIAPIPE_GPU_GPU_BUFFER_H_
 
+#include <algorithm>
+#include <functional>
 #include <memory>
 #include <utility>
 
+#include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/gpu/gpu_buffer_format.h"
 #include "mediapipe/gpu/gpu_buffer_storage.h"
@@ -56,8 +59,7 @@ class GpuBuffer {
   // Creates an empty buffer of a given size and format. It will be allocated
   // when a view is requested.
   GpuBuffer(int width, int height, Format format)
-      : GpuBuffer(std::make_shared<PlaceholderGpuBufferStorage>(width, height,
+      : holder_(std::make_shared<StorageHolder>(width, height, format)) {}
-                                                                format)) {}
 
   // Copy and move constructors and assignment operators are supported.
   GpuBuffer(const GpuBuffer& other) = default;
@@ -70,9 +72,8 @@ class GpuBuffer {
   // are not portable. Applications and calculators should normally obtain
   // GpuBuffers in a portable way from the framework, e.g. using
   // GpuBufferMultiPool.
-  explicit GpuBuffer(std::shared_ptr<internal::GpuBufferStorage> storage) {
+  explicit GpuBuffer(std::shared_ptr<internal::GpuBufferStorage> storage)
-    storages_.push_back(std::move(storage));
+      : holder_(std::make_shared<StorageHolder>(std::move(storage))) {}
-  }
 
 #if !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   // This is used to support backward-compatible construction of GpuBuffer from
@@ -84,9 +85,11 @@ class GpuBuffer {
       : GpuBuffer(internal::AsGpuBufferStorage(storage_convertible)) {}
 #endif  // !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 
-  int width() const { return current_storage().width(); }
+  int width() const { return holder_ ? holder_->width() : 0; }
-  int height() const { return current_storage().height(); }
+  int height() const { return holder_ ? holder_->height() : 0; }
-  GpuBufferFormat format() const { return current_storage().format(); }
+  GpuBufferFormat format() const {
+    return holder_ ? holder_->format() : GpuBufferFormat::kUnknown;
+  }
 
   // Converts to true iff valid.
   explicit operator bool() const { return operator!=(nullptr); }
@@ -122,31 +125,17 @@ class GpuBuffer {
   // using views.
   template <class T>
   std::shared_ptr<T> internal_storage() const {
-    for (const auto& s : storages_)
+    return holder_ ? holder_->internal_storage<T>() : nullptr;
-      if (s->down_cast<T>()) return std::static_pointer_cast<T>(s);
-    return nullptr;
   }
 
   std::string DebugString() const;
 
  private:
-  class PlaceholderGpuBufferStorage
-      : public internal::GpuBufferStorageImpl<PlaceholderGpuBufferStorage> {
-   public:
-    PlaceholderGpuBufferStorage(int width, int height, Format format)
-        : width_(width), height_(height), format_(format) {}
-    int width() const override { return width_; }
-    int height() const override { return height_; }
-    GpuBufferFormat format() const override { return format_; }
-
-   private:
-    int width_ = 0;
-    int height_ = 0;
-    GpuBufferFormat format_ = GpuBufferFormat::kUnknown;
-  };
-
   internal::GpuBufferStorage* GetStorageForView(TypeId view_provider_type,
-                                                bool for_writing) const;
+                                                bool for_writing) const {
+    return holder_ ? holder_->GetStorageForView(view_provider_type, for_writing)
+                   : nullptr;
+  }
 
   internal::GpuBufferStorage& GetStorageForViewOrDie(TypeId view_provider_type,
                                                      bool for_writing) const;
@@ -158,25 +147,49 @@ class GpuBuffer {
                 .template down_cast<VP>();
   }
 
-  std::shared_ptr<internal::GpuBufferStorage>& no_storage() const {
+  // This class manages a set of alternative storages for the contents of a
-    static auto placeholder =
+  // GpuBuffer. GpuBuffer was originally designed as a reference-type object,
-        std::static_pointer_cast<internal::GpuBufferStorage>(
+  // where a copy represents another reference to the same contents, so multiple
-            std::make_shared<PlaceholderGpuBufferStorage>(
+  // GpuBuffer instances can share the same StorageHolder.
-                0, 0, GpuBufferFormat::kUnknown));
+  class StorageHolder {
-    return placeholder;
+   public:
+    explicit StorageHolder(std::shared_ptr<internal::GpuBufferStorage> storage)
+        : StorageHolder(storage->width(), storage->height(),
+                        storage->format()) {
+      storages_.push_back(std::move(storage));
+    }
+    explicit StorageHolder(int width, int height, Format format)
+        : width_(width), height_(height), format_(format) {}
+
+    int width() const { return width_; }
+    int height() const { return height_; }
+    GpuBufferFormat format() const { return format_; }
+
+    internal::GpuBufferStorage* GetStorageForView(TypeId view_provider_type,
+                                                  bool for_writing) const;
+
+    template <class T>
+    std::shared_ptr<T> internal_storage() const {
+      absl::MutexLock lock(&mutex_);
+      for (const auto& s : storages_)
+        if (s->down_cast<T>()) return std::static_pointer_cast<T>(s);
+      return nullptr;
     }
 
-  const internal::GpuBufferStorage& current_storage() const {
+    std::string DebugString() const;
-    return storages_.empty() ? *no_storage() : *storages_[0];
-  }
-
-  internal::GpuBufferStorage& current_storage() {
-    return storages_.empty() ? *no_storage() : *storages_[0];
-  }
 
+   private:
+    int width_ = 0;
+    int height_ = 0;
+    GpuBufferFormat format_ = GpuBufferFormat::kUnknown;
     // This is mutable because view methods that do not change the contents may
     // still need to allocate new storages.
-  mutable std::vector<std::shared_ptr<internal::GpuBufferStorage>> storages_;
+    mutable absl::Mutex mutex_;
+    mutable std::vector<std::shared_ptr<internal::GpuBufferStorage>> storages_
+        ABSL_GUARDED_BY(mutex_);
+  };
+
+  std::shared_ptr<StorageHolder> holder_;
 
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   friend CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer);
@@ -184,15 +197,15 @@ class GpuBuffer {
 };
 
 inline bool GpuBuffer::operator==(std::nullptr_t other) const {
-  return storages_.empty();
+  return holder_ == other;
 }
 
 inline bool GpuBuffer::operator==(const GpuBuffer& other) const {
-  return storages_ == other.storages_;
+  return holder_ == other.holder_;
 }
 
 inline GpuBuffer& GpuBuffer::operator=(std::nullptr_t other) {
-  storages_.clear();
+  holder_ = other;
   return *this;
 }
 

mediapipe/gpu/gpu_buffer_test.cc

@@ -20,6 +20,7 @@
 #include "mediapipe/framework/port/gmock.h"
 #include "mediapipe/framework/port/gtest.h"
 #include "mediapipe/framework/tool/test_util.h"
+#include "mediapipe/gpu/gl_texture_buffer.h"
 #include "mediapipe/gpu/gl_texture_util.h"
 #include "mediapipe/gpu/gpu_buffer_storage_ahwb.h"
 #include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
@@ -228,5 +229,26 @@ TEST_F(GpuBufferTest, GlTextureViewRetainsWhatItNeeds) {
   EXPECT_TRUE(true);
 }
 
+TEST_F(GpuBufferTest, CopiesShareConversions) {
+  GpuBuffer buffer(300, 200, GpuBufferFormat::kBGRA32);
+  {
+    std::shared_ptr<ImageFrame> view = buffer.GetWriteView<ImageFrame>();
+    FillImageFrameRGBA(*view, 255, 0, 0, 255);
+  }
+
+  GpuBuffer other_handle = buffer;
+  RunInGlContext([&buffer] {
+    TempGlFramebuffer fb;
+    auto view = buffer.GetReadView<GlTextureView>(0);
+  });
+
+  // Check that other_handle also sees the same GlTextureBuffer as buffer.
+  // Note that this is deliberately written so that it still passes on platforms
+  // where we use another storage for GL textures (they will both be null).
+  // TODO: expose more accessors for testing?
+  EXPECT_EQ(other_handle.internal_storage<GlTextureBuffer>(),
+            buffer.internal_storage<GlTextureBuffer>());
+}
+
 }  // anonymous namespace
 }  // namespace mediapipe