diff --git a/mediapipe/framework/formats/tensor.h b/mediapipe/framework/formats/tensor.h index 9d3e90b6a..f5a99cde1 100644 --- a/mediapipe/framework/formats/tensor.h +++ b/mediapipe/framework/formats/tensor.h @@ -408,8 +408,8 @@ class Tensor { mutable std::function release_callback_; bool AllocateAHardwareBuffer(int size_alignment = 0) const; void CreateEglSyncAndFd() const; - // Use Ahwb for other views: OpenGL / CPU buffer. #endif // MEDIAPIPE_TENSOR_USE_AHWB + // Use Ahwb for other views: OpenGL / CPU buffer. static inline bool use_ahwb_ = false; // Expects the target SSBO to be already bound. bool AllocateAhwbMapToSsbo() const; diff --git a/mediapipe/framework/formats/tensor_ahwb.cc b/mediapipe/framework/formats/tensor_ahwb.cc index 3c3ec8b17..363c5efd0 100644 --- a/mediapipe/framework/formats/tensor_ahwb.cc +++ b/mediapipe/framework/formats/tensor_ahwb.cc @@ -212,9 +212,6 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferReadView() const { 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 target system."; bool transfer = !ahwb_; CHECK(AllocateAHardwareBuffer()) << "AHardwareBuffer is not supported on the target system."; @@ -315,7 +312,13 @@ void Tensor::MoveCpuOrSsboToAhwb() const { ahwb_, AHARDWAREBUFFER_USAGE_CPU_WRITE_RARELY, -1, nullptr, &dest); CHECK(error == 0) << "AHardwareBuffer_lock " << error; } - if (valid_ & kValidOpenGlBuffer) { + if (valid_ & kValidCpu) { + std::memcpy(dest, cpu_buffer_, bytes()); + // Free CPU memory because next time AHWB is mapped instead. + free(cpu_buffer_); + cpu_buffer_ = nullptr; + valid_ &= ~kValidCpu; + } else if (valid_ & kValidOpenGlBuffer) { gl_context_->Run([this, dest]() { glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_); const void* src = glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(), @@ -326,11 +329,9 @@ void Tensor::MoveCpuOrSsboToAhwb() const { }); opengl_buffer_ = GL_INVALID_INDEX; gl_context_ = nullptr; - } else if (valid_ & kValidCpu) { - std::memcpy(dest, cpu_buffer_, bytes()); - // Free CPU memory because next time AHWB is mapped instead. - free(cpu_buffer_); - cpu_buffer_ = nullptr; + // Reset OpenGL Buffer validness. The OpenGL buffer will be allocated on top + // of the Ahwb at the next request to the OpenGlBufferView. + valid_ &= ~kValidOpenGlBuffer; } else { LOG(FATAL) << "Can't convert tensor with mask " << valid_ << " into AHWB."; } diff --git a/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc b/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc index 7ccd9c7f5..a6ca00949 100644 --- a/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc +++ b/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc @@ -152,6 +152,36 @@ TEST_F(TensorAhwbGpuTest, TestReplacingCpuByAhwb) { { auto view = tensor.GetAHardwareBufferReadView(); EXPECT_NE(view.handle(), nullptr); + view.SetReadingFinishedFunc([](bool) { return true; }); + } + auto ptr = tensor.GetCpuReadView().buffer(); + EXPECT_NE(ptr, nullptr); + std::vector reference; + reference.resize(num_elements); + for (int i = 0; i < num_elements; i++) { + reference[i] = static_cast(i) / 10.0f; + } + EXPECT_THAT(absl::Span(ptr, num_elements), + testing::Pointwise(testing::FloatEq(), reference)); +} + +TEST_F(TensorAhwbGpuTest, TestReplacingGpuByAhwb) { + // Request the GPU view to get the ssbo allocated internally. + // 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})}; + RunInGlContext([&tensor] { + auto ssbo_view = tensor.GetOpenGlBufferWriteView(); + auto ssbo_name = ssbo_view.name(); + EXPECT_GT(ssbo_name, 0); + FillGpuBuffer(ssbo_name, tensor.shape().num_elements(), + tensor.element_type()); + }); + { + auto view = tensor.GetAHardwareBufferReadView(); + EXPECT_NE(view.handle(), nullptr); + view.SetReadingFinishedFunc([](bool) { return true; }); } auto ptr = tensor.GetCpuReadView().buffer(); EXPECT_NE(ptr, nullptr);