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:
parent
1924f1cdff
commit
2d9a969d10
|
@ -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",
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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, *)) {
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue
Block a user