Move flush hook to CvPixelBufferPoolWrapper constructor

This unifies the implementation of GpuBufferMultiPool::GetBufferFromSimplePool.

PiperOrigin-RevId: 488782173
This commit is contained in:
Camillo Lugaresi 2022-11-15 15:59:33 -08:00 committed by Copybara-Service
parent a4fe3eb094
commit 0c4522cb9f
3 changed files with 18 additions and 20 deletions

View File

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

View File

@ -32,8 +32,9 @@ namespace mediapipe {
class CvPixelBufferPoolWrapper { class CvPixelBufferPoolWrapper {
public: public:
CvPixelBufferPoolWrapper(int width, int height, GpuBufferFormat format, CvPixelBufferPoolWrapper(int width, int height, GpuBufferFormat format,
CFTimeInterval maxAge); CFTimeInterval maxAge,
CFHolder<CVPixelBufferRef> GetBuffer(std::function<void(void)> flush); std::function<void(void)> flush_texture_caches);
CFHolder<CVPixelBufferRef> GetBuffer();
int GetBufferCount() const { return count_; } int GetBufferCount() const { return count_; }
std::string GetDebugString() const; std::string GetDebugString() const;
@ -46,6 +47,7 @@ class CvPixelBufferPoolWrapper {
private: private:
CFHolder<CVPixelBufferPoolRef> pool_; CFHolder<CVPixelBufferPoolRef> pool_;
int count_ = 0; int count_ = 0;
std::function<void(void)> flush_texture_caches_;
}; };
} // namespace mediapipe } // namespace mediapipe

View File

@ -48,12 +48,8 @@ static constexpr int kRequestCountScrubInterval = 50;
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, kMaxInactiveBufferAge,
} flush_platform_caches_);
GpuBuffer GpuBufferMultiPool::GetBufferFromSimplePool(
BufferSpec spec, GpuBufferMultiPool::SimplePool& pool) {
return GpuBuffer(pool.GetBuffer(flush_platform_caches_));
} }
#else #else
@ -64,11 +60,6 @@ GpuBufferMultiPool::MakeSimplePool(const BufferSpec& spec) {
kKeepCount); kKeepCount);
} }
GpuBuffer GpuBufferMultiPool::GetBufferFromSimplePool(
BufferSpec spec, GpuBufferMultiPool::SimplePool& pool) {
return GpuBuffer(pool.GetBuffer());
}
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER #endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool( std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool(
@ -102,6 +93,11 @@ GpuBuffer GpuBufferMultiPool::GetBuffer(int width, int height,
} }
} }
GpuBuffer GpuBufferMultiPool::GetBufferFromSimplePool(
BufferSpec spec, GpuBufferMultiPool::SimplePool& pool) {
return GpuBuffer(pool.GetBuffer());
}
GpuBuffer GpuBufferMultiPool::GetBufferWithoutPool(const BufferSpec& spec) { GpuBuffer GpuBufferMultiPool::GetBufferWithoutPool(const BufferSpec& spec) {
return GpuBuffer(SimplePool::CreateBufferWithoutPool(spec.width, spec.height, return GpuBuffer(SimplePool::CreateBufferWithoutPool(spec.width, spec.height,
spec.format)); spec.format));