From c6ef465011bd395bd0a1a16b820cabfa585768b2 Mon Sep 17 00:00:00 2001
From: Guoyu Wang <62914304+gwang-msft@users.noreply.github.com>
Date: Wed, 26 Jan 2022 16:42:38 -0800
Subject: [PATCH] minor fix in node unit change (#10405)

---
 .../nnapi_builtin/builders/op_builder.cc      | 23 ++++++++++---------
 .../builders/op_support_checker.cc            |  2 +-
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc
index 63a83cfbc0..84022ea774 100644
--- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc
+++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc
@@ -857,10 +857,22 @@ void ReshapeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const
 // onnxruntime::nnapi::Model.
 /* static */ bool ReshapeOpBuilder::CanSkipReshape(const ModelBuilder& model_builder, const NodeUnit& node_unit,
                                                    size_t input_rank, size_t output_rank) {
+  // Since we know this is a Reshape NodeUnit, so we can safely assume there is only 1 output
+  // and the node_unit has only one output node.
   const auto& output_node_arg = node_unit.Outputs()[0].node_arg;
   const auto& output_name = output_node_arg.Name();
   const auto& output_node = *node_unit.GetOutputNodes()[0];
 
+  // Check if the Reshape output is a graph output, if so we cannot skip the Reshape
+  // We do not care the case where the Reshape output is a dead end
+  for (const auto* node_arg : model_builder.GetGraphViewer().GetOutputs()) {
+    if (node_arg == &output_node_arg) {
+      LOGS_DEFAULT(VERBOSE) << "Reshape/Flatten can not be skipped when the output is a graph output"
+                            << ", output name, " << output_name;
+      return false;
+    }
+  }
+
   // We will go through all the output edges
   for (auto it = output_node.OutputEdgesBegin(), end = output_node.OutputEdgesEnd(); it != end; ++it) {
     const auto& dest_node_unit = model_builder.GetNodeUnit(&it->GetNode());
@@ -897,17 +909,6 @@ void ReshapeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const
     }
   }
 
-  // If we reach here, we have all the Reshape outputs are used by gemm/matmul, or Reshape has no output edge
-  // Check if the Reshape output is a graph output, if so we cannot skip the Reshape
-  // We do not care the case where the Reshape output is a dead end
-  for (const auto* node_arg : model_builder.GetGraphViewer().GetOutputs()) {
-    if (node_arg == &output_node_arg) {
-      LOGS_DEFAULT(VERBOSE) << "Reshape/Flatten can not be skipped when the output is a graph output"
-                            << ", output name, " << output_name;
-      return false;
-    }
-  }
-
   LOGS_DEFAULT(VERBOSE) << "Skipping Reshape/Flatten node ["
                         << node_unit.Name() << "] with output, " << output_name;
   return true;
diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc
index 1b1c6e7990..5949e67d6b 100644
--- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc
+++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc
@@ -1236,7 +1236,7 @@ class ConcatOpSupportChecker : public BaseOpSupportChecker {
 bool ConcatOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& /* initializers */, const NodeUnit& node_unit,
                                                const OpSupportCheckParams& /* params */) const {
   Shape input_shape;
-  if (GetShape(node_unit.Inputs()[0].node_arg, input_shape))
+  if (!GetShape(node_unit.Inputs()[0].node_arg, input_shape))
     return false;
 
   const auto input_size = input_shape.size();