Project import generated by Copybara.
GitOrigin-RevId: 6e5aa035cd1f6a9333962df5d3ab97a05bd5744e
This commit is contained in:
MediaPipe Team
committed by
Sebastian Schmidt
Sebastian Schmidt
parent
4a20e9909d
commit
c688862570
@@ -0,0 +1,382 @@
|
||||
#include <cstdint>
|
||||
#include <utility>
|
||||
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
|
||||
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/port.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/gpu/gl_base.h"
|
||||
#include "third_party/GL/gl/include/EGL/egl.h"
|
||||
#include "third_party/GL/gl/include/EGL/eglext.h"
|
||||
#endif // MEDIAPIPE_TENSOR_USE_AHWB
|
||||
|
||||
namespace mediapipe {
|
||||
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
|
||||
|
||||
namespace {
|
||||
PFNGLBUFFERSTORAGEEXTERNALEXTPROC glBufferStorageExternalEXT;
|
||||
PFNEGLGETNATIVECLIENTBUFFERANDROIDPROC eglGetNativeClientBufferANDROID;
|
||||
PFNEGLDUPNATIVEFENCEFDANDROIDPROC eglDupNativeFenceFDANDROID;
|
||||
PFNEGLCREATESYNCKHRPROC eglCreateSyncKHR;
|
||||
PFNEGLWAITSYNCKHRPROC eglWaitSyncKHR;
|
||||
PFNEGLCLIENTWAITSYNCKHRPROC eglClientWaitSyncKHR;
|
||||
PFNEGLDESTROYSYNCKHRPROC eglDestroySyncKHR;
|
||||
|
||||
bool IsGlSupported() {
|
||||
static const bool extensions_allowed = [] {
|
||||
eglGetNativeClientBufferANDROID =
|
||||
reinterpret_cast<PFNEGLGETNATIVECLIENTBUFFERANDROIDPROC>(
|
||||
eglGetProcAddress("eglGetNativeClientBufferANDROID"));
|
||||
glBufferStorageExternalEXT =
|
||||
reinterpret_cast<PFNGLBUFFERSTORAGEEXTERNALEXTPROC>(
|
||||
eglGetProcAddress("glBufferStorageExternalEXT"));
|
||||
eglDupNativeFenceFDANDROID =
|
||||
reinterpret_cast<PFNEGLDUPNATIVEFENCEFDANDROIDPROC>(
|
||||
eglGetProcAddress("eglDupNativeFenceFDANDROID"));
|
||||
eglCreateSyncKHR = reinterpret_cast<PFNEGLCREATESYNCKHRPROC>(
|
||||
eglGetProcAddress("eglCreateSyncKHR"));
|
||||
eglWaitSyncKHR = reinterpret_cast<PFNEGLWAITSYNCKHRPROC>(
|
||||
eglGetProcAddress("eglWaitSyncKHR"));
|
||||
eglClientWaitSyncKHR = reinterpret_cast<PFNEGLCLIENTWAITSYNCKHRPROC>(
|
||||
eglGetProcAddress("eglClientWaitSyncKHR"));
|
||||
eglDestroySyncKHR = reinterpret_cast<PFNEGLDESTROYSYNCKHRPROC>(
|
||||
eglGetProcAddress("eglDestroySyncKHR"));
|
||||
return eglClientWaitSyncKHR && eglWaitSyncKHR &&
|
||||
eglGetNativeClientBufferANDROID && glBufferStorageExternalEXT &&
|
||||
eglCreateSyncKHR && eglDupNativeFenceFDANDROID && eglDestroySyncKHR;
|
||||
}();
|
||||
return extensions_allowed;
|
||||
}
|
||||
|
||||
absl::Status MapAHardwareBufferToGlBuffer(AHardwareBuffer* handle, size_t size,
|
||||
GLuint name) {
|
||||
if (!IsGlSupported()) {
|
||||
return absl::UnknownError(
|
||||
"No GL extension functions found to bind AHardwareBuffer and "
|
||||
"OpenGL buffer");
|
||||
}
|
||||
EGLClientBuffer native_buffer = eglGetNativeClientBufferANDROID(handle);
|
||||
if (!native_buffer) {
|
||||
return absl::UnknownError("Can't get native buffer");
|
||||
}
|
||||
glBufferStorageExternalEXT(GL_SHADER_STORAGE_BUFFER, 0, size, native_buffer,
|
||||
GL_MAP_READ_BIT | GL_MAP_WRITE_BIT |
|
||||
GL_MAP_COHERENT_BIT_EXT |
|
||||
GL_MAP_PERSISTENT_BIT_EXT);
|
||||
if (glGetError() == GL_NO_ERROR) {
|
||||
return absl::OkStatus();
|
||||
} else {
|
||||
return absl::InternalError("Error in glBufferStorageExternalEXT");
|
||||
}
|
||||
}
|
||||
|
||||
static inline int AlignedToPowerOf2(int value, int alignment) {
|
||||
// alignment must be a power of 2
|
||||
return ((value - 1) | (alignment - 1)) + 1;
|
||||
}
|
||||
|
||||
// This class keeps tensor's resources while the tensor is in use on GPU or TPU
|
||||
// but is already released on CPU. When a regular OpenGL buffer is bound to the
|
||||
// GPU queue for execution and released on client side then the buffer is still
|
||||
// not released because is being used by GPU. OpenGL driver keeps traking of
|
||||
// that. When OpenGL buffer is build on top of AHWB then the traking is done
|
||||
// with the DeleyedRelease which, actually, keeps record of all AHWBs allocated
|
||||
// and releases each of them if already used. EGL/GL fences are used to check
|
||||
// the status of a buffer.
|
||||
class DelayedReleaser {
|
||||
public:
|
||||
// Non-copyable
|
||||
DelayedReleaser(const DelayedReleaser&) = delete;
|
||||
DelayedReleaser& operator=(const DelayedReleaser&) = delete;
|
||||
// Non-movable
|
||||
DelayedReleaser(DelayedReleaser&&) = delete;
|
||||
DelayedReleaser& operator=(DelayedReleaser&&) = delete;
|
||||
|
||||
static void Add(AHardwareBuffer* ahwb, GLuint opengl_buffer,
|
||||
EGLSyncKHR ssbo_sync, GLsync ssbo_read,
|
||||
std::function<bool()>&& ahwb_written,
|
||||
std::shared_ptr<mediapipe::GlContext> gl_context,
|
||||
std::function<void()>&& callback) {
|
||||
static absl::Mutex mutex;
|
||||
absl::MutexLock lock(&mutex);
|
||||
// Using `new` to access a non-public constructor.
|
||||
to_release_.emplace_back(absl::WrapUnique(new DelayedReleaser(
|
||||
ahwb, opengl_buffer, ssbo_sync, ssbo_read, std::move(ahwb_written),
|
||||
gl_context, std::move(callback))));
|
||||
for (auto it = to_release_.begin(); it != to_release_.end();) {
|
||||
if ((*it)->IsSignaled()) {
|
||||
it = to_release_.erase(it);
|
||||
} else {
|
||||
++it;
|
||||
}
|
||||
}
|
||||
}
|
||||
~DelayedReleaser() {
|
||||
AHardwareBuffer_release(ahwb_);
|
||||
if (release_callback_) release_callback_();
|
||||
}
|
||||
|
||||
bool IsSignaled() {
|
||||
CHECK(!(ssbo_read_ && ahwb_written_))
|
||||
<< "ssbo_read_ and ahwb_written_ cannot both be set";
|
||||
if (ahwb_written_) {
|
||||
if (!ahwb_written_()) return false;
|
||||
gl_context_->Run([this]() {
|
||||
if (fence_sync_ != EGL_NO_SYNC_KHR && IsGlSupported()) {
|
||||
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
|
||||
if (egl_display != EGL_NO_DISPLAY) {
|
||||
eglDestroySyncKHR(egl_display, fence_sync_);
|
||||
}
|
||||
fence_sync_ = EGL_NO_SYNC_KHR;
|
||||
}
|
||||
glDeleteBuffers(1, &opengl_buffer_);
|
||||
opengl_buffer_ = GL_INVALID_INDEX;
|
||||
});
|
||||
return true;
|
||||
}
|
||||
|
||||
gl_context_->Run([this]() {
|
||||
if (ssbo_read_ != 0) {
|
||||
GLenum status = glClientWaitSync(ssbo_read_, 0,
|
||||
/* timeout ns = */ 0);
|
||||
if (status != GL_CONDITION_SATISFIED && status != GL_ALREADY_SIGNALED) {
|
||||
return;
|
||||
}
|
||||
glDeleteSync(ssbo_read_);
|
||||
ssbo_read_ = 0;
|
||||
|
||||
// Don't wait on ssbo_sync because it is ahead of ssbo_read_sync.
|
||||
if (fence_sync_ != EGL_NO_SYNC_KHR && IsGlSupported()) {
|
||||
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
|
||||
if (egl_display != EGL_NO_DISPLAY) {
|
||||
eglDestroySyncKHR(egl_display, fence_sync_);
|
||||
}
|
||||
}
|
||||
fence_sync_ = EGL_NO_SYNC_KHR;
|
||||
|
||||
glDeleteBuffers(1, &opengl_buffer_);
|
||||
opengl_buffer_ = GL_INVALID_INDEX;
|
||||
}
|
||||
});
|
||||
return opengl_buffer_ == GL_INVALID_INDEX;
|
||||
}
|
||||
|
||||
protected:
|
||||
AHardwareBuffer* ahwb_;
|
||||
GLuint opengl_buffer_;
|
||||
// TODO: use wrapper instead.
|
||||
EGLSyncKHR fence_sync_;
|
||||
// TODO: use wrapper instead.
|
||||
GLsync ssbo_read_;
|
||||
std::function<bool()> ahwb_written_;
|
||||
std::shared_ptr<mediapipe::GlContext> gl_context_;
|
||||
std::function<void()> release_callback_;
|
||||
static inline std::deque<std::unique_ptr<DelayedReleaser>> to_release_;
|
||||
|
||||
DelayedReleaser(AHardwareBuffer* ahwb, GLuint opengl_buffer,
|
||||
EGLSyncKHR fence_sync, GLsync ssbo_read,
|
||||
std::function<bool()>&& ahwb_written,
|
||||
std::shared_ptr<mediapipe::GlContext> gl_context,
|
||||
std::function<void()>&& callback)
|
||||
: ahwb_(ahwb),
|
||||
opengl_buffer_(opengl_buffer),
|
||||
fence_sync_(fence_sync),
|
||||
ssbo_read_(ssbo_read),
|
||||
ahwb_written_(std::move(ahwb_written)),
|
||||
gl_context_(gl_context),
|
||||
release_callback_(std::move(callback)) {}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
Tensor::AHardwareBufferView Tensor::GetAHardwareBufferReadView() const {
|
||||
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
|
||||
CHECK(valid_ != kValidNone) << "Tensor must be written prior to read from.";
|
||||
CHECK(!(valid_ & kValidOpenGlTexture2d))
|
||||
<< "Tensor conversion between OpenGL texture and AHardwareBuffer is not "
|
||||
"supported.";
|
||||
CHECK(ahwb_ || !(valid_ & kValidOpenGlBuffer))
|
||||
<< "Interoperability bettween OpenGL buffer and AHardwareBuffer is not "
|
||||
"supported on targe system.";
|
||||
CHECK(AllocateAHardwareBuffer())
|
||||
<< "AHardwareBuffer is not supported on the target system.";
|
||||
valid_ |= kValidAHardwareBuffer;
|
||||
if (valid_ & kValidOpenGlBuffer) CreateEglSyncAndFd();
|
||||
return {ahwb_,
|
||||
ssbo_written_,
|
||||
&fence_fd_, // The FD is created for SSBO -> AHWB synchronization.
|
||||
&ahwb_written_, // Filled by SetReadingFinishedFunc.
|
||||
&release_callback_,
|
||||
std::move(lock)};
|
||||
}
|
||||
|
||||
void Tensor::CreateEglSyncAndFd() const {
|
||||
gl_context_->Run([this]() {
|
||||
if (IsGlSupported()) {
|
||||
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
|
||||
if (egl_display != EGL_NO_DISPLAY) {
|
||||
fence_sync_ = eglCreateSyncKHR(egl_display,
|
||||
EGL_SYNC_NATIVE_FENCE_ANDROID, nullptr);
|
||||
if (fence_sync_ != EGL_NO_SYNC_KHR) {
|
||||
ssbo_written_ = eglDupNativeFenceFDANDROID(egl_display, fence_sync_);
|
||||
if (ssbo_written_ == -1) {
|
||||
eglDestroySyncKHR(egl_display, fence_sync_);
|
||||
fence_sync_ = EGL_NO_SYNC_KHR;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Can't use Sync object.
|
||||
if (fence_sync_ == EGL_NO_SYNC_KHR) glFinish();
|
||||
});
|
||||
}
|
||||
|
||||
Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
|
||||
int size_alignment) const {
|
||||
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
|
||||
CHECK(AllocateAHardwareBuffer(size_alignment))
|
||||
<< "AHardwareBuffer is not supported on the target system.";
|
||||
valid_ = kValidAHardwareBuffer;
|
||||
return {ahwb_,
|
||||
/*ssbo_written=*/-1,
|
||||
&fence_fd_, // For SetWritingFinishedFD.
|
||||
/*ahwb_written=*/nullptr, // The lifetime is managed by SSBO.
|
||||
&release_callback_,
|
||||
std::move(lock)};
|
||||
}
|
||||
|
||||
bool Tensor::AllocateAHardwareBuffer(int size_alignment) const {
|
||||
if (!use_ahwb_) return false;
|
||||
if (ahwb_ == nullptr) {
|
||||
AHardwareBuffer_Desc desc = {};
|
||||
if (size_alignment == 0) {
|
||||
desc.width = bytes();
|
||||
} else {
|
||||
// We expect allocations to be page-aligned, implicitly satisfying any
|
||||
// requirements from Edge TPU. No need to add a check for this,
|
||||
// since Edge TPU will check for us.
|
||||
desc.width = AlignedToPowerOf2(bytes(), size_alignment);
|
||||
}
|
||||
desc.height = 1;
|
||||
desc.layers = 1;
|
||||
desc.format = AHARDWAREBUFFER_FORMAT_BLOB;
|
||||
desc.usage = AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN |
|
||||
AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN |
|
||||
AHARDWAREBUFFER_USAGE_GPU_DATA_BUFFER;
|
||||
return AHardwareBuffer_allocate(&desc, &ahwb_) == 0;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Tensor::AllocateAhwbMapToSsbo() const {
|
||||
if (AllocateAHardwareBuffer()) {
|
||||
if (MapAHardwareBufferToGlBuffer(ahwb_, bytes(), opengl_buffer_).ok()) {
|
||||
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
|
||||
return true;
|
||||
}
|
||||
// Unable to make OpenGL <-> AHWB binding. Use regular SSBO instead.
|
||||
AHardwareBuffer_release(ahwb_);
|
||||
ahwb_ = nullptr;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// SSBO is created on top of AHWB. A fence is inserted into the GPU queue before
|
||||
// the GPU task that is going to read from the SSBO. When the writing into AHWB
|
||||
// is finished then the GPU reads from the SSBO.
|
||||
bool Tensor::InsertAhwbToSsboFence() const {
|
||||
if (!ahwb_) return false;
|
||||
if (fence_fd_ != -1) {
|
||||
// Can't wait for FD to be signaled on GPU.
|
||||
// TODO: wait on CPU instead.
|
||||
if (!IsGlSupported()) return true;
|
||||
|
||||
// Server-side fence.
|
||||
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
|
||||
if (egl_display == EGL_NO_DISPLAY) return true;
|
||||
EGLint sync_attribs[] = {EGL_SYNC_NATIVE_FENCE_FD_ANDROID,
|
||||
(EGLint)fence_fd_, EGL_NONE};
|
||||
fence_sync_ = eglCreateSyncKHR(egl_display, EGL_SYNC_NATIVE_FENCE_ANDROID,
|
||||
sync_attribs);
|
||||
if (fence_sync_ != EGL_NO_SYNC_KHR) {
|
||||
eglWaitSyncKHR(egl_display, fence_sync_, 0);
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void Tensor::MoveAhwbStuff(Tensor* src) {
|
||||
ahwb_ = std::exchange(src->ahwb_, nullptr);
|
||||
fence_sync_ = std::exchange(src->fence_sync_, EGL_NO_SYNC_KHR);
|
||||
ssbo_read_ = std::exchange(src->ssbo_read_, static_cast<GLsync>(0));
|
||||
ssbo_written_ = std::exchange(src->ssbo_written_, -1);
|
||||
fence_fd_ = std::exchange(src->fence_fd_, -1);
|
||||
ahwb_written_ = std::move(src->ahwb_written_);
|
||||
release_callback_ = std::move(src->release_callback_);
|
||||
}
|
||||
|
||||
void Tensor::ReleaseAhwbStuff() {
|
||||
if (fence_fd_ != -1) {
|
||||
close(fence_fd_);
|
||||
fence_fd_ = -1;
|
||||
}
|
||||
if (ahwb_) {
|
||||
if (ssbo_read_ != 0 || fence_sync_ != EGL_NO_SYNC_KHR) {
|
||||
if (ssbo_written_ != -1) close(ssbo_written_);
|
||||
DelayedReleaser::Add(ahwb_, opengl_buffer_, fence_sync_, ssbo_read_,
|
||||
std::move(ahwb_written_), gl_context_,
|
||||
std::move(release_callback_));
|
||||
opengl_buffer_ = GL_INVALID_INDEX;
|
||||
} else {
|
||||
AHardwareBuffer_release(ahwb_);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void* Tensor::MapAhwbToCpuRead() const {
|
||||
if (ahwb_) {
|
||||
if (!(valid_ & kValidCpu) && (valid_ & kValidOpenGlBuffer) &&
|
||||
ssbo_written_ == -1) {
|
||||
// EGLSync is failed. Use another synchronization method.
|
||||
// TODO: Use tflite::gpu::GlBufferSync and GlActiveSync.
|
||||
glFinish();
|
||||
}
|
||||
void* ptr;
|
||||
auto error =
|
||||
AHardwareBuffer_lock(ahwb_, AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN,
|
||||
ssbo_written_, nullptr, &ptr);
|
||||
CHECK(error == 0) << "AHardwareBuffer_lock " << error;
|
||||
close(ssbo_written_);
|
||||
ssbo_written_ = -1;
|
||||
return ptr;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void* Tensor::MapAhwbToCpuWrite() const {
|
||||
if (ahwb_) {
|
||||
// TODO: If previously acquired view is GPU write view then need to
|
||||
// be sure that writing is finished. That's a warning: two consequent write
|
||||
// views should be interleaved with read view.
|
||||
void* ptr;
|
||||
auto error = AHardwareBuffer_lock(
|
||||
ahwb_, AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN, -1, nullptr, &ptr);
|
||||
CHECK(error == 0) << "AHardwareBuffer_lock " << error;
|
||||
return ptr;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
#else // MEDIAPIPE_TENSOR_USE_AHWB
|
||||
|
||||
bool Tensor::AllocateAhwbMapToSsbo() const { return false; }
|
||||
bool Tensor::InsertAhwbToSsboFence() const { return false; }
|
||||
void Tensor::MoveAhwbStuff(Tensor* src) {}
|
||||
void Tensor::ReleaseAhwbStuff() {}
|
||||
void* Tensor::MapAhwbToCpuRead() const { return nullptr; }
|
||||
void* Tensor::MapAhwbToCpuWrite() const { return nullptr; }
|
||||
|
||||
#endif // MEDIAPIPE_TENSOR_USE_AHWB
|
||||
|
||||
} // namespace mediapipe
|
||||
Reference in New Issue
Block a user