cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Tensor: Cpu -> Ahwb storage transfer

Browse Source 
PiperOrigin-RevId: 494033280
This commit is contained in:
Nikolay Chirkov 2022-12-08 17:05:06 -08:00 committed by Copybara-Service
parent 05535db5f7
commit bea0caae65
4 changed files with 32 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mediapipe/framework/formats/tensor.cc
View File

@ -361,7 +361,7 @@ 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_);
if (!AllocateAhwbMapToSsbo()) {
if (!use_ahwb_ || !AllocateAhwbMapToSsbo()) {
glBufferData(GL_SHADER_STORAGE_BUFFER, bytes(), NULL, GL_STREAM_COPY);
}
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
@ -610,7 +610,7 @@ Tensor::CpuWriteView Tensor::GetCpuWriteView() const {
void Tensor::AllocateCpuBuffer() const {
if (!cpu_buffer_) {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
if (AllocateAHardwareBuffer()) return;
if (use_ahwb_ && AllocateAHardwareBuffer()) return;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_METAL_ENABLED
cpu_buffer_ = AllocateVirtualMemory(bytes());

2
mediapipe/framework/formats/tensor.h
View File

@ -409,8 +409,8 @@ class Tensor {
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
static inline bool use_ahwb_ = false;
// Expects the target SSBO to be already bound.
bool AllocateAhwbMapToSsbo() const;
bool InsertAhwbToSsboFence() const;

3
mediapipe/framework/formats/tensor_ahwb.cc
View File

@ -214,7 +214,7 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferReadView() const {
"supported.";
CHECK(ahwb_ || !(valid_ & kValidOpenGlBuffer))
<< "Interoperability bettween OpenGL buffer and AHardwareBuffer is not "
"supported on targe system.";
"supported on target system.";
bool transfer = !ahwb_;
CHECK(AllocateAHardwareBuffer())
<< "AHardwareBuffer is not supported on the target system.";
@ -268,7 +268,6 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
}
bool Tensor::AllocateAHardwareBuffer(int size_alignment) const {
if (!use_ahwb_) return false;
if (__builtin_available(android 26, *)) {
if (ahwb_ == nullptr) {
AHardwareBuffer_Desc desc = {};

28
mediapipe/framework/formats/tensor_ahwb_gpu_test.cc
View File

@ -136,6 +136,34 @@ TEST_F(TensorAhwbGpuTest, TestGpuToCpuFloat16) {
testing::Pointwise(NearWithPrecision(0.001), reference));
}
TEST_F(TensorAhwbGpuTest, TestReplacingCpuByAhwb) {
// Request the CPU view to get the memory to be allocated.
// Request Ahwb view then to transform the storage into Ahwb.
Tensor::SetPreferredStorageType(Tensor::StorageType::kDefault);
constexpr size_t num_elements = 20;
Tensor tensor{Tensor::ElementType::kFloat32, Tensor::Shape({num_elements})};
{
auto ptr = tensor.GetCpuWriteView().buffer<float>();
EXPECT_NE(ptr, nullptr);
for (int i = 0; i < num_elements; i++) {
ptr[i] = static_cast<float>(i) / 10.0f;
}
}
{
auto view = tensor.GetAHardwareBufferReadView();
EXPECT_NE(view.handle(), nullptr);
}
auto ptr = tensor.GetCpuReadView().buffer<float>();
EXPECT_NE(ptr, nullptr);
std::vector<float> reference;
reference.resize(num_elements);
for (int i = 0; i < num_elements; i++) {
reference[i] = static_cast<float>(i) / 10.0f;
}
EXPECT_THAT(absl::Span<const float>(ptr, num_elements),
testing::Pointwise(testing::FloatEq(), reference));
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
} // namespace