Store AHWB in NN canonical memory type -- runtime

Prior to this CL, the canonical memory type only held a SharedHandle,
which mirrors the behavior of native_handle_t/hidl_handle. This means
memory types including AHardwareBuffer were stored as this SharedHandle
type. With this CL, the canonical memory type is stored directly as
AHardwareBuffer to avoid using non-NDK AHardwareBuffer calls in the NN
runtime.

Bug: 179906132
Test: mma
Test: NeuralNetworksTest_static
Test: presubmit
Change-Id: Ic269a5297ac1cec17fc5c30821c240c2254eedab
Merged-In: Ic269a5297ac1cec17fc5c30821c240c2254eedab
(cherry picked from commit d0ed0309474dca1daabe6f6c73ed1d5f5e92e3fb)
diff --git a/common/MetaModel.cpp b/common/MetaModel.cpp
index 9f3baf2..68f0103 100644
--- a/common/MetaModel.cpp
+++ b/common/MetaModel.cpp
@@ -74,8 +74,14 @@
     // We shouldn't have to check whether the model is valid. However, it could
     // be invalid if there is an error in the slicing algorithm.
     auto maybeVersion = validate(model);
-    if (!maybeVersion.has_value() || maybeVersion.value() > version) {
-        LOG(WARNING) << "Sliced model fails validate()";
+    if (!maybeVersion.has_value()) {
+        LOG(WARNING) << "Sliced model fails validate(): " << maybeVersion.error();
+        CHECK(!strictSlicing);
+        return true;
+    }
+    if (maybeVersion.value() > version) {
+        LOG(WARNING) << "Sliced model fails validate(): insufficient version ("
+                     << maybeVersion.value() << " vs " << version << ")";
         CHECK(!strictSlicing);
         return true;
     }
diff --git a/common/SharedMemoryAndroid.cpp b/common/SharedMemoryAndroid.cpp
index 4cc19f6..4ef9985 100644
--- a/common/SharedMemoryAndroid.cpp
+++ b/common/SharedMemoryAndroid.cpp
@@ -43,7 +43,6 @@
 
 #ifndef NN_NO_AHWB
 #include <android/hardware_buffer.h>
-#include <vndk/hardware_buffer.h>
 #endif  // NN_NO_AHWB
 
 namespace android::nn {
@@ -75,44 +74,44 @@
     return maybeMemory;
 }
 
-GeneralResult<hardware::hidl_handle> hidlHandleFromSharedHandle(const SharedHandle& handle) {
-    if (handle == nullptr) {
-        return {};
-    }
-
+GeneralResult<hardware::hidl_handle> hidlHandleFromSharedHandle(const Handle& handle) {
     std::vector<base::unique_fd> fds;
-    fds.reserve(handle->fds.size());
-    for (const auto& fd : handle->fds) {
-        int dupFd = dup(fd);
+    fds.reserve(handle.fds.size());
+    for (const auto& fd : handle.fds) {
+        const int dupFd = dup(fd);
         if (dupFd == -1) {
             return NN_ERROR(ErrorStatus::GENERAL_FAILURE) << "Failed to dup the fd";
         }
         fds.emplace_back(dupFd);
     }
 
-    native_handle_t* nativeHandle = native_handle_create(handle->fds.size(), handle->ints.size());
+    constexpr size_t kMaxInt = std::numeric_limits<int>::max();
+    CHECK_LE(handle.fds.size(), kMaxInt);
+    CHECK_LE(handle.ints.size(), kMaxInt);
+    native_handle_t* nativeHandle = native_handle_create(static_cast<int>(handle.fds.size()),
+                                                         static_cast<int>(handle.ints.size()));
     if (nativeHandle == nullptr) {
         return NN_ERROR(ErrorStatus::GENERAL_FAILURE) << "Failed to create native_handle";
     }
     for (size_t i = 0; i < fds.size(); ++i) {
         nativeHandle->data[i] = fds[i].release();
     }
-    std::copy(handle->ints.begin(), handle->ints.end(), &nativeHandle->data[nativeHandle->numFds]);
+    std::copy(handle.ints.begin(), handle.ints.end(), &nativeHandle->data[nativeHandle->numFds]);
 
     hardware::hidl_handle hidlHandle;
     hidlHandle.setTo(nativeHandle, /*shouldOwn=*/true);
     return hidlHandle;
 }
 
-GeneralResult<SharedHandle> sharedHandleFromNativeHandle(const native_handle_t* handle) {
+GeneralResult<Handle> sharedHandleFromNativeHandle(const native_handle_t* handle) {
     if (handle == nullptr) {
-        return nullptr;
+        return NN_ERROR() << "sharedHandleFromNativeHandle was passed nullptr for handle";
     }
 
     std::vector<base::unique_fd> fds;
     fds.reserve(handle->numFds);
     for (int i = 0; i < handle->numFds; ++i) {
-        int dupFd = dup(handle->data[i]);
+        const int dupFd = dup(handle->data[i]);
         if (dupFd == -1) {
             return NN_ERROR(ErrorStatus::GENERAL_FAILURE) << "Failed to dup the fd";
         }
@@ -122,10 +121,7 @@
     std::vector<int> ints(&handle->data[handle->numFds],
                           &handle->data[handle->numFds + handle->numInts]);
 
-    return std::make_shared<const Handle>(Handle{
-            .fds = std::move(fds),
-            .ints = std::move(ints),
-    });
+    return Handle{.fds = std::move(fds), .ints = std::move(ints)};
 }
 
 GeneralResult<SharedMemory> createMemory(const hidl_memory& memory) {
@@ -138,7 +134,9 @@
 }
 
 GeneralResult<hidl_memory> createHidlMemory(const Memory& memory) {
-    return hidl_memory(memory.name, NN_TRY(hidlHandleFromSharedHandle(memory.handle)), memory.size);
+    return hidl_memory(memory.name,
+                       NN_TRY(hidlHandleFromSharedHandle(std::get<Handle>(memory.handle))),
+                       memory.size);
 }
 
 GeneralResult<Mapping> mapAshmem(const Memory& memory) {
@@ -173,10 +171,10 @@
 
 GeneralResult<Mapping> mapMemFd(const Memory& memory) {
     const size_t size = memory.size;
-    const SharedHandle& handle = memory.handle;
-    const int fd = handle->fds[0];
-    const int prot = handle->ints[0];
-    const size_t offset = getOffsetFromInts(handle->ints[1], handle->ints[2]);
+    const Handle& handle = std::get<Handle>(memory.handle);
+    const int fd = handle.fds[0];
+    const int prot = handle.ints[0];
+    const size_t offset = getOffsetFromInts(handle.ints[1], handle.ints[2]);
 
     std::shared_ptr<base::MappedFile> mapping = base::MappedFile::FromFd(fd, offset, size, prot);
     if (mapping == nullptr) {
@@ -190,62 +188,23 @@
 
 #ifndef NN_NO_AHWB
 
-static uint32_t roundUpToMultiple(uint32_t value, uint32_t multiple) {
-    return (value + multiple - 1) / multiple * multiple;
-}
+GeneralResult<Mapping> mapAhwbBlobMemory(const HardwareBufferHandle& memory) {
+    AHardwareBuffer_Desc desc;
+    AHardwareBuffer_describe(memory.get(), &desc);
 
-GeneralResult<Mapping> mapAhwbBlobMemory(const Memory& memory) {
-    const SharedHandle& handle = memory.handle;
-    const auto size = memory.size;
-    const auto format = AHARDWAREBUFFER_FORMAT_BLOB;
-    const auto usage = AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN | AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN;
-    const uint32_t width = size;
-    const uint32_t height = 1;  // height is always 1 for BLOB mode AHardwareBuffer.
-    const uint32_t layers = 1;  // layers is always 1 for BLOB mode AHardwareBuffer.
-
-    // AHardwareBuffer_createFromHandle() might fail because an allocator
-    // expects a specific stride value. In that case, we try to guess it by
-    // aligning the width to small powers of 2.
-    // TODO(b/174120849): Avoid stride assumptions.
-    AHardwareBuffer* hardwareBuffer = nullptr;
-    status_t status = UNKNOWN_ERROR;
-    for (uint32_t alignment : {1, 4, 32, 64, 128, 2, 8, 16}) {
-        const uint32_t stride = roundUpToMultiple(width, alignment);
-        AHardwareBuffer_Desc desc{
-                .width = width,
-                .height = height,
-                .layers = layers,
-                .format = format,
-                .usage = usage,
-                .stride = stride,
-        };
-        status = AHardwareBuffer_createFromHandle(&desc, NN_TRY(hidlHandleFromSharedHandle(handle)),
-                                                  AHARDWAREBUFFER_CREATE_FROM_HANDLE_METHOD_CLONE,
-                                                  &hardwareBuffer);
-        if (status == NO_ERROR) {
-            break;
-        }
-    }
-    if (status != NO_ERROR) {
-        return NN_ERROR(ErrorStatus::GENERAL_FAILURE)
-               << "Can't create AHardwareBuffer from handle. Error: " << status;
-    }
+    CHECK_EQ(desc.format, AHARDWAREBUFFER_FORMAT_BLOB);
+    const uint32_t size = desc.width;
 
     void* data = nullptr;
-    status = AHardwareBuffer_lock(hardwareBuffer, usage, -1, nullptr, &data);
+    const auto status = AHardwareBuffer_lock(memory.get(), desc.usage, -1, nullptr, &data);
     if (status != NO_ERROR) {
         return NN_ERROR(ErrorStatus::GENERAL_FAILURE)
                << "Can't lock the AHardwareBuffer. Error: " << status;
-        // TODO(b/169166682): do we need to call AHardwareBuffer_release?
     }
 
     // Create shared scoped object to munmap.
-    auto scoped = base::make_scope_guard([hardwareBuffer] {
-        AHardwareBuffer_unlock(hardwareBuffer, nullptr);
-        if (hardwareBuffer != nullptr) {
-            AHardwareBuffer_release(hardwareBuffer);
-        }
-    });
+    auto scoped = base::make_scope_guard(
+            [ahwb = memory.get()] { AHardwareBuffer_unlock(ahwb, nullptr); });
     auto sharedScoped = std::make_shared<decltype(scoped)>(std::move(scoped));
 
     return Mapping{.pointer = data, .size = size, .context = std::move(sharedScoped)};
@@ -256,14 +215,38 @@
            << "Unable to map non-BLOB AHardwareBuffer memory";
 }
 
+void freeHardwareBuffer(AHardwareBuffer* buffer) {
+    if (buffer) {
+        AHardwareBuffer_release(buffer);
+    }
+}
+
 #endif  // NN_NO_AHWB
 
+void freeNoop(AHardwareBuffer* /*buffer*/) {}
+
 }  // namespace
 
+#ifndef NN_NO_AHWB
+HardwareBufferHandle::HardwareBufferHandle(AHardwareBuffer* handle, bool takeOwnership)
+    : mHandle(handle, (takeOwnership ? freeHardwareBuffer : freeNoop)) {
+    CHECK(mHandle != nullptr);
+}
+#else
+HardwareBufferHandle::HardwareBufferHandle(AHardwareBuffer* handle, bool /*takeOwnership*/)
+    : mHandle(handle, freeNoop) {
+    CHECK(mHandle != nullptr);
+}
+#endif  // NN_NO_AHWB
+
+AHardwareBuffer* HardwareBufferHandle::get() const {
+    return mHandle.get();
+}
+
 GeneralResult<SharedMemory> createSharedMemory(size_t size) {
 #ifndef NN_COMPATIBILITY_LIBRARY_BUILD
     const auto memory = NN_TRY(allocateSharedMemory(size));
-    return createSharedMemoryFromHidlMemory(memory);
+    return createMemory(memory);
 #else
     (void)size;
     return NN_ERROR(ErrorStatus::INVALID_ARGUMENT) << "Unimplemented";
@@ -288,32 +271,22 @@
     const auto [lowOffsetBits, highOffsetBits] = getIntsFromOffset(offset);
     std::vector<int> ints = {prot, lowOffsetBits, highOffsetBits};
 
-    SharedHandle handle = std::make_shared<const Handle>(Handle{
-            .fds = std::move(fds),
-            .ints = std::move(ints),
-    });
+    auto handle = Handle{.fds = std::move(fds), .ints = std::move(ints)};
     return std::make_shared<const Memory>(
             Memory{.handle = std::move(handle), .size = size, .name = "mmap_fd"});
 }
 
-GeneralResult<SharedMemory> createSharedMemoryFromHidlMemory(const hardware::hidl_memory& memory) {
+GeneralResult<SharedMemory> createSharedMemoryFromAHWB(AHardwareBuffer* ahwb, bool takeOwnership) {
 #ifndef NN_NO_AHWB
-    return createMemory(memory);
-#else
-    (void)memory;
-    return NN_ERROR(ErrorStatus::INVALID_ARGUMENT) << "hidl_memory not supported";
-#endif  // NN_NO_AHWB
-}
+    CHECK(ahwb != nullptr);
+    auto handle = HardwareBufferHandle(ahwb, takeOwnership);
 
-GeneralResult<SharedMemory> createSharedMemoryFromAHWB(const AHardwareBuffer& ahwb) {
-#ifndef NN_NO_AHWB
     AHardwareBuffer_Desc bufferDesc;
-    AHardwareBuffer_describe(&ahwb, &bufferDesc);
-    const native_handle_t* handle = AHardwareBuffer_getNativeHandle(&ahwb);
+    AHardwareBuffer_describe(ahwb, &bufferDesc);
 
     if (bufferDesc.format == AHARDWAREBUFFER_FORMAT_BLOB) {
         return std::make_shared<const Memory>(Memory{
-                .handle = NN_TRY(sharedHandleFromNativeHandle(handle)),
+                .handle = std::move(handle),
                 .size = bufferDesc.width,
                 .name = "hardware_buffer_blob",
         });
@@ -321,12 +294,13 @@
 
     // memory size is not used for non-BLOB AHWB memory.
     return std::make_shared<const Memory>(Memory{
-            .handle = NN_TRY(sharedHandleFromNativeHandle(handle)),
+            .handle = std::move(handle),
             .size = 0,
             .name = "hardware_buffer",
     });
 #else
     (void)ahwb;
+    (void)takeOwnership;
     return NN_ERROR(ErrorStatus::INVALID_ARGUMENT)
            << "AHardwareBuffer memory not implemented for support library build yet";
 #endif  // NN_NO_AHWB
@@ -341,7 +315,7 @@
         return mapAshmem(*memory);
     }
     if (memory->name == "hardware_buffer_blob") {
-        return mapAhwbBlobMemory(*memory);
+        return mapAhwbBlobMemory(std::get<HardwareBufferHandle>(memory->handle));
     }
     if (memory->name == "hardware_buffer") {
         return mapAhwbMemory(*memory);
diff --git a/common/SharedMemoryHost.cpp b/common/SharedMemoryHost.cpp
index 9f7d2c6..00c9df9 100644
--- a/common/SharedMemoryHost.cpp
+++ b/common/SharedMemoryHost.cpp
@@ -38,7 +38,7 @@
     CHECK_LE(memory.size, std::numeric_limits<uint32_t>::max());
     const auto size = memory.size;
 
-    const int fd = memory.handle->fds[0];
+    const int fd = std::get<Handle>(memory.handle).fds[0];
     std::shared_ptr<base::MappedFile> mapping =
             base::MappedFile::FromFd(fd, /*offset=*/0, size, PROT_READ | PROT_WRITE);
     if (mapping == nullptr) {
@@ -56,7 +56,7 @@
 
 GeneralResult<Mapping> mapMemFd(const Memory& memory) {
     const size_t size = memory.size;
-    const SharedHandle& handle = memory.handle;
+    const Handle& handle = std::get<Handle>(memory.handle);
     const int fd = handle->fds[0];
     const int prot = handle->ints[0];
     const size_t offset = getOffsetFromInts(handle->ints[1], handle->ints[2]);
@@ -71,8 +71,19 @@
     return Mapping{.pointer = data, .size = size, .context = std::move(context)};
 }
 
+void freeNoop(AHardwareBuffer* /*buffer*/) {}
+
 }  // namespace
 
+HardwareBufferHandle::HardwareBufferHandle(AHardwareBuffer* handle, bool /*takeOwnership*/)
+    : mHandle(handle, freeNoop) {
+    CHECK(mHandle != nullptr);
+}
+
+AHardwareBuffer* HardwareBufferHandle::get() const {
+    return mHandle.get();
+}
+
 GeneralResult<SharedMemory> createSharedMemory(size_t size) {
     int fd = ashmem_create_region("NnapiAshmem", size);
     if (fd < 0) {
@@ -83,7 +94,7 @@
     std::vector<base::unique_fd> fds;
     fds.emplace_back(fd);
 
-    SharedHandle handle = std::make_shared<const Handle>(Handle{
+    auto handle = Handle{
             .fds = std::move(fds),
             .ints = {},
     });
@@ -109,18 +120,13 @@
     const auto [lowOffsetBits, highOffsetBits] = getIntsFromOffset(offset);
     std::vector<int> ints = {prot, lowOffsetBits, highOffsetBits};
 
-    SharedHandle handle = std::make_shared<const Handle>(Handle{
+    auto handle = Handle{
             .fds = std::move(fds),
             .ints = std::move(ints),
     });
     return SharedMemory{.handle = std::move(handle), .size = size, .name = "mmap_fd"};
 }
 
-GeneralResult<SharedMemory> createSharedMemoryFromHidlMemory(
-        const hardware::hidl_memory& /*memory*/) {
-    return NN_ERROR(ErrorStatus::INVALID_ARGUMENT) << "hidl_memory not supported on host";
-}
-
 GeneralResult<SharedMemory> createSharedMemoryFromAHWB(const AHardwareBuffer& /*ahwb*/) {
     return NN_ERROR(ErrorStatus::INVALID_ARGUMENT)
            << "AHardwareBuffer memory not supported on host";
diff --git a/common/TypeUtils.cpp b/common/TypeUtils.cpp
index ef19a41..693a0b1 100644
--- a/common/TypeUtils.cpp
+++ b/common/TypeUtils.cpp
@@ -704,8 +704,27 @@
               << ", .outputs=" << operation.outputs << "}";
 }
 
+static std::ostream& operator<<(std::ostream& os, const Handle& handle) {
+    return os << "<handle with " << handle.fds.size() << " fds and " << handle.ints.size()
+              << " ints>";
+}
+
 std::ostream& operator<<(std::ostream& os, const SharedHandle& handle) {
-    return os << (handle != nullptr ? "<non-empty handle>" : "<empty handle>");
+    if (handle == nullptr) {
+        return os << "<empty handle>";
+    }
+    return os << *handle;
+}
+
+static std::ostream& operator<<(std::ostream& os, const HardwareBufferHandle& handle) {
+    return os << (handle.get() != nullptr ? "<non-empty HardwareBufferHandle>"
+                                          : "<empty HardwareBufferHandle>");
+}
+
+static std::ostream& operator<<(std::ostream& os,
+                                const std::variant<Handle, HardwareBufferHandle>& handle) {
+    std::visit([&os](const auto& x) { os << x; }, handle);
+    return os;
 }
 
 std::ostream& operator<<(std::ostream& os, const Memory& memory) {
diff --git a/common/Validation.cpp b/common/Validation.cpp
index 069a661..cb1a3e1 100644
--- a/common/Validation.cpp
+++ b/common/Validation.cpp
@@ -713,34 +713,51 @@
     return version;
 }
 
-Result<Version> validateSharedHandle(const SharedHandle& handle) {
-    NN_VALIDATE(handle != nullptr);
-    NN_VALIDATE(std::all_of(handle->fds.begin(), handle->fds.end(),
+Result<Version> validateHandle(const Handle& handle) {
+    NN_VALIDATE(std::all_of(handle.fds.begin(), handle.fds.end(),
                             [](const base::unique_fd& fd) { return fd.ok(); }));
     return Version::ANDROID_OC_MR1;
 }
 
+Result<Version> validateSharedHandle(const SharedHandle& handle) {
+    NN_VALIDATE(handle != nullptr);
+    return validateHandle(*handle);
+}
+
 Result<Version> validateSharedMemory(const SharedMemory& memory) {
     NN_VALIDATE(memory != nullptr);
-    NN_TRY(validateSharedHandle(memory->handle));
 
     if (memory->name == "ashmem") {
         NN_VALIDATE_NE(memory->size, 0u);
+        NN_VALIDATE(std::holds_alternative<Handle>(memory->handle));
+        NN_TRY(validateHandle(std::get<Handle>(memory->handle)));
         return Version::ANDROID_OC_MR1;
     }
     if (memory->name == "mmap_fd") {
         NN_VALIDATE_NE(memory->size, 0u);
+        NN_VALIDATE(std::holds_alternative<Handle>(memory->handle));
+        NN_TRY(validateHandle(std::get<Handle>(memory->handle)));
         return Version::ANDROID_OC_MR1;
     }
     if (memory->name == "hardware_buffer_blob") {
         NN_VALIDATE_NE(memory->size, 0u);
+        NN_VALIDATE(std::holds_alternative<HardwareBufferHandle>(memory->handle));
+        NN_VALIDATE(std::get<HardwareBufferHandle>(memory->handle).get() != nullptr);
         return Version::ANDROID_Q;
     }
     if (memory->name == "hardware_buffer") {
         // For hardware_buffer memory, all size information is bound to the AHardwareBuffer, so
         // memory.size must be 0.
         NN_VALIDATE_EQ(memory->size, 0u);
-        return Version::ANDROID_Q;
+        // hardware_buffer can be represented by either Handle or HardwareBufferHandle.
+        if (const auto* handle = std::get_if<Handle>(&memory->handle)) {
+            NN_TRY(validateHandle(*handle));
+            return Version::ANDROID_Q;
+        }
+        if (const auto* handle = std::get_if<HardwareBufferHandle>(&memory->handle)) {
+            NN_VALIDATE(handle->get() != nullptr);
+            return Version::ANDROID_Q;
+        }
     }
 
     NN_VALIDATE_FAIL() << "Unknown memory type " << memory->name;
diff --git a/common/include/nnapi/SharedMemory.h b/common/include/nnapi/SharedMemory.h
index 0467433..c5a3a19 100644
--- a/common/include/nnapi/SharedMemory.h
+++ b/common/include/nnapi/SharedMemory.h
@@ -30,15 +30,23 @@
 // Forward declare AHardwareBuffer
 extern "C" typedef struct AHardwareBuffer AHardwareBuffer;
 
-// Forward declare hidl_memory
-namespace android::hardware {
-struct hidl_memory;
-}  // namespace android::hardware
-
 namespace android::nn {
 
+// RAII wrapper for AHardwareBuffer
+class HardwareBufferHandle {
+   public:
+    // Precondition: handle != nullptr
+    HardwareBufferHandle(AHardwareBuffer* handle, bool takeOwnership);
+
+    AHardwareBuffer* get() const;
+
+   private:
+    using Deleter = std::add_pointer_t<void(AHardwareBuffer*)>;
+    std::unique_ptr<AHardwareBuffer, Deleter> mHandle;
+};
+
 struct Memory {
-    SharedHandle handle;
+    std::variant<Handle, HardwareBufferHandle> handle;
     size_t size = 0;
     std::string name;
 };
@@ -81,9 +89,8 @@
 
 GeneralResult<SharedMemory> createSharedMemoryFromFd(size_t size, int prot, int fd, size_t offset);
 
-GeneralResult<SharedMemory> createSharedMemoryFromHidlMemory(const hardware::hidl_memory& memory);
-
-GeneralResult<SharedMemory> createSharedMemoryFromAHWB(const AHardwareBuffer& ahwb);
+// Precondition: ahwb != nullptr
+GeneralResult<SharedMemory> createSharedMemoryFromAHWB(AHardwareBuffer* ahwb, bool takeOwnership);
 
 struct Mapping {
     std::variant<void*, const void*> pointer;
diff --git a/driver/sample/Android.bp b/driver/sample/Android.bp
index 4db78c8..920ae44 100644
--- a/driver/sample/Android.bp
+++ b/driver/sample/Android.bp
@@ -149,6 +149,7 @@
         "SampleDriverSL.cpp",
     ],
     shared_libs: [
+        "libandroid",
         "liblog",
         "libnativewindow",
     ],
diff --git a/runtime/Memory.cpp b/runtime/Memory.cpp
index 511ffd3..e00c0f9 100644
--- a/runtime/Memory.cpp
+++ b/runtime/Memory.cpp
@@ -560,7 +560,8 @@
 MemoryFd::MemoryFd(SharedMemory memory) : RuntimeMemory(std::move(memory)) {}
 
 std::pair<int, std::unique_ptr<MemoryAHWB>> MemoryAHWB::create(const AHardwareBuffer& ahwb) {
-    auto memory = createSharedMemoryFromAHWB(ahwb);
+    auto memory = createSharedMemoryFromAHWB(const_cast<AHardwareBuffer*>(&ahwb),
+                                             /*takeOwnership=*/false);
     if (!memory.has_value()) {
         LOG(ERROR) << "Failed to create memory from AHWB: " << memory.error().message;
         return {convertErrorStatusToResultCode(memory.error().code), nullptr};
@@ -593,9 +594,8 @@
         LOG(ERROR) << "Failed to allocate BLOB mode AHWB.";
         return {ANEURALNETWORKS_OP_FAILED, nullptr};
     }
-    auto ahwbGuard = base::make_scope_guard([ahwb]() { AHardwareBuffer_release(ahwb); });
 
-    auto memory = createSharedMemoryFromAHWB(*ahwb);
+    auto memory = createSharedMemoryFromAHWB(ahwb, /*takeOWnership=*/true);
     if (!memory.has_value()) {
         LOG(ERROR) << "Failed to allocate BLOB mode AHWB: " << memory.error().message;
         return {convertErrorStatusToResultCode(memory.error().code), nullptr};
@@ -605,8 +605,8 @@
         LOG(ERROR) << "Failed to map BLOB mode AHWB: " << mapping.error().message;
         return {convertErrorStatusToResultCode(mapping.error().code), nullptr};
     }
-    auto memoryAHWB = std::make_unique<MemoryRuntimeAHWB>(
-            std::move(memory).value(), std::move(ahwbGuard), std::move(mapping).value());
+    auto memoryAHWB = std::make_unique<MemoryRuntimeAHWB>(std::move(memory).value(),
+                                                          std::move(mapping).value());
     return {ANEURALNETWORKS_NO_ERROR, std::move(memoryAHWB)};
 }
 
@@ -614,12 +614,8 @@
     return static_cast<uint8_t*>(std::get<void*>(kMapping.pointer));
 }
 
-MemoryRuntimeAHWB::MemoryRuntimeAHWB(SharedMemory memory,
-                                     base::ScopeGuard<std::function<void()>> ahwbScopeGuard,
-                                     Mapping mapping)
-    : RuntimeMemory(std::move(memory)),
-      kAhwbScopeGuard(std::move(ahwbScopeGuard)),
-      kMapping(std::move(mapping)) {}
+MemoryRuntimeAHWB::MemoryRuntimeAHWB(SharedMemory memory, Mapping mapping)
+    : RuntimeMemory(std::move(memory)), kMapping(std::move(mapping)) {}
 
 std::pair<int, std::unique_ptr<MemoryFromDevice>> MemoryFromDevice::create(SharedBuffer buffer) {
     if (buffer == nullptr) {
diff --git a/runtime/Memory.h b/runtime/Memory.h
index 494216e..f750035 100644
--- a/runtime/Memory.h
+++ b/runtime/Memory.h
@@ -347,11 +347,9 @@
     }
 
     // prefer using MemoryRuntimeAHWB::create
-    MemoryRuntimeAHWB(SharedMemory memory, base::ScopeGuard<std::function<void()>> ahwbScopeGuard,
-                      Mapping mapping);
+    MemoryRuntimeAHWB(SharedMemory memory, Mapping mapping);
 
    private:
-    const base::ScopeGuard<std::function<void()>> kAhwbScopeGuard;
     const Mapping kMapping;
 };