Change NNAPI Memory to ref-counted SharedMemory -- hal

Bug: 179906132
Test: mma
Test: NeuralNetworksTest_static
Test: presubmit
Change-Id: I6435db906a2efe4938da18149a1fcd6d24730a95
diff --git a/neuralnetworks/utils/common/include/nnapi/hal/CommonUtils.h b/neuralnetworks/utils/common/include/nnapi/hal/CommonUtils.h
index fef9d9c..b13785d 100644
--- a/neuralnetworks/utils/common/include/nnapi/hal/CommonUtils.h
+++ b/neuralnetworks/utils/common/include/nnapi/hal/CommonUtils.h
@@ -74,7 +74,7 @@
 std::vector<uint32_t> countNumberOfConsumers(size_t numberOfOperands,
                                              const std::vector<nn::Operation>& operations);
 
-nn::GeneralResult<nn::Memory> createSharedMemoryFromHidlMemory(const hidl_memory& memory);
+nn::GeneralResult<nn::SharedMemory> createSharedMemoryFromHidlMemory(const hidl_memory& memory);
 
 nn::GeneralResult<hidl_handle> hidlHandleFromSharedHandle(const nn::SharedHandle& handle);
 nn::GeneralResult<nn::SharedHandle> sharedHandleFromNativeHandle(const native_handle_t* handle);
diff --git a/neuralnetworks/utils/common/include/nnapi/hal/HandleError.h b/neuralnetworks/utils/common/include/nnapi/hal/HandleError.h
index 95a20a8..209b663 100644
--- a/neuralnetworks/utils/common/include/nnapi/hal/HandleError.h
+++ b/neuralnetworks/utils/common/include/nnapi/hal/HandleError.h
@@ -14,6 +14,9 @@
  * limitations under the License.
  */
 
+#ifndef ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_UTILS_COMMON_HANDLE_ERROR_H
+#define ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_UTILS_COMMON_HANDLE_ERROR_H
+
 #include <android/hidl/base/1.0/IBase.h>
 #include <hidl/HidlSupport.h>
 #include <nnapi/Result.h>
@@ -50,7 +53,8 @@
     })
 
 template <typename Type>
-nn::GeneralResult<Type> makeGeneralFailure(nn::Result<Type> result, nn::ErrorStatus status) {
+nn::GeneralResult<Type> makeGeneralFailure(
+        nn::Result<Type> result, nn::ErrorStatus status = nn::ErrorStatus::GENERAL_FAILURE) {
     if (!result.has_value()) {
         return nn::error(status) << std::move(result).error();
     }
@@ -75,7 +79,8 @@
 }
 
 template <typename Type>
-nn::ExecutionResult<Type> makeExecutionFailure(nn::Result<Type> result, nn::ErrorStatus status) {
+nn::ExecutionResult<Type> makeExecutionFailure(
+        nn::Result<Type> result, nn::ErrorStatus status = nn::ErrorStatus::GENERAL_FAILURE) {
     return makeExecutionFailure(makeGeneralFailure(result, status));
 }
 
@@ -86,4 +91,6 @@
     } else                                                              \
         return NN_ERROR(canonical)
 
-}  // namespace android::hardware::neuralnetworks::utils
\ No newline at end of file
+}  // namespace android::hardware::neuralnetworks::utils
+
+#endif  // ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_UTILS_COMMON_HANDLE_ERROR_H
diff --git a/neuralnetworks/utils/common/include/nnapi/hal/InvalidBuffer.h b/neuralnetworks/utils/common/include/nnapi/hal/InvalidBuffer.h
index 8c04b88..0e98c2e 100644
--- a/neuralnetworks/utils/common/include/nnapi/hal/InvalidBuffer.h
+++ b/neuralnetworks/utils/common/include/nnapi/hal/InvalidBuffer.h
@@ -31,9 +31,9 @@
   public:
     nn::Request::MemoryDomainToken getToken() const override;
 
-    nn::GeneralResult<void> copyTo(const nn::Memory& dst) const override;
+    nn::GeneralResult<void> copyTo(const nn::SharedMemory& dst) const override;
 
