Merge TP1A.220126.001

Change-Id: Ibf6bd2c20d9927fde8b2a05dde2b58bd8faea20f

3 files changed, 121 insertions, 31 deletions

diff --git a/neuralnetworks/utils/adapter/aidl/src/Burst.cpp b/neuralnetworks/utils/adapter/aidl/src/Burst.cpp
index 4fabb20635..a4a80faf2a 100644
--- a/neuralnetworks/utils/adapter/aidl/src/Burst.cpp
+++ b/neuralnetworks/utils/adapter/aidl/src/Burst.cpp
@@ -93,7 +93,8 @@ std::vector<nn::IBurst::OptionalCacheHold> ensureAllMemoriesAreCached(
 nn::ExecutionResult<ExecutionResult> executeSynchronously(
         const nn::IBurst& burst, const Burst::ThreadSafeMemoryCache& cache, const Request& request,
         const std::vector<int64_t>& memoryIdentifierTokens, bool measureTiming, int64_t deadlineNs,
-        int64_t loopTimeoutDurationNs) {
+        int64_t loopTimeoutDurationNs, const std::vector<TokenValuePair>& hints,
+        const std::vector<ExtensionNameAndPrefix>& extensionNameToPrefix) {
     if (request.pools.size() != memoryIdentifierTokens.size()) {
         return NN_ERROR(nn::ErrorStatus::INVALID_ARGUMENT)
                << "request.pools.size() != memoryIdentifierTokens.size()";
@@ -107,11 +108,13 @@ nn::ExecutionResult<ExecutionResult> executeSynchronously(
     const auto nnMeasureTiming = measureTiming ? nn::MeasureTiming::YES : nn::MeasureTiming::NO;
     const auto nnDeadline = NN_TRY(makeOptionalTimePoint(deadlineNs));
     const auto nnLoopTimeoutDuration = NN_TRY(makeOptionalDuration(loopTimeoutDurationNs));
+    auto nnHints = NN_TRY(convertInput(hints));
+    auto nnExtensionNameToPrefix = NN_TRY(convertInput(extensionNameToPrefix));
 
     const auto hold = ensureAllMemoriesAreCached(&nnRequest, memoryIdentifierTokens, burst, cache);
 
-    const auto result =
-            burst.execute(nnRequest, nnMeasureTiming, nnDeadline, nnLoopTimeoutDuration);
+    const auto result = burst.execute(nnRequest, nnMeasureTiming, nnDeadline, nnLoopTimeoutDuration,
+                                      nnHints, nnExtensionNameToPrefix);
 
     if (!result.ok() && result.error().code == nn::ErrorStatus::OUTPUT_INSUFFICIENT_SIZE) {
         const auto& [message, code, outputShapes] = result.error();
@@ -155,7 +158,24 @@ ndk::ScopedAStatus Burst::executeSynchronously(const Request& request,
                                                ExecutionResult* executionResult) {
     auto result =
             adapter::executeSynchronously(*kBurst, kMemoryCache, request, memoryIdentifierTokens,
-                                          measureTiming, deadlineNs, loopTimeoutDurationNs);
+                                          measureTiming, deadlineNs, loopTimeoutDurationNs, {}, {});
+    if (!result.has_value()) {
+        auto [message, code, _] = std::move(result).error();
+        const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
+        return ndk::ScopedAStatus::fromServiceSpecificErrorWithMessage(
+                static_cast<int32_t>(aidlCode), message.c_str());
+    }
+    *executionResult = std::move(result).value();
+    return ndk::ScopedAStatus::ok();
+}
+
+ndk::ScopedAStatus Burst::executeSynchronouslyWithConfig(
+        const Request& request, const std::vector<int64_t>& memoryIdentifierTokens,
+        const ExecutionConfig& config, int64_t deadlineNs, ExecutionResult* executionResult) {
+    auto result = adapter::executeSynchronously(
+            *kBurst, kMemoryCache, request, memoryIdentifierTokens, config.measureTiming,
+            deadlineNs, config.loopTimeoutDurationNs, config.executionHints,
+            config.extensionNameToPrefix);
     if (!result.has_value()) {
         auto [message, code, _] = std::move(result).error();
         const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
diff --git a/neuralnetworks/utils/adapter/aidl/src/Device.cpp b/neuralnetworks/utils/adapter/aidl/src/Device.cpp
index 763be7f3fa..84aaddbe9d 100644
--- a/neuralnetworks/utils/adapter/aidl/src/Device.cpp
+++ b/neuralnetworks/utils/adapter/aidl/src/Device.cpp
@@ -148,13 +148,14 @@ void notify(IPreparedModelCallback* callback, PrepareModelResult result) {
     }
 }
 
-nn::GeneralResult<void> prepareModel(const nn::SharedDevice& device, const Executor& executor,
-                                     const Model& model, ExecutionPreference preference,
-                                     Priority priority, int64_t deadlineNs,
-                                     const std::vector<ndk::ScopedFileDescriptor>& modelCache,
-                                     const std::vector<ndk::ScopedFileDescriptor>& dataCache,
-                                     const std::vector<uint8_t>& token,
-                                     const std::shared_ptr<IPreparedModelCallback>& callback) {
+nn::GeneralResult<void> prepareModel(
+        const nn::SharedDevice& device, const Executor& executor, const Model& model,
+        ExecutionPreference preference, Priority priority, int64_t deadlineNs,
+        const std::vector<ndk::ScopedFileDescriptor>& modelCache,
+        const std::vector<ndk::ScopedFileDescriptor>& dataCache, const std::vector<uint8_t>& token,
+        const std::vector<TokenValuePair>& hints,
+        const std::vector<ExtensionNameAndPrefix>& extensionNameToPrefix,
+        const std::shared_ptr<IPreparedModelCallback>& callback) {
     if (callback.get() == nullptr) {
         return NN_ERROR(nn::ErrorStatus::INVALID_ARGUMENT) << "Invalid callback";
     }
@@ -166,12 +167,16 @@ nn::GeneralResult<void> prepareModel(const nn::SharedDevice& device, const Execu
     auto nnModelCache = NN_TRY(convertInput(modelCache));
     auto nnDataCache = NN_TRY(convertInput(dataCache));
     const auto nnToken = NN_TRY(convertCacheToken(token));
+    auto nnHints = NN_TRY(convertInput(hints));
+    auto nnExtensionNameToPrefix = NN_TRY(convertInput(extensionNameToPrefix));
 
     Task task = [device, nnModel = std::move(nnModel), nnPreference, nnPriority, nnDeadline,
                  nnModelCache = std::move(nnModelCache), nnDataCache = std::move(nnDataCache),
-                 nnToken, callback] {
-        auto result = device->prepareModel(nnModel, nnPreference, nnPriority, nnDeadline,
-                                           nnModelCache, nnDataCache, nnToken);
+                 nnToken, nnHints = std::move(nnHints),
+                 nnExtensionNameToPrefix = std::move(nnExtensionNameToPrefix), callback] {
+        auto result =
+                device->prepareModel(nnModel, nnPreference, nnPriority, nnDeadline, nnModelCache,
+                                     nnDataCache, nnToken, nnHints, nnExtensionNameToPrefix);
         notify(callback.get(), std::move(result));
     };
     executor(std::move(task), nnDeadline);
@@ -273,8 +278,9 @@ ndk::ScopedAStatus Device::prepareModel(const Model& model, ExecutionPreference
                                         const std::vector<ndk::ScopedFileDescriptor>& dataCache,
                                         const std::vector<uint8_t>& token,
                                         const std::shared_ptr<IPreparedModelCallback>& callback) {
-    const auto result = adapter::prepareModel(kDevice, kExecutor, model, preference, priority,
-                                              deadlineNs, modelCache, dataCache, token, callback);
+    const auto result =
+            adapter::prepareModel(kDevice, kExecutor, model, preference, priority, deadlineNs,
+                                  modelCache, dataCache, token, {}, {}, callback);
     if (!result.has_value()) {
         const auto& [message, code] = result.error();
         const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
@@ -301,4 +307,21 @@ ndk::ScopedAStatus Device::prepareModelFromCache(
     return ndk::ScopedAStatus::ok();
 }
 
+ndk::ScopedAStatus Device::prepareModelWithConfig(
+        const Model& model, const PrepareModelConfig& config,
+        const std::shared_ptr<IPreparedModelCallback>& callback) {
+    const auto result = adapter::prepareModel(
+            kDevice, kExecutor, model, config.preference, config.priority, config.deadlineNs,
+            config.modelCache, config.dataCache, config.cacheToken, config.compilationHints,
+            config.extensionNameToPrefix, callback);
+    if (!result.has_value()) {
+        const auto& [message, code] = result.error();
+        const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
+        callback->notify(aidlCode, nullptr);
+        return ndk::ScopedAStatus::fromServiceSpecificErrorWithMessage(
+                static_cast<int32_t>(aidlCode), message.c_str());
+    }
+    return ndk::ScopedAStatus::ok();
+}
+
 }  // namespace aidl::android::hardware::neuralnetworks::adapter
diff --git a/neuralnetworks/utils/adapter/aidl/src/PreparedModel.cpp b/neuralnetworks/utils/adapter/aidl/src/PreparedModel.cpp
index 5cab62c625..790558fde4 100644
--- a/neuralnetworks/utils/adapter/aidl/src/PreparedModel.cpp
+++ b/neuralnetworks/utils/adapter/aidl/src/PreparedModel.cpp
@@ -118,17 +118,20 @@ nn::GeneralResult<nn::OptionalTimePoint> makeOptionalTimePoint(int64_t durationN
     return durationNs < 0 ? nn::OptionalTimePoint{} : nn::TimePoint(makeDuration(durationNs));
 }
 
-nn::ExecutionResult<ExecutionResult> executeSynchronously(const nn::IPreparedModel& preparedModel,
-                                                          const Request& request,
-                                                          bool measureTiming, int64_t deadlineNs,
-                                                          int64_t loopTimeoutDurationNs) {
+nn::ExecutionResult<ExecutionResult> executeSynchronously(
+        const nn::IPreparedModel& preparedModel, const Request& request, bool measureTiming,
+        int64_t deadlineNs, int64_t loopTimeoutDurationNs, const std::vector<TokenValuePair>& hints,
+        const std::vector<ExtensionNameAndPrefix>& extensionNameToPrefix) {
     const auto nnRequest = NN_TRY(convertInput(request));
     const auto nnMeasureTiming = measureTiming ? nn::MeasureTiming::YES : nn::MeasureTiming::NO;
     const auto nnDeadline = NN_TRY(makeOptionalTimePoint(deadlineNs));
     const auto nnLoopTimeoutDuration = NN_TRY(makeOptionalDuration(loopTimeoutDurationNs));
+    auto nnHints = NN_TRY(convertInput(hints));
+    auto nnExtensionNameToPrefix = NN_TRY(convertInput(extensionNameToPrefix));
 
     const auto result =
-            preparedModel.execute(nnRequest, nnMeasureTiming, nnDeadline, nnLoopTimeoutDuration);
+            preparedModel.execute(nnRequest, nnMeasureTiming, nnDeadline, nnLoopTimeoutDuration,
+                                  nnHints, nnExtensionNameToPrefix);
 
     if (!result.ok() && result.error().code == nn::ErrorStatus::OUTPUT_INSUFFICIENT_SIZE) {
         const auto& [message, code, outputShapes] = result.error();
@@ -147,16 +150,21 @@ nn::ExecutionResult<ExecutionResult> executeSynchronously(const nn::IPreparedMod
 nn::GeneralResult<FencedExecutionResult> executeFenced(
         const nn::IPreparedModel& preparedModel, const Request& request,
         const std::vector<ndk::ScopedFileDescriptor>& waitFor, bool measureTiming,
-        int64_t deadlineNs, int64_t loopTimeoutDurationNs, int64_t durationNs) {
+        int64_t deadlineNs, int64_t loopTimeoutDurationNs, int64_t durationNs,
+        const std::vector<TokenValuePair>& hints,
+        const std::vector<ExtensionNameAndPrefix>& extensionNameToPrefix) {
     const auto nnRequest = NN_TRY(convertInput(request));
     const auto nnWaitFor = NN_TRY(convertSyncFences(waitFor));
     const auto nnMeasureTiming = measureTiming ? nn::MeasureTiming::YES : nn::MeasureTiming::NO;
     const auto nnDeadline = NN_TRY(makeOptionalTimePoint(deadlineNs));
     const auto nnLoopTimeoutDuration = NN_TRY(makeOptionalDuration(loopTimeoutDurationNs));
     const auto nnDuration = NN_TRY(makeOptionalDuration(durationNs));
+    auto nnHints = NN_TRY(convertInput(hints));
+    auto nnExtensionNameToPrefix = NN_TRY(convertInput(extensionNameToPrefix));
 
     auto [syncFence, executeFencedInfoCallback] = NN_TRY(preparedModel.executeFenced(
-            nnRequest, nnWaitFor, nnMeasureTiming, nnDeadline, nnLoopTimeoutDuration, nnDuration));
+            nnRequest, nnWaitFor, nnMeasureTiming, nnDeadline, nnLoopTimeoutDuration, nnDuration,
+            nnHints, nnExtensionNameToPrefix));
 
     ndk::ScopedFileDescriptor fileDescriptor;
     if (syncFence.hasFd()) {
@@ -171,11 +179,16 @@ nn::GeneralResult<FencedExecutionResult> executeFenced(
 
 nn::GeneralResult<nn::SharedExecution> createReusableExecution(
         const nn::IPreparedModel& preparedModel, const Request& request, bool measureTiming,
-        int64_t loopTimeoutDurationNs) {
+        int64_t loopTimeoutDurationNs, const std::vector<TokenValuePair>& hints,
+        const std::vector<ExtensionNameAndPrefix>& extensionNameToPrefix) {
     const auto nnRequest = NN_TRY(convertInput(request));
     const auto nnMeasureTiming = measureTiming ? nn::MeasureTiming::YES : nn::MeasureTiming::NO;
     const auto nnLoopTimeoutDuration = NN_TRY(makeOptionalDuration(loopTimeoutDurationNs));
-    return preparedModel.createReusableExecution(nnRequest, nnMeasureTiming, nnLoopTimeoutDuration);
+    auto nnHints = NN_TRY(convertInput(hints));
+    auto nnExtensionNameToPrefix = NN_TRY(convertInput(extensionNameToPrefix));
+
+    return preparedModel.createReusableExecution(nnRequest, nnMeasureTiming, nnLoopTimeoutDuration,
+                                                 nnHints, nnExtensionNameToPrefix);
 }
 
 nn::ExecutionResult<ExecutionResult> executeSynchronously(const nn::IExecution& execution,
@@ -231,7 +244,7 @@ ndk::ScopedAStatus PreparedModel::executeSynchronously(const Request& request, b
                                                        int64_t loopTimeoutDurationNs,
                                                        ExecutionResult* executionResult) {
     auto result = adapter::executeSynchronously(*kPreparedModel, request, measureTiming, deadlineNs,
-                                                loopTimeoutDurationNs);
+                                                loopTimeoutDurationNs, {}, {});
     if (!result.has_value()) {
         const auto& [message, code, _] = result.error();
         const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
@@ -247,7 +260,41 @@ ndk::ScopedAStatus PreparedModel::executeFenced(
         bool measureTiming, int64_t deadlineNs, int64_t loopTimeoutDurationNs, int64_t durationNs,
         FencedExecutionResult* executionResult) {
     auto result = adapter::executeFenced(*kPreparedModel, request, waitFor, measureTiming,
-                                         deadlineNs, loopTimeoutDurationNs, durationNs);
+                                         deadlineNs, loopTimeoutDurationNs, durationNs, {}, {});
+    if (!result.has_value()) {
+        const auto& [message, code] = result.error();
+        const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
+        return ndk::ScopedAStatus::fromServiceSpecificErrorWithMessage(
+                static_cast<int32_t>(aidlCode), message.c_str());
+    }
+    *executionResult = std::move(result).value();
+    return ndk::ScopedAStatus::ok();
+}
+
+ndk::ScopedAStatus PreparedModel::executeSynchronouslyWithConfig(const Request& request,
+                                                                 const ExecutionConfig& config,
+                                                                 int64_t deadlineNs,
+                                                                 ExecutionResult* executionResult) {
+    auto result = adapter::executeSynchronously(
+            *kPreparedModel, request, config.measureTiming, deadlineNs,
+            config.loopTimeoutDurationNs, config.executionHints, config.extensionNameToPrefix);
+    if (!result.has_value()) {
+        const auto& [message, code, _] = result.error();
+        const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
+        return ndk::ScopedAStatus::fromServiceSpecificErrorWithMessage(
+                static_cast<int32_t>(aidlCode), message.c_str());
+    }
+    *executionResult = std::move(result).value();
+    return ndk::ScopedAStatus::ok();
+}
+
+ndk::ScopedAStatus PreparedModel::executeFencedWithConfig(
+        const Request& request, const std::vector<ndk::ScopedFileDescriptor>& waitFor,
+        const ExecutionConfig& config, int64_t deadlineNs, int64_t durationNs,
+        FencedExecutionResult* executionResult) {
+    auto result = adapter::executeFenced(*kPreparedModel, request, waitFor, config.measureTiming,
+                                         deadlineNs, config.loopTimeoutDurationNs, durationNs,
+                                         config.executionHints, config.extensionNameToPrefix);
     if (!result.has_value()) {
         const auto& [message, code] = result.error();
         const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);
@@ -275,11 +322,11 @@ nn::SharedPreparedModel PreparedModel::getUnderlyingPreparedModel() const {
 }
 
 ndk::ScopedAStatus PreparedModel::createReusableExecution(const Request& request,
-                                                          bool measureTiming,
-                                                          int64_t loopTimeoutDurationNs,
+                                                          const ExecutionConfig& config,
                                                           std::shared_ptr<IExecution>* execution) {
-    auto result = adapter::createReusableExecution(*kPreparedModel, request, measureTiming,
-                                                   loopTimeoutDurationNs);
+    auto result = adapter::createReusableExecution(
+            *kPreparedModel, request, config.measureTiming, config.loopTimeoutDurationNs,
+            config.executionHints, config.extensionNameToPrefix);
     if (!result.has_value()) {
         const auto& [message, code] = result.error();
         const auto aidlCode = utils::convert(code).value_or(ErrorStatus::GENERAL_FAILURE);