static bool DispatchSingleStep()

in nn-samples/sequence/src/main/cpp/sequence_model.cpp [562:668]


static bool DispatchSingleStep(ANeuralNetworksCompilation* compilation,
                              ANeuralNetworksMemory* sumIn,
                              uint32_t sumInLength,
                              ANeuralNetworksMemory* stateIn,
                              uint32_t stateInLength,
                              ANeuralNetworksMemory* sumOut,
                              uint32_t sumOutLength,
                              ANeuralNetworksMemory* stateOut,
                              uint32_t stateOutLength,
                              const ANeuralNetworksEvent* waitFor,
                              ANeuralNetworksEvent** event) {
  // Create an ANeuralNetworksExecution object from the compiled model.
  ANeuralNetworksExecution* execution;
  int32_t status = ANeuralNetworksExecution_create(compilation, &execution);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksExecution_create failed");
    return false;
  }

  // Set the memory for the sumIn tensor.
  // Note that the index "0" here means the first operand of the modelInputs
  // list {sumIn, stateIn}, which means sumIn.
  status = ANeuralNetworksExecution_setInputFromMemory(execution,
                                                       0,
                                                       nullptr,
                                                       sumIn,
                                                       0,
                                                       sumInLength
                                                           * sizeof(float));
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksExecution_setInputFromMemory failed for sumIn");
    return false;
  }

  // Set the memory for the stateIn tensor.
  // Note that the index "1" here means the first operand of the modelInputs
  // list {sumIn, stateIn}, which means stateIn.
  status = ANeuralNetworksExecution_setInputFromMemory(execution,
                                                       1,
                                                       nullptr,
                                                       stateIn,
                                                       0,
                                                       stateInLength
                                                           * sizeof(float));
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksExecution_setInputFromMemory failed for stateIn");
    return false;
  }

  // Set the sumOut tensor that will be filled by executing the model.
  status = ANeuralNetworksExecution_setOutputFromMemory(execution,
                                                        0,
                                                        nullptr,
                                                        sumOut,
                                                        0,
                                                        sumOutLength
                                                            * sizeof(float));
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksExecution_setOutputFromMemory failed for sumOut");
    return false;
  }

  // Set the stateOut tensor that will be filled by executing the model.
  status = ANeuralNetworksExecution_setOutputFromMemory(execution,
                                                        1,
                                                        nullptr,
                                                        stateOut,
                                                        0,
                                                        stateOutLength
                                                            * sizeof(float));
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksExecution_setOutputFromMemory failed for stateOut");
    return false;
  }

  // Dispatch the execution of the model.
  // Note that the execution here is asynchronous with dependencies.
  const ANeuralNetworksEvent* const* dependencies = nullptr;
  uint32_t numDependencies = 0;
  if (waitFor != nullptr) {
    dependencies = &waitFor;
    numDependencies = 1;
  }
  status = ANeuralNetworksExecution_startComputeWithDependencies(execution,
                                                                 dependencies,
                                                                 numDependencies,
                                                                 0,  // infinite timeout duration
                                                                 event);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR, LOG_TAG,
                        "ANeuralNetworksExecution_compute failed");
    return false;
  }

  ANeuralNetworksExecution_free(execution);
  return true;
}