Support dynamic output shape in 1.2 HAL interface.

Modify ExecutionCallback to support getting output shapes from device.
The sample driver is currently only returning an empty list, will return
the actual shape information in a subsequent CL.

Bug: 73506513
Test: NeuralNetworksTest_static
Change-Id: Ib0e176d92c7b548ee83c11a576976d8a536e51e6
Merged-In: Ib0e176d92c7b548ee83c11a576976d8a536e51e6
(cherry picked from commit 6e07734aa3af051609bf8f6e54f09ed3ddedb346)
diff --git a/common/include/HalInterfaces.h b/common/include/HalInterfaces.h
index 7f71b21..dc5e72e 100644
--- a/common/include/HalInterfaces.h
+++ b/common/include/HalInterfaces.h
@@ -57,6 +57,7 @@
 using ::android::hardware::neuralnetworks::V1_2::OperandType;
 using ::android::hardware::neuralnetworks::V1_2::Operation;
 using ::android::hardware::neuralnetworks::V1_2::OperationType;
+using ::android::hardware::neuralnetworks::V1_2::OutputShape;
 using ::android::hidl::allocator::V1_0::IAllocator;
 using ::android::hidl::memory::V1_0::IMemory;
 
diff --git a/driver/sample/SampleDriver.cpp b/driver/sample/SampleDriver.cpp
index 6da41ae..c171798 100644
--- a/driver/sample/SampleDriver.cpp
+++ b/driver/sample/SampleDriver.cpp
@@ -152,16 +152,18 @@
     return setRunTimePoolInfosFromHidlMemories(&mPoolInfos, mModel.pools);
 }
 
-static Return<void> notify(const sp<V1_0::IExecutionCallback>& callback,
-                           const ErrorStatus& status) {
+static Return<void> notify(const sp<V1_0::IExecutionCallback>& callback, const ErrorStatus& status,
+                           const hidl_vec<OutputShape>&) {
     return callback->notify(status);
 }
 
-static Return<void> notify(const sp<V1_2::IExecutionCallback>& callback,
-                           const ErrorStatus& status) {
-    return callback->notify_1_2(status);
+static Return<void> notify(const sp<V1_2::IExecutionCallback>& callback, const ErrorStatus& status,
+                           const hidl_vec<OutputShape>& outputShapes) {
+    return callback->notify_1_2(status, outputShapes);
 }
 
+// TODO(xusongw): Let callback notify actual output shape once dynamic output shape
+//                is supported in CpuExecutor.
 template <typename T_IExecutionCallback>
 void asyncExecute(const Request& request, const Model& model,
                   const std::vector<RunTimePoolInfo>& poolInfos,
@@ -170,7 +172,7 @@
                  "SampleDriver::asyncExecute");
     std::vector<RunTimePoolInfo> requestPoolInfos;
     if (!setRunTimePoolInfosFromHidlMemories(&requestPoolInfos, request.pools)) {
-        notify(callback, ErrorStatus::GENERAL_FAILURE);
+        notify(callback, ErrorStatus::GENERAL_FAILURE, {});
         return;
     }
 
@@ -181,7 +183,7 @@
     VLOG(DRIVER) << "executor.run returned " << n;
     ErrorStatus executionStatus =
             n == ANEURALNETWORKS_NO_ERROR ? ErrorStatus::NONE : ErrorStatus::GENERAL_FAILURE;
-    Return<void> returned = notify(callback, executionStatus);
+    Return<void> returned = notify(callback, executionStatus, {});
     if (!returned.isOk()) {
         LOG(ERROR) << " hidl callback failed to return properly: " << returned.description();
     }
@@ -198,7 +200,7 @@
         return ErrorStatus::INVALID_ARGUMENT;
     }
     if (!validateRequest(request, model)) {
-        notify(callback, ErrorStatus::INVALID_ARGUMENT);
+        notify(callback, ErrorStatus::INVALID_ARGUMENT, {});
         return ErrorStatus::INVALID_ARGUMENT;
     }
 
@@ -222,19 +224,22 @@
     return executeBase(request, mModel, mPoolInfos, callback);
 }
 
