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

Make GpuBuffer a shared_ptr to a storage collection

Browse Source 
PiperOrigin-RevId: 493519590
This commit is contained in:
Camillo Lugaresi 2022-12-06 23:54:11 -08:00 committed by Copybara-Service
parent 402834b4f2
commit 523d16dffa
4 changed files with 156 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mediapipe/gpu/BUILD
View File

@ -289,7 +289,9 @@ cc_library(
deps = [
":gpu_buffer_format",
":gpu_buffer_storage",
"@com_google_absl//absl/functional:bind_front",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/port:logging",
":gpu_buffer_storage_image_frame",

102
mediapipe/gpu/gpu_buffer.cc
View File

@ -3,6 +3,7 @@
#include <memory>
#include <utility>
#include "absl/functional/bind_front.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mediapipe/framework/port/logging.h"
@ -25,57 +26,101 @@ struct StorageTypeFormatter {
} // namespace
std::string GpuBuffer::DebugString() const {
return absl::StrCat("GpuBuffer[",
absl::StrJoin(storages_, ", ", StorageTypeFormatter()),
"]");
return holder_ ? absl::StrCat("GpuBuffer[", width(), "x", height(), " ",
format(), " as ", holder_->DebugString(), "]")
: "GpuBuffer[invalid]";
}
internal::GpuBufferStorage* GpuBuffer::GetStorageForView(
std::string GpuBuffer::StorageHolder::DebugString() const {
absl::MutexLock lock(&mutex_);
return absl::StrJoin(storages_, ", ", StorageTypeFormatter());
}
internal::GpuBufferStorage* GpuBuffer::StorageHolder::GetStorageForView(
TypeId view_provider_type, bool for_writing) const {
const std::shared_ptr<internal::GpuBufferStorage>* chosen_storage = nullptr;
std::shared_ptr<internal::GpuBufferStorage> chosen_storage;
std::function<std::shared_ptr<internal::GpuBufferStorage>()> conversion;
// First see if any current storage supports the view.
for (const auto& s : storages_) {
if (s->can_down_cast_to(view_provider_type)) {
chosen_storage = &s;
break;
}
}
// Then try to convert existing storages to one that does.
// TODO: choose best conversion.
if (!chosen_storage) {
{
absl::MutexLock lock(&mutex_);
// First see if any current storage supports the view.
for (const auto& s : storages_) {
if (auto converter = internal::GpuBufferStorageRegistry::Get()
.StorageConverterForViewProvider(
view_provider_type, s->storage_type())) {
if (auto new_storage = converter(s)) {
storages_.push_back(new_storage);
chosen_storage = &storages_.back();
if (s->can_down_cast_to(view_provider_type)) {
chosen_storage = s;
break;
}
}
// Then try to convert existing storages to one that does.
// TODO: choose best conversion.
if (!chosen_storage) {
for (const auto& s : storages_) {
if (auto converter = internal::GpuBufferStorageRegistry::Get()
.StorageConverterForViewProvider(
view_provider_type, s->storage_type())) {
conversion = absl::bind_front(converter, s);
break;
}
}
}
}
// Avoid invoking a converter or factory while holding the mutex.
// Two reasons:
// 1. Readers that don't need a conversion will not be blocked.
// 2. We use mutexes to make sure GL contexts are not used simultaneously on
// different threads, and we also rely on Mutex's deadlock detection
// heuristic, which enforces a consistent mutex acquisition order.
// This function is likely to be called within a GL context, and the
// conversion function may in turn use a GL context, and this may cause a
// false positive in the deadlock detector.
// TODO: we could use Mutex::ForgetDeadlockInfo instead.
if (conversion) {
auto new_storage = conversion();
absl::MutexLock lock(&mutex_);
// Another reader might have already completed and inserted the same
// conversion. TODO: prevent this?
for (const auto& s : storages_) {
if (s->can_down_cast_to(view_provider_type)) {
chosen_storage = s;
break;
}
}
if (!chosen_storage) {
storages_.push_back(std::move(new_storage));
chosen_storage = storages_.back();
}
}
if (for_writing) {
// This will temporarily hold storages to be released, and do so while the
// lock is not held (see above).
decltype(storages_) old_storages;
using std::swap;
if (chosen_storage) {
// Discard all other storages.
storages_ = {*chosen_storage};
chosen_storage = &storages_.back();
absl::MutexLock lock(&mutex_);
swap(old_storages, storages_);
storages_ = {chosen_storage};
} else {
// Allocate a new storage supporting the requested view.
if (auto factory =
internal::GpuBufferStorageRegistry::Get()
.StorageFactoryForViewProvider(view_provider_type)) {
if (auto new_storage = factory(width(), height(), format())) {
if (auto new_storage = factory(width_, height_, format_)) {
absl::MutexLock lock(&mutex_);
swap(old_storages, storages_);
storages_ = {std::move(new_storage)};
chosen_storage = &storages_.back();
chosen_storage = storages_.back();
}
}
}
}
return chosen_storage ? chosen_storage->get() : nullptr;
// It is ok to return a non-owning storage pointer here because this object
// ensures the storage's lifetime. Overwriting a GpuBuffer while readers are
// active would violate this, but it's not allowed in MediaPipe.
return chosen_storage ? chosen_storage.get() : nullptr;
}
internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
@ -84,8 +129,7 @@ internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
GpuBuffer::GetStorageForView(view_provider_type, for_writing);
CHECK(chosen_storage) << "no view provider found for requested view "
<< view_provider_type.name() << "; storages available: "
<< absl::StrJoin(storages_, ", ",
StorageTypeFormatter());
<< (holder_ ? holder_->DebugString() : "invalid");
DCHECK(chosen_storage->can_down_cast_to(view_provider_type));
return *chosen_storage;
}

105
mediapipe/gpu/gpu_buffer.h
View File

@ -15,9 +15,12 @@
#ifndef MEDIAPIPE_GPU_GPU_BUFFER_H_
#define MEDIAPIPE_GPU_GPU_BUFFER_H_
#include <algorithm>
#include <functional>
#include <memory>
#include <utility>
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/gpu/gpu_buffer_format.h"
#include "mediapipe/gpu/gpu_buffer_storage.h"
@ -56,8 +59,7 @@ class GpuBuffer {
// Creates an empty buffer of a given size and format. It will be allocated
// when a view is requested.
GpuBuffer(int width, int height, Format format)
: GpuBuffer(std::make_shared<PlaceholderGpuBufferStorage>(width, height,
format)) {}
: holder_(std::make_shared<StorageHolder>(width, height, format)) {}
// Copy and move constructors and assignment operators are supported.
GpuBuffer(const GpuBuffer& other) = default;
@ -70,9 +72,8 @@ class GpuBuffer {
// are not portable. Applications and calculators should normally obtain
// GpuBuffers in a portable way from the framework, e.g. using
// GpuBufferMultiPool.
explicit GpuBuffer(std::shared_ptr<internal::GpuBufferStorage> storage) {
storages_.push_back(std::move(storage));
}
explicit GpuBuffer(std::shared_ptr<internal::GpuBufferStorage> storage)
: holder_(std::make_shared<StorageHolder>(std::move(storage))) {}
#if !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
// This is used to support backward-compatible construction of GpuBuffer from
@ -84,9 +85,11 @@ class GpuBuffer {
: GpuBuffer(internal::AsGpuBufferStorage(storage_convertible)) {}
#endif // !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
int width() const { return current_storage().width(); }
int height() const { return current_storage().height(); }
GpuBufferFormat format() const { return current_storage().format(); }
int width() const { return holder_ ? holder_->width() : 0; }
int height() const { return holder_ ? holder_->height() : 0; }
GpuBufferFormat format() const {
return holder_ ? holder_->format() : GpuBufferFormat::kUnknown;
}
// Converts to true iff valid.
explicit operator bool() const { return operator!=(nullptr); }
@ -122,31 +125,17 @@ class GpuBuffer {
// using views.
template <class T>
std::shared_ptr<T> internal_storage() const {
for (const auto& s : storages_)
if (s->down_cast<T>()) return std::static_pointer_cast<T>(s);
return nullptr;
return holder_ ? holder_->internal_storage<T>() : nullptr;
}
std::string DebugString() const;
private:
class PlaceholderGpuBufferStorage
: public internal::GpuBufferStorageImpl<PlaceholderGpuBufferStorage> {
public:
PlaceholderGpuBufferStorage(int width, int height, Format format)
: width_(width), height_(height), format_(format) {}
int width() const override { return width_; }
int height() const override { return height_; }
GpuBufferFormat format() const override { return format_; }
private:
int width_ = 0;
int height_ = 0;
GpuBufferFormat format_ = GpuBufferFormat::kUnknown;
};
internal::GpuBufferStorage* GetStorageForView(TypeId view_provider_type,
bool for_writing) const;
bool for_writing) const {
return holder_ ? holder_->GetStorageForView(view_provider_type, for_writing)
: nullptr;
}
internal::GpuBufferStorage& GetStorageForViewOrDie(TypeId view_provider_type,
bool for_writing) const;
@ -158,25 +147,49 @@ class GpuBuffer {
.template down_cast<VP>();
}
std::shared_ptr<internal::GpuBufferStorage>& no_storage() const {
static auto placeholder =
std::static_pointer_cast<internal::GpuBufferStorage>(
std::make_shared<PlaceholderGpuBufferStorage>(
0, 0, GpuBufferFormat::kUnknown));
return placeholder;
}
// This class manages a set of alternative storages for the contents of a
// GpuBuffer. GpuBuffer was originally designed as a reference-type object,
// where a copy represents another reference to the same contents, so multiple
// GpuBuffer instances can share the same StorageHolder.
class StorageHolder {
public:
explicit StorageHolder(std::shared_ptr<internal::GpuBufferStorage> storage)
: StorageHolder(storage->width(), storage->height(),
storage->format()) {
storages_.push_back(std::move(storage));
}
explicit StorageHolder(int width, int height, Format format)
: width_(width), height_(height), format_(format) {}
const internal::GpuBufferStorage& current_storage() const {
return storages_.empty() ? *no_storage() : *storages_[0];
}
int width() const { return width_; }
int height() const { return height_; }
GpuBufferFormat format() const { return format_; }
internal::GpuBufferStorage& current_storage() {
return storages_.empty() ? *no_storage() : *storages_[0];
}
internal::GpuBufferStorage* GetStorageForView(TypeId view_provider_type,
bool for_writing) const;
// This is mutable because view methods that do not change the contents may
// still need to allocate new storages.
mutable std::vector<std::shared_ptr<internal::GpuBufferStorage>> storages_;
template <class T>
std::shared_ptr<T> internal_storage() const {
absl::MutexLock lock(&mutex_);
for (const auto& s : storages_)
if (s->down_cast<T>()) return std::static_pointer_cast<T>(s);
return nullptr;
}
std::string DebugString() const;
private:
int width_ = 0;
int height_ = 0;
GpuBufferFormat format_ = GpuBufferFormat::kUnknown;
// This is mutable because view methods that do not change the contents may
// still need to allocate new storages.
mutable absl::Mutex mutex_;
mutable std::vector<std::shared_ptr<internal::GpuBufferStorage>> storages_
ABSL_GUARDED_BY(mutex_);
};
std::shared_ptr<StorageHolder> holder_;
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
friend CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer);
@ -184,15 +197,15 @@ class GpuBuffer {
};
inline bool GpuBuffer::operator==(std::nullptr_t other) const {
return storages_.empty();
return holder_ == other;
}
inline bool GpuBuffer::operator==(const GpuBuffer& other) const {
return storages_ == other.storages_;
return holder_ == other.holder_;
}
inline GpuBuffer& GpuBuffer::operator=(std::nullptr_t other) {
storages_.clear();
holder_ = other;
return *this;
}

22
mediapipe/gpu/gpu_buffer_test.cc
View File

@ -20,6 +20,7 @@
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/tool/test_util.h"
#include "mediapipe/gpu/gl_texture_buffer.h"
#include "mediapipe/gpu/gl_texture_util.h"
#include "mediapipe/gpu/gpu_buffer_storage_ahwb.h"
#include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
@ -228,5 +229,26 @@ TEST_F(GpuBufferTest, GlTextureViewRetainsWhatItNeeds) {
EXPECT_TRUE(true);
}
TEST_F(GpuBufferTest, CopiesShareConversions) {
GpuBuffer buffer(300, 200, GpuBufferFormat::kBGRA32);
{
std::shared_ptr<ImageFrame> view = buffer.GetWriteView<ImageFrame>();
FillImageFrameRGBA(*view, 255, 0, 0, 255);
}
GpuBuffer other_handle = buffer;
RunInGlContext([&buffer] {
TempGlFramebuffer fb;
auto view = buffer.GetReadView<GlTextureView>(0);
});
// Check that other_handle also sees the same GlTextureBuffer as buffer.
// Note that this is deliberately written so that it still passes on platforms
// where we use another storage for GL textures (they will both be null).
// TODO: expose more accessors for testing?
EXPECT_EQ(other_handle.internal_storage<GlTextureBuffer>(),
buffer.internal_storage<GlTextureBuffer>());
}
} // anonymous namespace
} // namespace mediapipe