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
This commit is contained in:
Nikolay Chirkov 2022-12-28 16:07:09 -08:00 committed by Copybara-Service
parent 1924f1cdff
commit 2d9a969d10
4 changed files with 78 additions and 5 deletions

View File

@ -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",

View File

@ -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;

View File

@ -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, *)) {

View File

@ -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