Tensor1: memorize size_alignment when tracking the ahwb usage.

When CPU/GPU buffer allocated and the tracker selects Ahwb storage to be used then the properly recorded alignment must be used.

PiperOrigin-RevId: 498264759

mediapipe/framework/formats/BUILD

@@ -455,7 +455,7 @@ cc_library(
         ],
     }),
     deps = [
-        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/synchronization",
         "//mediapipe/framework:port",

mediapipe/framework/formats/tensor.h

@@ -24,7 +24,7 @@
 #include <utility>
 #include <vector>
 
-#include "absl/container/flat_hash_set.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/formats/tensor_internal.h"
 #include "mediapipe/framework/port.h"
@@ -434,8 +434,9 @@ class Tensor {
   mutable bool use_ahwb_ = false;
   mutable uint64_t ahwb_tracking_key_ = 0;
   // TODO: Tracks all unique tensors. Can grow to a large number. LRU
-  // can be more predicted.
-  static inline absl::flat_hash_set<uint64_t> ahwb_usage_track_;
+  // (Least Recently Used) can be more predicted.
+  // The value contains the size alignment parameter.
+  static inline absl::flat_hash_map<uint64_t, int> ahwb_usage_track_;
   // Expects the target SSBO to be already bound.
   bool AllocateAhwbMapToSsbo() const;
   bool InsertAhwbToSsboFence() const;

mediapipe/framework/formats/tensor_ahwb.cc

@@ -266,7 +266,12 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
 
 bool Tensor::AllocateAHardwareBuffer(int size_alignment) const {
   // Mark current tracking key as Ahwb-use.
-  ahwb_usage_track_.insert(ahwb_tracking_key_);
+  if (auto it = ahwb_usage_track_.find(ahwb_tracking_key_);
+      it != ahwb_usage_track_.end()) {
+    size_alignment = it->second;
+  } else if (ahwb_tracking_key_ != 0) {
+    ahwb_usage_track_.insert({ahwb_tracking_key_, size_alignment});
+  }
   use_ahwb_ = true;
 
   if (__builtin_available(android 26, *)) {

mediapipe/framework/formats/tensor_ahwb_test.cc

@@ -30,4 +30,71 @@ TEST(TensorAhwbTest, TestAHWBThenCpu) {
   }
 }
 
+TEST(TensorAhwbTest, TestAhwbAlignment) {
+  Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{5});
+  {
+    auto view = tensor.GetAHardwareBufferWriteView(16);
+    EXPECT_NE(view.handle(), nullptr);
+    if (__builtin_available(android 26, *)) {
+      AHardwareBuffer_Desc desc;
+      AHardwareBuffer_describe(view.handle(), &desc);
+      // sizeof(float) * 5 = 20, the closest aligned to 16 size is 32.
+      EXPECT_EQ(desc.width, 32);
+    }
+    view.SetWritingFinishedFD(-1, [](bool) { return true; });
+  }
+}
+
+// Tensor::GetCpuView uses source location mechanism that gives source file name
+// and line from where the method is called. The function is intended just to
+// have two calls providing the same source file name and line.
+auto GetCpuView(const Tensor &tensor) { return tensor.GetCpuWriteView(); }
+
+// The test checks the tracking mechanism: when a tensor's Cpu view is retrieved
+// for the first time then the source location is attached to the tensor. If the
+// Ahwb view is requested then from the tensor then the previously recorded Cpu
+// view request source location is marked for using Ahwb storage.
+// When a Cpu view with the same source location (but for the newly allocated
+// tensor) is requested and the location is marked to use Ahwb storage then the
+// Ahwb storage is allocated for the CpuView.
+TEST(TensorAhwbTest, TestTrackingAhwb) {
+  // Create first tensor and request Cpu and then Ahwb view to mark the source
+  // location for Ahwb storage.
+  {
+    Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{9});
+    {
+      auto view = GetCpuView(tensor);
+      EXPECT_NE(view.buffer<float>(), nullptr);
+    }
+    {
+      // Align size of the Ahwb by multiple of 16.
+      auto view = tensor.GetAHardwareBufferWriteView(16);
+      EXPECT_NE(view.handle(), nullptr);
+      view.SetReadingFinishedFunc([](bool) { return true; });
+    }
+  }
+  {
+    Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{9});
+    {
+      // The second tensor uses the same Cpu view source location so Ahwb
+      // storage is allocated internally.
+      auto view = GetCpuView(tensor);
+      EXPECT_NE(view.buffer<float>(), nullptr);
+    }
+    {
+      // Check the Ahwb size to be aligned to multiple of 16. The alignment is
+      // stored by previous requesting of the Ahwb view.
+      auto view = tensor.GetAHardwareBufferReadView();
+      EXPECT_NE(view.handle(), nullptr);
+      if (__builtin_available(android 26, *)) {
+        AHardwareBuffer_Desc desc;
+        AHardwareBuffer_describe(view.handle(), &desc);
+        // sizeof(float) * 9 = 36. The closest aligned size is 48.
+        EXPECT_EQ(desc.width, 48);
+      }
+      view.SetReadingFinishedFunc([](bool) { return true; });
+    }
+  }
+}
+
 }  // namespace mediapipe