Internal change.
PiperOrigin-RevId: 495483878
This commit is contained in:
Nikolay Chirkov
Copybara-Service
parent
e9e173f9fa
commit
d526b20e19
|
|
@ -1,71 +0,0 @@
|
||||||
#ifndef MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_HARDWARE_BUFFER_H_
|
|
||||||
#define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_HARDWARE_BUFFER_H_
|
|
||||||
|
|
||||||
#if !defined(MEDIAPIPE_NO_JNI) && \
|
|
||||||
(__ANDROID_API__ >= 26 || \
|
|
||||||
defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))
|
|
||||||
|
|
||||||
#include <android/hardware_buffer.h>
|
|
||||||
|
|
||||||
#include <cstdint>
|
|
||||||
|
|
||||||
#include "mediapipe/framework/formats/tensor_buffer.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_internal.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_v2.h"
|
|
||||||
|
|
||||||
namespace mediapipe {
|
|
||||||
|
|
||||||
// Supports:
|
|
||||||
// - float 16 and 32 bits
|
|
||||||
// - signed / unsigned integers 8,16,32 bits
|
|
||||||
class TensorHardwareBufferView;
|
|
||||||
struct TensorHardwareBufferViewDescriptor : public Tensor::ViewDescriptor {
|
|
||||||
using ViewT = TensorHardwareBufferView;
|
|
||||||
TensorBufferDescriptor buffer;
|
|
||||||
};
|
|
||||||
|
|
||||||
class TensorHardwareBufferView : public Tensor::View {
|
|
||||||
public:
|
|
||||||
TENSOR_UNIQUE_VIEW_TYPE_ID();
|
|
||||||
~TensorHardwareBufferView() = default;
|
|
||||||
|
|
||||||
const TensorHardwareBufferViewDescriptor& descriptor() const override {
|
|
||||||
return descriptor_;
|
|
||||||
}
|
|
||||||
AHardwareBuffer* handle() const { return ahwb_handle_; }
|
|
||||||
|
|
||||||
protected:
|
|
||||||
TensorHardwareBufferView(int access_capability, Tensor::View::Access access,
|
|
||||||
Tensor::View::State state,
|
|
||||||
const TensorHardwareBufferViewDescriptor& desc,
|
|
||||||
AHardwareBuffer* ahwb_handle)
|
|
||||||
: Tensor::View(kId, access_capability, access, state),
|
|
||||||
descriptor_(desc),
|
|
||||||
ahwb_handle_(ahwb_handle) {}
|
|
||||||
|
|
||||||
private:
|
|
||||||
bool MatchDescriptor(
|
|
||||||
uint64_t view_type_id,
|
|
||||||
const Tensor::ViewDescriptor& base_descriptor) const override {
|
|
||||||
if (!Tensor::View::MatchDescriptor(view_type_id, base_descriptor))
|
|
||||||
return false;
|
|
||||||
auto descriptor =
|
|
||||||
static_cast<const TensorHardwareBufferViewDescriptor&>(base_descriptor);
|
|
||||||
return descriptor.buffer.format == descriptor_.buffer.format &&
|
|
||||||
descriptor.buffer.size_alignment <=
|
|
||||||
descriptor_.buffer.size_alignment &&
|
|
||||||
descriptor_.buffer.size_alignment %
|
|
||||||
descriptor.buffer.size_alignment ==
|
|
||||||
0;
|
|
||||||
}
|
|
||||||
const TensorHardwareBufferViewDescriptor& descriptor_;
|
|
||||||
AHardwareBuffer* ahwb_handle_ = nullptr;
|
|
||||||
};
|
|
||||||
|
|
||||||
} // namespace mediapipe
|
|
||||||
|
|
||||||
#endif // !defined(MEDIAPIPE_NO_JNI) && \
|
|
||||||
(__ANDROID_API__ >= 26 || \
|
|
||||||
defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))
|
|
||||||
|
|
||||||
#endif // MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_HARDWARE_BUFFER_H_
|
|
||||||
|
|
@ -1,216 +0,0 @@
|
||||||
#if !defined(MEDIAPIPE_NO_JNI) && \
|
|
||||||
(__ANDROID_API__ >= 26 || \
|
|
||||||
defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))
|
|
||||||
|
|
||||||
#include <cstdint>
|
|
||||||
|
|
||||||
#include "absl/memory/memory.h"
|
|
||||||
#include "absl/status/status.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_backend.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_cpu_buffer.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_hardware_buffer.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_v2.h"
|
|
||||||
#include "util/task/status_macros.h"
|
|
||||||
|
|
||||||
namespace mediapipe {
|
|
||||||
namespace {
|
|
||||||
|
|
||||||
class TensorCpuViewImpl : public TensorCpuView {
|
|
||||||
public:
|
|
||||||
TensorCpuViewImpl(int access_capabilities, Tensor::View::Access access,
|
|
||||||
Tensor::View::State state,
|
|
||||||
const TensorCpuViewDescriptor& descriptor, void* pointer,
|
|
||||||
AHardwareBuffer* ahwb_handle)
|
|
||||||
: TensorCpuView(access_capabilities, access, state, descriptor, pointer),
|
|
||||||
ahwb_handle_(ahwb_handle) {}
|
|
||||||
~TensorCpuViewImpl() {
|
|
||||||
// If handle_ is null then this view is constructed in GetViews with no
|
|
||||||
// access.
|
|
||||||
if (ahwb_handle_) {
|
|
||||||
if (__builtin_available(android 26, *)) {
|
|
||||||
AHardwareBuffer_unlock(ahwb_handle_, nullptr);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
private:
|
|
||||||
AHardwareBuffer* ahwb_handle_;
|
|
||||||
};
|
|
||||||
|
|
||||||
class TensorHardwareBufferViewImpl : public TensorHardwareBufferView {
|
|
||||||
public:
|
|
||||||
TensorHardwareBufferViewImpl(
|
|
||||||
int access_capability, Tensor::View::Access access,
|
|
||||||
Tensor::View::State state,
|
|
||||||
const TensorHardwareBufferViewDescriptor& descriptor,
|
|
||||||
AHardwareBuffer* handle)
|
|
||||||
: TensorHardwareBufferView(access_capability, access, state, descriptor,
|
|
||||||
handle) {}
|
|
||||||
~TensorHardwareBufferViewImpl() = default;
|
|
||||||
};
|
|
||||||
|
|
||||||
class HardwareBufferCpuStorage : public TensorStorage {
|
|
||||||
public:
|
|
||||||
~HardwareBufferCpuStorage() {
|
|
||||||
if (!ahwb_handle_) return;
|
|
||||||
if (__builtin_available(android 26, *)) {
|
|
||||||
AHardwareBuffer_release(ahwb_handle_);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
static absl::Status CanProvide(
|
|
||||||
int access_capability, const Tensor::Shape& shape, uint64_t view_type_id,
|
|
||||||
const Tensor::ViewDescriptor& base_descriptor) {
|
|
||||||
// TODO: use AHardwareBuffer_isSupported for API >= 29.
|
|
||||||
static const bool is_ahwb_supported = [] {
|
|
||||||
if (__builtin_available(android 26, *)) {
|
|
||||||
AHardwareBuffer_Desc desc = {};
|
|
||||||
// Aligned to the largest possible virtual memory page size.
|
|
||||||
constexpr uint32_t kPageSize = 16384;
|
|
||||||
desc.width = kPageSize;
|
|
||||||
desc.height = 1;
|
|
||||||
desc.layers = 1;
|
|
||||||
desc.format = AHARDWAREBUFFER_FORMAT_BLOB;
|
|
||||||
desc.usage = AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN |
|
|
||||||
AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN;
|
|
||||||
AHardwareBuffer* handle;
|
|
||||||
if (AHardwareBuffer_allocate(&desc, &handle) != 0) return false;
|
|
||||||
AHardwareBuffer_release(handle);
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
return false;
|
|
||||||
}();
|
|
||||||
if (!is_ahwb_supported) {
|
|
||||||
return absl::UnavailableError(
|
|
||||||
"AHardwareBuffer is not supported on the platform.");
|
|
||||||
}
|
|
||||||
|
|
||||||
if (view_type_id != TensorCpuView::kId &&
|
|
||||||
view_type_id != TensorHardwareBufferView::kId) {
|
|
||||||
return absl::InvalidArgumentError(
|
|
||||||
"A view type is not supported by this storage.");
|
|
||||||
}
|
|
||||||
return absl::OkStatus();
|
|
||||||
}
|
|
||||||
|
|
||||||
std::vector<std::unique_ptr<Tensor::View>> GetViews(uint64_t latest_version) {
|
|
||||||
std::vector<std::unique_ptr<Tensor::View>> result;
|
|
||||||
auto update_state = latest_version == version_
|
|
||||||
? Tensor::View::State::kUpToDate
|
|
||||||
: Tensor::View::State::kOutdated;
|
|
||||||
if (ahwb_handle_) {
|
|
||||||
result.push_back(
|
|
||||||
std::unique_ptr<Tensor::View>(new TensorHardwareBufferViewImpl(
|
|
||||||
kAccessCapability, Tensor::View::Access::kNoAccess, update_state,
|
|
||||||
hw_descriptor_, ahwb_handle_)));
|
|
||||||
|
|
||||||
result.push_back(std::unique_ptr<Tensor::View>(new TensorCpuViewImpl(
|
|
||||||
kAccessCapability, Tensor::View::Access::kNoAccess, update_state,
|
|
||||||
cpu_descriptor_, nullptr, nullptr)));
|
|
||||||
}
|
|
||||||
return result;
|
|
||||||
}
|
|
||||||
|
|
||||||
absl::StatusOr<std::unique_ptr<Tensor::View>> GetView(
|
|
||||||
Tensor::View::Access access, const Tensor::Shape& shape,
|
|
||||||
uint64_t latest_version, uint64_t view_type_id,
|
|
||||||
const Tensor::ViewDescriptor& base_descriptor, int access_capability) {
|
|
||||||
MP_RETURN_IF_ERROR(
|
|
||||||
CanProvide(access_capability, shape, view_type_id, base_descriptor));
|
|
||||||
const auto& buffer_descriptor =
|
|
||||||
view_type_id == TensorHardwareBufferView::kId
|
|
||||||
? static_cast<const TensorHardwareBufferViewDescriptor&>(
|
|
||||||
base_descriptor)
|
|
||||||
.buffer
|
|
||||||
: static_cast<const TensorCpuViewDescriptor&>(base_descriptor)
|
|
||||||
.buffer;
|
|
||||||
if (!ahwb_handle_) {
|
|
||||||
if (__builtin_available(android 26, *)) {
|
|
||||||
AHardwareBuffer_Desc desc = {};
|
|
||||||
desc.width = TensorBufferSize(buffer_descriptor, shape);
|
|
||||||
desc.height = 1;
|
|
||||||
desc.layers = 1;
|
|
||||||
desc.format = AHARDWAREBUFFER_FORMAT_BLOB;
|
|
||||||
// TODO: Use access capabilities to set hints.
|
|
||||||
desc.usage = AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN |
|
|
||||||
AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN;
|
|
||||||
auto error = AHardwareBuffer_allocate(&desc, &ahwb_handle_);
|
|
||||||
if (error != 0) {
|
|
||||||
return absl::UnknownError(
|
|
||||||
absl::StrCat("Error allocating hardware buffer: ", error));
|
|
||||||
}
|
|
||||||
// Fill all possible views to provide it as proto views.
|
|
||||||
hw_descriptor_.buffer = buffer_descriptor;
|
|
||||||
cpu_descriptor_.buffer = buffer_descriptor;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if (buffer_descriptor.format != hw_descriptor_.buffer.format ||
|
|
||||||
buffer_descriptor.size_alignment >
|
|
||||||
hw_descriptor_.buffer.size_alignment ||
|
|
||||||
hw_descriptor_.buffer.size_alignment %
|
|
||||||
buffer_descriptor.size_alignment >
|
|
||||||
0) {
|
|
||||||
return absl::AlreadyExistsError(
|
|
||||||
"A view with different params is already allocated with this "
|
|
||||||
"storage");
|
|
||||||
}
|
|
||||||
|
|
||||||
absl::StatusOr<std::unique_ptr<Tensor::View>> result;
|
|
||||||
if (view_type_id == TensorHardwareBufferView::kId) {
|
|
||||||
result = GetAhwbView(access, shape, base_descriptor);
|
|
||||||
} else {
|
|
||||||
result = GetCpuView(access, shape, base_descriptor);
|
|
||||||
}
|
|
||||||
if (result.ok()) version_ = latest_version;
|
|
||||||
return result;
|
|
||||||
}
|
|
||||||
|
|
||||||
private:
|
|
||||||
absl::StatusOr<std::unique_ptr<Tensor::View>> GetAhwbView(
|
|
||||||
Tensor::View::Access access, const Tensor::Shape& shape,
|
|
||||||
const Tensor::ViewDescriptor& base_descriptor) {
|
|
||||||
return std::unique_ptr<Tensor::View>(new TensorHardwareBufferViewImpl(
|
|
||||||
kAccessCapability, access, Tensor::View::State::kUpToDate,
|
|
||||||
hw_descriptor_, ahwb_handle_));
|
|
||||||
}
|
|
||||||
|
|
||||||
absl::StatusOr<std::unique_ptr<Tensor::View>> GetCpuView(
|
|
||||||
Tensor::View::Access access, const Tensor::Shape& shape,
|
|
||||||
const Tensor::ViewDescriptor& base_descriptor) {
|
|
||||||
void* pointer = nullptr;
|
|
||||||
if (__builtin_available(android 26, *)) {
|
|
||||||
int error =
|
|
||||||
AHardwareBuffer_lock(ahwb_handle_,
|
|
||||||
access == Tensor::View::Access::kWriteOnly
|
|
||||||
? AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN
|
|
||||||
: AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN,
|
|
||||||
-1, nullptr, &pointer);
|
|
||||||
if (error != 0) {
|
|
||||||
return absl::UnknownError(
|
|
||||||
absl::StrCat("Error locking hardware buffer: ", error));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return std::unique_ptr<Tensor::View>(
|
|
||||||
new TensorCpuViewImpl(access == Tensor::View::Access::kWriteOnly
|
|
||||||
? Tensor::View::AccessCapability::kWrite
|
|
||||||
: Tensor::View::AccessCapability::kRead,
|
|
||||||
access, Tensor::View::State::kUpToDate,
|
|
||||||
cpu_descriptor_, pointer, ahwb_handle_));
|
|
||||||
}
|
|
||||||
|
|
||||||
static constexpr int kAccessCapability =
|
|
||||||
Tensor::View::AccessCapability::kRead |
|
|
||||||
Tensor::View::AccessCapability::kWrite;
|
|
||||||
TensorHardwareBufferViewDescriptor hw_descriptor_;
|
|
||||||
AHardwareBuffer* ahwb_handle_ = nullptr;
|
|
||||||
|
|
||||||
TensorCpuViewDescriptor cpu_descriptor_;
|
|
||||||
uint64_t version_ = 0;
|
|
||||||
};
|
|
||||||
TENSOR_REGISTER_STORAGE(HardwareBufferCpuStorage);
|
|
||||||
|
|
||||||
} // namespace
|
|
||||||
} // namespace mediapipe
|
|
||||||
|
|
||||||
#endif // !defined(MEDIAPIPE_NO_JNI) && (__ANDROID_API__ >= 26 ||
|
|
||||||
// defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))
|
|
||||||
|
|
@ -1,76 +0,0 @@
|
||||||
|
|
||||||
#if !defined(MEDIAPIPE_NO_JNI) && \
|
|
||||||
(__ANDROID_API__ >= 26 || \
|
|
||||||
defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))
|
|
||||||
#include <android/hardware_buffer.h>
|
|
||||||
|
|
||||||
#include <cstdint>
|
|
||||||
|
|
||||||
#include "mediapipe/framework/formats/tensor_cpu_buffer.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_hardware_buffer.h"
|
|
||||||
#include "mediapipe/framework/formats/tensor_v2.h"
|
|
||||||
#include "testing/base/public/gmock.h"
|
|
||||||
#include "testing/base/public/gunit.h"
|
|
||||||
|
|
||||||
namespace mediapipe {
|
|
||||||
|
|
||||||
namespace {
|
|
||||||
|
|
||||||
class TensorHardwareBufferTest : public ::testing::Test {
|
|
||||||
public:
|
|
||||||
TensorHardwareBufferTest() {}
|
|
||||||
~TensorHardwareBufferTest() override {}
|
|
||||||
};
|
|
||||||
|
|
||||||
TEST_F(TensorHardwareBufferTest, TestFloat32) {
|
|
||||||
Tensor tensor{Tensor::Shape({1})};
|
|
||||||
{
|
|
||||||
MP_ASSERT_OK_AND_ASSIGN(
|
|
||||||
auto view,
|
|
||||||
tensor.GetView<Tensor::View::Access::kWriteOnly>(
|
|
||||||
TensorHardwareBufferViewDescriptor{
|
|
||||||
.buffer = {.format =
|
|
||||||
TensorBufferDescriptor::Format::kFloat32}}));
|
|
||||||
EXPECT_NE(view->handle(), nullptr);
|
|
||||||
}
|
|
||||||
{
|
|
||||||
const auto& const_tensor = tensor;
|
|
||||||
MP_ASSERT_OK_AND_ASSIGN(
|
|
||||||
auto view,
|
|
||||||
const_tensor.GetView<Tensor::View::Access::kReadOnly>(
|
|
||||||
TensorCpuViewDescriptor{
|
|
||||||
.buffer = {.format =
|
|
||||||
TensorBufferDescriptor::Format::kFloat32}}));
|
|
||||||
EXPECT_NE(view->data<void>(), nullptr);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
TEST_F(TensorHardwareBufferTest, TestInt8Padding) {
|
|
||||||
Tensor tensor{Tensor::Shape({1})};
|
|
||||||
|
|
||||||
{
|
|
||||||
MP_ASSERT_OK_AND_ASSIGN(
|
|
||||||
auto view,
|
|
||||||
tensor.GetView<Tensor::View::Access::kWriteOnly>(
|
|
||||||
TensorHardwareBufferViewDescriptor{
|
|
||||||
.buffer = {.format = TensorBufferDescriptor::Format::kInt8,
|
|
||||||
.size_alignment = 4}}));
|
|
||||||
EXPECT_NE(view->handle(), nullptr);
|
|
||||||
}
|
|
||||||
{
|
|
||||||
const auto& const_tensor = tensor;
|
|
||||||
MP_ASSERT_OK_AND_ASSIGN(
|
|
||||||
auto view,
|
|
||||||
const_tensor.GetView<Tensor::View::Access::kReadOnly>(
|
|
||||||
TensorCpuViewDescriptor{
|
|
||||||
.buffer = {.format = TensorBufferDescriptor::Format::kInt8}}));
|
|
||||||
EXPECT_NE(view->data<void>(), nullptr);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
} // namespace
|
|
||||||
|
|
||||||
} // namespace mediapipe
|
|
||||||
|
|
||||||
#endif // !defined(MEDIAPIPE_NO_JNI) && (__ANDROID_API__ >= 26 ||
|
|
||||||
// defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))
|
|
||||||
Loading…
Reference in New Issue
Block a user