talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 6e5aa035cd1f6a9333962df5d3ab97a05bd5744e
This commit is contained in:
MediaPipe Team
2022-06-28 12:11:05 +00:00
committed by Sebastian Schmidt
parent 4a20e9909d
commit c688862570
144 changed files with 5772 additions and 2118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/formats/BUILD
+4
View File
@@ -448,6 +448,7 @@ cc_library(
srcs =
[
"tensor.cc",
"tensor_ahwb.cc",
],
hdrs = ["tensor.h"],
copts = select({
@@ -463,6 +464,9 @@ cc_library(
"-framework MetalKit",
],
"//conditions:default": [],
"//mediapipe:android": [
"-landroid",
],
}),
visibility = ["//visibility:public"],
deps = [
mediapipe/framework/formats/body_rig.proto
+3 -1
View File
@@ -19,7 +19,9 @@ package mediapipe;
// Joint of a 3D human model (e.g. elbow, knee, wrist). Contains 3D rotation of
// the joint and its visibility.
message Joint {
// Joint rotation in 6D contineous representation.
// Joint rotation in 6D contineous representation ordered as
// [a1, b1, a2, b2, a3, b3].
//
// Such representation is more sutable for NN model training and can be
// converted to quaternions and Euler angles if needed. Details can be found
// in https://arxiv.org/abs/1812.07035.
mediapipe/framework/formats/tensor.cc
+99 -37
View File
@@ -20,6 +20,9 @@
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/logging.h"
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#include "mediapipe/gpu/gl_base.h"
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#if MEDIAPIPE_METAL_ENABLED
#include <mach/mach_init.h>
@@ -319,28 +322,41 @@ void Tensor::AllocateOpenGlTexture2d() const {
Tensor::OpenGlBufferView Tensor::GetOpenGlBufferReadView() const {
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
LOG_IF(FATAL, !(valid_ & (kValidCpu | kValidOpenGlBuffer)))
<< "Tensor conversion between different GPU resources is not supported "
"yet.";
LOG_IF(FATAL, !(valid_ & (kValidCpu |
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
kValidAHardwareBuffer |
#endif // MEDIAPIPE_TENSOR_USE_AHWB
kValidOpenGlBuffer)))
<< "Tensor conversion between different GPU resources is not supported.";
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
AllocateOpenGlBuffer();
if (!(valid_ & kValidOpenGlBuffer)) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
void* ptr =
glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(),
GL_MAP_INVALIDATE_BUFFER_BIT | GL_MAP_WRITE_BIT);
std::memcpy(ptr, cpu_buffer_, bytes());
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
// If the call succeds then AHWB -> SSBO are synchronized so any usage of
// the SSBO is correct after this call.
if (!InsertAhwbToSsboFence()) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
void* ptr =
glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(),
GL_MAP_INVALIDATE_BUFFER_BIT | GL_MAP_WRITE_BIT);
std::memcpy(ptr, cpu_buffer_, bytes());
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
}
valid_ |= kValidOpenGlBuffer;
}
return {opengl_buffer_, std::move(lock)};
return {opengl_buffer_, std::move(lock),
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
&ssbo_read_
#else
nullptr
#endif // MEDIAPIPE_TENSOR_USE_AHWB
};
}
Tensor::OpenGlBufferView Tensor::GetOpenGlBufferWriteView() const {
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
AllocateOpenGlBuffer();
valid_ = kValidOpenGlBuffer;
return {opengl_buffer_, std::move(lock)};
return {opengl_buffer_, std::move(lock), nullptr};
}
void Tensor::AllocateOpenGlBuffer() const {
@@ -349,7 +365,10 @@ void Tensor::AllocateOpenGlBuffer() const {
LOG_IF(FATAL, !gl_context_) << "GlContext is not bound to the thread.";
glGenBuffers(1, &opengl_buffer_);
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
glBufferData(GL_SHADER_STORAGE_BUFFER, bytes(), NULL, GL_STREAM_COPY);
if (!AllocateAhwbMapToSsbo()) {
glBufferData(GL_SHADER_STORAGE_BUFFER, bytes(), NULL, GL_STREAM_COPY);
}
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
}
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
@@ -377,6 +396,8 @@ void Tensor::Move(Tensor* src) {
src->metal_buffer_ = nil;
#endif // MEDIAPIPE_METAL_ENABLED
MoveAhwbStuff(src);
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
gl_context_ = std::move(src->gl_context_);
frame_buffer_ = src->frame_buffer_;
@@ -395,27 +416,31 @@ void Tensor::Move(Tensor* src) {
Tensor::Tensor(ElementType element_type, const Shape& shape)
: element_type_(element_type), shape_(shape) {}
#if MEDIAPIPE_METAL_ENABLED
void Tensor::Invalidate() {
absl::MutexLock lock(&view_mutex_);
// If memory is allocated and not owned by the metal buffer.
// TODO: Re-design cpu buffer memory management.
if (cpu_buffer_ && !metal_buffer_) {
DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
}
metal_buffer_ = nil;
cpu_buffer_ = nullptr;
}
#else
void Tensor::Invalidate() {
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
GLuint cleanup_gl_tex = GL_INVALID_INDEX;
GLuint cleanup_gl_fb = GL_INVALID_INDEX;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
GLuint cleanup_gl_buf = GL_INVALID_INDEX;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
{
absl::MutexLock lock(&view_mutex_);
#if MEDIAPIPE_METAL_ENABLED
// If memory is allocated and not owned by the metal buffer.
// TODO: Re-design cpu buffer memory management.
if (cpu_buffer_ && !metal_buffer_) {
DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
}
metal_buffer_ = nil;
#else
if (cpu_buffer_) {
free(cpu_buffer_);
}
#endif // MEDIAPIPE_METAL_ENABLED
cpu_buffer_ = nullptr;
ReleaseAhwbStuff();
// Don't need to wait for the resource to be deleted bacause if will be
// released on last reference deletion inside the OpenGL driver.
@@ -429,28 +454,44 @@ void Tensor::Invalidate() {
}
// Do not hold the view mutex while invoking GlContext::RunWithoutWaiting,
// since that method may acquire the context's own lock.
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX ||
cleanup_gl_buf != GL_INVALID_INDEX)
gl_context_->RunWithoutWaiting([cleanup_gl_tex, cleanup_gl_fb
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
,
cleanup_gl_buf
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
]() {
if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX ||
cleanup_gl_buf != GL_INVALID_INDEX) {
gl_context_->RunWithoutWaiting(
[cleanup_gl_tex, cleanup_gl_fb, cleanup_gl_buf]() {
glDeleteTextures(1, &cleanup_gl_tex);
glDeleteFramebuffers(1, &cleanup_gl_fb);
glDeleteBuffers(1, &cleanup_gl_buf);
});
}
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX) {
gl_context_->RunWithoutWaiting([cleanup_gl_tex, cleanup_gl_fb]() {
glDeleteTextures(1, &cleanup_gl_tex);
glDeleteFramebuffers(1, &cleanup_gl_fb);
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
glDeleteBuffers(1, &cleanup_gl_buf);
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
});
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
if (cpu_buffer_) {
free(cpu_buffer_);
}
cpu_buffer_ = nullptr;
}
#endif // MEDIAPIPE_METAL_ENABLED
Tensor::CpuReadView Tensor::GetCpuReadView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
void* ptr = MapAhwbToCpuRead();
if (ptr) {
valid_ |= kValidCpu;
return {ptr, ahwb_, nullptr, std::move(lock)};
}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
AllocateCpuBuffer();
if (!(valid_ & kValidCpu)) {
// GPU-to-CPU synchronization and read-back.
@@ -512,18 +553,33 @@ Tensor::CpuReadView Tensor::GetCpuReadView() const {
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
valid_ |= kValidCpu;
}
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
return {cpu_buffer_, nullptr, nullptr, std::move(lock)};
#else
return {cpu_buffer_, std::move(lock)};
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
Tensor::CpuWriteView Tensor::GetCpuWriteView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
AllocateCpuBuffer();
valid_ = kValidCpu;
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
void* ptr = MapAhwbToCpuWrite();
if (ptr) {
return {ptr, ahwb_, &fence_fd_, std::move(lock)};
}
return {cpu_buffer_, nullptr, nullptr, std::move(lock)};
#else
return {cpu_buffer_, std::move(lock)};
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
void Tensor::AllocateCpuBuffer() const {
if (!cpu_buffer_) {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
if (AllocateAHardwareBuffer()) return;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_METAL_ENABLED
cpu_buffer_ = AllocateVirtualMemory(bytes());
#else
@@ -532,4 +588,10 @@ void Tensor::AllocateCpuBuffer() const {
}
}
void Tensor::SetPreferredStorageType(StorageType type) {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
use_ahwb_ = type == StorageType::kAhwb;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
} // namespace mediapipe
mediapipe/framework/formats/tensor.h
+144 -11
View File
@@ -30,6 +30,16 @@
#import <Metal/Metal.h>
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
#if __ANDROID_API__ < 26
#error MEDIAPIPE_TENSOR_USE_AHWB requires NDK version 26 or higher to be specified.
#endif // __ANDROID_API__ < 26
#include <android/hardware_buffer.h>
#include "third_party/GL/gl/include/EGL/egl.h"
#include "third_party/GL/gl/include/EGL/eglext.h"
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#include "mediapipe/gpu/gl_base.h"
#include "mediapipe/gpu/gl_context.h"
@@ -108,14 +118,37 @@ class Tensor {
return static_cast<typename std::tuple_element<
std::is_const<T>::value, std::tuple<P*, const P*> >::type>(buffer_);
}
CpuView(CpuView&& src) : View(std::move(src)), buffer_(src.buffer_) {
src.buffer_ = nullptr;
CpuView(CpuView&& src) : View(std::move(src)) {
buffer_ = std::exchange(src.buffer_, nullptr);
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
ahwb_ = std::exchange(src.ahwb_, nullptr);
fence_fd_ = std::exchange(src.fence_fd_, nullptr);
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
~CpuView() {
if (ahwb_) {
auto error = AHardwareBuffer_unlock(ahwb_, fence_fd_);
CHECK(error == 0) << "AHardwareBuffer_unlock " << error;
}
}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
protected:
friend class Tensor;
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
CpuView(T* buffer, AHardwareBuffer* ahwb, int* fence_fd,
std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)),
buffer_(buffer),
fence_fd_(fence_fd),
ahwb_(ahwb) {}
AHardwareBuffer* ahwb_;
int* fence_fd_;
#else
CpuView(T* buffer, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), buffer_(buffer) {}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
T* buffer_;
};
using CpuReadView = CpuView<const void>;
@@ -150,6 +183,60 @@ class Tensor {
MtlBufferView GetMtlBufferWriteView(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
class AHardwareBufferView : public View {
public:
AHardwareBuffer* handle() const { return handle_; }
AHardwareBufferView(AHardwareBufferView&& src) : View(std::move(src)) {
handle_ = std::exchange(src.handle_, nullptr);
file_descriptor_ = src.file_descriptor_;
fence_fd_ = std::exchange(src.fence_fd_, nullptr);
ahwb_written_ = std::exchange(src.ahwb_written_, nullptr);
release_callback_ = std::exchange(src.release_callback_, nullptr);
}
int file_descriptor() const { return file_descriptor_; }
void SetReadingFinishedFunc(std::function<bool()>&& func) {
CHECK(ahwb_written_)
<< "AHWB write view can't accept 'reading finished callback'";
*ahwb_written_ = std::move(func);
}
void SetWritingFinishedFD(int fd) {
CHECK(fence_fd_)
<< "AHWB read view can't accept 'writing finished file descriptor'";
*fence_fd_ = fd;
}
// The function is called when the tensor is released.
void SetReleaseCallback(std::function<void()> callback) {
*release_callback_ = std::move(callback);
}
protected:
friend class Tensor;
AHardwareBufferView(AHardwareBuffer* handle, int file_descriptor,
int* fence_fd, std::function<bool()>* ahwb_written,
std::function<void()>* release_callback,
std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)),
handle_(handle),
file_descriptor_(file_descriptor),
fence_fd_(fence_fd),
ahwb_written_(ahwb_written),
release_callback_(release_callback) {}
AHardwareBuffer* handle_;
int file_descriptor_;
// The view sets some Tensor's fields. The view is released prior to tensor.
int* fence_fd_;
std::function<bool()>* ahwb_written_;
std::function<void()>* release_callback_;
};
AHardwareBufferView GetAHardwareBufferReadView() const;
// size_alignment is an optional argument to tell the API to allocate
// a buffer that is padded to multiples of size_alignment bytes.
// size_alignment must be power of 2, i.e. 2, 4, 8, 16, 64, etc.
// If size_alignment is 0, then the buffer will not be padded.
AHardwareBufferView GetAHardwareBufferWriteView(int size_alignment = 0) const;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
// TODO: Use GlTextureView instead.
// Only float32 textures are supported with 1/2/3/4 depths.
@@ -188,16 +275,23 @@ class Tensor {
class OpenGlBufferView : public View {
public:
GLuint name() const { return name_; }
OpenGlBufferView(OpenGlBufferView&& src)
: View(std::move(src)), name_(src.name_) {
src.name_ = GL_INVALID_INDEX;
OpenGlBufferView(OpenGlBufferView&& src) : View(std::move(src)) {
name_ = std::exchange(src.name_, GL_INVALID_INDEX);
ssbo_read_ = std::exchange(src.ssbo_read_, nullptr);
}
~OpenGlBufferView() {
if (ssbo_read_) {
*ssbo_read_ = glFenceSync(GL_SYNC_GPU_COMMANDS_COMPLETE, 0);
}
}
protected:
friend class Tensor;
OpenGlBufferView(GLuint name, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), name_(name) {}
OpenGlBufferView(GLuint name, std::unique_ptr<absl::MutexLock>&& lock,
GLsync* ssbo_read)
: View(std::move(lock)), name_(name), ssbo_read_(ssbo_read) {}
GLuint name_;
GLsync* ssbo_read_;
};
// A valid OpenGL context must be bound to the calling thread due to possible
// GPU resource allocation.
@@ -223,16 +317,26 @@ class Tensor {
}
int bytes() const { return shape_.num_elements() * element_size(); }
bool ready_on_cpu() const { return valid_ & kValidCpu; }
bool ready_on_cpu() const {
return valid_ & (kValidAHardwareBuffer | kValidCpu);
}
bool ready_on_gpu() const {
return valid_ &
(kValidMetalBuffer | kValidOpenGlBuffer | kValidOpenGlTexture2d);
return valid_ & (kValidMetalBuffer | kValidOpenGlBuffer |
kValidAHardwareBuffer | kValidOpenGlTexture2d);
}
bool ready_as_metal_buffer() const { return valid_ & kValidMetalBuffer; }
bool ready_as_opengl_buffer() const { return valid_ & kValidOpenGlBuffer; }
bool ready_as_opengl_buffer() const {
return valid_ & (kValidAHardwareBuffer | kValidOpenGlBuffer);
}
bool ready_as_opengl_texture_2d() const {
return valid_ & kValidOpenGlTexture2d;
}
// Sets the type of underlying resource that is going to be allocated.
enum class StorageType {
kDefault,
kAhwb,
};
static void SetPreferredStorageType(StorageType type);
private:
void Move(Tensor*);
@@ -248,6 +352,7 @@ class Tensor {
kValidMetalBuffer = 1 << 1,
kValidOpenGlBuffer = 1 << 2,
kValidOpenGlTexture2d = 1 << 3,
kValidAHardwareBuffer = 1 << 5,
};
// A list of resource which are currently allocated and synchronized between
// each-other: valid_ = kValidCpu | kValidMetalBuffer;
@@ -264,6 +369,34 @@ class Tensor {
void AllocateMtlBuffer(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
mutable AHardwareBuffer* ahwb_ = nullptr;
// Signals when GPU finished writing into SSBO so AHWB can be used then. Or
// signals when writing into AHWB has been finished so GPU can read from SSBO.
// Sync and FD are bound together.
mutable EGLSyncKHR fence_sync_ = EGL_NO_SYNC_KHR;
// This FD signals when the writing into the SSBO has been finished.
mutable int ssbo_written_ = -1;
// An externally set FD that is wrapped with the EGL sync then to synchronize
// AHWB -> OpenGL SSBO.
mutable int fence_fd_ = -1;
// Reading from SSBO has been finished so SSBO can be released.
mutable GLsync ssbo_read_ = 0;
// An externally set function that signals when it is safe to release AHWB.
mutable std::function<bool()> ahwb_written_;
mutable std::function<void()> release_callback_;
bool AllocateAHardwareBuffer(int size_alignment = 0) const;
void CreateEglSyncAndFd() const;
// Use Ahwb for other views: OpenGL / CPU buffer.
static inline bool use_ahwb_ = false;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
bool AllocateAhwbMapToSsbo() const;
bool InsertAhwbToSsboFence() const;
void MoveAhwbStuff(Tensor* src);
void ReleaseAhwbStuff();
void* MapAhwbToCpuRead() const;
void* MapAhwbToCpuWrite() const;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
mutable std::shared_ptr<mediapipe::GlContext> gl_context_;
mutable GLuint opengl_texture2d_ = GL_INVALID_INDEX;
mediapipe/framework/formats/tensor_ahwb.cc
+382
View File
@@ -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
mediapipe/framework/formats/tensor_ahwb_test.cc
+59
View File
@@ -0,0 +1,59 @@
#include "mediapipe/gpu/gpu_test_base.h"
#include "testing/base/public/gmock.h"
#include "testing/base/public/gunit.h"
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
#include <android/hardware_buffer.h>
#include "mediapipe/framework/formats/tensor.h"
namespace mediapipe {
#if !MEDIAPIPE_DISABLE_GPU
class TensorAhwbTest : public mediapipe::GpuTestBase {
public:
};
TEST_F(TensorAhwbTest, TestCpuThenAHWB) {
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{1});
{
auto ptr = tensor.GetCpuWriteView().buffer<float>();
EXPECT_NE(ptr, nullptr);
}
{
auto ahwb = tensor.GetAHardwareBufferReadView().handle();
EXPECT_NE(ahwb, nullptr);
}
}
TEST_F(TensorAhwbTest, TestAHWBThenCpu) {
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{1});
{
auto ahwb = tensor.GetAHardwareBufferWriteView().handle();
EXPECT_NE(ahwb, nullptr);
}
{
auto ptr = tensor.GetCpuReadView().buffer<float>();
EXPECT_NE(ptr, nullptr);
}
}
TEST_F(TensorAhwbTest, TestCpuThenGl) {
RunInGlContext([] {
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{1});
{
auto ptr = tensor.GetCpuWriteView().buffer<float>();
EXPECT_NE(ptr, nullptr);
}
{
auto ssbo = tensor.GetOpenGlBufferReadView().name();
EXPECT_GT(ssbo, 0);
}
});
}
} // namespace mediapipe
#endif // !MEDIAPIPE_DISABLE_GPU
#endif // MEDIAPIPE_TENSOR_USE_AHWB
mediapipe/framework/formats/time_series_header.proto
+2
View File
@@ -21,6 +21,8 @@ syntax = "proto2";
package mediapipe;
option objc_class_prefix = "MediaPipe";
// Header for a uniformly sampled time series stream. Each Packet in
// the stream is a Matrix, and each column is a (vector-valued) sample of
// the series, i.e. each column corresponds to a distinct sample in time.