Cleanup NNAPI code (microsoft#5505)

* Cleanup NNAPI code

* Check return of GetNCHWInput

Author: guoyu-wang

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc

@@ -122,6 +122,7 @@ Status ModelBuilder::Prepare() {
   RETURN_STATUS_ON_ERROR(nnapi_->ANeuralNetworksModel_create(&nnapi_model_->model_));
   ORT_RETURN_IF_ERROR(GetTargetDevices());
   PreprocessInitializers();
+  PreprocessActivations();
   ORT_RETURN_IF_ERROR(RegisterInitializers());
   ORT_RETURN_IF_ERROR(RegisterModelInputs());
   ORT_RETURN_IF_ERROR(AddOperations());
@@ -190,6 +191,28 @@ void ModelBuilder::PreprocessInitializers() {
   }
 }
 
+void ModelBuilder::PreprocessActivations() {
+  const auto& node_indices = graph_viewer_.GetNodesInTopologicalOrder();
+  for (size_t i = 0; i < node_indices.size(); i++) {
+    const auto* node(graph_viewer_.GetNode(node_indices[i]));
+    const auto& op_type(node->OpType());
+
+    if (op_type == "Relu") {
+      activation_nodes_.emplace(node->Index(), ANEURALNETWORKS_FUSED_RELU);
+    } else if (op_type == "Clip") {  // Relu1 or Relu6
+      float min, max;
+      if (!GetClipMinMax(*this, *node, min, max))
+        continue;
+
+      if (min == -1.0f && max == 1.0f) {
+        activation_nodes_.emplace(node->Index(), ANEURALNETWORKS_FUSED_RELU1);
+      } else if (min == 0.0f && max == 6.0f) {
+        activation_nodes_.emplace(node->Index(), ANEURALNETWORKS_FUSED_RELU6);
+      }
+    }
+  }
+}
+
 // Help to get all quantized operators' input and the node(s) using the input
 std::unordered_map<std::string, vector<const Node*>> GetAllQuantizedOpInputs(const GraphViewer& graph_viewer) {
   std::unordered_map<std::string, vector<const Node*>> all_quantized_op_inputs;
@@ -554,9 +577,9 @@ int32_t ModelBuilder::FindActivation(const Node& node, const NodeArg& output) {
   for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) {
     const auto& dst_node = it->GetNode();
     const auto* dst_input = dst_node.InputDefs()[it->GetDstArgIndex()];
-    if (dst_node.OpType() == "Relu") {
+    if (Contains(activation_nodes_, dst_node.Index())) {
       if (&output == dst_input) {
-        fuse_code = ANEURALNETWORKS_FUSED_RELU;
+        fuse_code = activation_nodes_.at(dst_node.Index());
       }
     } else {
       // if there is any other non-relu node using the output

onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.h

@@ -132,6 +132,9 @@ class ModelBuilder {
   std::unordered_map<std::string, const ONNX_NAMESPACE::TensorProto&> initializers_;
   std::unordered_set<std::string> skipped_initializers_;
 
+  // All activation nodes (Relu, Relu1, Relu6) as a map <NodeIndex, activeation_code>
+  std::unordered_map<NodeIndex, int32_t> activation_nodes_;
+
   std::unordered_map<std::string, std::shared_ptr<IOpBuilder>> op_builders_;
 
   // Operands in nhwc
@@ -157,12 +160,18 @@ class ModelBuilder {
   Status Prepare() ORT_MUST_USE_RESULT;
 
   Status GetTargetDevices() ORT_MUST_USE_RESULT;
+  // Get names of all the initializers
   void GetAllInitializers();
+  // If a NNAPI operation will use initializers directly, we will add the initializers to the skip list
   void PreprocessInitializers();
+  // Preprocess all the activation nodes (Relu/Relu1/Relu6) for easy query later
+  void PreprocessActivations();
+  // Copy and process all the initializers to NNAPI model
   Status RegisterInitializers() ORT_MUST_USE_RESULT;
   Status RegisterModelInputs() ORT_MUST_USE_RESULT;
   Status AddOperations() ORT_MUST_USE_RESULT;
   Status RegisterModelOutputs() ORT_MUST_USE_RESULT;
+  // After constructing the NNAPI model, will set the shape inferencing record to the Model
   void RegisterModelShaper();
 
   Status SetOperandValue(uint32_t index, Model::NNMemory* memory,