-    nn::GeneralResult<void> copyFrom(const nn::Memory& src,
+    nn::GeneralResult<void> copyFrom(const nn::SharedMemory& src,
                                      const nn::Dimensions& dimensions) const override;
 };
 
diff --git a/neuralnetworks/utils/common/include/nnapi/hal/InvalidBurst.h b/neuralnetworks/utils/common/include/nnapi/hal/InvalidBurst.h
index 83e60b6..996858c 100644
--- a/neuralnetworks/utils/common/include/nnapi/hal/InvalidBurst.h
+++ b/neuralnetworks/utils/common/include/nnapi/hal/InvalidBurst.h
@@ -29,7 +29,7 @@
 
 class InvalidBurst final : public nn::IBurst {
   public:
-    OptionalCacheHold cacheMemory(const nn::Memory& memory) const override;
+    OptionalCacheHold cacheMemory(const nn::SharedMemory& memory) const override;
 
     nn::ExecutionResult<std::pair<std::vector<nn::OutputShape>, nn::Timing>> execute(
             const nn::Request& request, nn::MeasureTiming measure) const override;
diff --git a/neuralnetworks/utils/common/include/nnapi/hal/ResilientBuffer.h b/neuralnetworks/utils/common/include/nnapi/hal/ResilientBuffer.h
index d2c2469..c8ca6f2 100644
--- a/neuralnetworks/utils/common/include/nnapi/hal/ResilientBuffer.h
+++ b/neuralnetworks/utils/common/include/nnapi/hal/ResilientBuffer.h
@@ -46,9 +46,9 @@
 
     nn::Request::MemoryDomainToken getToken() const override;
 
-    nn::GeneralResult<void> copyTo(const nn::Memory& dst) const override;
+    nn::GeneralResult<void> copyTo(const nn::SharedMemory& dst) const override;
 
-    nn::GeneralResult<void> copyFrom(const nn::Memory& src,
+    nn::GeneralResult<void> copyFrom(const nn::SharedMemory& src,
                                      const nn::Dimensions& dimensions) const override;
 
   private:
diff --git a/neuralnetworks/utils/common/include/nnapi/hal/ResilientBurst.h b/neuralnetworks/utils/common/include/nnapi/hal/ResilientBurst.h
index 0df287f..3b87330 100644
--- a/neuralnetworks/utils/common/include/nnapi/hal/ResilientBurst.h
+++ b/neuralnetworks/utils/common/include/nnapi/hal/ResilientBurst.h
@@ -44,7 +44,7 @@
     nn::SharedBurst getBurst() const;
     nn::GeneralResult<nn::SharedBurst> recover(const nn::IBurst* failingBurst) const;
 
-    OptionalCacheHold cacheMemory(const nn::Memory& memory) const override;
+    OptionalCacheHold cacheMemory(const nn::SharedMemory& memory) const override;
 
     nn::ExecutionResult<std::pair<std::vector<nn::OutputShape>, nn::Timing>> execute(
             const nn::Request& request, nn::MeasureTiming measure) const override;
diff --git a/neuralnetworks/utils/common/src/CommonUtils.cpp b/neuralnetworks/utils/common/src/CommonUtils.cpp
index c04c8df..9021563 100644
--- a/neuralnetworks/utils/common/src/CommonUtils.cpp
+++ b/neuralnetworks/utils/common/src/CommonUtils.cpp
@@ -203,13 +203,13 @@
 nn::GeneralResult<void> unflushDataFromSharedToPointer(
         const nn::Request& request, const std::optional<nn::Request>& maybeRequestInShared) {
     if (!maybeRequestInShared.has_value() || maybeRequestInShared->pools.empty() ||
-        !std::holds_alternative<nn::Memory>(maybeRequestInShared->pools.back())) {
+        !std::holds_alternative<nn::SharedMemory>(maybeRequestInShared->pools.back())) {
         return {};
     }
     const auto& requestInShared = *maybeRequestInShared;
 
     // Map the memory.
-    const auto& outputMemory = std::get<nn::Memory>(requestInShared.pools.back());
+    const auto& outputMemory = std::get<nn::SharedMemory>(requestInShared.pools.back());
     const auto [pointer, size, context] = NN_TRY(map(outputMemory));
     const uint8_t* constantPointer =
             std::visit([](const auto& o) { return static_cast<const uint8_t*>(o); }, pointer);
diff --git a/neuralnetworks/utils/common/src/InvalidBuffer.cpp b/neuralnetworks/utils/common/src/InvalidBuffer.cpp
index c6f75d7..e73001d 100644
--- a/neuralnetworks/utils/common/src/InvalidBuffer.cpp
+++ b/neuralnetworks/utils/common/src/InvalidBuffer.cpp
@@ -30,11 +30,11 @@
     return nn::Request::MemoryDomainToken{};
 }
 
-nn::GeneralResult<void> InvalidBuffer::copyTo(const nn::Memory& /*dst*/) const {
+nn::GeneralResult<void> InvalidBuffer::copyTo(const nn::SharedMemory& /*dst*/) const {
     return NN_ERROR() << "InvalidBuffer";
 }
 
-nn::GeneralResult<void> InvalidBuffer::copyFrom(const nn::Memory& /*src*/,
+nn::GeneralResult<void> InvalidBuffer::copyFrom(const nn::SharedMemory& /*src*/,
                                                 const nn::Dimensions& /*dimensions*/) const {
     return NN_ERROR() << "InvalidBuffer";
 }
diff --git a/neuralnetworks/utils/common/src/InvalidBurst.cpp b/neuralnetworks/utils/common/src/InvalidBurst.cpp
index 4ca6603..81ca18d 100644
--- a/neuralnetworks/utils/common/src/InvalidBurst.cpp
+++ b/neuralnetworks/utils/common/src/InvalidBurst.cpp
@@ -26,7 +26,8 @@
 
 namespace android::hardware::neuralnetworks::utils {
 
-InvalidBurst::OptionalCacheHold InvalidBurst::cacheMemory(const nn::Memory& /*memory*/) const {
+InvalidBurst::OptionalCacheHold InvalidBurst::cacheMemory(
+        const nn::SharedMemory& /*memory*/) const {
     return nullptr;
 }
 
diff --git a/neuralnetworks/utils/common/src/ResilientBuffer.cpp b/neuralnetworks/utils/common/src/ResilientBuffer.cpp
index 47abbe2..1904375 100644
--- a/neuralnetworks/utils/common/src/ResilientBuffer.cpp
+++ b/neuralnetworks/utils/common/src/ResilientBuffer.cpp
@@ -99,12 +99,12 @@
     return getBuffer()->getToken();
 }
 
-nn::GeneralResult<void> ResilientBuffer::copyTo(const nn::Memory& dst) const {
+nn::GeneralResult<void> ResilientBuffer::copyTo(const nn::SharedMemory& dst) const {
     const auto fn = [&dst](const nn::IBuffer& buffer) { return buffer.copyTo(dst); };
     return protect(*this, fn);
 }
 
-nn::GeneralResult<void> ResilientBuffer::copyFrom(const nn::Memory& src,
+nn::GeneralResult<void> ResilientBuffer::copyFrom(const nn::SharedMemory& src,
                                                   const nn::Dimensions& dimensions) const {
     const auto fn = [&src, &dimensions](const nn::IBuffer& buffer) {
         return buffer.copyFrom(src, dimensions);
diff --git a/neuralnetworks/utils/common/src/ResilientBurst.cpp b/neuralnetworks/utils/common/src/ResilientBurst.cpp
index 0d3cb33..5ca868b 100644
--- a/neuralnetworks/utils/common/src/ResilientBurst.cpp
+++ b/neuralnetworks/utils/common/src/ResilientBurst.cpp
@@ -94,7 +94,8 @@
     return mBurst;
 }
 
-ResilientBurst::OptionalCacheHold ResilientBurst::cacheMemory(const nn::Memory& memory) const {
+ResilientBurst::OptionalCacheHold ResilientBurst::cacheMemory(
+        const nn::SharedMemory& memory) const {
     return getBurst()->cacheMemory(memory);
 }
 
diff --git a/neuralnetworks/utils/common/test/MockBuffer.h b/neuralnetworks/utils/common/test/MockBuffer.h
index c5405fb..59d5700 100644
--- a/neuralnetworks/utils/common/test/MockBuffer.h
+++ b/neuralnetworks/utils/common/test/MockBuffer.h
@@ -27,9 +27,9 @@
 class MockBuffer final : public IBuffer {
   public:
     MOCK_METHOD(Request::MemoryDomainToken, getToken, (), (const, override));
-    MOCK_METHOD(GeneralResult<void>, copyTo, (const Memory& dst), (const, override));
-    MOCK_METHOD(GeneralResult<void>, copyFrom, (const Memory& src, const Dimensions& dimensions),
-                (const, override));
+    MOCK_METHOD(GeneralResult<void>, copyTo, (const SharedMemory& dst), (const, override));
+    MOCK_METHOD(GeneralResult<void>, copyFrom,
+                (const SharedMemory& src, const Dimensions& dimensions), (const, override));
 };
 
 }  // namespace android::nn
diff --git a/neuralnetworks/utils/common/test/ResilientBufferTest.cpp b/neuralnetworks/utils/common/test/ResilientBufferTest.cpp
index deb9b7c..7afd020 100644
--- a/neuralnetworks/utils/common/test/ResilientBufferTest.cpp
+++ b/neuralnetworks/utils/common/test/ResilientBufferTest.cpp
@@ -15,9 +15,11 @@
  */
 
 #include <gmock/gmock.h>
+#include <nnapi/SharedMemory.h>
 #include <nnapi/TypeUtils.h>
 #include <nnapi/Types.h>
 #include <nnapi/hal/ResilientBuffer.h>
+#include <memory>
 #include <tuple>
 #include <utility>
 #include "MockBuffer.h"
@@ -113,7 +115,8 @@
     EXPECT_CALL(*mockBuffer, copyTo(_)).Times(1).WillOnce(Return(kNoError));
 
     // run test
-    const auto result = buffer->copyTo({});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyTo(memory);
 
     // verify result
     ASSERT_TRUE(result.has_value())
@@ -126,7 +129,8 @@
     EXPECT_CALL(*mockBuffer, copyTo(_)).Times(1).WillOnce(kReturnGeneralFailure);
 
     // run test
-    const auto result = buffer->copyTo({});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyTo(memory);
 
     // verify result
     ASSERT_FALSE(result.has_value());
@@ -140,7 +144,8 @@
     EXPECT_CALL(*mockBufferFactory, Call()).Times(1).WillOnce(kReturnGeneralFailure);
 
     // run test
-    const auto result = buffer->copyTo({});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyTo(memory);
 
     // verify result
     ASSERT_FALSE(result.has_value());
@@ -156,7 +161,8 @@
     EXPECT_CALL(*mockBufferFactory, Call()).Times(1).WillOnce(Return(recoveredMockBuffer));
 
     // run test
-    const auto result = buffer->copyTo({});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyTo(memory);
 
     // verify result
     ASSERT_TRUE(result.has_value())
@@ -169,7 +175,8 @@
     EXPECT_CALL(*mockBuffer, copyFrom(_, _)).Times(1).WillOnce(Return(kNoError));
 
     // run test
-    const auto result = buffer->copyFrom({}, {});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyFrom(memory, {});
 
     // verify result
     ASSERT_TRUE(result.has_value())
@@ -182,7 +189,8 @@
     EXPECT_CALL(*mockBuffer, copyFrom(_, _)).Times(1).WillOnce(kReturnGeneralFailure);
 
     // run test
-    const auto result = buffer->copyFrom({}, {});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyFrom(memory, {});
 
     // verify result
     ASSERT_FALSE(result.has_value());
@@ -196,7 +204,8 @@
     EXPECT_CALL(*mockBufferFactory, Call()).Times(1).WillOnce(kReturnGeneralFailure);
 
     // run test
-    const auto result = buffer->copyFrom({}, {});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyFrom(memory, {});
 
     // verify result
     ASSERT_FALSE(result.has_value());
@@ -212,7 +221,8 @@
     EXPECT_CALL(*mockBufferFactory, Call()).Times(1).WillOnce(Return(recoveredMockBuffer));
 
     // run test
-    const auto result = buffer->copyFrom({}, {});
+    const nn::SharedMemory memory = std::make_shared<const nn::Memory>();
+    const auto result = buffer->copyFrom(memory, {});
 
     // verify result
     ASSERT_TRUE(result.has_value())