From bea0caae6586343eb91e986b84aa00ef75fa67b1 Mon Sep 17 00:00:00 2001
From: Nikolay Chirkov <chirkov@google.com>
Date: Thu, 8 Dec 2022 17:05:06 -0800
Subject: [PATCH] Tensor: Cpu -> Ahwb storage transfer

PiperOrigin-RevId: 494033280
---
 mediapipe/framework/formats/tensor.cc         |  4 +--
 mediapipe/framework/formats/tensor.h          |  2 +-
 mediapipe/framework/formats/tensor_ahwb.cc    |  3 +-
 .../framework/formats/tensor_ahwb_gpu_test.cc | 28 +++++++++++++++++++
 4 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/mediapipe/framework/formats/tensor.cc b/mediapipe/framework/formats/tensor.cc
index 9e1406db..fdafbff5 100644
--- a/mediapipe/framework/formats/tensor.cc
+++ b/mediapipe/framework/formats/tensor.cc
@@ -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());
diff --git a/mediapipe/framework/formats/tensor.h b/mediapipe/framework/formats/tensor.h
index 151aa299..9d3e90b6 100644
--- a/mediapipe/framework/formats/tensor.h
+++ b/mediapipe/framework/formats/tensor.h
@@ -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;
diff --git a/mediapipe/framework/formats/tensor_ahwb.cc b/mediapipe/framework/formats/tensor_ahwb.cc
index 21bae959..3c3ec8b1 100644
--- a/mediapipe/framework/formats/tensor_ahwb.cc
+++ b/mediapipe/framework/formats/tensor_ahwb.cc
@@ -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 = {};
diff --git a/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc b/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc
index dd865a36..7ccd9c7f 100644
--- a/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc
+++ b/mediapipe/framework/formats/tensor_ahwb_gpu_test.cc
@@ -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