From 7c5c21665242ee3a2f3ce2875d6379e74ec9a22c Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Wed, 8 Nov 2023 20:04:12 -0800
Subject: [PATCH] Exposes a handle to AHardwareBuffers through a new GpuBuffer
 view

PiperOrigin-RevId: 580754933
---
 mediapipe/framework/formats/BUILD       |  9 +++++
 mediapipe/framework/formats/ahwb_view.h | 54 +++++++++++++++++++++++++
 mediapipe/framework/port.h              |  5 +++
 mediapipe/gpu/BUILD                     |  7 ++++
 4 files changed, 75 insertions(+)
 create mode 100644 mediapipe/framework/formats/ahwb_view.h

diff --git a/mediapipe/framework/formats/BUILD b/mediapipe/framework/formats/BUILD
index b36ea0211..047b95d32 100644
--- a/mediapipe/framework/formats/BUILD
+++ b/mediapipe/framework/formats/BUILD
@@ -124,6 +124,15 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ahwb_view",
+    hdrs = ["ahwb_view.h"],
+    deps = [
+        "//mediapipe/framework:port",
+        "//mediapipe/gpu:gpu_buffer_storage",
+    ],
+)
+
 cc_library(
     name = "affine_transform",
     srcs = ["affine_transform.cc"],
diff --git a/mediapipe/framework/formats/ahwb_view.h b/mediapipe/framework/formats/ahwb_view.h
new file mode 100644
index 000000000..0c8ad6323
--- /dev/null
+++ b/mediapipe/framework/formats/ahwb_view.h
@@ -0,0 +1,54 @@
+#ifndef MEDIAPIPE_FRAMEWORK_FORMATS_AHWB_VIEW_H_
+#define MEDIAPIPE_FRAMEWORK_FORMATS_AHWB_VIEW_H_
+
+#include "mediapipe/framework/port.h"
+#ifdef MEDIAPIPE_GPU_BUFFER_USE_AHWB
+#include <android/hardware_buffer.h>
+
+#include "mediapipe/gpu/gpu_buffer_storage.h"
+
+namespace mediapipe {
+
+// Wrapper to facilitate short lived access to Android Hardware Buffer objects.
+// Intended use cases:
+// - Extracting an AHWB for processing in another library after it's produced by
+// MediaPipe.
+// - Sending AHWBs to compute devices that are able to map the memory for their
+// own usage.
+// The AHWB abstractions in GpuBuffer and Tensor are likely more suitable for
+// other CPU/GPU uses of AHWBs.
+class AhwbView {
+ public:
+  explicit AhwbView(AHardwareBuffer* handle) : handle_(handle) {}
+  // Non-copyable
+  AhwbView(const AhwbView&) = delete;
+  AhwbView& operator=(const AhwbView&) = delete;
+  // Non-movable
+  AhwbView(AhwbView&&) = delete;
+
+  // Only supports synchronous usage. All users of GetHandle must finish
+  // accessing the buffer before this view object is destroyed to avoid race
+  // conditions.
+  // TODO: Support asynchronous usage.
+  const AHardwareBuffer* GetHandle() const { return handle_; }
+
+ private:
+  const AHardwareBuffer* handle_;
+};
+
+namespace internal {
+// Makes this class available as a GpuBuffer view.
+template <>
+class ViewProvider<AhwbView> {
+ public:
+  virtual ~ViewProvider() = default;
+  virtual const AhwbView GetReadView(types<AhwbView>) const = 0;
+  virtual AhwbView GetWriteView(types<AhwbView>) = 0;
+};
+
+}  // namespace internal
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_AHWB
+#endif  // MEDIAPIPE_FRAMEWORK_FORMATS_AHWB_VIEW_H_
diff --git a/mediapipe/framework/port.h b/mediapipe/framework/port.h
index e8b17e4d2..1bb4d4cdf 100644
--- a/mediapipe/framework/port.h
+++ b/mediapipe/framework/port.h
@@ -104,4 +104,9 @@
 #endif
 #endif  // MEDIAPIPE_HAS_RTTI
 
+// AHardware buffers are only available since Android API 26.
+#if (__ANDROID_API__ >= 26)
+#define MEDIAPIPE_GPU_BUFFER_USE_AHWB 1
+#endif
+
 #endif  // MEDIAPIPE_FRAMEWORK_PORT_H_
diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD
index f39b8d3f7..6ce4ec117 100644
--- a/mediapipe/gpu/BUILD
+++ b/mediapipe/gpu/BUILD
@@ -511,12 +511,19 @@ cc_library(
         ],
     }),
     deps = [
+        ":gl_base_hdr",
+        ":gl_context",
         ":gl_texture_buffer",
+        ":gl_texture_view",
         ":gpu_buffer_format",
         ":gpu_buffer_storage",
         ":image_frame_view",
+        "//mediapipe/framework:port",
+        "//mediapipe/framework/formats:ahwb_view",
         "//mediapipe/framework/formats:image_frame",
         "//mediapipe/framework/port:ret_check",
+        "//third_party/GL:EGL_headers",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/strings:str_format",
     ],
 )