Make GpuBuffer a shared_ptr to a storage collection
PiperOrigin-RevId: 493519590
This commit is contained in:
parent
402834b4f2
commit
523d16dffa
|
@ -289,7 +289,9 @@ cc_library(
|
||||||
deps = [
|
deps = [
|
||||||
":gpu_buffer_format",
|
":gpu_buffer_format",
|
||||||
":gpu_buffer_storage",
|
":gpu_buffer_storage",
|
||||||
|
"@com_google_absl//absl/functional:bind_front",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
|
"@com_google_absl//absl/synchronization",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/port:logging",
|
"//mediapipe/framework/port:logging",
|
||||||
":gpu_buffer_storage_image_frame",
|
":gpu_buffer_storage_image_frame",
|
||||||
|
|
|
@ -3,6 +3,7 @@
|
||||||
#include <memory>
|
#include <memory>
|
||||||
#include <utility>
|
#include <utility>
|
||||||
|
|
||||||
|
#include "absl/functional/bind_front.h"
|
||||||
#include "absl/strings/str_cat.h"
|
#include "absl/strings/str_cat.h"
|
||||||
#include "absl/strings/str_join.h"
|
#include "absl/strings/str_join.h"
|
||||||
#include "mediapipe/framework/port/logging.h"
|
#include "mediapipe/framework/port/logging.h"
|
||||||
|
@ -25,57 +26,101 @@ struct StorageTypeFormatter {
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
std::string GpuBuffer::DebugString() const {
|
std::string GpuBuffer::DebugString() const {
|
||||||
return absl::StrCat("GpuBuffer[",
|
return holder_ ? absl::StrCat("GpuBuffer[", width(), "x", height(), " ",
|
||||||
absl::StrJoin(storages_, ", ", StorageTypeFormatter()),
|
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 {
|
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_) {
|
absl::MutexLock lock(&mutex_);
|
||||||
if (s->can_down_cast_to(view_provider_type)) {
|
// First see if any current storage supports the view.
|
||||||
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_) {
|
for (const auto& s : storages_) {
|
||||||
if (auto converter = internal::GpuBufferStorageRegistry::Get()
|
if (s->can_down_cast_to(view_provider_type)) {
|
||||||
.StorageConverterForViewProvider(
|
chosen_storage = s;
|
||||||
view_provider_type, s->storage_type())) {
|
break;
|
||||||
if (auto new_storage = converter(s)) {
|
}
|
||||||
storages_.push_back(new_storage);
|
}
|
||||||
chosen_storage = &storages_.back();
|
|
||||||
|
// 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;
|
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) {
|
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) {
|
if (chosen_storage) {
|
||||||
// Discard all other storages.
|
// Discard all other storages.
|
||||||
storages_ = {*chosen_storage};
|
absl::MutexLock lock(&mutex_);
|
||||||
chosen_storage = &storages_.back();
|
swap(old_storages, storages_);
|
||||||
|
storages_ = {chosen_storage};
|
||||||
} else {
|
} else {
|
||||||
// Allocate a new storage supporting the requested view.
|
// Allocate a new storage supporting the requested view.
|
||||||
if (auto factory =
|
if (auto factory =
|
||||||
internal::GpuBufferStorageRegistry::Get()
|
internal::GpuBufferStorageRegistry::Get()
|
||||||
.StorageFactoryForViewProvider(view_provider_type)) {
|
.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)};
|
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(
|
internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
|
||||||
|
@ -84,8 +129,7 @@ internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
|
||||||
GpuBuffer::GetStorageForView(view_provider_type, for_writing);
|
GpuBuffer::GetStorageForView(view_provider_type, for_writing);
|
||||||
CHECK(chosen_storage) << "no view provider found for requested view "
|
CHECK(chosen_storage) << "no view provider found for requested view "
|
||||||
<< view_provider_type.name() << "; storages available: "
|
<< view_provider_type.name() << "; storages available: "
|
||||||
<< absl::StrJoin(storages_, ", ",
|
<< (holder_ ? holder_->DebugString() : "invalid");
|
||||||
StorageTypeFormatter());
|
|
||||||
DCHECK(chosen_storage->can_down_cast_to(view_provider_type));
|
DCHECK(chosen_storage->can_down_cast_to(view_provider_type));
|
||||||
return *chosen_storage;
|
return *chosen_storage;
|
||||||
}
|
}
|
||||||
|
|
|
@ -15,9 +15,12 @@
|
||||||
#ifndef MEDIAPIPE_GPU_GPU_BUFFER_H_
|
#ifndef MEDIAPIPE_GPU_GPU_BUFFER_H_
|
||||||
#define MEDIAPIPE_GPU_GPU_BUFFER_H_
|
#define MEDIAPIPE_GPU_GPU_BUFFER_H_
|
||||||
|
|
||||||
|
#include <algorithm>
|
||||||
|
#include <functional>
|
||||||
#include <memory>
|
#include <memory>
|
||||||
#include <utility>
|
#include <utility>
|
||||||
|
|
||||||
|
#include "absl/synchronization/mutex.h"
|
||||||
#include "mediapipe/framework/formats/image_frame.h"
|
#include "mediapipe/framework/formats/image_frame.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer_format.h"
|
#include "mediapipe/gpu/gpu_buffer_format.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer_storage.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
|
// Creates an empty buffer of a given size and format. It will be allocated
|
||||||
// when a view is requested.
|
// when a view is requested.
|
||||||
GpuBuffer(int width, int height, Format format)
|
GpuBuffer(int width, int height, Format format)
|
||||||
: GpuBuffer(std::make_shared<PlaceholderGpuBufferStorage>(width, height,
|
: holder_(std::make_shared<StorageHolder>(width, height, format)) {}
|
||||||
format)) {}
|
|
||||||
|
|
||||||
// Copy and move constructors and assignment operators are supported.
|
// Copy and move constructors and assignment operators are supported.
|
||||||
GpuBuffer(const GpuBuffer& other) = default;
|
GpuBuffer(const GpuBuffer& other) = default;
|
||||||
|
@ -70,9 +72,8 @@ class GpuBuffer {
|
||||||
// are not portable. Applications and calculators should normally obtain
|
// are not portable. Applications and calculators should normally obtain
|
||||||
// GpuBuffers in a portable way from the framework, e.g. using
|
// GpuBuffers in a portable way from the framework, e.g. using
|
||||||
// GpuBufferMultiPool.
|
// GpuBufferMultiPool.
|
||||||
explicit GpuBuffer(std::shared_ptr<internal::GpuBufferStorage> storage) {
|
explicit GpuBuffer(std::shared_ptr<internal::GpuBufferStorage> storage)
|
||||||
storages_.push_back(std::move(storage));
|
: holder_(std::make_shared<StorageHolder>(std::move(storage))) {}
|
||||||
}
|
|
||||||
|
|
||||||
#if !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
#if !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
||||||
// This is used to support backward-compatible construction of GpuBuffer from
|
// This is used to support backward-compatible construction of GpuBuffer from
|
||||||
|
@ -84,9 +85,11 @@ class GpuBuffer {
|
||||||
: GpuBuffer(internal::AsGpuBufferStorage(storage_convertible)) {}
|
: GpuBuffer(internal::AsGpuBufferStorage(storage_convertible)) {}
|
||||||
#endif // !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
#endif // !MEDIAPIPE_DISABLE_GPU && MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
||||||
|
|
||||||
int width() const { return current_storage().width(); }
|
int width() const { return holder_ ? holder_->width() : 0; }
|
||||||
int height() const { return current_storage().height(); }
|
int height() const { return holder_ ? holder_->height() : 0; }
|
||||||
GpuBufferFormat format() const { return current_storage().format(); }
|
GpuBufferFormat format() const {
|
||||||
|
return holder_ ? holder_->format() : GpuBufferFormat::kUnknown;
|
||||||
|
}
|
||||||
|
|
||||||
// Converts to true iff valid.
|
// Converts to true iff valid.
|
||||||
explicit operator bool() const { return operator!=(nullptr); }
|
explicit operator bool() const { return operator!=(nullptr); }
|
||||||
|
@ -122,31 +125,17 @@ class GpuBuffer {
|
||||||
// using views.
|
// using views.
|
||||||
template <class T>
|
template <class T>
|
||||||
std::shared_ptr<T> internal_storage() const {
|
std::shared_ptr<T> internal_storage() const {
|
||||||
for (const auto& s : storages_)
|
return holder_ ? holder_->internal_storage<T>() : nullptr;
|
||||||
if (s->down_cast<T>()) return std::static_pointer_cast<T>(s);
|
|
||||||
return nullptr;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
std::string DebugString() const;
|
std::string DebugString() const;
|
||||||
|
|
||||||
private:
|
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,
|
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,
|
internal::GpuBufferStorage& GetStorageForViewOrDie(TypeId view_provider_type,
|
||||||
bool for_writing) const;
|
bool for_writing) const;
|
||||||
|
@ -158,25 +147,49 @@ class GpuBuffer {
|
||||||
.template down_cast<VP>();
|
.template down_cast<VP>();
|
||||||
}
|
}
|
||||||
|
|
||||||
std::shared_ptr<internal::GpuBufferStorage>& no_storage() const {
|
// This class manages a set of alternative storages for the contents of a
|
||||||
static auto placeholder =
|
// GpuBuffer. GpuBuffer was originally designed as a reference-type object,
|
||||||
std::static_pointer_cast<internal::GpuBufferStorage>(
|
// where a copy represents another reference to the same contents, so multiple
|
||||||
std::make_shared<PlaceholderGpuBufferStorage>(
|
// GpuBuffer instances can share the same StorageHolder.
|
||||||
0, 0, GpuBufferFormat::kUnknown));
|
class StorageHolder {
|
||||||
return placeholder;
|
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 {
|
int width() const { return width_; }
|
||||||
return storages_.empty() ? *no_storage() : *storages_[0];
|
int height() const { return height_; }
|
||||||
}
|
GpuBufferFormat format() const { return format_; }
|
||||||
|
|
||||||
internal::GpuBufferStorage& current_storage() {
|
internal::GpuBufferStorage* GetStorageForView(TypeId view_provider_type,
|
||||||
return storages_.empty() ? *no_storage() : *storages_[0];
|
bool for_writing) const;
|
||||||
}
|
|
||||||
|
|
||||||
// This is mutable because view methods that do not change the contents may
|
template <class T>
|
||||||
// still need to allocate new storages.
|
std::shared_ptr<T> internal_storage() const {
|
||||||
mutable std::vector<std::shared_ptr<internal::GpuBufferStorage>> storages_;
|
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
|
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
||||||
friend CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer);
|
friend CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer);
|
||||||
|
@ -184,15 +197,15 @@ class GpuBuffer {
|
||||||
};
|
};
|
||||||
|
|
||||||
inline bool GpuBuffer::operator==(std::nullptr_t other) const {
|
inline bool GpuBuffer::operator==(std::nullptr_t other) const {
|
||||||
return storages_.empty();
|
return holder_ == other;
|
||||||
}
|
}
|
||||||
|
|
||||||
inline bool GpuBuffer::operator==(const GpuBuffer& other) const {
|
inline bool GpuBuffer::operator==(const GpuBuffer& other) const {
|
||||||
return storages_ == other.storages_;
|
return holder_ == other.holder_;
|
||||||
}
|
}
|
||||||
|
|
||||||
inline GpuBuffer& GpuBuffer::operator=(std::nullptr_t other) {
|
inline GpuBuffer& GpuBuffer::operator=(std::nullptr_t other) {
|
||||||
storages_.clear();
|
holder_ = other;
|
||||||
return *this;
|
return *this;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -20,6 +20,7 @@
|
||||||
#include "mediapipe/framework/port/gmock.h"
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
#include "mediapipe/framework/port/gtest.h"
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
#include "mediapipe/framework/tool/test_util.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/gl_texture_util.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer_storage_ahwb.h"
|
#include "mediapipe/gpu/gpu_buffer_storage_ahwb.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
|
#include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
|
||||||
|
@ -228,5 +229,26 @@ TEST_F(GpuBufferTest, GlTextureViewRetainsWhatItNeeds) {
|
||||||
EXPECT_TRUE(true);
|
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
|
} // anonymous namespace
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|
Loading…
Reference in New Issue
Block a user