-Return<ErrorStatus> SamplePreparedModel::executeSynchronously(const Request& request) {
+Return<void> SamplePreparedModel::executeSynchronously(const Request& request,
+                                                       executeSynchronously_cb cb) {
     NNTRACE_FULL(NNTRACE_LAYER_DRIVER, NNTRACE_PHASE_EXECUTION,
                  "SampleDriver::executeSynchronously");
     VLOG(DRIVER) << "executeSynchronously(" << SHOW_IF_DEBUG(toString(request)) << ")";
     if (!validateRequest(request, mModel)) {
-        return ErrorStatus::INVALID_ARGUMENT;
+        cb(ErrorStatus::INVALID_ARGUMENT, {});
+        return Void();
     }
 
     NNTRACE_FULL_SWITCH(NNTRACE_LAYER_DRIVER, NNTRACE_PHASE_INPUTS_AND_OUTPUTS,
                         "SampleDriver::executeSynchronously");
     std::vector<RunTimePoolInfo> requestPoolInfos;
     if (!setRunTimePoolInfosFromHidlMemories(&requestPoolInfos, request.pools)) {
-        return ErrorStatus::GENERAL_FAILURE;
+        cb(ErrorStatus::GENERAL_FAILURE, {});
+        return Void();
     }
 
     NNTRACE_FULL_SWITCH(NNTRACE_LAYER_DRIVER, NNTRACE_PHASE_EXECUTION,
@@ -242,7 +247,8 @@
     CpuExecutor executor;
     int n = executor.run(mModel, request, mPoolInfos, requestPoolInfos);
     VLOG(DRIVER) << "executor.run returned " << n;
-    return n == ANEURALNETWORKS_NO_ERROR ? ErrorStatus::NONE : ErrorStatus::GENERAL_FAILURE;
+    cb(n == ANEURALNETWORKS_NO_ERROR ? ErrorStatus::NONE : ErrorStatus::GENERAL_FAILURE, {});
+    return Void();
 }
 
 } // namespace sample_driver
diff --git a/driver/sample/SampleDriver.h b/driver/sample/SampleDriver.h
index ac05aee..dc968ca 100644
--- a/driver/sample/SampleDriver.h
+++ b/driver/sample/SampleDriver.h
@@ -66,7 +66,7 @@
                                 const sp<V1_0::IExecutionCallback>& callback) override;
     Return<ErrorStatus> execute_1_2(const Request& request,
                                     const sp<V1_2::IExecutionCallback>& callback) override;
-    Return<ErrorStatus> executeSynchronously(const Request& request) override;
+    Return<void> executeSynchronously(const Request& request, executeSynchronously_cb cb) override;
 
    private:
     Model mModel;
diff --git a/runtime/Callbacks.cpp b/runtime/Callbacks.cpp
index b08563d..93be77b 100644
--- a/runtime/Callbacks.cpp
+++ b/runtime/Callbacks.cpp
@@ -139,8 +139,10 @@
     return Void();
 }
 
-Return<void> ExecutionCallback::notify_1_2(ErrorStatus errorStatus) {
+Return<void> ExecutionCallback::notify_1_2(ErrorStatus errorStatus,
+                                           const hidl_vec<OutputShape>& outputShapes) {
     mErrorStatus = errorStatus;
+    mOutputShapes = outputShapes;
     CallbackBase::notify();
     return Void();
 }
@@ -150,6 +152,11 @@
     return mErrorStatus;
 }
 
+const std::vector<OutputShape>& ExecutionCallback::getOutputShapes() {
+    wait();
+    return mOutputShapes;
+}
+
 }  // namespace implementation
 }  // namespace V1_2
 }  // namespace neuralnetworks
