Generic MultiPool template

PiperOrigin-RevId: 488783176

mediapipe/gpu/BUILD

@@ -618,6 +618,7 @@ cc_library(
 cc_library(
     name = "multi_pool",
     hdrs = ["multi_pool.h"],
+    deps = ["//mediapipe/util:resource_cache"],
 )
 
 cc_library(

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc

@@ -71,12 +71,12 @@ std::string CvPixelBufferPoolWrapper::GetDebugString() const {
 void CvPixelBufferPoolWrapper::Flush() { CVPixelBufferPoolFlush(*pool_, 0); }
 
 CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::CreateBufferWithoutPool(
-    int width, int height, GpuBufferFormat format) {
+    const internal::GpuBufferSpec& spec) {
-  OSType cv_format = CVPixelFormatForGpuBufferFormat(format);
+  OSType cv_format = CVPixelFormatForGpuBufferFormat(spec.format);
   CHECK_NE(cv_format, -1) << "unsupported pixel format";
   CVPixelBufferRef buffer;
-  CVReturn err =
+  CVReturn err = CreateCVPixelBufferWithoutPool(spec.width, spec.height,
-      CreateCVPixelBufferWithoutPool(width, height, cv_format, &buffer);
+                                                cv_format, &buffer);
   CHECK(!err) << "Error creating pixel buffer: " << err;
   return MakeCFHolderAdopting(buffer);
 }

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h

@@ -38,11 +38,11 @@ class CvPixelBufferPoolWrapper {
                            CvTextureCacheManager* texture_caches);
 
   static std::shared_ptr<CvPixelBufferPoolWrapper> Create(
-      int width, int height, GpuBufferFormat format,
+      const internal::GpuBufferSpec& spec, const MultiPoolOptions& options,
-      const MultiPoolOptions& options,
       CvTextureCacheManager* texture_caches = nullptr) {
     return std::make_shared<CvPixelBufferPoolWrapper>(
-        width, height, format, options.max_inactive_buffer_age, texture_caches);
+        spec.width, spec.height, spec.format, options.max_inactive_buffer_age,
+        texture_caches);
   }
 
   CFHolder<CVPixelBufferRef> GetBuffer();
@@ -53,7 +53,7 @@ class CvPixelBufferPoolWrapper {
   void Flush();
 
   static CFHolder<CVPixelBufferRef> CreateBufferWithoutPool(
-      int width, int height, GpuBufferFormat format);
+      const internal::GpuBufferSpec& spec);
 
  private:
   CFHolder<CVPixelBufferPoolRef> pool_;

mediapipe/gpu/gl_texture_buffer_pool.h

@@ -42,9 +42,8 @@ class GlTextureBufferPool
   }
 
   static std::shared_ptr<GlTextureBufferPool> Create(
-      int width, int height, GpuBufferFormat format,
+      const internal::GpuBufferSpec& spec, const MultiPoolOptions& options) {
-      const MultiPoolOptions& options) {
+    return Create(spec.width, spec.height, spec.format, options.keep_count);
-    return Create(width, height, format, options.keep_count);
   }
 
   // Obtains a buffers. May either be reused or created anew.
@@ -59,8 +58,8 @@ class GlTextureBufferPool
   std::pair<int, int> GetInUseAndAvailableCounts();
 
   static GlTextureBufferSharedPtr CreateBufferWithoutPool(
-      int width, int height, GpuBufferFormat format) {
+      const internal::GpuBufferSpec& spec) {
-    return GlTextureBuffer::Create(width, height, format);
+    return GlTextureBuffer::Create(spec.width, spec.height, spec.format);
   }
 
  private:

mediapipe/gpu/gpu_buffer_format.h

@@ -153,6 +153,34 @@ inline GpuBufferFormat GpuBufferFormatForCVPixelFormat(OSType format) {
 
 #endif  // __APPLE__
 
+namespace internal {
+
+struct GpuBufferSpec {
+  GpuBufferSpec(int w, int h, GpuBufferFormat f)
+      : width(w), height(h), format(f) {}
+
+  template <typename H>
+  friend H AbslHashValue(H h, const GpuBufferSpec& spec) {
+    return H::combine(std::move(h), spec.width, spec.height,
+                      static_cast<uint32_t>(spec.format));
+  }
+
+  int width;
+  int height;
+  GpuBufferFormat format;
+};
+
+// BufferSpec equality operators
+inline bool operator==(const GpuBufferSpec& lhs, const GpuBufferSpec& rhs) {
+  return lhs.width == rhs.width && lhs.height == rhs.height &&
+         lhs.format == rhs.format;
+}
+inline bool operator!=(const GpuBufferSpec& lhs, const GpuBufferSpec& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+}  // namespace internal
+
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_GPU_GPU_BUFFER_FORMAT_H_

mediapipe/gpu/gpu_buffer_multi_pool.cc

@@ -16,51 +16,7 @@
 
 #include <tuple>
 
-#include "absl/memory/memory.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/port/logging.h"
-#include "mediapipe/gpu/gpu_shared_data_internal.h"
 
-namespace mediapipe {
+namespace mediapipe {}  // namespace mediapipe
-
-std::shared_ptr<GpuBufferMultiPool::SimplePool>
-GpuBufferMultiPool::DefaultMakeSimplePool(
-    const GpuBufferMultiPool::BufferSpec& spec,
-    const MultiPoolOptions& options) {
-  return SimplePool::Create(spec.width, spec.height, spec.format, options);
-}
-
-std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool(
-    const BufferSpec& spec) {
-  std::shared_ptr<SimplePool> pool;
-  std::vector<std::shared_ptr<SimplePool>> evicted;
-  {
-    absl::MutexLock lock(&mutex_);
-    pool =
-        cache_.Lookup(spec, [this](const BufferSpec& spec, int request_count) {
-          return (request_count >= options_.min_requests_before_pool)
-                     ? create_simple_pool_(spec, options_)
-                     : nullptr;
-        });
-    evicted = cache_.Evict(options_.max_pool_count,
-                           options_.request_count_scrub_interval);
-  }
-  // Evicted pools, and their buffers, will be released without holding the
-  // lock.
-  return pool;
-}
-
-GpuBuffer GpuBufferMultiPool::GetBuffer(int width, int height,
-                                        GpuBufferFormat format) {
-  BufferSpec key(width, height, format);
-  std::shared_ptr<SimplePool> pool = RequestPool(key);
-  if (pool) {
-    // Note: we release our multipool lock before accessing the simple pool.
-    return GpuBuffer(pool->GetBuffer());
-  } else {
-    return GpuBuffer(
-        SimplePool::CreateBufferWithoutPool(width, height, format));
-  }
-}
-
-}  // namespace mediapipe

mediapipe/gpu/gpu_buffer_multi_pool.h

@@ -22,15 +22,9 @@
 #ifndef MEDIAPIPE_GPU_GPU_BUFFER_MULTI_POOL_H_
 #define MEDIAPIPE_GPU_GPU_BUFFER_MULTI_POOL_H_
 
-#include "absl/hash/hash.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/gpu/gpu_buffer.h"
 #include "mediapipe/gpu/multi_pool.h"
-#include "mediapipe/util/resource_cache.h"
-
-#ifdef __APPLE__
-#include "mediapipe/gpu/pixel_buffer_pool_util.h"
-#endif  // __APPLE__
 
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 #include "mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h"
@@ -40,77 +34,24 @@
 
 namespace mediapipe {
 
-struct GpuSharedData;
 class CvPixelBufferPoolWrapper;
 
-class GpuBufferMultiPool {
+class GpuBufferMultiPool : public MultiPool<
- public:
-  GpuBufferMultiPool(MultiPoolOptions options = kDefaultMultiPoolOptions)
-      : options_(options) {}
-
-  // Obtains a buffer. May either be reused or created anew.
-  GpuBuffer GetBuffer(int width, int height,
-                      GpuBufferFormat format = GpuBufferFormat::kBGRA32);
-
-  // This class is not intended as part of the public api of this class. It is
-  // public only because it is used as a map key type, and the map
-  // implementation needs access to, e.g., the equality operator.
-  struct BufferSpec {
-    BufferSpec(int w, int h, mediapipe::GpuBufferFormat f)
-        : width(w), height(h), format(f) {}
-
-    template <typename H>
-    friend H AbslHashValue(H h, const BufferSpec& spec) {
-      return H::combine(std::move(h), spec.width, spec.height,
-                        static_cast<uint32_t>(spec.format));
-    }
-
-    int width;
-    int height;
-    mediapipe::GpuBufferFormat format;
-  };
-
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-  using SimplePool = CvPixelBufferPoolWrapper;
+                               CvPixelBufferPoolWrapper,
 #else
-  using SimplePool = GlTextureBufferPool;
+                               GlTextureBufferPool,
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+                               internal::GpuBufferSpec, GpuBuffer> {
+ public:
+  using MultiPool::MultiPool;
 
-  using SimplePoolFactory = std::function<std::shared_ptr<SimplePool>(
+  GpuBuffer GetBuffer(int width, int height,
-      const BufferSpec& spec, const MultiPoolOptions& options)>;
+                      GpuBufferFormat format = GpuBufferFormat::kBGRA32) {
-
+    return Get(internal::GpuBufferSpec(width, height, format));
-  void SetSimplePoolFactory(SimplePoolFactory create_simple_pool) {
-    create_simple_pool_ = create_simple_pool;
   }
-
- private:
-  static std::shared_ptr<SimplePool> DefaultMakeSimplePool(
-      const GpuBufferMultiPool::BufferSpec& spec,
-      const MultiPoolOptions& options);
-
-  // Requests a simple buffer pool for the given spec. This may return nullptr
-  // if we have not yet reached a sufficient number of requests to allocate a
-  // pool, in which case the caller should invoke CreateBufferWithoutPool.
-  std::shared_ptr<SimplePool> RequestPool(const BufferSpec& spec);
-
-  MultiPoolOptions options_;
-  absl::Mutex mutex_;
-  mediapipe::ResourceCache<BufferSpec, std::shared_ptr<SimplePool>> cache_
-      ABSL_GUARDED_BY(mutex_);
-  SimplePoolFactory create_simple_pool_ = DefaultMakeSimplePool;
 };
 
-// BufferSpec equality operators
-inline bool operator==(const GpuBufferMultiPool::BufferSpec& lhs,
-                       const GpuBufferMultiPool::BufferSpec& rhs) {
-  return lhs.width == rhs.width && lhs.height == rhs.height &&
-         lhs.format == rhs.format;
-}
-inline bool operator!=(const GpuBufferMultiPool::BufferSpec& lhs,
-                       const GpuBufferMultiPool::BufferSpec& rhs) {
-  return !operator==(lhs, rhs);
-}
-
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_GPU_GPU_BUFFER_MULTI_POOL_H_

mediapipe/gpu/gpu_shared_data_internal.cc

@@ -80,18 +80,20 @@ GpuResources::StatusOrGpuResources GpuResources::Create(
   return gpu_resources;
 }
 
-GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context) {
+GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context)
+#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+    : texture_caches_(std::make_shared<CvTextureCacheManager>()),
+      gpu_buffer_pool_(
+          [tc = texture_caches_](const internal::GpuBufferSpec& spec,
+                                 const MultiPoolOptions& options) {
+            return CvPixelBufferPoolWrapper::Create(spec, options, tc.get());
+          })
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+{
   gl_key_context_[SharedContextKey()] = gl_context;
   named_executors_[kGpuExecutorName] =
       std::make_shared<GlContextExecutor>(gl_context.get());
 #if __APPLE__
-  texture_caches_ = std::make_shared<CvTextureCacheManager>();
-  gpu_buffer_pool().SetSimplePoolFactory(
-      [tc = texture_caches_](const GpuBufferMultiPool::BufferSpec& spec,
-                             const MultiPoolOptions& options) {
-        return CvPixelBufferPoolWrapper::Create(spec.width, spec.height,
-                                                spec.format, options, tc.get());
-      });
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   texture_caches_->RegisterTextureCache(gl_context->cv_texture_cache());
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER

mediapipe/gpu/gpu_shared_data_internal.h

@@ -87,13 +87,15 @@ class GpuResources {
   std::map<std::string, std::string> node_key_;
   std::map<std::string, std::shared_ptr<GlContext>> gl_key_context_;
 
+#ifdef MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+  std::shared_ptr<CvTextureCacheManager> texture_caches_;
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+
   // The pool must be destructed before the gl_context, but after the
   // ios_gpu_data, so the declaration order is important.
   GpuBufferMultiPool gpu_buffer_pool_;
 
 #ifdef __APPLE__
-  std::shared_ptr<CvTextureCacheManager> texture_caches_;
-
   // Note that this is an Objective-C object.
   MPPGraphGPUData* ios_gpu_data_;
 #endif  // defined(__APPLE__)

mediapipe/gpu/multi_pool.h

@@ -12,16 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// This class lets calculators allocate GpuBuffers of various sizes, caching
-// and reusing them as needed. It does so by automatically creating and using
-// platform-specific buffer pools for the requested sizes.
-//
-// This class is not meant to be used directly by calculators, but is instead
-// used by GlCalculatorHelper to allocate buffers.
-
 #ifndef MEDIAPIPE_GPU_MULTI_POOL_H_
 #define MEDIAPIPE_GPU_MULTI_POOL_H_
 
+#include "mediapipe/util/resource_cache.h"
+
 namespace mediapipe {
 
 struct MultiPoolOptions {
@@ -42,6 +37,81 @@ struct MultiPoolOptions {
 
 static constexpr MultiPoolOptions kDefaultMultiPoolOptions;
 
+// MultiPool is a generic class for vending reusable resources of type Item,
+// which are assumed to be relatively expensive to create, so that reusing them
+// is beneficial.
+// Items are classified by Spec; when an item with a given Spec is requested,
+// an old Item with the same Spec can be reused, if available; otherwise a new
+// Item will be created. When user code is done with an Item, it is returned
+// to the pool for reuse.
+// In order to manage this, a MultiPool contains a map of Specs to SimplePool;
+// each SimplePool manages Items with the same Spec, which are thus considered
+// interchangeable.
+// Item retention and eviction policies are controlled by options.
+// A concrete example would be a pool of GlTextureBuffer, grouped by dimensions
+// and format.
+template <class SimplePool, class Spec, class Item>
+class MultiPool {
+ public:
+  using SimplePoolFactory = std::function<std::shared_ptr<SimplePool>(
+      const Spec& spec, const MultiPoolOptions& options)>;
+
+  MultiPool(SimplePoolFactory factory = DefaultMakeSimplePool,
+            MultiPoolOptions options = kDefaultMultiPoolOptions)
+      : create_simple_pool_(factory), options_(options) {}
+
+  // Obtains an item. May either be reused or created anew.
+  Item Get(const Spec& spec);
+
+ private:
+  static std::shared_ptr<SimplePool> DefaultMakeSimplePool(
+      const Spec& spec, const MultiPoolOptions& options) {
+    return SimplePool::Create(spec, options);
+  }
+
+  // Requests a simple buffer pool for the given spec. This may return nullptr
+  // if we have not yet reached a sufficient number of requests to allocate a
+  // pool, in which case the caller should invoke CreateBufferWithoutPool.
+  std::shared_ptr<SimplePool> RequestPool(const Spec& spec);
+
+  absl::Mutex mutex_;
+  mediapipe::ResourceCache<Spec, std::shared_ptr<SimplePool>> cache_
+      ABSL_GUARDED_BY(mutex_);
+  SimplePoolFactory create_simple_pool_ = DefaultMakeSimplePool;
+  MultiPoolOptions options_;
+};
+
+template <class SimplePool, class Spec, class Item>
+std::shared_ptr<SimplePool> MultiPool<SimplePool, Spec, Item>::RequestPool(
+    const Spec& spec) {
+  std::shared_ptr<SimplePool> pool;
+  std::vector<std::shared_ptr<SimplePool>> evicted;
+  {
+    absl::MutexLock lock(&mutex_);
+    pool = cache_.Lookup(spec, [this](const Spec& spec, int request_count) {
+      return (request_count >= options_.min_requests_before_pool)
+                 ? create_simple_pool_(spec, options_)
+                 : nullptr;
+    });
+    evicted = cache_.Evict(options_.max_pool_count,
+                           options_.request_count_scrub_interval);
+  }
+  // Evicted pools, and their buffers, will be released without holding the
+  // lock.
+  return pool;
+}
+
+template <class SimplePool, class Spec, class Item>
+Item MultiPool<SimplePool, Spec, Item>::Get(const Spec& spec) {
+  std::shared_ptr<SimplePool> pool = RequestPool(spec);
+  if (pool) {
+    // Note: we release our multipool lock before accessing the simple pool.
+    return Item(pool->GetBuffer());
+  } else {
+    return Item(SimplePool::CreateBufferWithoutPool(spec));
+  }
+}
+
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_GPU_MULTI_POOL_H_