Allow customizing MultiPool options

These don't need to be constants.

PiperOrigin-RevId: 488782713
This commit is contained in:
Camillo Lugaresi 2022-11-15 16:01:56 -08:00 committed by Copybara-Service
parent f13903b7c5
commit 7ef3185ecb
2 changed files with 26 additions and 19 deletions

View File

@ -23,26 +23,12 @@
namespace mediapipe { namespace mediapipe {
// Keep this many buffers allocated for a given frame size.
static constexpr int kKeepCount = 2;
// The maximum size of the GpuBufferMultiPool. When the limit is reached, the
// oldest BufferSpec will be dropped.
static constexpr int kMaxPoolCount = 10;
// Time in seconds after which an inactive buffer can be dropped from the pool.
// Currently only used with CVPixelBufferPool.
static constexpr float kMaxInactiveBufferAge = 0.25;
// Skip allocating a buffer pool until at least this many requests have been
// made for a given BufferSpec.
static constexpr int kMinRequestsBeforePool = 2;
// Do a deeper flush every this many requests.
static constexpr int kRequestCountScrubInterval = 50;
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
std::shared_ptr<GpuBufferMultiPool::SimplePool> std::shared_ptr<GpuBufferMultiPool::SimplePool>
GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec) { GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec) {
return std::make_shared<CvPixelBufferPoolWrapper>( return std::make_shared<CvPixelBufferPoolWrapper>(
spec.width, spec.height, spec.format, kMaxInactiveBufferAge, spec.width, spec.height, spec.format, options_.max_inactive_buffer_age,
flush_platform_caches_); flush_platform_caches_);
} }
@ -51,7 +37,7 @@ GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec) {
std::shared_ptr<GpuBufferMultiPool::SimplePool> std::shared_ptr<GpuBufferMultiPool::SimplePool>
GpuBufferMultiPool::MakeSimplePool(const BufferSpec& spec) { GpuBufferMultiPool::MakeSimplePool(const BufferSpec& spec) {
return GlTextureBufferPool::Create(spec.width, spec.height, spec.format, return GlTextureBufferPool::Create(spec.width, spec.height, spec.format,
kKeepCount); options_.keep_count);
} }
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER #endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
@ -64,11 +50,12 @@ std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool(
absl::MutexLock lock(&mutex_); absl::MutexLock lock(&mutex_);
pool = pool =
cache_.Lookup(spec, [this](const BufferSpec& spec, int request_count) { cache_.Lookup(spec, [this](const BufferSpec& spec, int request_count) {
return (request_count >= kMinRequestsBeforePool) return (request_count >= options_.min_requests_before_pool)
? MakeSimplePool(spec) ? MakeSimplePool(spec)
: nullptr; : nullptr;
}); });
evicted = cache_.Evict(kMaxPoolCount, kRequestCountScrubInterval); evicted = cache_.Evict(options_.max_pool_count,
options_.request_count_scrub_interval);
} }
// Evicted pools, and their buffers, will be released without holding the // Evicted pools, and their buffers, will be released without holding the
// lock. // lock.

View File

@ -42,9 +42,28 @@ namespace mediapipe {
struct GpuSharedData; struct GpuSharedData;
class CvPixelBufferPoolWrapper; class CvPixelBufferPoolWrapper;
struct MultiPoolOptions {
// Keep this many buffers allocated for a given frame size.
int keep_count = 2;
// The maximum size of the GpuBufferMultiPool. When the limit is reached, the
// oldest BufferSpec will be dropped.
int max_pool_count = 10;
// Time in seconds after which an inactive buffer can be dropped from the
// pool. Currently only used with CVPixelBufferPool.
float max_inactive_buffer_age = 0.25;
// Skip allocating a buffer pool until at least this many requests have been
// made for a given BufferSpec.
int min_requests_before_pool = 2;
// Do a deeper flush every this many requests.
int request_count_scrub_interval = 50;
};
static constexpr MultiPoolOptions kDefaultMultiPoolOptions;
class GpuBufferMultiPool { class GpuBufferMultiPool {
public: public:
GpuBufferMultiPool() {} GpuBufferMultiPool(MultiPoolOptions options = kDefaultMultiPoolOptions)
: options_(options) {}
// Obtains a buffer. May either be reused or created anew. // Obtains a buffer. May either be reused or created anew.
GpuBuffer GetBuffer(int width, int height, GpuBuffer GetBuffer(int width, int height,
@ -85,6 +104,7 @@ class GpuBufferMultiPool {
// pool, in which case the caller should invoke CreateBufferWithoutPool. // pool, in which case the caller should invoke CreateBufferWithoutPool.
std::shared_ptr<SimplePool> RequestPool(const BufferSpec& spec); std::shared_ptr<SimplePool> RequestPool(const BufferSpec& spec);
MultiPoolOptions options_;
absl::Mutex mutex_; absl::Mutex mutex_;
mediapipe::ResourceCache<BufferSpec, std::shared_ptr<SimplePool>> cache_ mediapipe::ResourceCache<BufferSpec, std::shared_ptr<SimplePool>> cache_
ABSL_GUARDED_BY(mutex_); ABSL_GUARDED_BY(mutex_);