diff --git a/runtime/Callbacks.h b/runtime/Callbacks.h
index 5acedda..4cd9aa2 100644
--- a/runtime/Callbacks.h
+++ b/runtime/Callbacks.h
@@ -282,8 +282,9 @@
      * Either IExecutionCallback::notify or IExecutionCallback::notify_1_2 must
      * be called exactly once on a given ExecutionCallback object.
      *
-     * @param status Error status returned from asynchronously preparing the
-     *               model; will be:
+     * @param status Error status returned from launching the asynchronous task
+     *               (if the launch fails) or from the asynchronous task itself
+     *               (if the launch succeeds). Must be:
      *               - NONE if the asynchronous execution was successful
      *               - DEVICE_UNAVAILABLE if driver is offline or busy
      *               - GENERAL_FAILURE if there is an unspecified error
@@ -292,27 +293,78 @@
      *               - INVALID_ARGUMENT if the input request is invalid
      */
     Return<void> notify(ErrorStatus status) override;
-    Return<void> notify_1_2(ErrorStatus status) override;
+
+    /**
+     * Similar to IExecutionCallback::notify, but for V1_2::IPreparedModel to
+     * also notify output shapes along with error status.
+     *
+     * @param status Error status returned from launching the asynchronous task
+     *               (if the launch fails) or from the asynchronous task itself
+     *               (if the launch succeeds). Must be:
+     *               - NONE if the asynchronous execution was successful
+     *               - DEVICE_UNAVAILABLE if driver is offline or busy
+     *               - GENERAL_FAILURE if the asynchronous task resulted in an
+     *                 unspecified error
+     *               - OUTPUT_INSUFFICIENT_SIZE if at least one output
+     *                 operand buffer is not large enough to store the
+     *                 corresponding output
+     *               - INVALID_ARGUMENT if one of the input arguments to
+     *                 prepareModel is invalid
+     * @param outputShapes A list of shape information of model output operands.
+     *                     The index into "outputShapes" corresponds to the index
+     *                     of the output operand in the Request outputs vector.
+     *                     outputShapes must be empty unless the status is either
+     *                     NONE or OUTPUT_INSUFFICIENT_SIZE.
+     */
+    Return<void> notify_1_2(ErrorStatus status, const hidl_vec<OutputShape>& outputShapes) override;
+
+    // An overload of the latest notify interface to hide the version from ExecutionBuilder.
+    Return<void> notify(ErrorStatus status, const hidl_vec<OutputShape>& outputShapes) {
+        return notify_1_2(status, outputShapes);
+    }
 
     /**
      * Retrieves the error status returned from the asynchronous task launched
-     * by IPreparedModel::execute. If IPreparedModel::execute has not finished
+     * by either IPreparedModel::execute or IPreparedModel::execute_1_2. If
+     * IPreparedModel::execute or IPreparedModel::execute_1_2 has not finished
      * asynchronously executing, this call will block until the asynchronous task
      * notifies the object.
      *
-     * @return status Error status returned from asynchronously preparing the
-     *                model; will be:
+     * @return status Error status returned from launching the asynchronous task
+     *                (if the launch fails) or from the asynchronous task itself
+     *                (if the launch succeeds). Must be:
      *                - NONE if the asynchronous execution was successful
      *                - DEVICE_UNAVAILABLE if driver is offline or busy
-     *                - GENERAL_FAILURE if there is an unspecified error
-     *                - OUTPUT_INSUFFICIENT_SIZE if provided output buffer is
-     *                  not large enough to store the resultant values
-     *                - INVALID_ARGUMENT if the input request is invalid
+     *                - GENERAL_FAILURE if the asynchronous task resulted in an
+     *                  unspecified error
+     *                - OUTPUT_INSUFFICIENT_SIZE if at least one output
+     *                  operand buffer is not large enough to store the
+     *                  corresponding output
+     *                - INVALID_ARGUMENT if one of the input arguments to
+     *                  prepareModel is invalid
      */
     ErrorStatus getStatus();
 
