From db3cb68d919693adb729437d1223d29c30736f27 Mon Sep 17 00:00:00 2001 From: Nikolay Chirkov Date: Fri, 9 Dec 2022 07:27:11 -0800 Subject: [PATCH] Internal change. PiperOrigin-RevId: 494166776 --- .../formats/tensor_hardware_buffer.h | 71 ++++++ .../tensor_hardware_buffer_cpu_storage.cc | 216 ++++++++++++++++++ ...tensor_hardware_buffer_cpu_storage_test.cc | 76 ++++++ 3 files changed, 363 insertions(+) create mode 100644 mediapipe/framework/formats/tensor_hardware_buffer.h create mode 100644 mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage.cc create mode 100644 mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage_test.cc diff --git a/mediapipe/framework/formats/tensor_hardware_buffer.h b/mediapipe/framework/formats/tensor_hardware_buffer.h new file mode 100644 index 000000000..fa0241bde --- /dev/null +++ b/mediapipe/framework/formats/tensor_hardware_buffer.h @@ -0,0 +1,71 @@ +#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 + +#include + +#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(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_ diff --git a/mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage.cc b/mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage.cc new file mode 100644 index 000000000..9c223ce2c --- /dev/null +++ b/mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage.cc @@ -0,0 +1,216 @@ +#if !defined(MEDIAPIPE_NO_JNI) && \ + (__ANDROID_API__ >= 26 || \ + defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__)) + +#include + +#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> GetViews(uint64_t latest_version) { + std::vector> result; + auto update_state = latest_version == version_ + ? Tensor::View::State::kUpToDate + : Tensor::View::State::kOutdated; + if (ahwb_handle_) { + result.push_back( + std::unique_ptr(new TensorHardwareBufferViewImpl( + kAccessCapability, Tensor::View::Access::kNoAccess, update_state, + hw_descriptor_, ahwb_handle_))); + + result.push_back(std::unique_ptr(new TensorCpuViewImpl( + kAccessCapability, Tensor::View::Access::kNoAccess, update_state, + cpu_descriptor_, nullptr, nullptr))); + } + return result; + } + + absl::StatusOr> 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( + base_descriptor) + .buffer + : static_cast(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> 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> GetAhwbView( + Tensor::View::Access access, const Tensor::Shape& shape, + const Tensor::ViewDescriptor& base_descriptor) { + return std::unique_ptr(new TensorHardwareBufferViewImpl( + kAccessCapability, access, Tensor::View::State::kUpToDate, + hw_descriptor_, ahwb_handle_)); + } + + absl::StatusOr> 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( + 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__)) diff --git a/mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage_test.cc b/mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage_test.cc new file mode 100644 index 000000000..0afa9899f --- /dev/null +++ b/mediapipe/framework/formats/tensor_hardware_buffer_cpu_storage_test.cc @@ -0,0 +1,76 @@ + +#if !defined(MEDIAPIPE_NO_JNI) && \ + (__ANDROID_API__ >= 26 || \ + defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__)) +#include + +#include + +#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( + 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( + TensorCpuViewDescriptor{ + .buffer = {.format = + TensorBufferDescriptor::Format::kFloat32}})); + EXPECT_NE(view->data(), nullptr); + } +} + +TEST_F(TensorHardwareBufferTest, TestInt8Padding) { + Tensor tensor{Tensor::Shape({1})}; + + { + MP_ASSERT_OK_AND_ASSIGN( + auto view, + tensor.GetView( + 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( + TensorCpuViewDescriptor{ + .buffer = {.format = TensorBufferDescriptor::Format::kInt8}})); + EXPECT_NE(view->data(), nullptr); + } +} + +} // namespace + +} // namespace mediapipe + +#endif // !defined(MEDIAPIPE_NO_JNI) && (__ANDROID_API__ >= 26 || + // defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__))