Internal change.
PiperOrigin-RevId: 495483878
This commit is contained in:
		
							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