- private:
+    /**
+     * Retrieves the output shapes returned from the asynchronous task launched
+     * by IPreparedModel::execute_1_2. If IPreparedModel::execute_1_2 has not finished
+     * asynchronously executing, this call will block until the asynchronous task
+     * notifies the object.
+     *
+     * If the asynchronous task was launched by IPreparedModel::execute, an empty vector
+     * will be returned.
+     *
+     * @return outputShapes A list of shape information of model output operands.
+     *                      The index into "outputShapes" corresponds to the index
+     *                      of the output operand in the Request outputs vector.
+     *                      outputShapes must be empty unless the status is either
+     *                      NONE or OUTPUT_INSUFFICIENT_SIZE.
+     */
+    const std::vector<OutputShape>& getOutputShapes();
+
+   private:
     ErrorStatus mErrorStatus;
+    std::vector<OutputShape> mOutputShapes;
 };
 
 
diff --git a/runtime/ExecutionBuilder.cpp b/runtime/ExecutionBuilder.cpp
index fb18a06..4a6b288 100644
--- a/runtime/ExecutionBuilder.cpp
+++ b/runtime/ExecutionBuilder.cpp
@@ -616,10 +616,8 @@
     if (DeviceManager::get()->syncExecHal()) {
         VLOG(EXECUTION) << "Before mPreparedModel->executeSynchronously() "
                         << SHOW_IF_DEBUG(toString(request));
-        Return<ErrorStatus> syncCallbackStatus = mPreparedModel->executeSynchronously(request);
-        executionCallback->notify(syncCallbackStatus.isOk()
-                                          ? static_cast<ErrorStatus>(syncCallbackStatus)
-                                          : ErrorStatus::GENERAL_FAILURE);
+        auto syncExecuteResult = mPreparedModel->executeSynchronously(request);
+        executionCallback->notify(syncExecuteResult.first, syncExecuteResult.second);
     } else {
         VLOG(EXECUTION) << "Before mPreparedModel->execute() " << SHOW_IF_DEBUG(toString(request));
         // Execute.
@@ -679,7 +677,7 @@
     NNTRACE_RT(NNTRACE_PHASE_EXECUTION, "computeOnCpu");
     CpuExecutor executor;
     int err = executor.run(model, request, modelPoolInfos, requestPoolInfos);
-    executionCallback->notify_1_2(convertResultCodeToErrorStatus(err));
+    executionCallback->notify_1_2(convertResultCodeToErrorStatus(err), {});
 }
 
 int StepExecutor::startComputeOnCpu(sp<ExecutionCallback>* synchronizationCallback) {
diff --git a/runtime/VersionedInterfaces.cpp b/runtime/VersionedInterfaces.cpp
index 9896c61..e669540 100644
--- a/runtime/VersionedInterfaces.cpp
+++ b/runtime/VersionedInterfaces.cpp
@@ -49,23 +49,29 @@
     }
 }
 
