diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD index 1efe75b52..747d131ba 100644 --- a/mediapipe/gpu/BUILD +++ b/mediapipe/gpu/BUILD @@ -607,6 +607,7 @@ cc_library( ":gpu_buffer", ":gpu_shared_data_header", ":multi_pool", + ":reusable_pool", "//mediapipe/framework:calculator_context", "//mediapipe/framework:calculator_node", "//mediapipe/framework/port:logging", @@ -615,6 +616,16 @@ cc_library( ], ) +cc_library( + name = "reusable_pool", + hdrs = ["reusable_pool.h"], + deps = [ + ":multi_pool", + "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/synchronization", + ], +) + cc_library( name = "multi_pool", hdrs = ["multi_pool.h"], diff --git a/mediapipe/gpu/gl_texture_buffer.h b/mediapipe/gpu/gl_texture_buffer.h index 124a0ec2f..a770163b5 100644 --- a/mediapipe/gpu/gl_texture_buffer.h +++ b/mediapipe/gpu/gl_texture_buffer.h @@ -71,6 +71,11 @@ class GlTextureBuffer // Create a texture with a copy of the data in image_frame. static std::unique_ptr Create(const ImageFrame& image_frame); + static std::unique_ptr Create( + const internal::GpuBufferSpec& spec) { + return Create(spec.width, spec.height, spec.format); + } + // Wraps an existing texture, but does not take ownership of it. // deletion_callback is invoked when the GlTextureBuffer is released, so // the caller knows that the texture is no longer in use. diff --git a/mediapipe/gpu/gl_texture_buffer_pool.cc b/mediapipe/gpu/gl_texture_buffer_pool.cc index 3d5a8cdaa..599381a34 100644 --- a/mediapipe/gpu/gl_texture_buffer_pool.cc +++ b/mediapipe/gpu/gl_texture_buffer_pool.cc @@ -16,79 +16,4 @@ #include "absl/synchronization/mutex.h" -namespace mediapipe { - -GlTextureBufferPool::GlTextureBufferPool(int width, int height, - GpuBufferFormat format, int keep_count) - : width_(width), - height_(height), - format_(format), - keep_count_(keep_count) {} - -GlTextureBufferSharedPtr GlTextureBufferPool::GetBuffer() { - std::unique_ptr buffer; - bool reuse = false; - - { - absl::MutexLock lock(&mutex_); - if (available_.empty()) { - buffer = GlTextureBuffer::Create(width_, height_, format_); - if (!buffer) return nullptr; - } else { - buffer = std::move(available_.back()); - available_.pop_back(); - reuse = true; - } - - ++in_use_count_; - } - - // This needs to wait on consumer sync points, therefore it should not be - // done while holding the mutex. - if (reuse) { - buffer->Reuse(); - } - - // Return a shared_ptr with a custom deleter that adds the buffer back - // to our available list. - std::weak_ptr weak_pool(shared_from_this()); - return std::shared_ptr( - buffer.release(), [weak_pool](GlTextureBuffer* buf) { - auto pool = weak_pool.lock(); - if (pool) { - pool->Return(absl::WrapUnique(buf)); - } else { - delete buf; - } - }); -} - -std::pair GlTextureBufferPool::GetInUseAndAvailableCounts() { - absl::MutexLock lock(&mutex_); - return {in_use_count_, available_.size()}; -} - -void GlTextureBufferPool::Return(std::unique_ptr buf) { - std::vector> trimmed; - { - absl::MutexLock lock(&mutex_); - --in_use_count_; - available_.emplace_back(std::move(buf)); - TrimAvailable(&trimmed); - } - // The trimmed buffers will be released without holding the lock. -} - -void GlTextureBufferPool::TrimAvailable( - std::vector>* trimmed) { - int keep = std::max(keep_count_ - in_use_count_, 0); - if (available_.size() > keep) { - auto trim_it = std::next(available_.begin(), keep); - if (trimmed) { - std::move(trim_it, available_.end(), std::back_inserter(*trimmed)); - } - available_.erase(trim_it, available_.end()); - } -} - -} // namespace mediapipe +namespace mediapipe {} // namespace mediapipe diff --git a/mediapipe/gpu/gl_texture_buffer_pool.h b/mediapipe/gpu/gl_texture_buffer_pool.h index 29fc3c01c..726d0528d 100644 --- a/mediapipe/gpu/gl_texture_buffer_pool.h +++ b/mediapipe/gpu/gl_texture_buffer_pool.h @@ -24,11 +24,11 @@ #include "absl/synchronization/mutex.h" #include "mediapipe/gpu/gl_texture_buffer.h" #include "mediapipe/gpu/multi_pool.h" +#include "mediapipe/gpu/reusable_pool.h" namespace mediapipe { -class GlTextureBufferPool - : public std::enable_shared_from_this { +class GlTextureBufferPool : public ReusablePool { public: // Creates a pool. This pool will manage buffers of the specified dimensions, // and will keep keep_count buffers around for reuse. @@ -37,52 +37,32 @@ class GlTextureBufferPool static std::shared_ptr Create(int width, int height, GpuBufferFormat format, int keep_count) { - return std::shared_ptr( - new GlTextureBufferPool(width, height, format, keep_count)); + return Create({width, height, format}, {.keep_count = keep_count}); } static std::shared_ptr Create( const internal::GpuBufferSpec& spec, const MultiPoolOptions& options) { - return Create(spec.width, spec.height, spec.format, options.keep_count); + return std::shared_ptr( + new GlTextureBufferPool(spec, options)); } - // Obtains a buffers. May either be reused or created anew. - // A GlContext must be current when this is called. - GlTextureBufferSharedPtr GetBuffer(); - - int width() const { return width_; } - int height() const { return height_; } - GpuBufferFormat format() const { return format_; } - - // This method is meant for testing. - std::pair GetInUseAndAvailableCounts(); + int width() const { return spec_.width; } + int height() const { return spec_.height; } + GpuBufferFormat format() const { return spec_.format; } static GlTextureBufferSharedPtr CreateBufferWithoutPool( const internal::GpuBufferSpec& spec) { - return GlTextureBuffer::Create(spec.width, spec.height, spec.format); + return GlTextureBuffer::Create(spec); } - private: - GlTextureBufferPool(int width, int height, GpuBufferFormat format, - int keep_count); + protected: + GlTextureBufferPool(const internal::GpuBufferSpec& spec, + const MultiPoolOptions& options) + : ReusablePool( + [this] { return GlTextureBuffer::Create(spec_); }, options), + spec_(spec) {} - // Return a buffer to the pool. - void Return(std::unique_ptr buf); - - // If the total number of buffers is greater than keep_count, destroys any - // surplus buffers that are no longer in use. - void TrimAvailable(std::vector>* trimmed) - ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); - - const int width_; - const int height_; - const GpuBufferFormat format_; - const int keep_count_; - - absl::Mutex mutex_; - int in_use_count_ ABSL_GUARDED_BY(mutex_) = 0; - std::vector> available_ - ABSL_GUARDED_BY(mutex_); + const internal::GpuBufferSpec spec_; }; } // namespace mediapipe diff --git a/mediapipe/gpu/reusable_pool.h b/mediapipe/gpu/reusable_pool.h new file mode 100644 index 000000000..ddeaa5ba7 --- /dev/null +++ b/mediapipe/gpu/reusable_pool.h @@ -0,0 +1,145 @@ +// Copyright 2019 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Consider this file an implementation detail. None of this is part of the +// public API. + +#ifndef MEDIAPIPE_GPU_REUSABLE_POOL_H_ +#define MEDIAPIPE_GPU_REUSABLE_POOL_H_ + +#include +#include + +#include "absl/functional/any_invocable.h" +#include "absl/synchronization/mutex.h" +#include "mediapipe/gpu/multi_pool.h" + +namespace mediapipe { + +template +class ReusablePool : public std::enable_shared_from_this> { + public: + using ItemFactory = absl::AnyInvocable() const>; + + // Creates a pool. This pool will manage buffers of the specified dimensions, + // and will keep keep_count buffers around for reuse. + // We enforce creation as a shared_ptr so that we can use a weak reference in + // the buffers' deleters. + static std::shared_ptr> Create( + ItemFactory item_factory, const MultiPoolOptions& options) { + return std::shared_ptr>( + new ReusablePool(std::move(item_factory), options)); + } + + // Obtains a buffer. May either be reused or created anew. + // A GlContext must be current when this is called. + std::shared_ptr GetBuffer(); + + // This method is meant for testing. + std::pair GetInUseAndAvailableCounts(); + + protected: + ReusablePool(ItemFactory item_factory, const MultiPoolOptions& options) + : item_factory_(std::move(item_factory)), + keep_count_(options.keep_count) {} + + private: + // Return a buffer to the pool. + void Return(std::unique_ptr buf); + + // If the total number of buffers is greater than keep_count, destroys any + // surplus buffers that are no longer in use. + void TrimAvailable(std::vector>* trimmed) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + const ItemFactory item_factory_; + const int keep_count_; + + absl::Mutex mutex_; + int in_use_count_ ABSL_GUARDED_BY(mutex_) = 0; + std::vector> available_ ABSL_GUARDED_BY(mutex_); +}; + +template +inline std::shared_ptr ReusablePool::GetBuffer() { + std::unique_ptr buffer; + bool reuse = false; + + { + absl::MutexLock lock(&mutex_); + if (available_.empty()) { + buffer = item_factory_(); + if (!buffer) return nullptr; + } else { + buffer = std::move(available_.back()); + available_.pop_back(); + reuse = true; + } + + ++in_use_count_; + } + + // This needs to wait on consumer sync points, therefore it should not be + // done while holding the mutex. + if (reuse) { + buffer->Reuse(); + } + + // Return a shared_ptr with a custom deleter that adds the buffer back + // to our available list. + std::weak_ptr> weak_pool(this->shared_from_this()); + return std::shared_ptr(buffer.release(), [weak_pool](Item* buf) { + auto pool = weak_pool.lock(); + if (pool) { + pool->Return(absl::WrapUnique(buf)); + } else { + delete buf; + } + }); +} + +template +inline std::pair ReusablePool::GetInUseAndAvailableCounts() { + absl::MutexLock lock(&mutex_); + return {in_use_count_, available_.size()}; +} + +template +void ReusablePool::Return(std::unique_ptr buf) { + std::vector> trimmed; + { + absl::MutexLock lock(&mutex_); + --in_use_count_; + available_.emplace_back(std::move(buf)); + TrimAvailable(&trimmed); + } + // The trimmed buffers will be released without holding the lock. +} + +template +void ReusablePool::TrimAvailable( + std::vector>* trimmed) { + int keep = std::max(keep_count_ - in_use_count_, 0); + if (available_.size() > keep) { + auto trim_it = std::next(available_.begin(), keep); + if (trimmed) { + std::move(trim_it, available_.end(), std::back_inserter(*trimmed)); + } + available_.erase(trim_it, available_.end()); + } +} + +} // namespace mediapipe + +#endif // MEDIAPIPE_GPU_REUSABLE_POOL_H_