From 9f942e1a3e6eae55653205b0acc048ed406e96af Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Thu, 30 Mar 2023 15:39:43 -0700 Subject: [PATCH] Graph transformer to ensure unique DQ nodes for QDQ node units (#15145) ### Description Add required graph transformer to duplicate DQ nodes to ensure that QDQ node units have unique DQ nodes. This condition is necessary for QDQ node unit processing. ### Motivation and Context There is an existing Python utility that does this: https://github.com/microsoft/onnxruntime/blob/c7ced7a5e9c5ba1245b03f84edd537f45432b2e9/tools/python/util/qdq_helpers/qdq_model_utils.py#L77 This PR implements it as a graph transformer so it is integrated into ORT and does not require a separate step to update the model. There are also tests to ensure that its effects are not undone by basic level graph optimizations. --- cmake/onnxruntime_unittests.cmake | 6 + .../core/common/inlined_containers_fwd.h | 10 +- include/onnxruntime/core/graph/graph_nodes.h | 4 +- .../core/optimizer/graph_transformer.h | 11 +- onnxruntime/core/graph/graph.cc | 13 +- .../core/optimizer/graph_transformer_utils.cc | 31 ++- .../ensure_unique_dq_for_node_unit.cc | 171 ++++++++++++++ .../ensure_unique_dq_for_node_unit.h | 37 +++ .../selectors_actions/qdq_selectors.cc | 41 ++-- .../selectors_actions/qdq_selectors.h | 3 - .../selectors_actions/shared/utils.cc | 24 ++ .../selectors_actions/shared/utils.h | 7 + .../nnapi/nnapi_builtin/nnapi_api_helper.cc | 4 +- .../nnapi/nnapi_builtin/nnapi_api_helper.h | 5 +- .../nnapi_builtin/nnapi_execution_provider.h | 2 +- .../providers/shared/node_unit/node_unit.cc | 5 +- onnxruntime/core/session/inference_session.cc | 62 +++-- onnxruntime/core/session/inference_session.h | 5 +- .../kernel_type_str_resolver_utils_test.cc | 3 +- .../ensure_unique_dq_for_node_unit_test.cc | 222 ++++++++++++++++++ .../optimizer/graph_transform_test_builder.cc | 14 +- .../test/optimizer/qdq_transformer_test.cc | 85 ++++++- .../optimizer/transpose_optimizer_test.cc | 22 +- .../test/providers/nnapi/nnapi_basic_test.cc | 2 - ...qdq_with_multi_consumer_dq_nodes.fixed.txt | 10 +- .../util/qdq_helpers/optimize_qdq_model.py | 16 +- .../util/qdq_helpers/qdq_model_utils.py | 108 --------- .../python/util/qdq_helpers/test/__init__.py | 0 .../qdq_helpers/test/test_qdq_model_utils.py | 32 --- 29 files changed, 689 insertions(+), 266 deletions(-) create mode 100644 onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.cc create mode 100644 onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h create mode 100644 onnxruntime/test/optimizer/ensure_unique_dq_for_node_unit_test.cc delete mode 100644 tools/python/util/qdq_helpers/qdq_model_utils.py delete mode 100644 tools/python/util/qdq_helpers/test/__init__.py delete mode 100644 tools/python/util/qdq_helpers/test/test_qdq_model_utils.py diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index e26660963a..930db84ef4 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -41,8 +41,14 @@ function(AddTest) if (MSVC) target_compile_options(${_UT_TARGET} PRIVATE "/wd6330") endif() + set_target_properties(${_UT_TARGET} PROPERTIES FOLDER "ONNXRuntimeTest") + if (MSVC) + # set VS debugger working directory to the test program's directory + set_target_properties(${_UT_TARGET} PROPERTIES VS_DEBUGGER_WORKING_DIRECTORY $) + endif() + if (_UT_DEPENDS) add_dependencies(${_UT_TARGET} ${_UT_DEPENDS}) endif(_UT_DEPENDS) diff --git a/include/onnxruntime/core/common/inlined_containers_fwd.h b/include/onnxruntime/core/common/inlined_containers_fwd.h index 9e5215c2e4..21a55f9b31 100644 --- a/include/onnxruntime/core/common/inlined_containers_fwd.h +++ b/include/onnxruntime/core/common/inlined_containers_fwd.h @@ -57,6 +57,7 @@ namespace onnxruntime { #ifndef DISABLE_ABSEIL /// Inspired by LLVM SmallVector with ONNX Runtime adjustments for abseil. +/// https://github.com/llvm/llvm-project/blob/a85b37d0ca819776c6034c2dbda2b21e54e3393a/llvm/include/llvm/ADT/SmallVector.h#L1128-L1179 /// /// Helper class for calculating the default number of inline elements for /// `InlinedVector`. @@ -77,6 +78,9 @@ struct CalculateInlinedVectorDefaultInlinedElements { // it contradicts 1. static constexpr size_t kPreferredInlinedVectorSizeof = 64; + // Largest allowed element size for default element count calculation. + static constexpr size_t kElementSizeCutoff = 256; + // static_assert that sizeof(T) is not "too big". // // Because the InlinedVector must have at least one inlined element, it is possible @@ -98,7 +102,7 @@ struct CalculateInlinedVectorDefaultInlinedElements { // happens on a 32-bit host and then fails due to sizeof(T) *increasing* on a // 64-bit host, is expected to be very rare. static_assert( - sizeof(absl::InlinedVector) <= kPreferredInlinedVectorSizeof, + sizeof(T) <= kElementSizeCutoff, "You are trying to use a default number of inlined elements for " "`InlinedVector` but `sizeof(T)` is really big! Please use an " "explicit number of inlined elements with `InlinedVector` to make " @@ -106,8 +110,8 @@ struct CalculateInlinedVectorDefaultInlinedElements { // Discount the size of the header itself when calculating the maximum inline // bytes. - static constexpr size_t PreferredInlineBytes = - kPreferredInlinedVectorSizeof - (sizeof(absl::InlinedVector) - sizeof(T)); + static constexpr size_t InlinedVectorHeaderSize = sizeof(absl::InlinedVector) - sizeof(T); + static constexpr size_t PreferredInlineBytes = kPreferredInlinedVectorSizeof - InlinedVectorHeaderSize; static constexpr size_t NumElementsThatFit = PreferredInlineBytes / sizeof(T); static constexpr size_t value = NumElementsThatFit == 0 ? 1 : NumElementsThatFit; diff --git a/include/onnxruntime/core/graph/graph_nodes.h b/include/onnxruntime/core/graph/graph_nodes.h index 8a119e3302..4fa2848a1d 100644 --- a/include/onnxruntime/core/graph/graph_nodes.h +++ b/include/onnxruntime/core/graph/graph_nodes.h @@ -134,13 +134,13 @@ class ValidNodes { } /** Return the current Node&. This will be const if the iterator was returned from a const GraphNodes instance. */ - reference operator*() { + reference operator*() const { // if iterator is valid we always have a non-nullptr node // if this is a nullptr we're at end_ and this shouldn't be being called return **current_; } - pointer operator->() { + pointer operator->() const { return current_->get(); } diff --git a/include/onnxruntime/core/optimizer/graph_transformer.h b/include/onnxruntime/core/optimizer/graph_transformer.h index 579f1445a9..f61eb2dab9 100644 --- a/include/onnxruntime/core/optimizer/graph_transformer.h +++ b/include/onnxruntime/core/optimizer/graph_transformer.h @@ -38,13 +38,13 @@ class GraphTransformer { @param[out] modified Set to true if the Graph was modified. @returns Status with success or error information. */ - common::Status Apply(Graph& graph, bool& modified, const logging::Logger& logger) const; + Status Apply(Graph& graph, bool& modified, const logging::Logger& logger) const; virtual bool ShouldOnlyApplyOnce() const { return false; } protected: /** Helper method to call ApplyImpl on any subgraphs in the Node. */ - common::Status Recurse(Node& node, bool& modified, int graph_level, const logging::Logger& logger) const { + Status Recurse(Node& node, bool& modified, int graph_level, const logging::Logger& logger) const { int subgraph_level = ++graph_level; for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) { auto& subgraph = *entry.second; @@ -62,10 +62,9 @@ class GraphTransformer { // You MUST call Recurse for all valid Nodes in the graph to ensure any subgraphs in control flow nodes // (Scan/If/Loop) are processed as well. // You should avoid calling Graph::Resolve in ApplyImpl unless you are 100% sure it's required. In most cases - // the call to Graph::Resolve in Apply prior to ApplyImpl being called, and after ApplyImpl fore the main graph - // completes (if 'modified' is true) should suffice. - virtual common::Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) - const = 0; + // the call to Graph::Resolve in GraphTransformer::Apply after the call to ApplyImpl (if 'modified' is true) + // should suffice. + virtual Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const = 0; const std::string name_; const InlinedHashSet compatible_provider_types_; diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index a7cacbeb4c..082fd7e015 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -898,7 +898,7 @@ void Node::Init(const std::string& name, } } } -#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) +#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) Node::Definitions& Node::MutableDefinitions() noexcept { @@ -1625,11 +1625,12 @@ Status Graph::BuildConnections(std::unordered_set& outer_scope_node Node& output_node = *entry->second.first; AddEdge(output_node.Index(), node->Index(), entry->second.second, input_slot_index); - // If this Graph was built manually, remove the implicit input from the graph outputs - // if it is present there and not explicitly listed in the ordered graph outputs - // (as that implies we should leave it as an output). - // If the Graph was loaded from a GraphProto, honor the explicit graph outputs and leave as is. - if (!is_loaded_from_model_file_) { + // If this Graph was built manually and the outputs were not manually set, remove the implicit input from + // the graph outputs if it is present there. + // + // Otherwise, if the Graph was loaded from a GraphProto or the outputs were manually set, honor the + // explicit graph outputs and leave as is. + if (!is_loaded_from_model_file_ && !graph_outputs_manually_set_) { graph_outputs_.erase(std::remove(graph_outputs_.begin(), graph_outputs_.end(), node_arg), graph_outputs_.end()); } diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 7fbef8a9b4..8a46537c97 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -10,18 +10,13 @@ #include "core/optimizer/nhwc_transformer.h" #include "core/optimizer/qdq_transformer/qdq_final_cleanup.h" #include "core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.h" -#include "core/optimizer/rocm_blas_alt_impl.h" #include "core/optimizer/selectors_actions/selector_action_transformer_apply_contexts.h" #include "core/session/onnxruntime_session_options_config_keys.h" -#include "core/optimizer/conv_add_act_fusion.h" #if !defined(ORT_MINIMAL_BUILD) #include "core/mlas/inc/mlas.h" #include "core/optimizer/attention_fusion.h" -#ifdef MLAS_TARGET_AMD64_IX86 -#include "core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.h" -#endif #include "core/optimizer/bias_dropout_fusion.h" #include "core/optimizer/bias_gelu_fusion.h" #include "core/optimizer/bias_softmax_fusion.h" @@ -29,6 +24,7 @@ #include "core/optimizer/common_subexpression_elimination.h" #include "core/optimizer/constant_folding.h" #include "core/optimizer/constant_sharing.h" +#include "core/optimizer/conv_add_act_fusion.h" #include "core/optimizer/conv_add_fusion.h" #include "core/optimizer/conv_bn_fusion.h" #include "core/optimizer/conv_mul_fusion.h" @@ -56,21 +52,26 @@ #include "core/optimizer/nchwc_transformer.h" #include "core/optimizer/noop_elimination.h" #include "core/optimizer/not_where_fusion.h" +#ifdef MLAS_TARGET_AMD64_IX86 +#include "core/optimizer/qdq_transformer/avx2_weight_s8_to_u8.h" +#endif #include "core/optimizer/qdq_transformer/clip_quantizelinear.h" +#include "core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h" #include "core/optimizer/qdq_transformer/qdq_propagation.h" #include "core/optimizer/qdq_transformer/qdq_s8_to_u8.h" #include "core/optimizer/qdq_transformer/relu_quantizelinear.h" #include "core/optimizer/quick_gelu_fusion.h" #include "core/optimizer/relu_clip_fusion.h" #include "core/optimizer/reshape_fusion.h" +#include "core/optimizer/rocm_blas_alt_impl.h" #include "core/optimizer/rule_based_graph_transformer.h" #include "core/optimizer/skip_layer_norm_fusion.h" #include "core/optimizer/slice_elimination.h" #include "core/optimizer/transpose_optimizer/ort_transpose_optimizer.h" #include "core/optimizer/unsqueeze_elimination.h" #ifdef ENABLE_TRAINING_CORE -#include "orttraining/core/optimizer/bitmask_dropout_replacement.h" #include "orttraining/core/optimizer/bias_softmax_dropout_fusion.h" +#include "orttraining/core/optimizer/bitmask_dropout_replacement.h" #include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h" #endif #ifdef ENABLE_TRAINING @@ -212,22 +213,28 @@ InlinedVector> GenerateTransformers( if (!disable_quant_qdq) { transformers.emplace_back(std::make_unique()); + + // EnsureUniqueDQForNodeUnit is actually a required graph transformation. The unique DQ per QDQ node unit input + // condition that it ensures is important for the partitioning that happens after Level1 optimizers are run. + // It runs unconditionally in InferenceSession::TransformGraph() prior to Level1 optimizers. + // We also put it here with other Level1 optimizers so that it can fix things up after their changes. + transformers.emplace_back(std::make_unique()); } + // add __backwardpass attribute to nodes after YieldOp, ROCm-only + const InlinedHashSet rocm_ep = {onnxruntime::kRocmExecutionProvider}; + transformers.emplace_back(std::make_unique(rocm_ep)); + // run TransposeOptimizer last as it works in a slightly different way by moving Transpose nodes around. // shouldn't affect the end result - just easier to debug any issue if it's last. auto cpu_allocator = cpu_execution_provider.GetAllocator(OrtMemTypeDefault); transformers.emplace_back(std::make_unique(std::move(cpu_allocator))); - - // add __backwardpass attribute to nodes after YieldOp, ROCm-only - const InlinedHashSet rocm_ep = {onnxruntime::kRocmExecutionProvider}; - transformers.emplace_back(std::make_unique(rocm_ep)); } break; case TransformerLevel::Level2: { const bool enable_quant_qdq_cleanup = session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsEnableQuantQDQCleanup, "0") == "1"; -#ifndef DISABLE_CONTRIB_OPS +#if !defined(DISABLE_CONTRIB_OPS) const bool qdq_is_int8_allowed = session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsQDQIsInt8Allowed, QDQIsInt8Allowed() ? "1" : "0") == "1"; @@ -309,7 +316,7 @@ InlinedVector> GenerateTransformers( } #endif -#endif +#endif // !defined(DISABLE_CONTRIB_OPS) // The QDQFinalCleanupTransformer must run AFTER other transformers that fuse Q/DQ nodes. Otherwise, their // fusions might be prevented if this one removes a Q/DQ node too early. transformers.emplace_back(std::make_unique(enable_quant_qdq_cleanup)); diff --git a/onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.cc b/onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.cc new file mode 100644 index 0000000000..e50efd8aa1 --- /dev/null +++ b/onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.cc @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h" + +#include + +#include "core/common/common.h" +#include "core/common/narrow.h" +#include "core/graph/graph_utils.h" +#include "core/graph/graph_viewer.h" +#include "core/optimizer/qdq_transformer/qdq_util.h" + +namespace onnxruntime { + +namespace { + +constexpr auto* kTransformerName = "EnsureUniqueDQForNodeUnit"; + +// Given `original_dq_output_edge`, an edge from DQ to an input of Y, duplicate DQ to a new node, DQ'. +// After duplication, the input of Y from `original_dq_output_edge` is the only consumer of DQ'. +// +// Convert this: -> DQ -> ... +// | +// +---> Y +// +// To this: -> DQ -> ... +// | +// +-------> DQ' -> Y +// +// Note: may be graph inputs/initializers or nodes. +Status DuplicateDQForOutputEdge(const graph_utils::GraphEdge& original_dq_output_edge, Graph& graph) { + // DQ + Node* original_dq_node_ptr = graph.GetNode(original_dq_output_edge.src_node); + assert(original_dq_node_ptr != nullptr); + Node& original_dq_node = *original_dq_node_ptr; + + // Y + Node* dst_node_ptr = graph.GetNode(original_dq_output_edge.dst_node); + assert(dst_node_ptr != nullptr); + Node& dst_node = *dst_node_ptr; + + // set up new NodeArg + auto& new_dq_output_nodearg = + graph.GetOrCreateNodeArg(graph.GenerateNodeArgName(original_dq_output_edge.arg_name + "/duplicated"), nullptr); + + // set up new Node + const auto& dq_inputs = original_dq_node.MutableInputDefs(); + + // DQ' + Node& new_dq_node = graph.AddNode(graph.GenerateNodeName(original_dq_node.Name() + "/duplicated"), + QDQ::DQOpName, + MakeString("Added by ", kTransformerName), + dq_inputs, + {&new_dq_output_nodearg}); + + // set up edges + // remove DQ -> Y + graph_utils::GraphEdge::RemoveGraphEdges(graph, {original_dq_output_edge}); + + // add DQ_input -> DQ' if DQ_input is a node + const auto original_dq_input_edges = graph_utils::GraphEdge::GetNodeInputEdges(original_dq_node); + for (const auto& original_dq_input_edge : original_dq_input_edges) { + // create edge from the original DQ node's input node + graph.AddEdge(original_dq_input_edge.src_node, new_dq_node.Index(), + original_dq_input_edge.src_arg_index, original_dq_input_edge.dst_arg_index); + } + + // add DQ' -> Y + dst_node.MutableInputDefs()[original_dq_output_edge.dst_arg_index] = &new_dq_output_nodearg; + graph.AddEdge(new_dq_node.Index(), original_dq_output_edge.dst_node, 0, original_dq_output_edge.dst_arg_index); + + return Status::OK(); +} + +Status EnsureUniqueDQForEachExplicitOutputEdge(const Node& node, Graph& graph, bool& modified) { + if (!QDQ::MatchDQNode(node)) { + return Status::OK(); + } + + const auto& dq_node = node; + + // QDQ node units are only formed by nodes in the current graph, and not between nodes in the current graph and a + // subgraph. Consequently, we only duplicate DQ's for edges to explicit inputs and skip edges to implicit (subgraph) + // inputs. + + const bool produces_graph_output = graph.NodeProducesGraphOutput(node); + + auto dq_output_edges = graph_utils::GraphEdge::GetNodeOutputEdges(dq_node, 0); + + // Check for common cases where we don't need to do anything: + // 1. DQ has no output edges. + // + // It may have a graph output. Either way, there is no same-graph consumer node input to consider. + // + // 2. DQ has exactly one output edge and no graph output. + // + // If the output edge is to an implicit input, there are only consumer nodes in a subgraph which we can ignore. + // Otherwise, the output edge is to an explicit input of the only same-graph consumer node so the DQ is already + // unique. + if (dq_output_edges.empty() || + (dq_output_edges.size() == 1 && !produces_graph_output)) { + return Status::OK(); + } + + // Remove edges to implicit inputs from consideration. + const auto dq_output_edges_to_explicit_inputs_end = std::remove_if( + dq_output_edges.begin(), dq_output_edges.end(), + [&const_graph = std::as_const(graph)](const graph_utils::GraphEdge& dq_output_edge) { + const Node* consumer_node_ptr = const_graph.GetNode(dq_output_edge.dst_node); + assert(consumer_node_ptr != nullptr); + const Node& consumer_node = *consumer_node_ptr; + const auto consumer_explicit_input_defs_count = consumer_node.InputDefs().size(); + return narrow(dq_output_edge.dst_arg_index) >= consumer_explicit_input_defs_count; + }); + + const bool has_subgraph_consumer = dq_output_edges_to_explicit_inputs_end != dq_output_edges.end(); + if (has_subgraph_consumer) { + dq_output_edges.erase(dq_output_edges_to_explicit_inputs_end, dq_output_edges.end()); + } + + // If the original DQ produces a graph output or has a subgraph consumer node, we preserve any of those output + // relationships and duplicate new unique DQ's for each of the same-graph consumer nodes. It is important that the + // original DQ node does not become part of a QDQ node unit as it has an additional consumer. + // + // Otherwise, where the original DQ has only same-graph consumer nodes, we can reuse the original DQ as the unique + // DQ for the first same-graph consumer node and duplicate new unique DQ's for the rest of them. + const bool must_preserve_original_dq_outside_of_node_unit = produces_graph_output || has_subgraph_consumer; + + auto edge_to_process = dq_output_edges.begin(); + if (!must_preserve_original_dq_outside_of_node_unit && edge_to_process != dq_output_edges.end()) { + // reuse original DQ for the first edge, start duplicating from the next edge + ++edge_to_process; + } + + for (; edge_to_process != dq_output_edges.end(); ++edge_to_process) { + ORT_RETURN_IF_ERROR(DuplicateDQForOutputEdge(*edge_to_process, graph)); + modified = true; + } + + return Status::OK(); +} + +} // namespace + +Status EnsureUniqueDQForNodeUnit::ApplyImpl(Graph& graph, + bool& modified, + int graph_level, + const logging::Logger& logger) const { + const GraphViewer graph_viewer{graph}; + const auto& node_indices = graph_viewer.GetNodesInTopologicalOrder(); + for (const auto node_idx : node_indices) { + auto* node_ptr = graph.GetNode(node_idx); + if (!node_ptr) { + continue; + } + + Node& node = *node_ptr; + ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger)); + + ORT_RETURN_IF_ERROR(EnsureUniqueDQForEachExplicitOutputEdge(node, graph, modified)); + } + + return Status::OK(); +} + +EnsureUniqueDQForNodeUnit::EnsureUniqueDQForNodeUnit() + : GraphTransformer{kTransformerName, {}} { +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h b/onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h new file mode 100644 index 0000000000..d4a47007a7 --- /dev/null +++ b/onnxruntime/core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +/** + * Graph transformer that duplicates DQ nodes in order to ensure that each potential QDQ node unit has unique DQ nodes + * for its inputs, which is necessary for QDQ node unit processing. + * + * In particular, it ensures that each output edge from a DQ to an explicit input (i.e., to a node in the same graph) + * will have a unique DQ of which that explicit input is the only consumer. + * + * Before: + * + * DQ -> X + * | + * +---> Y + * + * After: + * + * DQ -> X + * + * DQ' -> Y (DQ' is a duplicate of DQ) + */ +class EnsureUniqueDQForNodeUnit : public GraphTransformer { + public: + EnsureUniqueDQForNodeUnit(); + + private: + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc index 2e063011ad..f9d070044f 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc @@ -8,6 +8,7 @@ #include "core/graph/graph.h" #include "core/optimizer/initializer.h" #include "core/optimizer/qdq_transformer/qdq_util.h" +#include "core/optimizer/qdq_transformer/selectors_actions/shared/utils.h" #include "core/optimizer/utils.h" namespace onnxruntime { @@ -19,9 +20,8 @@ int NumActualValues(const Node& node, bool input) { return gsl::narrow_cast(std::count_if(defs.cbegin(), defs.cend(), [](const NodeArg* def) { return def && def->Exists(); })); } -} // namespace -static std::vector FindQDQNodes(const GraphViewer& graph_viewer, const Node& node, bool find_dq_nodes) { +std::vector FindQDQNodes(const GraphViewer& graph_viewer, const Node& node, bool find_dq_nodes) { // First get all the upstream (DQ) or downstream (Q) nodes std::vector nodes = find_dq_nodes ? graph_utils::FindParentsByType(node, QDQ::DQOpName) @@ -36,6 +36,7 @@ static std::vector FindQDQNodes(const GraphViewer& graph_viewer, co return nodes; } +} // namespace bool NodeGroupSelector::CheckQDQNodes(const GraphViewer& graph_viewer, const Node& node, const std::vector& dq_nodes, @@ -51,11 +52,8 @@ bool NodeGroupSelector::CheckQDQNodes(const GraphViewer& graph_viewer, const Nod return false; } - auto does_node_produce_graph_output = [&graph_viewer](const Node* node_ptr) { - return graph_viewer.NodeProducesGraphOutput(*node_ptr); - }; - - if (std::any_of(dq_nodes.begin(), dq_nodes.end(), does_node_produce_graph_output)) { + if (const auto dq_validation_status = QDQ::ValidateNodeGroupDQNodes(graph_viewer, node, dq_nodes); + !dq_validation_status.IsOK()) { return false; } @@ -96,8 +94,8 @@ std::optional BaseSelector::Select(const GraphViewer& gr // TODO(edgchen1) update NodeGroup to use InlinedVector builder.input_nodes.assign(qdq_group->dq_nodes.begin(), qdq_group->dq_nodes.end()); builder.output_nodes.assign(qdq_group->q_nodes.begin(), qdq_group->q_nodes.end()); - //builder.input_nodes = qdq_group->dq_nodes; - //builder.output_nodes = qdq_group->q_nodes; + // builder.input_nodes = qdq_group->dq_nodes; + // builder.output_nodes = qdq_group->q_nodes; builder.target_node = qdq_group->target_node; UpdateBuilder(builder); @@ -122,17 +120,17 @@ bool DropQDQNodeGroupSelector::Check(const GraphViewer& graph_viewer, return IsQDQPairSupported(q_node, dq_node, get_const_initializer, graph_viewer.ModelPath()); } -bool DropDQNodeGroupSelector::CheckDQNodes(const Node& node, const std::vector& dq_nodes) const { - int num_dq_inputs = NumActualValues(node, true); - - return num_dq_inputs == gsl::narrow_cast(dq_nodes.size()); -} - bool DropDQNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& node, const std::vector& dq_nodes, const std::vector& q_nodes) const { - if (!CheckDQNodes(node, dq_nodes)) { + int num_dq_inputs = NumActualValues(node, true); + if (num_dq_inputs != gsl::narrow_cast(dq_nodes.size())) { + return false; + } + + if (const auto dq_validation_status = QDQ::ValidateNodeGroupDQNodes(graph_viewer, node, dq_nodes); + !dq_validation_status.IsOK()) { return false; } @@ -317,20 +315,19 @@ void GemmSelector::UpdateBuilder(NodesToOptimizeIndicesBuilder& builder) const { builder.input_nodes.resize(3, NodesToOptimizeIndices::kEmptyNodeIndex); } -bool WhereNodeGroupSelector::Check(const GraphViewer &graph_viewer, const Node &node, - const std::vector &dq_nodes, - const std::vector &q_nodes) const { +bool WhereNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& node, + const std::vector& dq_nodes, + const std::vector& q_nodes) const { // Where has 1 boolean input and 2 dq inputs - if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes,2)) { + if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes, 2)) { return false; } - const int32_t dt_input_1 = dq_nodes[0]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); + const int32_t dt_input_1 = dq_nodes[0]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); const int32_t dt_input_2 = dq_nodes[1]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); const int32_t dt_output = q_nodes[0]->OutputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); return dt_input_1 == dt_input_2 && dt_input_1 == dt_output; - } bool InstanceNormalizationNodeGroupSelector::Check(const GraphViewer& graph_viewer, diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h index acdbd07eb1..2efc9419ed 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h @@ -59,9 +59,6 @@ class DropQDQNodeGroupSelector : public NodeGroupSelector { // Single DQ -> node. class DropDQNodeGroupSelector : public NodeGroupSelector { - // base check that we have the expected number of DQ inputs. - bool CheckDQNodes(const Node& node, const std::vector& dq_nodes) const; - bool Check(const GraphViewer& graph_viewer, const Node& node, const std::vector& dq_nodes, const std::vector& q_nodes) const override; diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc index 66e5789755..6bc40344ef 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.cc @@ -208,6 +208,30 @@ std::vector SelectorManager::GetQDQSelections(const GraphViewer& grap return qdq_selections; } +Status ValidateNodeGroupDQNodes(const GraphViewer& graph_viewer, + const Node& target_node, + gsl::span dq_nodes) { + // Within a QDQ node group, a target node input is the only consumer of each DQ. + // This should have been ensured by the EnsureUniqueDQForNodeUnit graph transformer, but other graph modifications + // may have happened since. Verify that this is still true. + for (const auto* dq_node : dq_nodes) { + const bool dq_produces_graph_output = graph_viewer.NodeProducesGraphOutput(*dq_node); + ORT_RETURN_IF(dq_produces_graph_output, + "QDQ node group cannot have DQ node that produces a graph output. DQ node: ", dq_node->Name(), + ", target node: ", target_node.Name()); + + const bool dq_has_single_output_edge_to_target = + dq_node->GetOutputEdgesCount() == 1 && + dq_node->OutputEdgesBegin()->GetNode().Index() == target_node.Index(); + ORT_RETURN_IF_NOT(dq_has_single_output_edge_to_target, + "QDQ node group cannot have DQ that doesn't have a single output edge to the target node. " + "DQ node: ", + dq_node->Name(), ", target node: ", target_node.Name()); + } + + return Status::OK(); +} + } // namespace QDQ } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.h b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.h index e148987263..c64f40d0dd 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.h +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/shared/utils.h @@ -5,6 +5,7 @@ #include #include "core/common/common.h" +#include "core/common/gsl.h" #include "core/common/inlined_containers.h" #include "core/graph/basic_types.h" @@ -77,5 +78,11 @@ class SelectorManager { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(SelectorManager); }; +// Checks whether the provided DQ nodes are valid for forming a QDQ node group with the provided target node. +// Returns successful status if so, failed status with reason otherwise. +Status ValidateNodeGroupDQNodes(const GraphViewer& graph_viewer, + const Node& target_node, + gsl::span dq_nodes); + } // namespace QDQ } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.cc index 350078ac78..5e64f532e7 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.cc @@ -66,7 +66,7 @@ static int32_t GetDeviceFeatureLevelInternal(const NnApi& nnapi_handle, gsl::spa * */ Status GetTargetDevices(const NnApi& nnapi_handle, TargetDeviceOption target_device_option, - InlinedVector& devices) { + DeviceWrapperVector& devices) { // GetTargetDevices is only supported when NNAPI runtime feature level >= ANEURALNETWORKS_FEATURE_LEVEL_3 if (GetNNAPIRuntimeFeatureLevel(nnapi_handle) < ANEURALNETWORKS_FEATURE_LEVEL_3) return Status::OK(); @@ -134,7 +134,7 @@ std::string GetDevicesDescription(gsl::span devices) { // Return -1 if failed. // It's not necessary to handle the error here, because level=-1 will refuse all ops in NNAPI EP. int32_t GetNNAPIEffectiveFeatureLevelFromTargetDeviceOption(const NnApi& nnapi_handle, TargetDeviceOption target_device_option) { - InlinedVector nnapi_target_devices; + DeviceWrapperVector nnapi_target_devices; if (auto st = GetTargetDevices(nnapi_handle, target_device_option, nnapi_target_devices); !st.IsOK()) { LOGS_DEFAULT(WARNING) << "GetTargetDevices failed for:" << st.ErrorMessage(); return -1; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.h index 0f8a22a208..a5ab7f92f8 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_api_helper.h @@ -54,12 +54,13 @@ constexpr const char* const kNnapiCpuDeviceName = "nnapi-reference"; */ int32_t GetNNAPIEffectiveFeatureLevel(const NnApi& nnapi_handle, gsl::span device_handles); +using DeviceWrapperVector = InlinedVector; + /** * Get all target devices specified by target_device_option. - * */ Status GetTargetDevices(const NnApi& nnapi_handle, TargetDeviceOption target_device_option, - InlinedVector& nnapi_target_devices); + DeviceWrapperVector& nnapi_target_devices); int32_t GetNNAPIEffectiveFeatureLevelFromTargetDeviceOption(const NnApi& nnapi_handle, TargetDeviceOption target_device_option); diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_execution_provider.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_execution_provider.h index 1df9802fa3..e4911511e6 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_execution_provider.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/nnapi_execution_provider.h @@ -46,7 +46,7 @@ class NnapiExecutionProvider : public IExecutionProvider { // For Android NNAPI and stub implementation. const NnApi* nnapi_handle_ = nullptr; - InlinedVector nnapi_target_devices_; + nnapi::DeviceWrapperVector nnapi_target_devices_; nnapi::TargetDeviceOption target_device_option_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/shared/node_unit/node_unit.cc b/onnxruntime/core/providers/shared/node_unit/node_unit.cc index 9ba122c4cd..10dd58ba28 100644 --- a/onnxruntime/core/providers/shared/node_unit/node_unit.cc +++ b/onnxruntime/core/providers/shared/node_unit/node_unit.cc @@ -146,18 +146,19 @@ std::vector GetQDQIODefs(const Node& target_node, const QDQ::Node } // namespace NodeUnit::NodeUnit(const Node& node) - :target_node_(node), + : target_node_(node), type_(Type::SingleNode) { InitForSingleNode(); } NodeUnit::NodeUnit(const GraphViewer& graph_viewer, const QDQ::NodeGroup& node_group) - :q_nodes_{GetQDQIONodes(graph_viewer, node_group, false /* is_input */)}, + : q_nodes_{GetQDQIONodes(graph_viewer, node_group, false /* is_input */)}, dq_nodes_{GetQDQIONodes(graph_viewer, node_group, true /* is_input */)}, target_node_(*graph_viewer.GetNode(node_group.target_node)), type_(Type::QDQGroup), inputs_{GetQDQIODefs(target_node_, node_group, true /* is_input */)}, outputs_{GetQDQIODefs(target_node_, node_group, false /* is_input */)} { + ORT_THROW_IF_ERROR(QDQ::ValidateNodeGroupDQNodes(graph_viewer, target_node_, dq_nodes_)); } const std::string& NodeUnit::Domain() const noexcept { return target_node_.Domain(); } diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index aec4ab4b03..e9899e38ad 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -40,6 +40,7 @@ #include "core/optimizer/graph_transformer_utils.h" #include "core/optimizer/graph_transformer.h" #include "core/optimizer/insert_cast_transformer.h" +#include "core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h" #include "core/optimizer/rule_based_graph_transformer.h" #include "core/optimizer/selectors_actions/selector_action_transformer_apply_contexts.h" #include "core/optimizer/transformer_memcpy.h" @@ -368,7 +369,6 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const : #if !defined(ORT_MINIMAL_BUILD) graph_transformer_mgr_(session_options.max_num_graph_transformation_steps), - insert_cast_transformer_("CastFloat16Transformer"), #endif logging_manager_(session_env.GetLoggingManager()), environment_(session_env) { @@ -383,7 +383,6 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, : #if !defined(ORT_MINIMAL_BUILD) graph_transformer_mgr_(session_options.max_num_graph_transformation_steps), - insert_cast_transformer_("CastFloat16Transformer"), #endif logging_manager_(session_env.GetLoggingManager()), external_intra_op_thread_pool_(external_intra_op_thread_pool), @@ -398,7 +397,6 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const const PathString& model_uri) : model_location_(model_uri), graph_transformer_mgr_(session_options.max_num_graph_transformation_steps), - insert_cast_transformer_("CastFloat16Transformer"), logging_manager_(session_env.GetLoggingManager()), environment_(session_env) { auto status = Model::Load(model_location_, model_proto_); @@ -420,7 +418,6 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, InferenceSession::InferenceSession(const SessionOptions& session_options, const Environment& session_env, std::istream& model_istream) : graph_transformer_mgr_(session_options.max_num_graph_transformation_steps), - insert_cast_transformer_("CastFloat16Transformer"), logging_manager_(session_env.GetLoggingManager()), environment_(session_env) { Status st = Model::Load(model_istream, &model_proto_); @@ -433,7 +430,6 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const InferenceSession::InferenceSession(const SessionOptions& session_options, const Environment& session_env, const void* model_data, int model_data_len) : graph_transformer_mgr_(session_options.max_num_graph_transformation_steps), - insert_cast_transformer_("CastFloat16Transformer"), logging_manager_(session_env.GetLoggingManager()), environment_(session_env) { const bool result = model_proto_.ParseFromArray(model_data, model_data_len); @@ -597,7 +593,7 @@ common::Status InferenceSession::SaveToOrtFormat(const PathString& filepath) con ORT_RETURN_IF_ERROR(kernel_type_str_resolver.RegisterGraphNodeOpSchemas(model_->MainGraph())); ORT_RETURN_IF_ERROR(standalone::RegisterCustomOpNodeSchemas(kernel_type_str_resolver, model_->MainGraph())); - for (const auto* op_schema : saved_runtime_optimization_produced_node_op_schemas_) { + for (const auto op_schema : saved_runtime_optimization_produced_node_op_schemas_) { ORT_RETURN_IF_ERROR(kernel_type_str_resolver.RegisterOpSchema(*op_schema)); } @@ -869,15 +865,31 @@ common::Status InferenceSession::Load() { common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool saving_model_in_ort_format) { // The transformer order: - // 1. run level 1 optimizations. these only use ONNX operators. - // 2. partition nodes based on EP capabilities. EPs may fuse nodes during this process. - // 3. run all optimizations. level 2 and 3 optimizations use contrib ops. - // 4. insert cast nodes - // 5. insert copy nodes + // 1. ensure potential QDQ node units have unique DQ nodes (required transformer). + // - This is a required transformer as the ORT code has a hard requirement there are no overlapping QDQ node units. + // - We run it here in case optimizers are disabled. + // 2. run level 1 optimizations. these only use ONNX operators. + // 3. partition nodes based on EP capabilities. EPs may fuse nodes during this process. + // 4. run level 2+ optimizations. level 2 and 3 optimizations use contrib ops. + // 5. insert cast nodes (required transformer). + // 6. insert copy nodes (required transformer). - // first apply execution provider independent level 1 graph optimizations. - ORT_RETURN_IF_ERROR_SESSIONID_( - graph_transformer_mgr_.ApplyTransformers(graph, TransformerLevel::Level1, *session_logger_)); + auto apply_transformer_once = [](const GraphTransformer& transformer, const logging::Logger& logger, + Graph& graph) { + bool modified = false; + return transformer.Apply(graph, modified, logger); + }; + + // ensure potential QDQ node units have unique DQ nodes + if (const bool disable_quant_qdq = + session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsDisableQuantQDQ, "0") == "1"; + !disable_quant_qdq) { + EnsureUniqueDQForNodeUnit ensure_unique_dq_for_node_unit{}; + ORT_RETURN_IF_ERROR_SESSIONID_(apply_transformer_once(ensure_unique_dq_for_node_unit, *session_logger_, graph)); + } + + // apply execution provider independent level 1 graph optimizations. + ORT_RETURN_IF_ERROR_SESSIONID_(graph_transformer_mgr_.ApplyTransformers(graph, TransformerLevel::Level1, *session_logger_)); // if saving model to ORT format we only assign nodes a custom EP can handle and don't compile them. // we do this to preserve the original nodes in the model but prevent optimizers from changing them. @@ -950,20 +962,24 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool graph_transformer_mgr_.ApplyTransformers(graph, static_cast(i), *session_logger_)); } - bool modified = false; // Insert cast node/s. - ORT_RETURN_IF_ERROR_SESSIONID_(insert_cast_transformer_.Apply(graph, modified, *session_logger_)); - - std::vector provider_types; - for (auto& provider_ptr : execution_providers_) { - provider_types.push_back(provider_ptr->Type()); + { + InsertCastTransformer insert_cast_transformer{"CastFloat16Transformer"}; + ORT_RETURN_IF_ERROR_SESSIONID_(apply_transformer_once(insert_cast_transformer, *session_logger_, graph)); } // Insert copy node/s. - MemcpyTransformer copy_transformer{provider_types, kernel_registry_manager_}; - ORT_RETURN_IF_ERROR_SESSIONID_(copy_transformer.Apply(graph, modified, *session_logger_)); + { + std::vector provider_types; + for (auto& provider_ptr : execution_providers_) { + provider_types.push_back(provider_ptr->Type()); + } - return common::Status::OK(); + MemcpyTransformer copy_transformer{provider_types, kernel_registry_manager_}; + ORT_RETURN_IF_ERROR_SESSIONID_(apply_transformer_once(copy_transformer, *session_logger_, graph)); + } + + return Status::OK(); } #endif // !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index c2313a852a..473c38c1af 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -658,10 +658,7 @@ class InferenceSession { onnxruntime::GraphTransformerManager graph_transformer_mgr_; - InsertCastTransformer insert_cast_transformer_; - - // assuming that OpSchema* elements are not null. our version of gsl::not_null doesn't specialize std::hash. - InlinedHashSet saved_runtime_optimization_produced_node_op_schemas_; + InlinedHashSet> saved_runtime_optimization_produced_node_op_schemas_; #endif // Any GraphTransformer/RewriteRule name in this set will not be enabled. InlinedHashSet optimizers_to_disable_; diff --git a/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc b/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc index 69dd81b0fc..ac213f70b1 100644 --- a/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc +++ b/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc @@ -53,7 +53,7 @@ TEST(KernelTypeStrResolverUtilsTest, VerifyLayoutTransformationRequiredOpsResolv TEST(KernelTypeStrResolverUtilsTest, DISABLED_PrintExpectedLayoutTransformationRequiredOpsResolverByteArray) { #if defined(DISABLE_CONTRIB_OPS) FAIL() << "Contrib ops must be enabled."; -#endif // defined(DISABLE_CONTRIB_OPS) +#else // defined(DISABLE_CONTRIB_OPS) KernelTypeStrResolver expected_resolver; ASSERT_STATUS_OK(LoadLayoutTransformationRequiredOpsFromOpSchemas(expected_resolver)); @@ -75,6 +75,7 @@ TEST(KernelTypeStrResolverUtilsTest, DISABLED_PrintExpectedLayoutTransformationR os << "\n };\n"; std::cout << os.str(); +#endif // defined(DISABLE_CONTRIB_OPS) } } // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/ensure_unique_dq_for_node_unit_test.cc b/onnxruntime/test/optimizer/ensure_unique_dq_for_node_unit_test.cc new file mode 100644 index 0000000000..86fae32c12 --- /dev/null +++ b/onnxruntime/test/optimizer/ensure_unique_dq_for_node_unit_test.cc @@ -0,0 +1,222 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/qdq_transformer/ensure_unique_dq_for_node_unit.h" + +#include "gtest/gtest.h" + +#include "test/optimizer/graph_transform_test_builder.h" +#include "test/test_environment.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime::test { + +namespace { + +struct GraphConfig { + size_t num_explicit_consumer_nodes{1}; + size_t num_inputs_per_explicit_consumer_node{1}; + bool has_graph_output{false}; + bool has_subgraph_consumer{false}; +}; + +auto GetGraphBuilder(GraphConfig config) { + return [=](ModelTestBuilder& builder) { + const auto input_shape = std::vector{1, 2, 4}; + constexpr float scale = 0.5f; + constexpr uint8_t zero_point = 0; + + auto* dq_input = builder.MakeInput(input_shape, uint8_t{0}, uint8_t{255}); + auto* dq_output = config.has_graph_output ? builder.MakeOutput() : builder.MakeIntermediate(); + builder.AddDequantizeLinearNode(dq_input, scale, zero_point, dq_output); + + for (size_t i = 0; i < config.num_explicit_consumer_nodes; ++i) { + // use Concat for the explicit consumer node as it supports a variadic number of inputs + auto* explicit_consumer_output = builder.MakeOutput(); + std::vector explicit_consumer_inputs(config.num_inputs_per_explicit_consumer_node, dq_output); + auto& concat = builder.AddNode("Concat", explicit_consumer_inputs, {explicit_consumer_output}); + concat.AddAttribute("axis", int64_t{-1}); + } + + if (config.has_subgraph_consumer) { + auto create_if_subgraph = [&](bool condition) { + auto model = Model{"model for generating graph proto", true, DefaultLoggingManager().DefaultLogger()}; + auto& graph = model.MainGraph(); + + const auto& dq_output_name = dq_output->Name(); + + ONNX_NAMESPACE::TypeProto type_proto{}; + type_proto.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + + NodeArg& local_dq_output = graph.GetOrCreateNodeArg(dq_output_name, &type_proto); + graph.AddOuterScopeNodeArg(local_dq_output.Name()); + + NodeArg& out = graph.GetOrCreateNodeArg(condition ? "output_then" : "output_else", &type_proto); + graph.AddNode("identity", "Identity", "pass through identity", {&local_dq_output}, {&out}); + + graph.SetOutputs({&out}); + + ORT_THROW_IF_ERROR(graph.Resolve()); + + return graph.ToGraphProto(); + }; + + auto* if_input = builder.MakeInitializerBool({}, {true}); + auto* if_output = builder.MakeOutput(); + Node& if_node = builder.AddNode("If", {if_input}, {if_output}); + if_node.AddAttribute("then_branch", create_if_subgraph(true)); + if_node.AddAttribute("else_branch", create_if_subgraph(false)); + } + }; +} + +void RunEnsureUniqueDQForNodeUnitTest(std::function graph_builder_fn, + int expected_dq_count) { + constexpr int opset_version = 12; + + { + SCOPED_TRACE("test with standalone transformer"); + + auto post_transform_check_fn = [expected_dq_count](const Graph& graph) { + const auto op_counts = CountOpsInGraph(graph); + const auto actual_dq_count = OpCount(op_counts, "DequantizeLinear"); + ORT_RETURN_IF_NOT(actual_dq_count == expected_dq_count, + "Expected DQ count: ", expected_dq_count, ", actual: ", actual_dq_count); + return Status::OK(); + }; + + EXPECT_STATUS_OK(TestGraphTransformer( + graph_builder_fn, + opset_version, + DefaultLoggingManager().DefaultLogger(), + std::make_unique(), + TransformerLevel::Level1, + 5, + {}, + post_transform_check_fn)); + } + + { + SCOPED_TRACE("test with basic transformers"); + + auto post_transform_check_fn = [expected_dq_count](const InferenceSessionWrapper& session) { + const auto& graph = session.GetGraph(); + const auto op_counts = CountOpsInGraph(graph); + ASSERT_EQ(OpCount(op_counts, "DequantizeLinear"), expected_dq_count); + }; + + TransformerTester( + graph_builder_fn, + post_transform_check_fn, + TransformerLevel::Default, + TransformerLevel::Level1, + opset_version); + } +} + +} // namespace + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodes) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 3; + config.num_inputs_per_explicit_consumer_node = 1; + + // expected count = one for each explicit consumer node (3), reusing the original one = 3 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 3); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodesWithGraphOutput) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 3; + config.num_inputs_per_explicit_consumer_node = 1; + config.has_graph_output = true; + + // expected count = preserved original (1) + one for each explicit consumer node (3) = 4 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 4); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodesWithSubgraphConsumer) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 3; + config.num_inputs_per_explicit_consumer_node = 1; + config.has_subgraph_consumer = true; + + // expected count = preserved original (1) + one for each explicit consumer node (3) = 4 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 4); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodesWithSubgraphConsumerAndGraphOutput) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 3; + config.num_inputs_per_explicit_consumer_node = 1; + config.has_graph_output = true; + config.has_subgraph_consumer = true; + + // expected count = preserved original (1) + one for each explicit consumer node (3) = 4 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 4); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodeInputs) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 2; + config.num_inputs_per_explicit_consumer_node = 5; + + // expected count = one for each explicit consumer node input (2 * 5), reusing the original one = 10 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 10); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodeInputsWithGraphOutput) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 2; + config.num_inputs_per_explicit_consumer_node = 5; + config.has_graph_output = true; + + // expected count = preserved original (1) + one for each explicit consumer node input (2 * 5) = 11 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 11); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodeInputsWithSubgraphConsumer) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 2; + config.num_inputs_per_explicit_consumer_node = 5; + config.has_subgraph_consumer = true; + + // expected count = preserved original (1) + one for each explicit consumer node input (2 * 5) = 11 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 11); +} + +TEST(EnsureUniqueDQForNodeUnitTests, DQSharedAmongNodeInputsWithSubgraphConsumerAndGraphOutput) { + GraphConfig config{}; + config.num_explicit_consumer_nodes = 2; + config.num_inputs_per_explicit_consumer_node = 5; + config.has_graph_output = true; + config.has_subgraph_consumer = true; + + // expected count = preserved original (1) + one for each explicit consumer node input (2 * 5) = 11 + RunEnsureUniqueDQForNodeUnitTest(GetGraphBuilder(config), 11); +} + +TEST(EnsureUniqueDQForNodeUnitTests, QDQWithMultiConsumerDQNodes) { + constexpr auto model_uri = ORT_TSTR("testdata/qdq_with_multi_consumer_dq_nodes.onnx"); + + SessionOptions session_options{}; + // test interaction with level 1 transformers + session_options.graph_optimization_level = TransformerLevel::Level1; + + InferenceSessionWrapper session{session_options, GetEnvironment()}; + + ASSERT_STATUS_OK(session.Load(model_uri)); + + const auto op_count_before = CountOpsInGraph(session.GetGraph()); + + ASSERT_STATUS_OK(session.Initialize()); + + const auto op_count_after = CountOpsInGraph(session.GetGraph()); + + // there are 3 DQ nodes with 2 consumers (an earlier Conv and later Add) + // additionally the last one also provides a graph output + // based on that there should be 3 new DQ nodes for the internal consumers and 1 new one for the graph output + EXPECT_EQ(OpCount(op_count_before, "DequantizeLinear") + 4, OpCount(op_count_after, "DequantizeLinear")); +} + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/graph_transform_test_builder.cc b/onnxruntime/test/optimizer/graph_transform_test_builder.cc index 80f17fdda3..32edbaf507 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_builder.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_builder.cc @@ -97,8 +97,9 @@ void TransformerTester(const std::function& buil &fetches)); if (level == target_level) { - ASSERT_TRUE(check_transformed_graph); - check_transformed_graph(session); + if (check_transformed_graph) { + check_transformed_graph(session); + } } }; @@ -148,12 +149,17 @@ Status TestGraphTransformer(const std::function& domain_to_version, {}, logger); Graph& graph = model.MainGraph(); ModelTestBuilder helper(graph); + ORT_RETURN_IF_NOT(build_test_case, "build_test_case must be provided"); build_test_case(helper); helper.SetGraphOutputs(); ORT_RETURN_IF_ERROR(graph.Resolve()); - ORT_RETURN_IF_ERROR(pre_graph_checker(graph)); + if (pre_graph_checker) { + ORT_RETURN_IF_ERROR(pre_graph_checker(graph)); + } ORT_RETURN_IF_ERROR(graph_transformation_mgr.ApplyTransformers(graph, level, logger)); - ORT_RETURN_IF_ERROR(post_graph_checker(graph)); + if (post_graph_checker) { + ORT_RETURN_IF_ERROR(post_graph_checker(graph)); + } } return Status::OK(); diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc index e2dcc7fac2..1f42018360 100644 --- a/onnxruntime/test/optimizer/qdq_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc @@ -913,7 +913,7 @@ TEST(QDQTransformerTests, DoubleQDQ_Without_Last_Node_Being_Output) { TransformerLevel::Level1); }; test_case(0, 2, 2); - test_case(1, 2, 2); + test_case(1, 2, 3); // EnsureUniqueDQForNodeUnit will duplicate first DQ, so expect one more (3) test_case(2, 2, 2); test_case(3, 1, 1); } @@ -2768,12 +2768,15 @@ TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicQDQCleanup) { }; // if we block removal of the DQ node the Q node in the pair will not be removed either - int expected_qdq_count = 0 + (block_removal_of_first_dq ? 1 : 0) + (block_removal_of_last_dq ? 1 : 0); + const int expected_qdq_count = 0 + (block_removal_of_first_dq ? 1 : 0) + (block_removal_of_last_dq ? 1 : 0); + // blocking removal of DQ by adding an additional edge will cause EnsureUniqueDQForNodeUnit to duplicate the DQ, + // so we expect twice as many DQ's as original QDQ pairs + const int expected_dq_count = expected_qdq_count * 2; - auto check_graph = [expected_qdq_count](InferenceSessionWrapper& session) { + auto check_graph = [expected_qdq_count, expected_dq_count](InferenceSessionWrapper& session) { auto op_to_count = CountOpsInGraph(session.GetGraph()); EXPECT_EQ(op_to_count["QuantizeLinear"], expected_qdq_count); - EXPECT_EQ(op_to_count["DequantizeLinear"], expected_qdq_count); + EXPECT_EQ(op_to_count["DequantizeLinear"], expected_dq_count); EXPECT_EQ(op_to_count["Concat"], 1); }; @@ -2898,8 +2901,8 @@ TEST(QDQTransformerTests, QDQFinalCleanupTransformer_GraphInputToOutput) { test_case(false); } -// Not fuse if DQ produces graph output -TEST(QDQTransformerTests, DQ_Produce_Graph_Output) { +#if !defined(DISABLE_CONTRIB_OPS) +TEST(QDQTransformerTests, QDQSoftmaxWithDQProducingGraphOutput) { auto test_case = [&](const std::vector& input_shape, int64_t axis) { auto build_test_case = [&](ModelTestBuilder& builder) { auto* input_arg = builder.MakeInput(input_shape, -5.f, 5.f); @@ -2935,10 +2938,12 @@ TEST(QDQTransformerTests, DQ_Produce_Graph_Output) { auto check_graph = [&](InferenceSessionWrapper& session) { auto op_to_count = CountOpsInGraph(session.GetGraph()); - EXPECT_EQ(op_to_count["com.microsoft.QLinearSoftmax"], 0); - EXPECT_EQ(op_to_count["Softmax"], 1); - EXPECT_EQ(op_to_count["QuantizeLinear"], 2); - EXPECT_EQ(op_to_count["DequantizeLinear"], 2); + + // expect fusion because DQ duplication ensures that the node unit has unique DQ nodes + EXPECT_EQ(op_to_count["com.microsoft.QLinearSoftmax"], 1); + EXPECT_EQ(op_to_count["Softmax"], 0); + EXPECT_EQ(op_to_count["QuantizeLinear"], 1); + EXPECT_EQ(op_to_count["DequantizeLinear"], 2); // duplicate of first DQ and original second DQ }; TransformerTester(build_test_case, @@ -2947,12 +2952,68 @@ TEST(QDQTransformerTests, DQ_Produce_Graph_Output) { TransformerLevel::Level2, 12 /*opset_version*/, 0.01 /*per_sample_tolerance*/, - 0.01 /*relative_per_sample_tolerance*/, - std::make_unique(QDQIsInt8Allowed())); + 0.01 /*relative_per_sample_tolerance*/); }; test_case({1, 12, 37}, -1); } +// DQ produces graph output - special case for DropDQ path where there is only a DQ -> Node with no trailing Q +TEST(QDQTransformerTests, DropDQSelectorWithDQProducingGraphOutput) { + auto test_case = [&](const std::vector& input_shape, int64_t axis, bool dq_produces_graph_output) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput(input_shape, -5.f, 5.f); + auto* output_arg = builder.MakeOutput(); + + // add input QDQ + auto* input_q_output = builder.MakeIntermediate(); + auto* dq_output_arg = dq_produces_graph_output ? builder.MakeOutput() : builder.MakeIntermediate(); + + builder.AddQuantizeLinearNode(input_arg, .105f, 127, input_q_output); + builder.AddDequantizeLinearNode(input_q_output, .105f, 127, dq_output_arg); + + // add ArgMax + auto* argmax_output = builder.MakeIntermediate(); + auto& argmax_node = builder.AddNode("ArgMax", {dq_output_arg}, {argmax_output}); + argmax_node.AddAttribute("axis", axis); + + // add output Identity + builder.AddNode("Identity", {argmax_output}, {output_arg}); + }; + + auto check_graph = [&](InferenceSessionWrapper& session) { + const Graph& graph = session.GetGraph(); + + auto op_to_count = CountOpsInGraph(graph); + const auto expected_dq_count = + dq_produces_graph_output + ? 1 // EnsureUniqueDQForNodeUnit duplicates one DQ and DropDQ drops one DQ + : 0; // DropDQ drops one DQ + EXPECT_EQ(op_to_count["DequantizeLinear"], expected_dq_count); + + const auto& nodes = graph.Nodes(); + const auto argmax_node_it = std::find_if(nodes.cbegin(), + nodes.cend(), + [](const Node& node) { return node.OpType() == "ArgMax"; }); + ASSERT_NE(argmax_node_it, nodes.cend()); + + // the DQ from Q -> DQ -> ArgMax should have been dropped, look for the Q -> ArgMax edge + ASSERT_EQ(argmax_node_it->GetInputEdgesCount(), static_cast(1)); + EXPECT_EQ(argmax_node_it->InputEdgesBegin()->GetNode().OpType(), "QuantizeLinear"); + }; + + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 12 /*opset_version*/); + }; + + // test with and without the DQ producing a graph output to validate the test hits DropDQ + test_case({1, 4, 8}, -1, false); + test_case({1, 4, 8}, -1, true); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/transpose_optimizer_test.cc b/onnxruntime/test/optimizer/transpose_optimizer_test.cc index 10928dec34..f95033771b 100644 --- a/onnxruntime/test/optimizer/transpose_optimizer_test.cc +++ b/onnxruntime/test/optimizer/transpose_optimizer_test.cc @@ -14,7 +14,6 @@ #include "test/test_environment.h" #include "test/optimizer/graph_transform_test_builder.h" -#include "test/optimizer/qdq_test_utils.h" #include "test/providers/internal_testing/internal_testing_execution_provider.h" #include "test/util/include/asserts.h" #include "test/util/include/inference_session_wrapper.h" @@ -3735,13 +3734,22 @@ TEST(TransposeOptimizerTests, TestDequantizeLinearTransposePropagation) { }; auto check_graph = [&](InferenceSessionWrapper& session) { - std::vector expected_op_types_in_order{ - "DequantizeLinear", - "Transpose", - "Transpose"}; + const auto& graph = session.GetGraph(); - const auto op_types_in_order = GetNodeOpTypesInTopologicalOrder(session.GetGraph()); - EXPECT_EQ(op_types_in_order, expected_op_types_in_order); + const auto op_count = CountOpsInGraph(graph); + decltype(op_count) expected_op_count{ + {"DequantizeLinear", 2}, // EnsureUniqueDQForNodeUnit should duplicate the original DQ + {"Transpose", 2}, + }; + ASSERT_EQ(op_count, expected_op_count); + + // Transposes should be pushed, so check for Transpose -> DQ edges + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Transpose") { + ASSERT_EQ(node.GetOutputEdgesCount(), static_cast(1)); + ASSERT_EQ(node.OutputEdgesBegin()->GetNode().OpType(), "DequantizeLinear"); + } + } }; TransformerTester(build_test_case_1, diff --git a/onnxruntime/test/providers/nnapi/nnapi_basic_test.cc b/onnxruntime/test/providers/nnapi/nnapi_basic_test.cc index aa2f678942..6d2a6c2b43 100644 --- a/onnxruntime/test/providers/nnapi/nnapi_basic_test.cc +++ b/onnxruntime/test/providers/nnapi/nnapi_basic_test.cc @@ -327,13 +327,11 @@ TEST(NnapiExecutionProviderTest, TestQDQResizeNCHW) { // NNAPI EP does not support the default setting of Resize Op // Use bi-linear and asymmetric for NNAPI EP only auto Mode = ExpectedEPNodeAssignment::None; -#if defined(__ANDROID__) const auto* nnapi_handle = NnApiImplementation(); if (nnapi_handle && nnapi::GetNNAPIEffectiveFeatureLevelFromTargetDeviceOption( *nnapi_handle, nnapi::TargetDeviceOption::ALL_DEVICES) >= ANEURALNETWORKS_FEATURE_LEVEL_3) { Mode = ExpectedEPNodeAssignment::All; } -#endif RunQDQModelTest(BuildQDQResizeTestCase({1, 3, 64, 64} /* input_shape */, {1, 3, 32, 32} /* sizes_data */, "linear" /* mode */, diff --git a/onnxruntime/test/testdata/transform/qdq_with_multi_consumer_dq_nodes.fixed.txt b/onnxruntime/test/testdata/transform/qdq_with_multi_consumer_dq_nodes.fixed.txt index ec0e17dda0..11c28e4bec 100644 --- a/onnxruntime/test/testdata/transform/qdq_with_multi_consumer_dq_nodes.fixed.txt +++ b/onnxruntime/test/testdata/transform/qdq_with_multi_consumer_dq_nodes.fixed.txt @@ -1,6 +1,8 @@ -Model was created by running +Model was created by running -python -m onnxruntime.tools.qdq_helpers.optimize_qdq_model \onnxruntime\test\testdata\qdq_with_multi_consumer_dq_nodes.onnx \onnxruntime\test\testdata\transform -qdq_with_multi_consumer_dq_nodes.fixed.onnx +python -m onnxruntime.tools.qdq_helpers.optimize_qdq_model \onnxruntime\test\testdata\qdq_with_multi_consumer_dq_nodes.onnx \onnxruntime\test\testdata\transform\qdq_with_multi_consumer_dq_nodes.fixed.onnx -This results in a model that has duplicated DequantizeLinear (DQ) nodes so that each QDQ node group has no shared nodes. The CommonSubexpressionElimination optimizer could potentially combine these duplicated DQ nodes, which would break the QDQ handling. Due to this we ignore DQ nodes in that optimizer. \ No newline at end of file +This results in a model that has duplicated DequantizeLinear (DQ) nodes so that each QDQ node group has no shared nodes. The CommonSubexpressionElimination optimizer could potentially combine these duplicated DQ nodes, which would break the QDQ handling. Due to this we ignore DQ nodes in that optimizer. + +Note: The onnxruntime.tools.qdq_helpers.optimize_qdq_model tool duplicates DQ nodes from ORT 1.12 to 1.14. +In later versions, the DQ duplication is done by an ORT graph transformer. diff --git a/tools/python/util/qdq_helpers/optimize_qdq_model.py b/tools/python/util/qdq_helpers/optimize_qdq_model.py index 7691f407a8..0733db298f 100644 --- a/tools/python/util/qdq_helpers/optimize_qdq_model.py +++ b/tools/python/util/qdq_helpers/optimize_qdq_model.py @@ -8,16 +8,11 @@ import pathlib import onnx -from .qdq_model_utils import fix_dq_nodes_with_multiple_consumers - def optimize_qdq_model(): parser = argparse.ArgumentParser( os.path.basename(__file__), - description=""" - Update a QDQ format ONNX model to ensure optimal performance when executed using - ONNX Runtime. - """, + description="Update a QDQ format ONNX model to ensure optimal performance when executed using ONNX Runtime.", ) parser.add_argument("input_model", type=pathlib.Path, help="Provide path to ONNX model to update.") @@ -27,8 +22,13 @@ def optimize_qdq_model(): model = onnx.load(str(args.input_model.resolve(strict=True))) - # there's just one utility to run currently but we expect that will grow - fix_dq_nodes_with_multiple_consumers(model) + # run QDQ model optimizations here + + # Originally, the fixing up of DQ nodes with multiple consumers was implemented as one such optimization. + # That was moved to an ORT graph transformer. + print("As of ORT 1.15, the fixing up of DQ nodes with multiple consumers is done by an ORT graph transformer.") + + # There are no optimizations being run currently but we expect that there may be in the future. onnx.save(model, str(args.output_model.resolve())) diff --git a/tools/python/util/qdq_helpers/qdq_model_utils.py b/tools/python/util/qdq_helpers/qdq_model_utils.py deleted file mode 100644 index d15c9d65d2..0000000000 --- a/tools/python/util/qdq_helpers/qdq_model_utils.py +++ /dev/null @@ -1,108 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -import onnx - -from ..onnx_model_utils import get_producer_consumer_maps, iterate_graph_per_graph_func - - -def _duplicate_dq_nodes_with_multiple_consumers(graph: onnx.GraphProto, **kwargs): - updated_graphs = kwargs["updated_graphs"] - node_to_consumers = kwargs["node_to_consumers"] - validate_updates = kwargs["validate_updates"] - - nodes_to_update = [] - for node in filter(lambda node: node.op_type == "DequantizeLinear", graph.node): - # node providing graph output won't have consumer nodes - consumers = node_to_consumers[node] if node in node_to_consumers else [] - if len(consumers) > 1: - if not all(consumer in graph.node for consumer in consumers): - # TODO: If this does ever occur, as long as it's only consumed in one subgraph we could leave that - # value as is (no need to handle recursing into the subgraph) and update the consumers in this - # graph only - raise IndexError( - "DequantizeLinear node output is consumed by a subgraph. " "This is not currently supported." - ) - - nodes_to_update.append(node) - - if validate_updates: - if nodes_to_update: - # internal error. we somehow missed an update in the first pass when validate_upates was false - raise ValueError("Graph still has DequantizeLinear nodes with multiple consumers.") - - return - - if nodes_to_update: - dup_idx = 0 - new_graph = onnx.GraphProto() - graph_outputs = {output.name for output in graph.output} - for node in graph.node: - new_graph.node.append(node) - if node in nodes_to_update: - is_graph_output = node.output[0] in graph_outputs - # create duplicate DQ nodes as needed so that there is one consumer per node. - # this allows us to cleanly create a QDQ node group with no DQ nodes shared with other QDQ node groups. - # if the node produces a graph output we need a duplicate DQ node for every consumer node. - # if not, we can leave the first consumer as is and create duplicate nodes for the other consumers. - start_idx = 0 if is_graph_output else 1 - consumers = list(node_to_consumers[node])[start_idx:] - - for idx, consumer in enumerate(consumers): - # create duplicate DQ node - duplicate = onnx.NodeProto() - duplicate.CopyFrom(node) - # update node name for debugging. use the global dup idx for node duplication - duplicate.name += f"/qdq_utils_dup_{dup_idx}" - - # update output. use the local idx for value duplication - orig_output = node.output[0] - new_output = f"{orig_output}/qdq_utils_dup_{idx}" - duplicate.output[0] = new_output - - # update input on the consumer node. - for input_idx, input_name in enumerate(consumer.input): - if input_name == orig_output: - consumer.input[input_idx] = new_output - - new_graph.node.append(duplicate) - dup_idx += 1 - - # replace nodes - del graph.node[:] - graph.node.extend(new_graph.node) - updated_graphs.append(graph) - - -def fix_dq_nodes_with_multiple_consumers(model): - """ - Update a model if any DequantizeLinear nodes have multiple consumers. - The QDQ node unit processing is overly complicated if this is the case, as the DQ node would be in multiple units, - and the units may end up in different partitions at runtime. - :param model: QDQ model to update - """ - node_to_producers, node_to_consumers = get_producer_consumer_maps(model.graph) - - updated_graphs = [] # list of GraphProto instances that were updated_graphs - iterate_graph_per_graph_func( - model.graph, - _duplicate_dq_nodes_with_multiple_consumers, - node_to_consumers=node_to_consumers, - validate_updates=False, - updated_graphs=updated_graphs, - ) - - if updated_graphs: - updated_graphs = [] - node_to_producers, node_to_consumers = get_producer_consumer_maps(model.graph) - iterate_graph_per_graph_func( - model.graph, - _duplicate_dq_nodes_with_multiple_consumers, - node_to_consumers=node_to_consumers, - validate_updates=True, - updated_graphs=updated_graphs, - ) - - # validate with check and by running shape inference. - onnx.checker.check_model(model) - _ = onnx.shape_inference.infer_shapes(model) diff --git a/tools/python/util/qdq_helpers/test/__init__.py b/tools/python/util/qdq_helpers/test/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tools/python/util/qdq_helpers/test/test_qdq_model_utils.py b/tools/python/util/qdq_helpers/test/test_qdq_model_utils.py deleted file mode 100644 index d7cd7ac200..0000000000 --- a/tools/python/util/qdq_helpers/test/test_qdq_model_utils.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -import pathlib -import unittest - -import onnx - -from ..qdq_model_utils import fix_dq_nodes_with_multiple_consumers - -script_dir = pathlib.Path(__file__).parent -ort_root = script_dir.parents[4] - -# example usage from /tools/python -# python -m unittest util/qdq_helpers/test/test_qdq_model_utils.py -# NOTE: at least on Windows you must use that as the working directory for all the imports to be happy - - -class TestQDQUtils(unittest.TestCase): - def test_fix_DQ_with_multiple_consumers(self): # noqa: N802 - """ """ - model_path = ort_root / "onnxruntime" / "test" / "testdata" / "qdq_with_multi_consumer_dq_nodes.onnx" - model = onnx.load(str(model_path)) - - orig_dq_nodes = [n for n in model.graph.node if n.op_type == "DequantizeLinear"] - fix_dq_nodes_with_multiple_consumers(model) - new_dq_nodes = [n for n in model.graph.node if n.op_type == "DequantizeLinear"] - - # there are 3 DQ nodes with 2 consumers (an earlier Conv and later Add) - # additionally the last one also provides a graph output - # based on that there should be 3 new DQ nodes for the internal consumers and 1 new one for the graph output - self.assertEqual(len(orig_dq_nodes) + 4, len(new_dq_nodes))