Finish implementing single-partition graphs.

This includes wiring up the partitioning algorithm
to the NN API compilation and execution phases under
debug.nn.partition.test=2: When we have a single-
partition graph, we rely on the device selection and
compilation performed by the partitioning algorithm,
rather than doing it at execution time as we were
before.

Bug: 63905942
Test: mma (user)
      mma (userdebug)
      ml/nn/runtime/tests (userdebug)
      (with debug.nn.partition.test 0, 1, and 2;
       no new failures, logcat looks plausible,
       compilation happens at the expected times)

Change-Id: I6cd8b53625588d5aca81c87cc8600735e384b8e0

1 files changed, 79 insertions, 12 deletions

diff --git a/nn/runtime/ExecutionPlan.h b/nn/runtime/ExecutionPlan.h
index 2b01b4e20..b586cc0dd 100644
--- a/nn/runtime/ExecutionPlan.h
+++ b/nn/runtime/ExecutionPlan.h
@@ -61,7 +61,7 @@ public:
         mSubModelOutputs.insert(std::make_pair(fromModelIndex, it->second));
     }
 
-    void finishSubModel();
+    int finishSubModel();
 
     void dump() const;
 private:
@@ -73,6 +73,7 @@ private:
     uint32_t mIndex;  // index of step within plan
     std::shared_ptr<ModelBuilder> mSubModel;
     std::shared_ptr<Device> mDevice;  // nullptr signifies CPU
+    sp<IPreparedModel> mPreparedSubModel;  // not used for CPU
 
     // Inputs of original model that are also inputs of this submodel:
     //     (fromModel index, subModel index)
@@ -92,29 +93,95 @@ private:
 
 class ExecutionPlan {
 public:
-    std::shared_ptr<ExecutionStep> newStep(const std::shared_ptr<Device> device);
+    ExecutionPlan(const ExecutionPlan&) = delete;
+    ExecutionPlan& operator=(const ExecutionPlan&) = delete;
 
-    void finishSubModels();
+    ExecutionPlan() { }
+    ~ExecutionPlan() { delete mBody; }
+
+    std::shared_ptr<ExecutionStep> createNewStep(const std::shared_ptr<Device> device);
+
+    void becomeSingleStep(const std::shared_ptr<Device> device,
+                          const ModelBuilder* model);
+
+    int finish();
 
     void recordTemporaryDef(uint32_t fromModelIndex, uint32_t stepIndex) {
-        nnAssert(mTemporaryToDefiningStep.count(fromModelIndex) == 0);
-        mTemporaryToDefiningStep.insert(std::make_pair(fromModelIndex, stepIndex));
+        auto& temporaryToDefiningStep = compound()->mTemporaryToDefiningStep;
+        nnAssert(temporaryToDefiningStep.count(fromModelIndex) == 0);
+        temporaryToDefiningStep.insert(std::make_pair(fromModelIndex, stepIndex));
     }
 
     void dump() const;
 
+    // TODO: This member function is only temporary, until we finish
+    // fully integrating ExecutionPlan with the compilation and
+    // execution phases of the NN API.  The return value is as follows:
+    //
+    //     NO_ERROR
+    //         There's exactly one partition, and it was successfully compiled.
+    //
+    //         If device is not nullptr, *device has been set (set to nullptr
+    //         in the case of CPU execution).
+    //
+    //         If preparedModel is not nullptr, *preparedModel has been set
+    //         (set to nullptr in the case of CPU execution).
+    //
+    //     OP_FAILED
+    //         There's exactly one partition, but it was not successfully compiled.
+    //
+    //     BAD_STATE
+    //         There are zero or multiple partitions.
+    //
+    int getSimplePlan(std::shared_ptr<Device>* device = nullptr,
+                      sp<IPreparedModel>* preparedModel = nullptr) const;
+
 private:
     void findSubModelOutputs();
 
-    // TODO: Some of the data is working state information that
-    // shouldn't be needed after we've constructed but not executed
-    // the plan.
+    struct Body {
+        virtual ~Body() {}
+        virtual void dump() const = 0;
+        virtual int finish() = 0;
+    };
+
+    struct SimpleBody : Body {
+        SimpleBody(std::shared_ptr<Device> device, const ModelBuilder* model) :
+                mDevice(device), mModel(model) {}
 
-    std::vector<std::shared_ptr<ExecutionStep>> mSteps;
+        void dump() const override;
+        int finish() override;
 
-    // Map from original operand index to defining step index.
-    // Used for all (and only) TEMPORARY_VARIABLEs.
-    std::unordered_map<uint32_t, uint32_t> mTemporaryToDefiningStep;
+        std::shared_ptr<Device> mDevice;  // nullptr signifies CPU
+        const ModelBuilder* mModel;
+        sp<IPreparedModel> mPreparedModel;  // not used for CPU
+        bool mSuccessfulFinish = false;
+    };
+
+    struct CompoundBody : Body {
+        void dump() const override;
+        int finish() override;
+
+        // TODO: Some of the data is working state information that
+        // shouldn't be needed after we've constructed but not
+        // executed the plan.
+
+        std::vector<std::shared_ptr<ExecutionStep>> mSteps;
+
+        // Map from original operand index to defining step index.
+        // Used for all (and only) TEMPORARY_VARIABLEs.
+        std::unordered_map<uint32_t, uint32_t> mTemporaryToDefiningStep;
+
+    private:
+        void findSubModelOutputs();
+    };
+
+    enum { EMPTY, SIMPLE, COMPOUND } mState = EMPTY;
+    Body* mBody = nullptr;
+    CompoundBody* compound() {
+        nnAssert(mState == COMPOUND);
+        return static_cast<CompoundBody*>(mBody);
+    }
 };
 
 }  // namespace nn