bool SimpleSequenceModel::CreateOpaqueMemories()

in nn-samples/sequence/src/main/cpp/sequence_model.cpp [410:557]


bool SimpleSequenceModel::CreateOpaqueMemories() {
  int32_t status;

  // Create opaque memories for sum tensors.
  // We start from creating a memory descriptor and describing all of the
  // intended memory usages.
  ANeuralNetworksMemoryDesc* sumDesc = nullptr;
  status = ANeuralNetworksMemoryDesc_create(&sumDesc);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemoryDesc_create failed");
    return false;
  }

  // Specify that the state memory will be used as the first input (sumIn)
  // of the compilation. Note that the index "0" here means the first operand
  // of the modelInputs list {sumIn, stateIn}, which means sumIn.
  status =
      ANeuralNetworksMemoryDesc_addInputRole(sumDesc, compilation_, 0, 1.0f);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR, LOG_TAG,
                        "ANeuralNetworksMemoryDesc_addInputRole failed");
    ANeuralNetworksMemoryDesc_free(sumDesc);
    return false;
  }

  // Specify that the state memory will also be used as the first output
  // (sumOut) of the compilation. Note that the index "0" here means the
  // first operand of the modelOutputs list {sumOut, stateOut}, which means
  // sumOut.
  status =
      ANeuralNetworksMemoryDesc_addOutputRole(sumDesc, compilation_, 0, 1.0f);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR, LOG_TAG,
                        "ANeuralNetworksMemoryDesc_addOutputRole failed");
    ANeuralNetworksMemoryDesc_free(sumDesc);
    return false;
  }

  // Finish the memory descriptor.
  status = ANeuralNetworksMemoryDesc_finish(sumDesc);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemoryDesc_finish failed");
    ANeuralNetworksMemoryDesc_free(sumDesc);
    return false;
  }

  // Create two opaque memories from the finished descriptor: one for input
  // and one for output. We will swap the two memories after each single
  // execution step.
  status = ANeuralNetworksMemory_createFromDesc(sumDesc, &memoryOpaqueSumIn_);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemory_createFromDesc failed for sum memory #1");
    ANeuralNetworksMemoryDesc_free(sumDesc);
    return false;
  }
  status = ANeuralNetworksMemory_createFromDesc(sumDesc, &memoryOpaqueSumOut_);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemory_createFromDesc failed for sum memory #2");
    ANeuralNetworksMemoryDesc_free(sumDesc);
    return false;
  }

  // It is safe to free the memory descriptor once all of the memories have
  // been created.
  ANeuralNetworksMemoryDesc_free(sumDesc);

  // Create opaque memories for state tensors.
  // We start from creating a memory descriptor and describing all of the
  // intended memory usages.
  ANeuralNetworksMemoryDesc* stateDesc = nullptr;
  status = ANeuralNetworksMemoryDesc_create(&stateDesc);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemoryDesc_create failed");
    return false;
  }

  // Specify that the state memory will be used as the second input (stateIn)
  // of the compilation. Note that the index "1" here means the second operand
  // of the modelInputs list {sumIn, stateIn}, which means stateIn.
  status =
      ANeuralNetworksMemoryDesc_addInputRole(stateDesc, compilation_, 1, 1.0f);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR, LOG_TAG,
                        "ANeuralNetworksMemoryDesc_addInputRole failed");
    ANeuralNetworksMemoryDesc_free(stateDesc);
    return false;
  }

  // Specify that the state memory will also be used as the second output
  // (stateOut) of the compilation. Note that the index "1" here means the
  // second operand of the modelOutputs list {sumOut, stateOut}, which means
  // stateOut.
  status =
      ANeuralNetworksMemoryDesc_addOutputRole(stateDesc, compilation_, 1, 1.0f);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR, LOG_TAG,
                        "ANeuralNetworksMemoryDesc_addOutputRole failed");
    ANeuralNetworksMemoryDesc_free(stateDesc);
    return false;
  }

  // Finish the memory descriptor.
  status = ANeuralNetworksMemoryDesc_finish(stateDesc);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemoryDesc_finish failed");
    ANeuralNetworksMemoryDesc_free(stateDesc);
    return false;
  }

  // Create two opaque memories from the finished descriptor: one for input
  // and one for output. We will swap the two memories after each single
  // execution step.
  status =
      ANeuralNetworksMemory_createFromDesc(stateDesc, &memoryOpaqueStateIn_);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemory_createFromDesc failed for state memory #1");
    ANeuralNetworksMemoryDesc_free(stateDesc);
    return false;
  }
  status =
      ANeuralNetworksMemory_createFromDesc(stateDesc, &memoryOpaqueStateOut_);
  if (status != ANEURALNETWORKS_NO_ERROR) {
    __android_log_print(ANDROID_LOG_ERROR,
                        LOG_TAG,
                        "ANeuralNetworksMemory_createFromDesc failed for state memory #2");
    ANeuralNetworksMemoryDesc_free(stateDesc);
    return false;
  }

  // It is safe to free the memory descriptor once all of the memories have
  // been created.
  ANeuralNetworksMemoryDesc_free(stateDesc);
  return true;
}