Change prepareModel to return model by value

Bug: N/A
Test: mma
Test: NeuralNetworksTest_static
Test: CtsNNAPITestCases
Change-Id: I17663d6c2e2c82c54a18dc6cdedb87b309f9af71
Merged-In: I17663d6c2e2c82c54a18dc6cdedb87b309f9af71
(cherry picked from commit fab58a7a71578ced5c9dd66d395f8d29cdf89ead)

4 files changed, 85 insertions, 76 deletions

diff --git a/nn/runtime/ExecutionBuilder.cpp b/nn/runtime/ExecutionBuilder.cpp
index b0ead8af8..87bc77f45 100644
--- a/nn/runtime/ExecutionBuilder.cpp
+++ b/nn/runtime/ExecutionBuilder.cpp
@@ -725,9 +725,11 @@ int StepExecutor::startComputeOnCpuFallback(sp<ExecutionCallback>* synchronizati
     mPreparedModel = nullptr;
     // TODO: Propagate user preference to this point instead of using default value of
     // ANEURALNETWORKS_PREFER_FAST_SINGLE_ANSWER.
-    ExecutionPreference preference =
+    const ExecutionPreference preference =
             static_cast<ExecutionPreference>(ANEURALNETWORKS_PREFER_FAST_SINGLE_ANSWER);
-    NN_RETURN_IF_ERROR(mDevice->prepareModel(model, preference, {}, {}, {}, &mPreparedModel));
+    const auto [n, preparedModel] = mDevice->prepareModel(model, preference, {}, {}, {});
+    mPreparedModel = preparedModel;
+    NN_RETURN_IF_ERROR(n);
     return startCompute(synchronizationCallback, /*burstController=*/nullptr);
 }
 
