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

Tensor: Make tensor not requiring "-x objective-c++" option.

Browse Source 
In this case tensor.h is compiled differently for C++ and Objective-C++ that violates ODR (once definition rule). Tensor has no virtual methods conditionally compiled but some Metal-related data members. Instead, unique_ptr to MtlResources that is declared as forward structure is unconditionally defined in the tensor class. MtlResources is defined differently in cc-file only that compiled just once per project so no ODR violation is here.

PiperOrigin-RevId: 504029286
This commit is contained in:
Nikolay Chirkov
2023-01-23 11:00:35 -08:00
committed by Copybara-Service
parent 921b6a6bef
commit 1124569c29
12 changed files with 177 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/formats/BUILD
+4 -1
View File
@@ -431,7 +431,10 @@ cc_library(
hdrs = [
"tensor.h",
"//mediapipe/framework/formats/tensor:internal.h",
],
] + select({
"//mediapipe:ios": ["tensor_mtl_buffer_view.h"],
"//conditions:default": [],
}),
copts = select({
"//mediapipe:apple": [
"-x objective-c++",
mediapipe/framework/formats/tensor.cc
+70 -55
View File
@@ -25,8 +25,11 @@
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#if MEDIAPIPE_METAL_ENABLED
#import <Metal/Metal.h>
#include <mach/mach_init.h>
#include <mach/vm_map.h>
#include "mediapipe/framework/formats/tensor_mtl_buffer_view.h"
#else
#include <cstdlib>
#endif // MEDIAPIPE_METAL_ENABLED
@@ -61,6 +64,12 @@ int BhwcDepthFromShape(const Tensor::Shape& shape) {
// 3) pad/"unpad" the bitmap after transfer CPU <-> GPU
#if MEDIAPIPE_METAL_ENABLED
// No ODR violation here because this file compiled just once per project.
struct MtlResources {
id<MTLCommandBuffer> command_buffer = nil;
id<MTLDevice> device = nil;
id<MTLBuffer> metal_buffer = nil;
};
namespace {
// MTLBuffer can use existing properly aligned and allocated CPU memory.
size_t AlignToPageSize(size_t size) {
@@ -83,52 +92,56 @@ void DeallocateVirtualMemory(void* pointer, size_t size) {
}
} // namespace
Tensor::MtlBufferView Tensor::GetMtlBufferReadView(
id<MTLCommandBuffer> command_buffer) const {
LOG_IF(FATAL, valid_ == kValidNone)
void MtlBufferView::AllocateMtlBuffer(const Tensor& tensor,
id<MTLDevice> device) {
tensor.mtl_resources_->device = device;
if (!tensor.cpu_buffer_) {
// It also means that the metal buffer is not allocated yet.
tensor.cpu_buffer_ = AllocateVirtualMemory(tensor.bytes());
}
if (!tensor.mtl_resources_->metal_buffer) {
tensor.mtl_resources_->metal_buffer = [tensor.mtl_resources_->device
newBufferWithBytesNoCopy:tensor.cpu_buffer_
length:AlignToPageSize(tensor.bytes())
options:MTLResourceStorageModeShared |
MTLResourceCPUCacheModeDefaultCache
deallocator:^(void* pointer, NSUInteger length) {
DeallocateVirtualMemory(pointer, length);
}];
}
}
MtlBufferView MtlBufferView::GetReadView(const Tensor& tensor,
id<MTLCommandBuffer> command_buffer) {
LOG_IF(FATAL, tensor.valid_ == Tensor::kValidNone)
<< "Tensor must be written prior to read from.";
LOG_IF(FATAL, !(valid_ & (kValidCpu | kValidMetalBuffer)))
LOG_IF(FATAL,
!(tensor.valid_ & (Tensor::kValidCpu | Tensor::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)};
auto lock(absl::make_unique<absl::MutexLock>(&tensor.view_mutex_));
tensor.valid_ |= Tensor::kValidMetalBuffer;
AllocateMtlBuffer(tensor, [command_buffer device]);
return {tensor.mtl_resources_->metal_buffer, std::move(lock)};
}
Tensor::MtlBufferView Tensor::GetMtlBufferWriteView(
id<MTLCommandBuffer> command_buffer) const {
MtlBufferView MtlBufferView::GetWriteView(const Tensor& tensor,
id<MTLCommandBuffer> command_buffer) {
// 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.mtl_resources_->command_buffer = command_buffer;
return GetWriteView(tensor, [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);
}];
}
MtlBufferView MtlBufferView::GetWriteView(const Tensor& tensor,
id<MTLDevice> device) {
auto lock(absl::make_unique<absl::MutexLock>(&tensor.view_mutex_));
tensor.valid_ = Tensor::kValidMetalBuffer;
AllocateMtlBuffer(tensor, device);
return {tensor.mtl_resources_->metal_buffer, std::move(lock)};
}
#else
struct MtlResources {};
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
@@ -379,6 +392,9 @@ Tensor& Tensor::operator=(Tensor&& src) {
return *this;
}
Tensor::Tensor(Tensor&& src) { Move(&src); }
Tensor::~Tensor() { Invalidate(); }
void Tensor::Move(Tensor* src) {
valid_ = src->valid_;
src->valid_ = kValidNone;
@@ -388,15 +404,7 @@ void Tensor::Move(Tensor* src) {
cpu_buffer_ = src->cpu_buffer_;
src->cpu_buffer_ = nullptr;
ahwb_tracking_key_ = src->ahwb_tracking_key_;
#if MEDIAPIPE_METAL_ENABLED
device_ = src->device_;
src->device_ = nil;
command_buffer_ = src->command_buffer_;
src->command_buffer_ = nil;
metal_buffer_ = src->metal_buffer_;
src->metal_buffer_ = nil;
#endif // MEDIAPIPE_METAL_ENABLED
mtl_resources_ = std::move(src->mtl_resources_);
MoveAhwbStuff(src);
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
@@ -415,12 +423,15 @@ void Tensor::Move(Tensor* src) {
}
Tensor::Tensor(ElementType element_type, const Shape& shape)
: element_type_(element_type), shape_(shape) {}
: element_type_(element_type),
shape_(shape),
mtl_resources_(std::make_unique<MtlResources>()) {}
Tensor::Tensor(ElementType element_type, const Shape& shape,
const QuantizationParameters& quantization_parameters)
: element_type_(element_type),
shape_(shape),
quantization_parameters_(quantization_parameters) {}
quantization_parameters_(quantization_parameters),
mtl_resources_(std::make_unique<MtlResources>()) {}
#if MEDIAPIPE_METAL_ENABLED
void Tensor::Invalidate() {
@@ -432,13 +443,16 @@ 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_) {
if (cpu_buffer_ && !mtl_resources_->metal_buffer) {
DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
}
metal_buffer_ = nil;
command_buffer_ = nil;
device_ = nil;
cpu_buffer_ = nullptr;
// This becomes NULL if the tensor is moved.
if (mtl_resources_) {
mtl_resources_->metal_buffer = nil;
mtl_resources_->command_buffer = nil;
mtl_resources_->device = nil;
}
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
// Don't need to wait for the resource to be deleted bacause if will be
// released on last reference deletion inside the OpenGL driver.
@@ -532,10 +546,11 @@ Tensor::CpuReadView Tensor::GetCpuReadView() const {
// 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];
LOG_IF(FATAL, !mtl_resources_->command_buffer)
<< "Metal -> CPU synchronization "
"requires MTLCommandBuffer to be set.";
if (mtl_resources_->command_buffer) {
[mtl_resources_->command_buffer waitUntilCompleted];
}
}
#endif // MEDIAPIPE_METAL_ENABLED
mediapipe/framework/formats/tensor.h
+7 -39
View File
@@ -29,9 +29,6 @@
#include "mediapipe/framework/formats/tensor/internal.h"
#include "mediapipe/framework/port.h"
#if MEDIAPIPE_METAL_ENABLED
#import <Metal/Metal.h>
#endif // MEDIAPIPE_METAL_ENABLED
#ifndef MEDIAPIPE_NO_JNI
#if __ANDROID_API__ >= 26 || defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__)
#define MEDIAPIPE_TENSOR_USE_AHWB 1
@@ -66,7 +63,6 @@
#endif
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.
@@ -91,6 +87,7 @@ namespace mediapipe {
// float* pointer = view.buffer<float>();
// ...reading the cpu memory...
struct MtlResources;
class Tensor {
class View {
public:
@@ -144,9 +141,9 @@ class Tensor {
Tensor(const Tensor&) = delete;
Tensor& operator=(const Tensor&) = delete;
// Move-only.
Tensor(Tensor&& src) { Move(&src); }
Tensor(Tensor&& src);
Tensor& operator=(Tensor&&);
~Tensor() { Invalidate(); }
~Tensor();
template <typename T>
class CpuView : public View {
@@ -182,33 +179,6 @@ class Tensor {
uint64_t source_location_hash =
tensor_internal::FnvHash64(builtin_FILE(), builtin_LINE())) 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
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
using FinishingFunc = std::function<bool(bool)>;
class AHardwareBufferView : public View {
@@ -372,6 +342,7 @@ class Tensor {
}
private:
friend class MtlBufferView;
void Move(Tensor*);
void Invalidate();
@@ -396,12 +367,9 @@ class Tensor {
mutable void* cpu_buffer_ = nullptr;
void AllocateCpuBuffer() const;
#if MEDIAPIPE_METAL_ENABLED
mutable id<MTLCommandBuffer> command_buffer_ = nil;
mutable id<MTLDevice> device_ = nil;
mutable id<MTLBuffer> metal_buffer_ = nil;
void AllocateMtlBuffer(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
// Forward declaration of the MtlResources provides compile-time verification
// of ODR if this header includes any actual code that uses MtlResources.
mutable std::unique_ptr<MtlResources> mtl_resources_;
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
mutable AHardwareBuffer* ahwb_ = nullptr;
mediapipe/framework/formats/tensor_mtl_buffer_view.h
+61
View File
@@ -0,0 +1,61 @@
// 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_MTL_BUFFER_VIEW_H_
#define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_MTL_BUFFER_VIEW_H_
#import <Metal/Metal.h>
#include <algorithm>
#include <functional>
#include <initializer_list>
#include <numeric>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port.h"
namespace mediapipe {
class MtlBufferView : public Tensor::View {
public:
// The command buffer status is checked for completeness if GPU-to-CPU
// synchronization is required.
static MtlBufferView GetReadView(const Tensor& tensor,
id<MTLCommandBuffer> command_buffer);
static MtlBufferView GetWriteView(const Tensor& tensor,
id<MTLCommandBuffer> command_buffer);
static MtlBufferView GetWriteView(const Tensor& tensor, id<MTLDevice> device);
id<MTLBuffer> buffer() const { return buffer_; }
MtlBufferView(MtlBufferView&& src)
: Tensor::View(std::move(src)), buffer_(src.buffer_) {
src.buffer_ = nil;
}
protected:
friend class Tensor;
static void AllocateMtlBuffer(const Tensor& tensor, id<MTLDevice> device);
MtlBufferView(id<MTLBuffer> buffer, std::unique_ptr<absl::MutexLock>&& lock)
: Tensor::View(std::move(lock)), buffer_(buffer) {}
id<MTLBuffer> buffer_;
};
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_MTL_BUFFER_VIEW_H_