minor fix in node unit change (#10405)

This commit is contained in:
Guoyu Wang 2022-01-26 16:42:38 -08:00 committed by GitHub
parent ea9c8a7cdc
commit c6ef465011
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 12 deletions

View file

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

View file

@ -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();