cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Generic MultiPool template

Browse Source 
PiperOrigin-RevId: 488783176
This commit is contained in:
Camillo Lugaresi 2022-11-15 16:03:41 -08:00 committed by Copybara-Service
parent b9fa2e3496
commit 53d015af08
10 changed files with 142 additions and 143 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mediapipe/gpu/BUILD
View File

@ -618,6 +618,7 @@ cc_library(
cc_library(
name = "multi_pool",
hdrs = ["multi_pool.h"],
deps = ["//mediapipe/util:resource_cache"],
)
cc_library(

8
mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc
View File

@ -71,12 +71,12 @@ std::string CvPixelBufferPoolWrapper::GetDebugString() const {
void CvPixelBufferPoolWrapper::Flush() { CVPixelBufferPoolFlush(*pool_, 0); }
CFHolder<CVPixelBufferRef> 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);
}

8
mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h
View File

@ -38,11 +38,11 @@ class CvPixelBufferPoolWrapper {
CvTextureCacheManager* texture_caches);
static std::shared_ptr<CvPixelBufferPoolWrapper> 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<CvPixelBufferPoolWrapper>(
width, height, format, options.max_inactive_buffer_age, texture_caches);
spec.width, spec.height, spec.format, options.max_inactive_buffer_age,
texture_caches);
}
CFHolder<CVPixelBufferRef> GetBuffer();
@ -53,7 +53,7 @@ class CvPixelBufferPoolWrapper {
void Flush();
static CFHolder<CVPixelBufferRef> CreateBufferWithoutPool(
int width, int height, GpuBufferFormat format);
const internal::GpuBufferSpec& spec);
private:
CFHolder<CVPixelBufferPoolRef> pool_;

9
mediapipe/gpu/gl_texture_buffer_pool.h
View File

@ -42,9 +42,8 @@ class GlTextureBufferPool
}
static std::shared_ptr<GlTextureBufferPool> 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<int, int> 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:

28
mediapipe/gpu/gpu_buffer_format.h
View File

@ -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 <typename H>
friend H AbslHashValue(H h, const GpuBufferSpec& spec) {
return H::combine(std::move(h), spec.width, spec.height,
static_cast<uint32_t>(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_

46
mediapipe/gpu/gpu_buffer_multi_pool.cc
View File

@ -16,51 +16,7 @@
#include <tuple>
#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::SimplePool>
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(
const BufferSpec& spec) {
std::shared_ptr<SimplePool> pool;
std::vector<std::shared_ptr<SimplePool>> 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<SimplePool> 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

77
mediapipe/gpu/gpu_buffer_multi_pool.h
View File

@ -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 <typename H>
friend H AbslHashValue(H h, const BufferSpec& spec) {
return H::combine(std::move(h), spec.width, spec.height,
static_cast<uint32_t>(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<std::shared_ptr<SimplePool>(
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<SimplePool> 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<SimplePool> RequestPool(const BufferSpec& spec);
MultiPoolOptions options_;
absl::Mutex mutex_;
mediapipe::ResourceCache<BufferSpec, std::shared_ptr<SimplePool>> 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_

18
mediapipe/gpu/gpu_shared_data_internal.cc
View File

@ -80,18 +80,20 @@ GpuResources::StatusOrGpuResources GpuResources::Create(
return gpu_resources;
}
GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context) {
GpuResources::GpuResources(std::shared_ptr<GlContext> gl_context)
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
: texture_caches_(std::make_shared<CvTextureCacheManager>()),
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<GlContextExecutor>(gl_context.get());
#if __APPLE__
texture_caches_ = std::make_shared<CvTextureCacheManager>();
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

6
mediapipe/gpu/gpu_shared_data_internal.h
View File

@ -87,13 +87,15 @@ class GpuResources {
std::map<std::string, std::string> node_key_;
std::map<std::string, std::shared_ptr<GlContext>> gl_key_context_;
#ifdef MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
std::shared_ptr<CvTextureCacheManager> 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<CvTextureCacheManager> texture_caches_;
// Note that this is an Objective-C object.
MPPGraphGPUData* ios_gpu_data_;
#endif // defined(__APPLE__)

84
mediapipe/gpu/multi_pool.h
View File

@ -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 SimplePool, class Spec, class Item>
class MultiPool {
public:
using SimplePoolFactory = std::function<std::shared_ptr<SimplePool>(
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<SimplePool> 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<SimplePool> RequestPool(const Spec& spec);
absl::Mutex mutex_;
mediapipe::ResourceCache<Spec, std::shared_ptr<SimplePool>> cache_
ABSL_GUARDED_BY(mutex_);
SimplePoolFactory create_simple_pool_ = DefaultMakeSimplePool;
MultiPoolOptions options_;
};
template <class SimplePool, class Spec, class Item>
std::shared_ptr<SimplePool> MultiPool<SimplePool, Spec, Item>::RequestPool(
const Spec& spec) {
std::shared_ptr<SimplePool> pool;
std::vector<std::shared_ptr<SimplePool>> 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 <class SimplePool, class Spec, class Item>
Item MultiPool<SimplePool, Spec, Item>::Get(const Spec& spec) {
std::shared_ptr<SimplePool> 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_