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;
}