Make it possible to override the SimplePool factory used by MultiPool

This means MultiPool no longer needs a SetFlushPlatformCaches method, which was too specific to the CVPixelBufferPool.

PiperOrigin-RevId: 488783003

mediapipe/gpu/BUILD

@@ -368,6 +368,7 @@ cc_library(
         ],
     }),
     deps = [
+        ":cv_texture_cache_manager",
         ":gpu_buffer_format",
         ":multi_pool",
         ":pixel_buffer_pool_util",

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc

@@ -25,13 +25,13 @@ namespace mediapipe {
 
 CvPixelBufferPoolWrapper::CvPixelBufferPoolWrapper(
     int width, int height, GpuBufferFormat format, CFTimeInterval maxAge,
-    std::function<void(void)> flush_texture_caches) {
+    CvTextureCacheManager* texture_caches) {
   OSType cv_format = CVPixelFormatForGpuBufferFormat(format);
   CHECK_NE(cv_format, -1) << "unsupported pixel format";
   pool_ = MakeCFHolderAdopting(
       /* keep count is 0 because the age param keeps buffers around anyway */
       CreateCVPixelBufferPool(width, height, cv_format, 0, maxAge));
-  flush_texture_caches_ = std::move(flush_texture_caches);
+  texture_caches_ = texture_caches;
 }
 
 CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::GetBuffer() {
@@ -47,12 +47,12 @@ CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::GetBuffer() {
         kCFAllocatorDefault, *pool_, (__bridge CFDictionaryRef)auxAttributes,
         &buffer);
     if (err != kCVReturnWouldExceedAllocationThreshold) break;
-    if (flush_texture_caches_ && !tried_flushing) {
+    if (texture_caches_ && !tried_flushing) {
       // Call the flush function to potentially release old holds on buffers
       // and try again to create a pixel buffer.
       // This is used to flush CV texture caches, which may retain buffers until
       // flushed.
-      flush_texture_caches_();
+      texture_caches_->FlushTextureCaches();
       tried_flushing = true;
     } else {
       ++threshold;

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h

@@ -23,6 +23,7 @@
 #define MEDIAPIPE_GPU_CV_PIXEL_BUFFER_POOL_WRAPPER_H_
 
 #include "CoreFoundation/CFBase.h"
+#include "mediapipe/gpu/cv_texture_cache_manager.h"
 #include "mediapipe/gpu/gpu_buffer_format.h"
 #include "mediapipe/gpu/multi_pool.h"
 #include "mediapipe/gpu/pixel_buffer_pool_util.h"
@@ -34,15 +35,14 @@ class CvPixelBufferPoolWrapper {
  public:
   CvPixelBufferPoolWrapper(int width, int height, GpuBufferFormat format,
                            CFTimeInterval maxAge,
-                           std::function<void(void)> flush_texture_caches);
+                           CvTextureCacheManager* texture_caches);
 
   static std::shared_ptr<CvPixelBufferPoolWrapper> Create(
       int width, int height, GpuBufferFormat format,
       const MultiPoolOptions& options,
-      std::function<void(void)> flush_texture_caches = nullptr) {
+      CvTextureCacheManager* texture_caches = nullptr) {
     return std::make_shared<CvPixelBufferPoolWrapper>(
-        width, height, format, options.max_inactive_buffer_age,
-        flush_texture_caches);
+        width, height, format, options.max_inactive_buffer_age, texture_caches);
   }
 
   CFHolder<CVPixelBufferRef> GetBuffer();
@@ -58,7 +58,7 @@ class CvPixelBufferPoolWrapper {
  private:
   CFHolder<CVPixelBufferPoolRef> pool_;
   int count_ = 0;
-  std::function<void(void)> flush_texture_caches_;
+  CvTextureCacheManager* texture_caches_;
 };
 
 }  // namespace mediapipe

mediapipe/gpu/gpu_buffer_multi_pool.cc

@@ -24,15 +24,10 @@
 namespace mediapipe {
 
 std::shared_ptr<GpuBufferMultiPool::SimplePool>
-GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec,
-                                   const MultiPoolOptions& options) {
-#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-  return CvPixelBufferPoolWrapper::Create(spec.width, spec.height, spec.format,
-                                          options, flush_platform_caches_);
-#else
-  return GlTextureBufferPool::Create(spec.width, spec.height, spec.format,
-                                     options);
-#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+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(
@@ -44,7 +39,7 @@ std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool(
     pool =
         cache_.Lookup(spec, [this](const BufferSpec& spec, int request_count) {
           return (request_count >= options_.min_requests_before_pool)
-                     ? MakeSimplePool(spec, options_)
+                     ? create_simple_pool_(spec, options_)
                      : nullptr;
         });
     evicted = cache_.Evict(options_.max_pool_count,

mediapipe/gpu/gpu_buffer_multi_pool.h

@@ -52,10 +52,6 @@ class GpuBufferMultiPool {
   GpuBuffer GetBuffer(int width, int height,
                       GpuBufferFormat format = GpuBufferFormat::kBGRA32);
 
-  void SetFlushPlatformCaches(std::function<void(void)> flush_platform_caches) {
-    flush_platform_caches_ = flush_platform_caches;
-  }
-
   // 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.
@@ -74,14 +70,21 @@ class GpuBufferMultiPool {
     mediapipe::GpuBufferFormat format;
   };
 
- private:
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   using SimplePool = CvPixelBufferPoolWrapper;
 #else
   using SimplePool = GlTextureBufferPool;
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 
-  std::shared_ptr<SimplePool> MakeSimplePool(
+  using SimplePoolFactory = std::function<std::shared_ptr<SimplePool>(
+      const BufferSpec& spec, const MultiPoolOptions& options)>;
+
+  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);
 
@@ -94,8 +97,7 @@ class GpuBufferMultiPool {
   absl::Mutex mutex_;
   mediapipe::ResourceCache<BufferSpec, std::shared_ptr<SimplePool>> cache_
       ABSL_GUARDED_BY(mutex_);
-  // This is used to hook up the TextureCacheManager on Apple platforms.
-  std::function<void(void)> flush_platform_caches_;
+  SimplePoolFactory create_simple_pool_ = DefaultMakeSimplePool;
 };
 
 // BufferSpec equality operators

mediapipe/gpu/gpu_shared_data_internal.cc

@@ -86,8 +86,12 @@ GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context) {
       std::make_shared<GlContextExecutor>(gl_context.get());
 #if __APPLE__
   texture_caches_ = std::make_shared<CvTextureCacheManager>();
-  gpu_buffer_pool().SetFlushPlatformCaches(
-      [tc = texture_caches_] { tc->FlushTextureCaches(); });
+  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