-ErrorStatus VersionedIPreparedModel::executeSynchronously(const Request& request) {
+std::pair<ErrorStatus, hidl_vec<OutputShape>> VersionedIPreparedModel::executeSynchronously(
+        const Request& request) {
     if (mPreparedModelV1_2 != nullptr) {
-        Return<ErrorStatus> ret = mPreparedModelV1_2->executeSynchronously(request);
+        std::pair<ErrorStatus, hidl_vec<OutputShape>> result;
+        Return<void> ret = mPreparedModelV1_2->executeSynchronously(
+                request, [&result](ErrorStatus error, const hidl_vec<OutputShape>& outputShapes) {
+                    result = {error, outputShapes};
+                });
         if (!ret.isOk()) {
             LOG(ERROR) << "executeSynchronously failure: " << ret.description();
-            return ErrorStatus::GENERAL_FAILURE;
+            return {ErrorStatus::GENERAL_FAILURE, {}};
         }
-        return static_cast<ErrorStatus>(ret);
+        return result;
     } else {
         // Simulate synchronous execution.
         sp<ExecutionCallback> callback = new ExecutionCallback();
         ErrorStatus ret = execute(request, callback);
         if (ret != ErrorStatus::NONE) {
-            return ret;
+            return {ret, {}};
         }
         callback->wait();
-        return callback->getStatus();
+        // callback->getOutputShapes() will always return an empty hidl vector.
+        return {callback->getStatus(), callback->getOutputShapes()};
     }
 }
 
diff --git a/runtime/VersionedInterfaces.h b/runtime/VersionedInterfaces.h
index d2900a3..1c4311f 100644
--- a/runtime/VersionedInterfaces.h
+++ b/runtime/VersionedInterfaces.h
@@ -336,12 +336,18 @@
      *                - NONE if execution is performed successfully
      *                - DEVICE_UNAVAILABLE if driver is offline or busy
      *                - GENERAL_FAILURE if there is an unspecified error
-     *                - OUTPUT_INSUFFICIENT_SIZE if provided output buffer is
-     *                  not large enough to store the resultant values
+     *                - OUTPUT_INSUFFICIENT_SIZE if at least one output
+     *                  operand buffer is not large enough to store the
+     *                  corresponding output
      *                - INVALID_ARGUMENT if one of the input arguments is
      *                  invalid
+     * @return outputShapes A list of shape information of model output operands.
+     *                      The index into "outputShapes" corresponds with the index
+     *                      of the output operand in the Request outputs vector.
+     *                      outputShapes nust be empty unless the status is either
+     *                      NONE or OUTPUT_INSUFFICIENT_SIZE.
      */
-    ErrorStatus executeSynchronously(const Request& request);
+    std::pair<ErrorStatus, hidl_vec<OutputShape>> executeSynchronously(const Request& request);
 
     /**
      * Returns whether this handle to an IPreparedModel object is valid or not.
diff --git a/runtime/test/TestExecution.cpp b/runtime/test/TestExecution.cpp
index dcbfc07..7abef57 100644
--- a/runtime/test/TestExecution.cpp
+++ b/runtime/test/TestExecution.cpp
@@ -78,17 +78,21 @@
         if (mErrorStatus == ErrorStatus::NONE) {
             return mPreparedModelV1_2->execute_1_2(request, callback);
         } else {
-            callback->notify_1_2(mErrorStatus);
+            callback->notify_1_2(mErrorStatus, {});
             return ErrorStatus::NONE;
         }
     }
 
-    Return<ErrorStatus> executeSynchronously(const Request& request) override {
+    Return<void> executeSynchronously(const Request& request, executeSynchronously_cb cb) override {
         CHECK(mPreparedModelV1_2 != nullptr) << "V1_2 prepared model is nullptr.";
         if (mErrorStatus == ErrorStatus::NONE) {
-            return mPreparedModelV1_2->executeSynchronously(request);
+            return mPreparedModelV1_2->executeSynchronously(
+                    request, [&cb](ErrorStatus error, const hidl_vec<OutputShape>& outputShapes) {
+                        cb(error, outputShapes);
+                    });
         } else {
-            return mErrorStatus;
+            cb(mErrorStatus, {});
+            return Void();
         }
     }
 
diff --git a/runtime/test/TestPartitioning.cpp b/runtime/test/TestPartitioning.cpp
index d1c668f..eba0756 100644
--- a/runtime/test/TestPartitioning.cpp
+++ b/runtime/test/TestPartitioning.cpp
@@ -215,8 +215,9 @@
      Return<ErrorStatus> execute_1_2(const Request&, const sp<V1_2::IExecutionCallback>&) override {
          return ErrorStatus::DEVICE_UNAVAILABLE;
      }
-     Return<ErrorStatus> executeSynchronously(const Request&) override {
-         return ErrorStatus::DEVICE_UNAVAILABLE;
+     Return<void> executeSynchronously(const Request&, executeSynchronously_cb cb) override {
+         cb(ErrorStatus::DEVICE_UNAVAILABLE, {});
+         return Void();
      }
     };
 public: