Merge "Consolidate Manager's prepareModel methods"
diff --git a/nn/runtime/ExecutionBuilder.cpp b/nn/runtime/ExecutionBuilder.cpp
index 2dccb9a..97e847b 100644
--- a/nn/runtime/ExecutionBuilder.cpp
+++ b/nn/runtime/ExecutionBuilder.cpp
@@ -771,14 +771,14 @@
 int StepExecutor::startComputeOnCpuFallback(sp<ExecutionCallback>* synchronizationCallback) {
     NNTRACE_RT(NNTRACE_PHASE_EXECUTION, "StepExecutor::startComputeOnCpuFallback");
     VLOG(EXECUTION) << "Re-compile the model on CPU";
-    const Model model = mModel->makeHidlModel();
     mDevice = DeviceManager::getCpuDevice();
     mPreparedModel = nullptr;
+    const ModelFactory makeModel = [this] { return mModel->makeHidlModel(); };
     // TODO: Propagate user preference to this point instead of using default value of
     // ANEURALNETWORKS_PREFER_FAST_SINGLE_ANSWER.
     const ExecutionPreference preference =
             static_cast<ExecutionPreference>(ANEURALNETWORKS_PREFER_FAST_SINGLE_ANSWER);
-    const auto [n, preparedModel] = mDevice->prepareModel(model, preference, {}, {});
+    const auto [n, preparedModel] = mDevice->prepareModel(makeModel, 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 0a51680..5f656e1 100644
--- a/nn/runtime/ExecutionPlan.cpp
+++ b/nn/runtime/ExecutionPlan.cpp
@@ -58,23 +58,6 @@
 
 using namespace hal;
 
-std::pair<int, std::shared_ptr<PreparedModel>> compile(
-        const Device& device, const ModelFactory& makeModel, ExecutionPreference preference,
-        const std::string& cacheDir, const std::optional<CacheToken>& maybeToken) {
-    // Attempt to compile from cache if token is present.
-    if (maybeToken.has_value()) {
-        const auto [n, preparedModel] = device.prepareModelFromCache(cacheDir, *maybeToken);
-        if (n == ANEURALNETWORKS_NO_ERROR) {
-            return {n, preparedModel};
-        }
-    }
-
-    // Fallback to full compilation (possibly with token) if
-    // prepareModelFromCache could not be used or failed.
-    const Model model = makeModel();
-    return device.prepareModel(model, preference, cacheDir, maybeToken);
-}
-
 // Compiles the model on device.
 // If compilation caching is available, depending on ExecutionPlan::mState, the token may only have
 // been initialized by the user provided token (SIMPLE body), or is already re-hashed by the
@@ -97,7 +80,7 @@
     const ModelFactory makeModel = [&model] { return model.makeHidlModel(); };
     const ExecutionPreference preference = static_cast<ExecutionPreference>(executionPreference);
     const auto [n, returnedPreparedModel] =
-            compile(device, makeModel, preference, cacheDir, cacheToken);
+            device.prepareModel(makeModel, preference, cacheDir, cacheToken);
     *preparedModel = returnedPreparedModel;
     return n;
 }
diff --git a/nn/runtime/Manager.cpp b/nn/runtime/Manager.cpp
index d430566..6a3882d 100644
--- a/nn/runtime/Manager.cpp
+++ b/nn/runtime/Manager.cpp
@@ -77,12 +77,17 @@
     }
 
     std::pair<int, std::shared_ptr<PreparedModel>> prepareModel(
-            const Model& model, ExecutionPreference preference, const std::string& cacheDir,
+            const ModelFactory& makeModel, ExecutionPreference preference,
+            const std::string& cacheDir,
             const std::optional<CacheToken>& maybeToken) const override;
-    std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCache(
-            const std::string& cacheDir, const CacheToken& token) const override;
 
    private:
+    std::pair<int, std::shared_ptr<PreparedModel>> prepareModelInternal(
+            const Model& model, ExecutionPreference preference, const std::string& cacheDir,
+            const std::optional<CacheToken>& maybeToken) const;
+    std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCacheInternal(
+            const std::string& cacheDir, const CacheToken& token) const;
+
     std::string mName;
     std::string mVersionString;
     const std::shared_ptr<VersionedIDevice> mInterface;
@@ -331,7 +336,7 @@
     return {ANEURALNETWORKS_NO_ERROR, std::make_shared<DriverPreparedModel>(preparedModel)};
 }
 
-std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModel(
+std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModelInternal(
         const Model& model, ExecutionPreference preference, const std::string& cacheDir,
         const std::optional<CacheToken>& maybeToken) const {
     // Note that some work within VersionedIDevice will be subtracted from the IPC layer
@@ -353,7 +358,7 @@
     return prepareModelCheck(status, preparedModel, "prepareModel", getName());
 }
 
-std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModelFromCache(
+std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModelFromCacheInternal(
         const std::string& cacheDir, 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");
@@ -371,6 +376,23 @@
     return prepareModelCheck(status, preparedModel, "prepareModelFromCache", getName());
 }
 
+std::pair<int, std::shared_ptr<PreparedModel>> DriverDevice::prepareModel(
+        const ModelFactory& makeModel, ExecutionPreference preference, const std::string& cacheDir,
+        const std::optional<CacheToken>& maybeToken) const {
+    // Attempt to compile from cache if token is present.
+    if (maybeToken.has_value()) {
+        const auto [n, preparedModel] = prepareModelFromCacheInternal(cacheDir, *maybeToken);
+        if (n == ANEURALNETWORKS_NO_ERROR) {
+            return {n, preparedModel};
+        }
+    }
+
+    // Fallback to full compilation (possibly with token) if
+    // prepareModelFromCache could not be used or failed.
+    const Model model = makeModel();
+    return prepareModelInternal(model, preference, cacheDir, maybeToken);
+}
+
 // Convert ModelArgumentInfo to HIDL RequestArgument. For pointer arguments, use the location
 // information in ptrArgsLocations.
 static void setRequestArgumentArray(const std::vector<ModelArgumentInfo>& argumentInfos,
@@ -566,13 +588,9 @@
     bool isCachingSupported() const override { return false; }
 
     std::pair<int, std::shared_ptr<PreparedModel>> prepareModel(
-            const Model& model, ExecutionPreference preference, const std::string& cacheDir,
+            const ModelFactory& makeModel, ExecutionPreference preference,
+            const std::string& cacheDir,
             const std::optional<CacheToken>& maybeToken) const override;
-    std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCache(
-            const std::string& /*cacheDir*/, const CacheToken& /*token*/) const override {
-        CHECK(false) << "Should never call prepareModelFromCache on CpuDevice";
-        return {ANEURALNETWORKS_OP_FAILED, nullptr};
-    }
 
    private:
     CpuDevice() = default;
@@ -628,11 +646,12 @@
 }
 
 std::pair<int, std::shared_ptr<PreparedModel>> CpuDevice::prepareModel(
-        const Model& model, ExecutionPreference preference, const std::string& /*cacheDir*/,
-        const std::optional<CacheToken>& maybeToken) const {
+        const ModelFactory& makeModel, ExecutionPreference preference,
+        const std::string& /*cacheDir*/, const std::optional<CacheToken>& maybeToken) const {
     CHECK(!maybeToken.has_value())
             << "Should never call prepareModel with cache information on CpuDevice";
 
+    const Model model = makeModel();
     if (!validateModel(model) || !validateExecutionPreference(preference)) {
         return {ANEURALNETWORKS_OP_FAILED, nullptr};
     }
diff --git a/nn/runtime/Manager.h b/nn/runtime/Manager.h
index fe2b03e..82042d1 100644
--- a/nn/runtime/Manager.h
+++ b/nn/runtime/Manager.h
@@ -84,11 +84,9 @@
     virtual bool isCachingSupported() const = 0;
 
     virtual std::pair<int, std::shared_ptr<PreparedModel>> prepareModel(
-            const hal::Model& model, hal::ExecutionPreference preference,
+            const hal::ModelFactory& makeModel, hal::ExecutionPreference preference,
             const std::string& cacheDir,
             const std::optional<hal::CacheToken>& maybeToken) const = 0;
-    virtual std::pair<int, std::shared_ptr<PreparedModel>> prepareModelFromCache(
-            const std::string& cacheDir, const hal::CacheToken& token) const = 0;
 };
 
 // Manages the NN HAL devices.  Only one instance of this class will exist.