From 53d015af08c96d39ecee97bdfa11cc5b5a882cec Mon Sep 17 00:00:00 2001 From: Camillo Lugaresi Date: Tue, 15 Nov 2022 16:03:41 -0800 Subject: [PATCH] Generic MultiPool template PiperOrigin-RevId: 488783176 --- mediapipe/gpu/BUILD | 1 + mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc | 8 +- mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h | 8 +- mediapipe/gpu/gl_texture_buffer_pool.h | 9 +- mediapipe/gpu/gpu_buffer_format.h | 28 +++++++ mediapipe/gpu/gpu_buffer_multi_pool.cc | 46 +--------- mediapipe/gpu/gpu_buffer_multi_pool.h | 77 ++--------------- mediapipe/gpu/gpu_shared_data_internal.cc | 18 ++-- mediapipe/gpu/gpu_shared_data_internal.h | 6 +- mediapipe/gpu/multi_pool.h | 84 +++++++++++++++++-- 10 files changed, 142 insertions(+), 143 deletions(-) diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD index 36527736b..1efe75b52 100644 --- a/mediapipe/gpu/BUILD +++ b/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( diff --git a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc index d8155f5cf..6e077ae6e 100644 --- a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc +++ b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc @@ -71,12 +71,12 @@ std::string CvPixelBufferPoolWrapper::GetDebugString() const { void CvPixelBufferPoolWrapper::Flush() { CVPixelBufferPoolFlush(*pool_, 0); } CFHolder CvPixelBufferPoolWrapper::CreateBufferWithoutPool( - int width, int height, GpuBufferFormat format) { - OSType cv_format = CVPixelFormatForGpuBufferFormat(format); + const internal::GpuBufferSpec& spec) { + OSType cv_format = CVPixelFormatForGpuBufferFormat(spec.format); CHECK_NE(cv_format, -1) << "unsupported pixel format"; CVPixelBufferRef buffer; - CVReturn err = - CreateCVPixelBufferWithoutPool(width, height, cv_format, &buffer); + CVReturn err = CreateCVPixelBufferWithoutPool(spec.width, spec.height, + cv_format, &buffer); CHECK(!err) << "Error creating pixel buffer: " << err; return MakeCFHolderAdopting(buffer); } diff --git a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h index 7d0aec4eb..4d71adbf2 100644 --- a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h +++ b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h @@ -38,11 +38,11 @@ class CvPixelBufferPoolWrapper { CvTextureCacheManager* texture_caches); static std::shared_ptr Create( - int width, int height, GpuBufferFormat format, - const MultiPoolOptions& options, + const internal::GpuBufferSpec& spec, const MultiPoolOptions& options, CvTextureCacheManager* texture_caches = nullptr) { return std::make_shared( - width, height, format, options.max_inactive_buffer_age, texture_caches); + spec.width, spec.height, spec.format, options.max_inactive_buffer_age, + texture_caches); } CFHolder GetBuffer(); @@ -53,7 +53,7 @@ class CvPixelBufferPoolWrapper { void Flush(); static CFHolder CreateBufferWithoutPool( - int width, int height, GpuBufferFormat format); + const internal::GpuBufferSpec& spec); private: CFHolder pool_; diff --git a/mediapipe/gpu/gl_texture_buffer_pool.h b/mediapipe/gpu/gl_texture_buffer_pool.h index fee46915e..29fc3c01c 100644 --- a/mediapipe/gpu/gl_texture_buffer_pool.h +++ b/mediapipe/gpu/gl_texture_buffer_pool.h @@ -42,9 +42,8 @@ class GlTextureBufferPool } static std::shared_ptr Create( - int width, int height, GpuBufferFormat format, - const MultiPoolOptions& options) { - return Create(width, height, format, options.keep_count); + const internal::GpuBufferSpec& spec, const MultiPoolOptions& options) { + return Create(spec.width, spec.height, spec.format, options.keep_count); } // Obtains a buffers. May either be reused or created anew. @@ -59,8 +58,8 @@ class GlTextureBufferPool std::pair GetInUseAndAvailableCounts(); static GlTextureBufferSharedPtr CreateBufferWithoutPool( - int width, int height, GpuBufferFormat format) { - return GlTextureBuffer::Create(width, height, format); + const internal::GpuBufferSpec& spec) { + return GlTextureBuffer::Create(spec.width, spec.height, spec.format); } private: diff --git a/mediapipe/gpu/gpu_buffer_format.h b/mediapipe/gpu/gpu_buffer_format.h index 45f054d31..06c5a0439 100644 --- a/mediapipe/gpu/gpu_buffer_format.h +++ b/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 + friend H AbslHashValue(H h, const GpuBufferSpec& spec) { + return H::combine(std::move(h), spec.width, spec.height, + static_cast(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_ diff --git a/mediapipe/gpu/gpu_buffer_multi_pool.cc b/mediapipe/gpu/gpu_buffer_multi_pool.cc index 744ccea2d..e2ed523e4 100644 --- a/mediapipe/gpu/gpu_buffer_multi_pool.cc +++ b/mediapipe/gpu/gpu_buffer_multi_pool.cc @@ -16,51 +16,7 @@ #include -#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 { - -std::shared_ptr -GpuBufferMultiPool::DefaultMakeSimplePool( - const GpuBufferMultiPool::BufferSpec& spec, - const MultiPoolOptions& options) { - return SimplePool::Create(spec.width, spec.height, spec.format, options); -} - -std::shared_ptr GpuBufferMultiPool::RequestPool( - const BufferSpec& spec) { - std::shared_ptr pool; - std::vector> 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 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 +namespace mediapipe {} // namespace mediapipe diff --git a/mediapipe/gpu/gpu_buffer_multi_pool.h b/mediapipe/gpu/gpu_buffer_multi_pool.h index 88428d053..827cf514a 100644 --- a/mediapipe/gpu/gpu_buffer_multi_pool.h +++ b/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 { - 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 - friend H AbslHashValue(H h, const BufferSpec& spec) { - return H::combine(std::move(h), spec.width, spec.height, - static_cast(spec.format)); - } - - int width; - int height; - mediapipe::GpuBufferFormat format; - }; - +class GpuBufferMultiPool : public MultiPool< #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( - const BufferSpec& spec, const MultiPoolOptions& options)>; - - void SetSimplePoolFactory(SimplePoolFactory create_simple_pool) { - create_simple_pool_ = create_simple_pool; + GpuBuffer GetBuffer(int width, int height, + GpuBufferFormat format = GpuBufferFormat::kBGRA32) { + return Get(internal::GpuBufferSpec(width, height, format)); } - - private: - static std::shared_ptr 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 RequestPool(const BufferSpec& spec); - - MultiPoolOptions options_; - absl::Mutex mutex_; - mediapipe::ResourceCache> 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_ diff --git a/mediapipe/gpu/gpu_shared_data_internal.cc b/mediapipe/gpu/gpu_shared_data_internal.cc index 6633c2f00..52db88633 100644 --- a/mediapipe/gpu/gpu_shared_data_internal.cc +++ b/mediapipe/gpu/gpu_shared_data_internal.cc @@ -80,18 +80,20 @@ GpuResources::StatusOrGpuResources GpuResources::Create( return gpu_resources; } -GpuResources::GpuResources(std::shared_ptr gl_context) { +GpuResources::GpuResources(std::shared_ptr gl_context) +#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER + : texture_caches_(std::make_shared()), + 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(gl_context.get()); #if __APPLE__ - texture_caches_ = std::make_shared(); - 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 diff --git a/mediapipe/gpu/gpu_shared_data_internal.h b/mediapipe/gpu/gpu_shared_data_internal.h index 12a7a1296..4fe6ba04e 100644 --- a/mediapipe/gpu/gpu_shared_data_internal.h +++ b/mediapipe/gpu/gpu_shared_data_internal.h @@ -87,13 +87,15 @@ class GpuResources { std::map node_key_; std::map> gl_key_context_; +#ifdef MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER + std::shared_ptr 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 texture_caches_; - // Note that this is an Objective-C object. MPPGraphGPUData* ios_gpu_data_; #endif // defined(__APPLE__) diff --git a/mediapipe/gpu/multi_pool.h b/mediapipe/gpu/multi_pool.h index e504fc820..8a3cf6be0 100644 --- a/mediapipe/gpu/multi_pool.h +++ b/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 MultiPool { + public: + using SimplePoolFactory = std::function( + 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 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 RequestPool(const Spec& spec); + + absl::Mutex mutex_; + mediapipe::ResourceCache> cache_ + ABSL_GUARDED_BY(mutex_); + SimplePoolFactory create_simple_pool_ = DefaultMakeSimplePool; + MultiPoolOptions options_; +}; + +template +std::shared_ptr MultiPool::RequestPool( + const Spec& spec) { + std::shared_ptr pool; + std::vector> 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 +Item MultiPool::Get(const Spec& spec) { + std::shared_ptr 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_