diff --git a/nn/runtime/ExecutionPlan.cpp b/nn/runtime/ExecutionPlan.cpp
index fe7e55d7b..049905d56 100644
--- a/nn/runtime/ExecutionPlan.cpp
+++ b/nn/runtime/ExecutionPlan.cpp
@@ -124,32 +124,44 @@ bool getCacheHandles(const std::string& cacheDir, const uint8_t* token,
 bool compileFromCache(const std::shared_ptr<Device>& device, const std::string& cacheDir,
                       const uint8_t* token, std::shared_ptr<PreparedModel>* preparedModel) {
     CHECK(token != nullptr && device != nullptr);
-    VLOG(COMPILATION) << "compileFromCache";
+    CHECK(preparedModel != nullptr);
     *preparedModel = nullptr;
+    VLOG(COMPILATION) << "compileFromCache";
+
     CacheToken cacheToken(token);
     hidl_vec<hidl_handle> modelCache, dataCache;
     NN_RET_CHECK(getCacheHandles(cacheDir, token, device->getNumberOfCacheFilesNeeded(),
                                  /*createIfNotExist=*/false, &modelCache, &dataCache));
-    int ret = device->prepareModelFromCache(modelCache, dataCache, cacheToken, preparedModel);
-    return ret == ANEURALNETWORKS_NO_ERROR;
+
+    const auto [n, returnedPreparedModel] =
+            device->prepareModelFromCache(modelCache, dataCache, cacheToken);
+    *preparedModel = returnedPreparedModel;
+    return n == ANEURALNETWORKS_NO_ERROR;
 }
 
 int compileModelAndCache(const std::shared_ptr<Device>& device, const ModelBuilder* model,
                          int32_t executionPreference, const std::string& cacheDir,
                          const uint8_t* token, std::shared_ptr<PreparedModel>* preparedModel) {
     CHECK(device != nullptr);
+    CHECK(preparedModel != nullptr);
     *preparedModel = nullptr;
+
     uint8_t dummyToken[ANEURALNETWORKS_BYTE_SIZE_OF_CACHE_TOKEN] = {0};
     CacheToken cacheToken(token == nullptr ? dummyToken : token);
+
     hidl_vec<hidl_handle> modelCache, dataCache;
     if (token == nullptr || !getCacheHandles(cacheDir, token, device->getNumberOfCacheFilesNeeded(),
                                              /*createIfNotExist=*/true, &modelCache, &dataCache)) {
         modelCache.resize(0);
         dataCache.resize(0);
     }
-    return device->prepareModel(model->makeHidlModel(),
-                                static_cast<ExecutionPreference>(executionPreference), modelCache,
-                                dataCache, cacheToken, preparedModel);
+
+    const Model hidlModel = model->makeHidlModel();
+    const ExecutionPreference preference = static_cast<ExecutionPreference>(executionPreference);
+    const auto [n, returnedPreparedModel] =
+            device->prepareModel(hidlModel, preference, modelCache, dataCache, cacheToken);
+    *preparedModel = returnedPreparedModel;
+    return n;
 }
 
 // Compiles the model on device.
@@ -161,6 +173,9 @@ int compile(std::shared_ptr<Device> device, const ModelBuilder* model, int32_t e
             const std::string& cacheDir, TokenHasher* token,
             std::shared_ptr<PreparedModel>* preparedModel) {
     CHECK(device != nullptr);
+    CHECK(preparedModel != nullptr);
+    *preparedModel = nullptr;
+
     const uint8_t* tokenData = nullptr;
     if (device->isCachingSupported() && token->ok() && token->updateFromString(device->getName()) &&
         token->updateFromString(device->getVersionString()) &&
@@ -170,6 +185,7 @@ int compile(std::shared_ptr<Device> device, const ModelBuilder* model, int32_t e
     if (tokenData != nullptr && compileFromCache(device, cacheDir, tokenData, preparedModel)) {
         return ANEURALNETWORKS_NO_ERROR;
     }
+
     return compileModelAndCache(device, model, executionPreference, cacheDir, tokenData,
                                 preparedModel);
 }
diff --git a/nn/runtime/Manager.cpp b/nn/runtime/Manager.cpp
index 26b372ff1..c081572b5 100644
--- a/nn/runtime/Manager.cpp
+++ b/nn/runtime/Manager.cpp
@@ -81,13 +81,13 @@ class DriverDevice : public Device {
         return mNumCacheFiles;
     }
 
-    int prepareModel(const Model& hidlModel, ExecutionPreference executionPreference,
-                     const hidl_vec<hidl_handle>& modelCache,
-                     const hidl_vec<hidl_handle>& dataCache, const CacheToken& token,
-                     std::shared_ptr<PreparedModel>* preparedModel) const override;
-    int prepareModelFromCache(const hidl_vec<hidl_handle>& modelCache,
-                              const hidl_vec<hidl_handle>& dataCache, const CacheToken& token,
-                              std::shared_ptr<PreparedModel>* preparedModel) const override;
+    std::pair<int, std::shared_ptr<PreparedModel>> prepareModel(
+            const Model& hidlModel, ExecutionPreference executionPreference,
+            const hidl_vec<hidl_handle>& modelCache, const hidl_vec<hidl_handle>& dataCache,
+            const CacheToken& token) const override;
+    std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCache(
+            const hidl_vec<hidl_handle>& modelCache, const hidl_vec<hidl_handle>& dataCache,
+            const CacheToken& token) const override;
 
    private:
     std::string mName;
@@ -257,52 +257,45 @@ PerformanceInfo DriverDevice::getPerformance(OperandType type) const {
     return lookup(mCapabilities.operandPerformance, type);
 }
 
-static int prepareModelCheck(ErrorStatus status,
-                             const std::shared_ptr<VersionedIPreparedModel>& preparedModel,
-                             const char* prepareName, const char* serviceName,
-                             std::shared_ptr<PreparedModel>* preparedModelOut) {
-    CHECK(preparedModelOut != nullptr) << "prepareModelCheck -- preparedModelOut must be non-null";
-    *preparedModelOut = nullptr;
-
+static std::pair<int, std::shared_ptr<PreparedModel>> prepareModelCheck(
+        ErrorStatus status, const std::shared_ptr<VersionedIPreparedModel>& preparedModel,
+        const char* prepareName, const char* serviceName) {
     if (status != ErrorStatus::NONE) {
         LOG(ERROR) << prepareName << " on " << serviceName << " failed: "
                    << "prepareReturnStatus=" << toString(status);
-        return ANEURALNETWORKS_OP_FAILED;
+        return {ANEURALNETWORKS_OP_FAILED, nullptr};
     }
     if (preparedModel == nullptr) {
         LOG(ERROR) << prepareName << " on " << serviceName << " failed: preparedModel is nullptr";
-        return ANEURALNETWORKS_OP_FAILED;
+        return {ANEURALNETWORKS_OP_FAILED, nullptr};
     }
 
-    *preparedModelOut = std::make_shared<DriverPreparedModel>(preparedModel);
-    return ANEURALNETWORKS_NO_ERROR;
+    return {ANEURALNETWORKS_NO_ERROR, std::make_shared<DriverPreparedModel>(preparedModel)};
 }
 
-int DriverDevice::prepareModel(const Model& hidlModel, ExecutionPreference executionPreference,
-                               const hidl_vec<hidl_handle>& modelCache,
-                               const hidl_vec<hidl_handle>& dataCache, const CacheToken& token,
-                               std::shared_ptr<PreparedModel>* preparedModel) const {
+std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModel(
+        const Model& hidlModel, ExecutionPreference executionPreference,
+        const hidl_vec<hidl_handle>& modelCache, const hidl_vec<hidl_handle>& dataCache,
+        const CacheToken& token) const {
     // Note that some work within VersionedIDevice will be subtracted from the IPC layer
     NNTRACE_FULL(NNTRACE_LAYER_IPC, NNTRACE_PHASE_COMPILATION, "prepareModel");
 
-    const auto [status, localPreparedModel] =
+    const auto [status, preparedModel] =
             mInterface->prepareModel(hidlModel, executionPreference, modelCache, dataCache, token);
 
-    return prepareModelCheck(status, localPreparedModel, "prepareModel", getName(), preparedModel);
+    return prepareModelCheck(status, preparedModel, "prepareModel", getName());
 }
 
-int DriverDevice::prepareModelFromCache(const hidl_vec<hidl_handle>& modelCache,
-                                        const hidl_vec<hidl_handle>& dataCache,
-                                        const CacheToken& token,
-                                        std::shared_ptr<PreparedModel>* preparedModel) const {
+std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModelFromCache(
+        const hidl_vec<hidl_handle>& modelCache, const hidl_vec<hidl_handle>& dataCache,
+        const CacheToken& token) const {
     // Note that some work within VersionedIDevice will be subtracted from the IPC layer
     NNTRACE_FULL(NNTRACE_LAYER_IPC, NNTRACE_PHASE_COMPILATION, "prepareModelFromCache");
 
-    const auto [status, localPreparedModel] =
+    const auto [status, preparedModel] =
             mInterface->prepareModelFromCache(modelCache, dataCache, token);
 
-    return prepareModelCheck(status, localPreparedModel, "prepareModelFromCache", getName(),
-                             preparedModel);
+    return prepareModelCheck(status, preparedModel, "prepareModelFromCache", getName());
 }
 
 // Convert ModelArgumentInfo to HIDL RequestArgument. For pointer arguments, use the location
@@ -501,14 +494,15 @@ class CpuDevice : public Device {
         return kNumCacheFiles;
     }
 
-    int prepareModel(const Model& hidlModel, ExecutionPreference executionPreference,
-                     const hidl_vec<hidl_handle>& modelCache,
-                     const hidl_vec<hidl_handle>& dataCache, const CacheToken&,
-                     std::shared_ptr<PreparedModel>* preparedModel) const override;
-    int prepareModelFromCache(const hidl_vec<hidl_handle>&, const hidl_vec<hidl_handle>&,
-                              const CacheToken&, std::shared_ptr<PreparedModel>*) const override {
+    std::pair<int, std::shared_ptr<PreparedModel>> prepareModel(
+            const Model& hidlModel, ExecutionPreference executionPreference,
+            const hidl_vec<hidl_handle>& modelCache, const hidl_vec<hidl_handle>& dataCache,
+            const CacheToken&) const override;
+    std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCache(
+            const hidl_vec<hidl_handle>&, const hidl_vec<hidl_handle>&,
+            const CacheToken&) const override {
         CHECK(false) << "Should never call prepareModelFromCache on CpuDevice";
-        return ANEURALNETWORKS_OP_FAILED;
+        return {ANEURALNETWORKS_OP_FAILED, nullptr};
     }
 
    private:
@@ -527,9 +521,10 @@ class CpuDevice : public Device {
 // A special abstracted PreparedModel for the CPU, constructed by CpuDevice.
 class CpuPreparedModel : public PreparedModel {
    public:
-    // Factory method for CpuPreparedModel. Returns ANEURALNETWORKS_NO_ERROR if
-    // successfully created.
-    static int create(Model hidlModel, std::shared_ptr<PreparedModel>* preparedModel);
+    // Factory method for CpuPreparedModel. Returns ANEURALNETWORKS_NO_ERROR and
+    // a prepared model object if successfully created. Returns an error code
+    // and nullptr otherwise.
+    static std::pair<int, std::shared_ptr<PreparedModel>> create(Model hidlModel);
 
     std::tuple<int, std::vector<OutputShape>, Timing> execute(
             const std::vector<ModelArgumentInfo>& inputs,
@@ -541,10 +536,11 @@ class CpuPreparedModel : public PreparedModel {
         return nullptr;
     }
 
-   private:
+    // Prefer to use CpuPreparedModel::create.
     CpuPreparedModel(Model model, std::vector<RunTimePoolInfo> poolInfos)
         : mModel(std::move(model)), mModelPoolInfos(std::move(poolInfos)) {}
 
+   private:
     const Model mModel;
     const std::vector<RunTimePoolInfo> mModelPoolInfos;
 };
@@ -565,34 +561,29 @@ void CpuDevice::getSupportedOperations(const MetaModel& metaModel,
     *supportedOperations = std::move(result);
 }
 
-int CpuDevice::prepareModel(const Model& hidlModel, ExecutionPreference executionPreference,
-                            const hidl_vec<hidl_handle>& modelCache,
-                            const hidl_vec<hidl_handle>& dataCache, const CacheToken&,
-                            std::shared_ptr<PreparedModel>* preparedModel) const {
+std::pair<int, std::shared_ptr<PreparedModel>> CpuDevice::prepareModel(
+        const Model& hidlModel, ExecutionPreference executionPreference,
+        const hidl_vec<hidl_handle>& modelCache, const hidl_vec<hidl_handle>& dataCache,
+        const CacheToken&) const {
     CHECK(modelCache.size() == 0 && dataCache.size() == 0)
             << "Should never call prepareModel with cache information on CpuDevice";
-    CHECK(preparedModel != nullptr) << "CpuDevice::prepareModel -- preparedModel must be non-null";
-    *preparedModel = nullptr;
 
     if (!validateModel(hidlModel) || !validateExecutionPreference(executionPreference)) {
-        return ANEURALNETWORKS_OP_FAILED;
+        return {ANEURALNETWORKS_OP_FAILED, nullptr};
     }
 
-    return CpuPreparedModel::create(hidlModel, preparedModel);
+    return CpuPreparedModel::create(hidlModel);
 }
 
-int CpuPreparedModel::create(Model hidlModel, std::shared_ptr<PreparedModel>* preparedModel) {
-    CHECK(preparedModel != nullptr);
-    *preparedModel = nullptr;
-
+std::pair<int, std::shared_ptr<PreparedModel>> CpuPreparedModel::create(Model hidlModel) {
     std::vector<RunTimePoolInfo> poolInfos;
     if (!setRunTimePoolInfosFromHidlMemories(&poolInfos, hidlModel.pools)) {
-        return ANEURALNETWORKS_UNMAPPABLE;
+        return {ANEURALNETWORKS_UNMAPPABLE, nullptr};
     }
 
-    *preparedModel = std::shared_ptr<CpuPreparedModel>(
-            new CpuPreparedModel(std::move(hidlModel), std::move(poolInfos)));
-    return ANEURALNETWORKS_NO_ERROR;
+    std::shared_ptr<PreparedModel> preparedModel =
+            std::make_shared<CpuPreparedModel>(std::move(hidlModel), std::move(poolInfos));
+    return {ANEURALNETWORKS_NO_ERROR, std::move(preparedModel)};
 }
 
 static std::tuple<int, std::vector<OutputShape>, Timing> computeOnCpu(
diff --git a/nn/runtime/Manager.h b/nn/runtime/Manager.h
index 54a49217c..26af7e08e 100644
--- a/nn/runtime/Manager.h
+++ b/nn/runtime/Manager.h
@@ -22,6 +22,7 @@
 #include <map>
 #include <memory>
 #include <string>
+#include <tuple>
 #include <unordered_set>
 #include <utility>
 #include <vector>
@@ -83,16 +84,15 @@ class Device {
     virtual std::pair<uint32_t, uint32_t> getNumberOfCacheFilesNeeded() const = 0;
     bool isCachingSupported() const;
 
-    virtual int prepareModel(const hal::Model& hidlModel,
-                             hal::ExecutionPreference executionPreference,
-                             const hal::hidl_vec<hal::hidl_handle>& modelCache,
-                             const hal::hidl_vec<hal::hidl_handle>& dataCache,
-                             const hal::CacheToken& token,
-                             std::shared_ptr<PreparedModel>* preparedModel) const = 0;
-    virtual int prepareModelFromCache(const hal::hidl_vec<hal::hidl_handle>& modelCache,
-                                      const hal::hidl_vec<hal::hidl_handle>& dataCache,
-                                      const hal::CacheToken& token,
-                                      std::shared_ptr<PreparedModel>* preparedModel) const = 0;
+    virtual std::pair<int, std::shared_ptr<PreparedModel>> prepareModel(
+            const hal::Model& hidlModel, hal::ExecutionPreference executionPreference,
+            const hal::hidl_vec<hal::hidl_handle>& modelCache,
+            const hal::hidl_vec<hal::hidl_handle>& dataCache,
+            const hal::CacheToken& token) const = 0;
+    virtual std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCache(
+            const hal::hidl_vec<hal::hidl_handle>& modelCache,
+            const hal::hidl_vec<hal::hidl_handle>& dataCache,
+            const hal::CacheToken& token) const = 0;
 };
 
 // Manages the NN HAL devices.  Only one instance of this class will exist.