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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: f7d09ed033907b893638a8eb4148efa11c0f09a6
This commit is contained in:
MediaPipe Team
2020-11-04 19:09:58 -05:00
committed by chuoling
parent a8d6ce95c4
commit f96eadd6df
250 changed files with 15261 additions and 4620 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/formats/BUILD
+60
View File
@@ -48,6 +48,18 @@ mediapipe_proto_library(
visibility = ["//visibility:public"],
)
mediapipe_register_type(
base_name = "classification",
include_headers = ["mediapipe/framework/formats/classification.pb.h"],
types = [
"::mediapipe::Classification",
"::mediapipe::ClassificationList",
"::std::vector<::mediapipe::Classification>",
"::std::vector<::mediapipe::ClassificationList>",
],
deps = [":classification_cc_proto"],
)
mediapipe_proto_library(
name = "image_format_proto",
srcs = ["image_format.proto"],
@@ -289,3 +301,51 @@ cc_test(
"@com_google_absl//absl/memory",
],
)
cc_library(
name = "tensor",
srcs = ["tensor.cc"],
hdrs = ["tensor.h"],
copts = select({
"//mediapipe:apple": [
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
"//conditions:default": [],
}),
linkopts = select({
"//mediapipe:ios": [
"-framework CoreVideo",
"-framework MetalKit",
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
"@com_google_absl//absl/memory",
"@com_google_absl//absl/synchronization",
"//mediapipe/framework:port",
"//mediapipe/framework/port:logging",
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_base",
"//mediapipe/gpu:gl_context",
],
}),
)
cc_test(
name = "tensor_test",
srcs = ["tensor_test.cc"],
deps = [
":tensor",
"//mediapipe/framework/port:gtest_main",
] + select({
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer_format",
],
"//mediapipe/gpu:disable_gpu": [],
}),
)
mediapipe/framework/formats/tensor.cc
+431
View File
@@ -0,0 +1,431 @@
// Copyright 2020 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mediapipe/framework/formats/tensor.h"
#include <cstdint>
#include <utility>
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/logging.h"
#if MEDIAPIPE_METAL_ENABLED
#include <mach/mach_init.h>
#include <mach/vm_map.h>
#else
#include <cstdlib>
#endif // MEDIAPIPE_METAL_ENABLED
namespace mediapipe {
int BhwcBatchFromShape(const Tensor::Shape& shape) {
LOG_IF(FATAL, shape.dims.empty())
<< "Tensor::Shape must be non-empty to retrieve a named dimension";
return shape.dims[0];
}
int BhwcHeightFromShape(const Tensor::Shape& shape) {
LOG_IF(FATAL, shape.dims.empty())
<< "Tensor::Shape must be non-empty to retrieve a named dimension";
return shape.dims.size() < 4 ? 1 : shape.dims[shape.dims.size() - 3];
}
int BhwcWidthFromShape(const Tensor::Shape& shape) {
LOG_IF(FATAL, shape.dims.empty())
<< "Tensor::Shape must be non-empty to retrieve a named dimension";
return shape.dims.size() < 3 ? 1 : shape.dims[shape.dims.size() - 2];
}
int BhwcDepthFromShape(const Tensor::Shape& shape) {
LOG_IF(FATAL, shape.dims.empty())
<< "Tensor::Shape must be non-empty to retrieve a named dimension";
return shape.dims.size() < 2 ? 1 : shape.dims[shape.dims.size() - 1];
}
// TODO: Match channels count and padding for Texture2D:
// 1) support 1/2/4 channesl texture for 1/2/3-4 depth.
// 2) Allocate cpu_buffer_ with padded amount of memory
// 3) pad/"unpad" the bitmap after transfer CPU <-> GPU
#if MEDIAPIPE_METAL_ENABLED
namespace {
// MTLBuffer can use existing properly aligned and allocated CPU memory.
size_t AlignToPageSize(size_t size) {
auto page_size = getpagesize();
return (size + page_size - 1) / page_size * page_size;
}
void* AllocateVirtualMemory(size_t size) {
vm_address_t data;
auto error = vm_allocate(mach_task_self(), &data, AlignToPageSize(size),
VM_FLAGS_ANYWHERE);
LOG_IF(FATAL, error != KERN_SUCCESS)
<< "Can't allocate virtual memory for Tensor.";
return reinterpret_cast<void*>(data);
}
void DeallocateVirtualMemory(void* pointer, size_t size) {
vm_deallocate(mach_task_self(), reinterpret_cast<vm_address_t>(pointer),
size);
}
} // namespace
Tensor::MtlBufferView Tensor::GetMtlBufferReadView(
id<MTLCommandBuffer> command_buffer) const {
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
LOG_IF(FATAL, !(valid_ & (kValidCpu | kValidMetalBuffer)))
<< "Tensor conversion between different GPU resources is not supported "
"yet.";
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
valid_ |= kValidMetalBuffer;
AllocateMtlBuffer([command_buffer device]);
return {metal_buffer_, std::move(lock)};
}
Tensor::MtlBufferView Tensor::GetMtlBufferWriteView(
id<MTLCommandBuffer> command_buffer) const {
// Don't overwrite command buffer at which the metal buffer has been written
// so we can wait until completed.
command_buffer_ = command_buffer;
return GetMtlBufferWriteView([command_buffer device]);
}
Tensor::MtlBufferView Tensor::GetMtlBufferWriteView(
id<MTLDevice> device) const {
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
valid_ = kValidMetalBuffer;
AllocateMtlBuffer(device);
return {metal_buffer_, std::move(lock)};
}
void Tensor::AllocateMtlBuffer(id<MTLDevice> device) const {
device_ = device;
if (!cpu_buffer_) {
// It also means that the metal buffer is not allocated yet.
cpu_buffer_ = AllocateVirtualMemory(bytes());
}
if (!metal_buffer_) {
metal_buffer_ =
[device_ newBufferWithBytesNoCopy:cpu_buffer_
length:AlignToPageSize(bytes())
options:MTLResourceStorageModeShared |
MTLResourceCPUCacheModeDefaultCache
deallocator:^(void* pointer, NSUInteger length) {
DeallocateVirtualMemory(pointer, length);
}];
}
}
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
Tensor::OpenGlTexture2dView Tensor::GetOpenGlTexture2dReadView() const {
LOG_IF(FATAL, BhwcDepthFromShape(shape_) > 4)
<< "OpenGlTexture2d supports depth <= 4. Current depth is "
<< BhwcDepthFromShape(shape_);
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
LOG_IF(FATAL, !(valid_ & (kValidCpu | kValidOpenGlTexture2d)))
<< "Tensor conversion between different GPU resources is not supported "
"yet.";
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
AllocateOpenGlTexture2d();
if (!(valid_ & kValidOpenGlTexture2d)) {
uint8_t* buffer;
std::unique_ptr<uint8_t[]> temp_buffer;
if (BhwcDepthFromShape(shape_) == 4) {
buffer = reinterpret_cast<uint8_t*>(cpu_buffer_);
} else {
const int padded_depth = 4;
const int padded_depth_size = padded_depth * element_size();
const int padded_size = BhwcBatchFromShape(shape_) *
BhwcHeightFromShape(shape_) *
BhwcWidthFromShape(shape_) * padded_depth_size;
temp_buffer = absl::make_unique<uint8_t[]>(padded_size);
buffer = temp_buffer.get();
uint8_t* src_buffer = reinterpret_cast<uint8_t*>(cpu_buffer_);
const int actual_depth_size = BhwcDepthFromShape(shape_) * element_size();
for (int e = 0;
e < BhwcBatchFromShape(shape_) * BhwcHeightFromShape(shape_) *
BhwcWidthFromShape(shape_);
e++) {
std::memcpy(buffer, src_buffer, actual_depth_size);
src_buffer += actual_depth_size;
buffer += padded_depth_size;
}
}
// Transfer from CPU memory into GPU memory.
glBindTexture(GL_TEXTURE_2D, opengl_texture2d_);
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, BhwcWidthFromShape(shape_),
BhwcHeightFromShape(shape_), GL_RGBA, GL_FLOAT, buffer);
glBindTexture(GL_TEXTURE_2D, 0);
valid_ |= kValidOpenGlTexture2d;
}
return {opengl_texture2d_, std::move(lock)};
}
Tensor::OpenGlTexture2dView Tensor::GetOpenGlTexture2dWriteView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
AllocateOpenGlTexture2d();
valid_ = kValidOpenGlTexture2d;
return {opengl_texture2d_, std::move(lock)};
}
void Tensor::AllocateOpenGlTexture2d() const {
if (opengl_texture2d_ == GL_INVALID_INDEX) {
gl_context_ = mediapipe::GlContext::GetCurrent();
LOG_IF(FATAL, !gl_context_) << "GlContext is not bound to the thread.";
glGenTextures(1, &opengl_texture2d_);
glBindTexture(GL_TEXTURE_2D, opengl_texture2d_);
// Texture2D represents a buffer with computable data so should be fetched
// but not sampled - can affect performance. Also on GLES2.0 sampling is not
// supported from floating point textures.
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
glTexStorage2D(GL_TEXTURE_2D, 1, GL_RGBA32F, BhwcWidthFromShape(shape_),
BhwcHeightFromShape(shape_));
glBindTexture(GL_TEXTURE_2D, 0);
}
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
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.";
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);
valid_ |= kValidOpenGlBuffer;
}
return {opengl_buffer_, std::move(lock)};
}
Tensor::OpenGlBufferView Tensor::GetOpenGlBufferWriteView() const {
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
AllocateOpenGlBuffer();
valid_ = kValidOpenGlBuffer;
return {opengl_buffer_, std::move(lock)};
}
void Tensor::AllocateOpenGlBuffer() const {
if (opengl_buffer_ == GL_INVALID_INDEX) {
gl_context_ = mediapipe::GlContext::GetCurrent();
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);
}
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
Tensor& Tensor::operator=(Tensor&& src) {
if (this != &src) {
Invalidate();
Move(&src);
}
return *this;
}
void Tensor::Move(Tensor* src) {
valid_ = src->valid_;
src->valid_ = kValidNone;
shape_ = src->shape();
element_type_ = src->element_type();
src->element_type_ = ElementType::kNone; // Mark as invalidated.
cpu_buffer_ = src->cpu_buffer_;
src->cpu_buffer_ = nullptr;
#if MEDIAPIPE_METAL_ENABLED
device_ = src->device_;
command_buffer_ = src->command_buffer_;
metal_buffer_ = src->metal_buffer_;
src->metal_buffer_ = nil;
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
gl_context_ = std::move(src->gl_context_);
opengl_texture2d_ = src->opengl_texture2d_;
src->opengl_texture2d_ = GL_INVALID_INDEX;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
opengl_buffer_ = src->opengl_buffer_;
src->opengl_buffer_ = GL_INVALID_INDEX;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
}
Tensor::Tensor(ElementType element_type, const Shape& shape)
: element_type_(element_type), shape_(shape) {}
void Tensor::Invalidate() {
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;
// Don't need to wait for the resource to be deleted bacause if will be
// released on last reference deletion inside the OpenGL driver.
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
if (opengl_texture2d_ != GL_INVALID_INDEX) {
GLuint opengl_texture2d = opengl_texture2d_;
gl_context_->RunWithoutWaiting(
[opengl_texture2d]() { glDeleteTextures(1, &opengl_texture2d); });
opengl_texture2d_ = GL_INVALID_INDEX;
}
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
if (opengl_buffer_ != GL_INVALID_INDEX) {
GLuint opengl_buffer = opengl_buffer_;
gl_context_->RunWithoutWaiting(
[opengl_buffer]() { glDeleteBuffers(1, &opengl_buffer); });
opengl_buffer_ = GL_INVALID_INDEX;
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
}
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.";
AllocateCpuBuffer();
if (!(valid_ & kValidCpu)) {
// GPU-to-CPU synchronization and read-back.
#if MEDIAPIPE_METAL_ENABLED
if (valid_ & kValidMetalBuffer) {
LOG_IF(FATAL, !command_buffer_) << "Metal -> CPU synchronization "
"requires MTLCommandBuffer to be set.";
if (command_buffer_) {
[command_buffer_ waitUntilCompleted];
}
}
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
if (valid_ & kValidOpenGlBuffer) {
gl_context_->Run([this]() {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
const void* ptr = glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(),
GL_MAP_READ_BIT);
std::memcpy(cpu_buffer_, ptr, bytes());
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
});
} else
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
// Transfer data from texture if not transferred from SSBO/MTLBuffer
// yet.
if (valid_ & kValidOpenGlTexture2d) {
gl_context_->Run([this]() {
GLint current_fbo;
glGetIntegerv(GL_FRAMEBUFFER_BINDING, &current_fbo);
uint8_t* buffer;
std::unique_ptr<uint8_t[]> temp_buffer;
if (BhwcDepthFromShape(shape_) == 4) {
buffer = reinterpret_cast<uint8_t*>(cpu_buffer_);
} else {
const int padded_depth = (BhwcDepthFromShape(shape_) + 3) / 4 * 4;
const int padded_size =
BhwcBatchFromShape(shape_) * BhwcHeightFromShape(shape_) *
BhwcWidthFromShape(shape_) * padded_depth * element_size();
temp_buffer = absl::make_unique<uint8_t[]>(padded_size);
buffer = temp_buffer.get();
}
GLint color_attachment_name;
glGetFramebufferAttachmentParameteriv(
GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
GL_FRAMEBUFFER_ATTACHMENT_OBJECT_NAME, &color_attachment_name);
if (color_attachment_name != opengl_texture2d_) {
// Save the viewport. Note that we assume that the color attachment is
// a GL_TEXTURE_2D texture.
GLint viewport[4];
glGetIntegerv(GL_VIEWPORT, viewport);
// Set the data from GLTexture object.
glViewport(0, 0, BhwcWidthFromShape(shape_),
BhwcHeightFromShape(shape_));
glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
GL_TEXTURE_2D, opengl_texture2d_, 0);
glReadPixels(0, 0, BhwcWidthFromShape(shape_),
BhwcHeightFromShape(shape_), GL_RGBA, GL_FLOAT, buffer);
// Restore from the saved viewport and color attachment name.
glViewport(viewport[0], viewport[1], viewport[2], viewport[3]);
glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
GL_TEXTURE_2D, color_attachment_name, 0);
} else {
glReadPixels(0, 0, BhwcWidthFromShape(shape_),
BhwcHeightFromShape(shape_), GL_RGBA, GL_FLOAT, buffer);
}
if (BhwcDepthFromShape(shape_) < 4) {
uint8_t* dest_buffer = reinterpret_cast<uint8_t*>(cpu_buffer_);
const int actual_depth_size =
BhwcDepthFromShape(shape_) * element_size();
const int padded_depth_size = 4 * element_size();
for (int e = 0;
e < BhwcBatchFromShape(shape_) * BhwcHeightFromShape(shape_) *
BhwcWidthFromShape(shape_);
e++) {
std::memcpy(dest_buffer, buffer, actual_depth_size);
dest_buffer += actual_depth_size;
buffer += padded_depth_size;
}
}
});
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
valid_ |= kValidCpu;
}
return {cpu_buffer_, std::move(lock)};
}
Tensor::CpuWriteView Tensor::GetCpuWriteView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
AllocateCpuBuffer();
valid_ = kValidCpu;
return {cpu_buffer_, std::move(lock)};
}
void Tensor::AllocateCpuBuffer() const {
if (!cpu_buffer_) {
#if MEDIAPIPE_METAL_ENABLED
cpu_buffer_ = AllocateVirtualMemory(bytes());
#else
cpu_buffer_ = malloc(bytes());
#endif // MEDIAPIPE_METAL_ENABLED
}
}
} // namespace mediapipe
mediapipe/framework/formats/tensor.h
+266
View File
@@ -0,0 +1,266 @@
// Copyright 2020 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_H_
#define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_H_
#include <algorithm>
#include <initializer_list>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/port.h"
#if MEDIAPIPE_METAL_ENABLED
#import <Metal/Metal.h>
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#include "mediapipe/gpu/gl_base.h"
#include "mediapipe/gpu/gl_context.h"
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
namespace mediapipe {
// Tensor is a container of multi-dimensional data that supports sharing the
// content across different backends and APIs, currently: CPU / Metal / OpenGL.
// Texture2DView is limited to 4 dimensions.
// The content is accessible through requesting device specific views.
// Acquiring a view guarantees that the content is not changed by another thread
// until the view is released.
//
// Tensor::MtlBufferView view = tensor.GetMtlBufferWriteView(mtl_device);
// mtl_device is used to create MTLBuffer
// id<MTLBuffer> buffer = view.buffer();
// For OpenGL the code below must be called by a thread with valid OpenGL ES
// context bound:
// GLuint buffer = view.buffer();
// Then the buffer can be bound to the GPU command buffer.
// ...binding the buffer to the command buffer...
// ...commiting command buffer and releasing the view...
//
// The following request for the CPU view will be blocked until the GPU view is
// released and the GPU task is finished.
//
// auto view = tensor.GetCpuReadView();
// float* pointer = view.buffer<float>();
// ...reading the cpu memory...
class Tensor {
class View {
public:
// Non-copyable.
View(const View&) = delete;
View& operator=(const View&) = delete;
View(View&& src) = default;
protected:
View(std::unique_ptr<absl::MutexLock>&& lock) : lock_(std::move(lock)) {}
std::unique_ptr<absl::MutexLock> lock_;
};
public:
// No resources are allocated here.
enum class ElementType { kNone, kFloat16, kFloat32 };
struct Shape {
Shape() = default;
Shape(std::initializer_list<int> dimensions) : dims(dimensions) {}
Shape(const std::vector<int>& dimensions) : dims(dimensions) {}
int num_elements() const {
int res = dims.empty() ? 0 : 1;
std::for_each(dims.begin(), dims.end(), [&res](int i) { res *= i; });
return res;
}
std::vector<int> dims;
};
Tensor(ElementType element_type, const Shape& shape);
// Non-copyable.
Tensor(const Tensor&) = delete;
Tensor& operator=(const Tensor&) = delete;
// Move-only.
Tensor(Tensor&& src) { Move(&src); }
Tensor& operator=(Tensor&&);
~Tensor() { Invalidate(); }
template <typename T>
class CpuView : public View {
public:
template <typename P>
auto buffer() const {
// const and non-const return type selection.
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;
}
protected:
friend class Tensor;
CpuView(T* buffer, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), buffer_(buffer) {}
T* buffer_;
};
using CpuReadView = CpuView<const void>;
CpuReadView GetCpuReadView() const;
using CpuWriteView = CpuView<void>;
CpuWriteView GetCpuWriteView() const;
#if MEDIAPIPE_METAL_ENABLED
// TODO: id<MTLBuffer> vs. MtlBufferView.
class MtlBufferView : public View {
public:
id<MTLBuffer> buffer() const { return buffer_; }
MtlBufferView(MtlBufferView&& src)
: View(std::move(src)), buffer_(src.buffer_) {
src.buffer_ = nil;
}
protected:
friend class Tensor;
MtlBufferView(id<MTLBuffer> buffer, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), buffer_(buffer) {}
id<MTLBuffer> buffer_;
};
// The command buffer status is checked for completeness if GPU-to-CPU
// synchronization is required.
// TODO: Design const and non-const view acquiring.
MtlBufferView GetMtlBufferReadView(id<MTLCommandBuffer> command_buffer) const;
MtlBufferView GetMtlBufferWriteView(
id<MTLCommandBuffer> command_buffer) const;
// Allocate new buffer.
// TODO: GPU-to-CPU design considerations.
MtlBufferView GetMtlBufferWriteView(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
// TODO: Use GlTextureView instead.
// Only float32 textures are supported with 1/2/3/4 depths.
// OpenGlTexture2dView currently only supports BHWC memory layout.
class OpenGlTexture2dView : public View {
public:
GLuint name() const { return name_; }
OpenGlTexture2dView(OpenGlTexture2dView&& src)
: View(std::move(src)), name_(src.name_) {
src.name_ = GL_INVALID_INDEX;
}
protected:
friend class Tensor;
OpenGlTexture2dView(GLuint name, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), name_(name) {}
GLuint name_;
};
// A valid OpenGL context must be bound to the calling thread due to possible
// GPU resource allocation.
OpenGlTexture2dView GetOpenGlTexture2dReadView() const;
OpenGlTexture2dView GetOpenGlTexture2dWriteView() const;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
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;
}
protected:
friend class Tensor;
OpenGlBufferView(GLuint name, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), name_(name) {}
GLuint name_;
};
// A valid OpenGL context must be bound to the calling thread due to possible
// GPU resource allocation.
OpenGlBufferView GetOpenGlBufferReadView() const;
OpenGlBufferView GetOpenGlBufferWriteView() const;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
const Shape& shape() const { return shape_; }
ElementType element_type() const { return element_type_; }
int element_size() const {
switch (element_type_) {
case ElementType::kNone:
return 0;
case ElementType::kFloat16:
return 2;
case ElementType::kFloat32:
return sizeof(float);
}
}
int bytes() const { return shape_.num_elements() * element_size(); }
bool ready_on_cpu() const { return valid_ & kValidCpu; }
bool ready_on_gpu() const {
return valid_ &
(kValidMetalBuffer | kValidOpenGlBuffer | kValidOpenGlTexture2d);
}
bool ready_as_metal_buffer() const { return valid_ & kValidMetalBuffer; }
bool ready_as_opengl_buffer() const { return valid_ & kValidOpenGlBuffer; }
bool ready_as_opengl_texture_2d() const {
return valid_ & kValidOpenGlTexture2d;
}
private:
void Move(Tensor*);
void Invalidate();
ElementType element_type_;
Shape shape_;
// The flags describe the current source of truth resource type.
enum {
kValidNone = 0,
kValidCpu = 1 << 0,
kValidMetalBuffer = 1 << 1,
kValidOpenGlBuffer = 1 << 2,
kValidOpenGlTexture2d = 1 << 3,
};
// A list of resource which are currently allocated and synchronized between
// each-other: valid_ = kValidCpu | kValidMetalBuffer;
mutable int valid_ = 0;
// The mutex is locked by Get*View and is kept by all Views.
mutable absl::Mutex view_mutex_;
mutable void* cpu_buffer_ = nullptr;
void AllocateCpuBuffer() const;
#if MEDIAPIPE_METAL_ENABLED
mutable id<MTLCommandBuffer> command_buffer_;
mutable id<MTLDevice> device_;
mutable id<MTLBuffer> metal_buffer_;
void AllocateMtlBuffer(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
mutable std::shared_ptr<mediapipe::GlContext> gl_context_;
mutable GLuint opengl_texture2d_ = GL_INVALID_INDEX;
void AllocateOpenGlTexture2d() const;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
mutable GLuint opengl_buffer_ = GL_INVALID_INDEX;
void AllocateOpenGlBuffer() const;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
};
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_H_
mediapipe/framework/formats/tensor_test.cc
+62
View File
@@ -0,0 +1,62 @@
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer_format.h"
#endif
namespace mediapipe {
TEST(General, TestDimensions) {
Tensor t1(Tensor::ElementType::kFloat32, Tensor::Shape{1, 2, 3, 4});
EXPECT_EQ(t1.shape().num_elements(), 1 * 2 * 3 * 4);
Tensor t2(Tensor::ElementType::kFloat16, Tensor::Shape{4, 3, 2, 3});
EXPECT_EQ(t2.shape().num_elements(), 4 * 3 * 2 * 3);
}
TEST(General, TestDataTypes) {
Tensor t1(Tensor::ElementType::kFloat32, Tensor::Shape{1, 2, 3, 4});
EXPECT_EQ(t1.bytes(), t1.shape().num_elements() * sizeof(float));
Tensor t2(Tensor::ElementType::kFloat16, Tensor::Shape{4, 3, 2, 3});
EXPECT_EQ(t2.bytes(), t2.shape().num_elements() * 2);
}
TEST(Cpu, TestMemoryAllocation) {
Tensor t1(Tensor::ElementType::kFloat32, Tensor::Shape{4, 3, 2, 3});
auto v1 = t1.GetCpuWriteView();
float* f1 = v1.buffer<float>();
EXPECT_NE(f1, nullptr);
}
TEST(Cpu, TestTensorMove) {
Tensor t1(Tensor::ElementType::kFloat32, Tensor::Shape{4, 3, 2, 3});
void* p1 = t1.GetCpuWriteView().buffer<float>();
EXPECT_NE(p1, nullptr);
Tensor t2(std::move(t1));
EXPECT_NE(t2.bytes(), 0);
EXPECT_EQ(t1.bytes(), 0); // NOLINT
void* p2 = t2.GetCpuWriteView().buffer<float>();
EXPECT_EQ(p1, p2);
}
TEST(Cpu, TestViewMove) {
Tensor t(Tensor::ElementType::kFloat32, Tensor::Shape{4, 3, 2, 3});
auto v1 = t.GetCpuWriteView();
auto p1 = v1.buffer<float>();
EXPECT_NE(p1, nullptr);
Tensor::CpuWriteView v2(std::move(v1));
auto p2 = v2.buffer<float>();
EXPECT_EQ(p1, p2);
EXPECT_EQ(v1.buffer<float>(), nullptr); // NOLINT
}
} // namespace mediapipe
int main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
mediapipe/framework/legacy_calculator_support.h
+1 -1
View File
@@ -61,7 +61,7 @@ class LegacyCalculatorSupport {
// platforms.
#ifndef __APPLE__
ABSL_CONST_INIT
#endif // !__APPLE__
#endif // !__APPLE__
static thread_local C* current_; // NOLINT
};
};
mediapipe/framework/port.h
+4 -1
View File
@@ -73,10 +73,13 @@
#elif defined(MEDIAPIPE_OSX)
#define MEDIAPIPE_OPENGL_ES_VERSION 0
#define MEDIAPIPE_METAL_ENABLED 1
#else
#elif defined(__EMSCRIPTEN__)
// WebGL config.
#define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_30
#define MEDIAPIPE_METAL_ENABLED 0
#else
#define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_31
#define MEDIAPIPE_METAL_ENABLED 0
#endif
#endif
mediapipe/framework/tool/BUILD
+1
View File
@@ -132,6 +132,7 @@ cc_library(
visibility = ["//mediapipe/framework:mediapipe_internal"],
deps = [
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:input_stream_shard",
"//mediapipe/framework:packet",
"//mediapipe/framework:packet_generator_cc_proto",
"//mediapipe/framework:packet_set",
mediapipe/framework/tool/options_util.h
+30 -8
View File
@@ -16,6 +16,7 @@
#define MEDIAPIPE_FRAMEWORK_TOOL_OPTIONS_UTIL_H_
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/input_stream_shard.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/packet_generator.pb.h"
#include "mediapipe/framework/packet_set.h"
@@ -96,21 +97,42 @@ void GetNodeOptions(const CalculatorGraphConfig::Node& node_config, T* result) {
// packet can hold either the specified options type T or CalculatorOptions.
// Fields are either replaced or merged depending on field merge_fields.
template <typename T>
inline T RetrieveOptions(const T& base, const PacketSet& packet_set,
const std::string& tag_name) {
if (packet_set.HasTag(tag_name)) {
const Packet& packet = packet_set.Tag(tag_name);
inline T RetrieveOptions(const T& base, const Packet& options_packet) {
if (!options_packet.IsEmpty()) {
T packet_options;
if (packet.ValidateAsType<T>().ok()) {
packet_options = packet.Get<T>();
} else if (packet.ValidateAsType<CalculatorOptions>().ok()) {
GetExtension<T>(packet.Get<CalculatorOptions>(), &packet_options);
if (options_packet.ValidateAsType<T>().ok()) {
packet_options = options_packet.Get<T>();
} else if (options_packet.ValidateAsType<CalculatorOptions>().ok()) {
GetExtension<T>(options_packet.Get<CalculatorOptions>(), &packet_options);
}
return tool::MergeOptions(base, packet_options);
}
return base;
}
// Combine a base options message with an optional side packet from
// a PacketSet such as a calculator's input-side-packets.
template <typename T>
inline T RetrieveOptions(const T& base, const PacketSet& packet_set,
const std::string& tag_name = "OPTIONS") {
if (packet_set.HasTag(tag_name)) {
return tool::RetrieveOptions(base, packet_set.Tag(tag_name));
}
return base;
}
// Combine a base options message with an optional input packet from
// an InputStreamShardSet such as a calculator's input streams.
template <typename T>
inline T RetrieveOptions(const T& base, const InputStreamShardSet& stream_set,
const std::string& tag_name = "OPTIONS") {
if (stream_set.HasTag(tag_name)) {
Packet options_packet = stream_set.Tag(tag_name).Value();
return tool::RetrieveOptions(base, options_packet);
}
return base;
}
// Extracts the options message of a specified type from a
// CalculatorGraphConfig::Node.
class OptionsMap {