diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 19c43df7c6..f956d4f6fe 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -185,13 +185,15 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique()); transformers.emplace_back(std::make_unique( session_options.free_dimension_overrides)); - auto cpu_allocator = cpu_execution_provider.GetAllocator(0, OrtMemTypeDefault); - transformers.emplace_back(std::make_unique(std::move(cpu_allocator))); if (!disable_quant_qdq) { transformers.emplace_back(std::make_unique()); } + // 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(0, OrtMemTypeDefault); + transformers.emplace_back(std::make_unique(std::move(cpu_allocator))); } break; case TransformerLevel::Level2: { diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index cd0b4c4d11..be5e7aa39f 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -69,7 +69,7 @@ Status NhwcTransformer::ApplyImpl(Graph& graph, bool& modified, int graph_level, } if (modified) { - Optimize(*api_graph, /*allow_extended_ops*/ true); + Optimize(*api_graph, /*allow_extended_ops*/ true, kCpuExecutionProvider); } return Status::OK(); diff --git a/onnxruntime/core/optimizer/transpose_optimizer/ort_transpose_optimizer.h b/onnxruntime/core/optimizer/transpose_optimizer/ort_transpose_optimizer.h index 9535e23a18..0d8a82cee9 100644 --- a/onnxruntime/core/optimizer/transpose_optimizer/ort_transpose_optimizer.h +++ b/onnxruntime/core/optimizer/transpose_optimizer/ort_transpose_optimizer.h @@ -21,6 +21,13 @@ class TransposeOptimizer : public GraphTransformer { : GraphTransformer("TransposeOptimizer"), cpu_allocator_(std::move(cpu_allocator)) {} Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; + + // One run should be sufficient. + // The second phase of optimization may swap a DequantizeLinear -> Transpose back, so multiple runs would + // keep swapping the order of the nodes in the first and second phases, leading to always returning true for + // modified. + // see https://github.com/microsoft/onnxruntime/blob/e3a2d5cca8bcefe064f83d57e46ea51ddb2b16e8/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc#L1917-L1921 + bool ShouldOnlyApplyOnce() const override { return true; } }; } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc b/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc index c082b583fc..f60e39a638 100644 --- a/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc +++ b/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc @@ -1876,12 +1876,31 @@ OptimizeResult OptimizeImpl(OptimizerCtx& ctx) { } bool changed = false; + bool have_dq = false; + // if nodes are assigned we only process those that match the EP in the context + bool ignore_assigned_nodes = ctx.provider_type.empty(); + // Optimize graph. Nodes will be modified during iteration, but nodes are never deleted before we reach them. // New transpose nodes are inserted, but always as an input to an existing node. for (size_t i = 0; i < nodes.size(); ++i) { api::NodeRef& node = *nodes[i]; + if (node.OpType() == "DequantizeLinear") { + have_dq = true; + } + + // it's not clear how we should handle assignment of new Transpose nodes created during optimization, so ignore. + // e.g. we may need to transpose the input of a node we move a Transpose past. if that node is assigned to + // an EP that doesn't support Transpose, the new node should use the CPU EP. but if that node is assigned to + // an EP that does support the Transpose we should assign the new node to that EP. + // as we do not know what each EP supports, it's safer to not optimize in order to maintain the EP assignments + // made during partitioning. + if (ignore_assigned_nodes && !node.GetExecutionProviderType().empty()) { + continue; + } + if (ctx.mode == OptimizerMode::OPTIMIZE_LAYOUT_TRANSFORM && - ctx.layout_sensitive_ops.count(node.OpType()) && node.GetExecutionProviderType() != ctx.provider_type) { + ctx.layout_sensitive_ops.count(node.OpType()) && + node.GetExecutionProviderType() != ctx.provider_type) { // If the current op is layout sensitive and it is not assigned to the given provider // then do not process transpose. continue; @@ -1907,21 +1926,27 @@ OptimizeResult OptimizeImpl(OptimizerCtx& ctx) { } } - // Currently limiting the second optimization pass to layout transform mode - // TODO: Enable this for both the modes. - if (ctx.mode == OptimizerMode::OPTIMIZE_TRANSPOSE) { + if (!have_dq) { result.graph_modified = changed; return result; } // Run second optimization pass. - // If any transpose succeeds a DQ node, move it above the DQ node. - // In case of QDQ models this helps to preserve the QDQ node unit + // If any transpose succeeds a DQ node, move it above the DQ node if it's not part of a QDQ node group. + // In QDQ models this helps to preserve the QDQ node group when a Transpose was pushed across a DQ into + // an existing QDQ node group. // In all other scenarios this is beneficial as well because moving transpose above DQ node is more efficient as // transpose node now handles less data. auto graph_nodes = ctx.graph.Nodes(); for (size_t i = 1; i < graph_nodes.size(); i++) { - if (graph_nodes[i]->OpType() == "Transpose") { + const auto& node = *graph_nodes[i]; + + // TODO: if we want to handle this we need to propagate the assigned EP to the new Transpose node. + if (ignore_assigned_nodes && !node.GetExecutionProviderType().empty()) { + continue; + } + + if (node.OpType() == "Transpose") { auto& transpose_node = *graph_nodes[i]; auto dq_node = ctx.graph.GetNodeProducingOutput(transpose_node.Inputs()[0]); if (!dq_node || dq_node->OpType() != "DequantizeLinear") { @@ -1929,11 +1954,11 @@ OptimizeResult OptimizeImpl(OptimizerCtx& ctx) { } auto consumers = ctx.graph.GetValueConsumers(transpose_node.Outputs()[0]); - bool is_part_of_qdq_unit = std::find_if(consumers->nodes.cbegin(), consumers->nodes.cend(), - [](const std::unique_ptr& node) { - return node->OpType() == "QuantizeLinear"; - }) != consumers->nodes.cend(); - if (is_part_of_qdq_unit) { + bool is_part_of_qdq_group = std::find_if(consumers->nodes.cbegin(), consumers->nodes.cend(), + [](const std::unique_ptr& node) { + return node->OpType() == "QuantizeLinear"; + }) != consumers->nodes.cend(); + if (is_part_of_qdq_group) { continue; } @@ -1942,9 +1967,11 @@ OptimizeResult OptimizeImpl(OptimizerCtx& ctx) { if (!perm.has_value()) { continue; } + if (!HandleQuantizeDequantizeScale(ctx.graph, *perm, *dq_node, ctx.opset)) { continue; } + TransposeFirstInput(ctx, *dq_node, *perm); // remove existing transpose node diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index e09d6947a9..6aaaa28c95 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -920,13 +920,13 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, SessionState& session_state, bool saving_model_in_ort_format) { // The transformer order: - // 1. built-in graph rewriter - // 2. each execution provider's transformer - // 3. do node placement according to kernel definition - // 4. insert copy nodes - // 5. insert cast nodes. + // 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 - // first apply global(execution provider independent), level 1(default/system/basic) graph to graph optimizations + // first apply execution provider independent level 1 graph optimizations. ORT_RETURN_IF_ERROR_SESSIONID_( graph_transformer_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *session_logger_)); @@ -954,15 +954,15 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, auto mode = saving_model_in_ort_format ? GraphPartitioner::Mode::kAssignOnly : GraphPartitioner::Mode::kNormal; - // Do partitioning based on execution providers' capability. + // Do partitioning based on execution providers' capabilities. GraphPartitioner partitioner(kernel_registry_manager, providers); ORT_RETURN_IF_ERROR_SESSIONID_(partitioner.Partition(graph, session_state.GetMutableFuncMgr(), layout_transformer::TransformLayoutForCompilingEP, mode)); - // apply transformers except default transformers - // Default transformers are required for correctness and they are owned and run by inference session - for (int i = static_cast(TransformerLevel::Level1); i <= static_cast(TransformerLevel::MaxLevel); i++) { + // apply Level2 and higher transformers. + // we do not run Level 1 again as those transformers assume partitioning will run later to do node assignment. + for (int i = static_cast(TransformerLevel::Level2); i <= static_cast(TransformerLevel::MaxLevel); i++) { ORT_RETURN_IF_ERROR_SESSIONID_( graph_transformer_mgr.ApplyTransformers(graph, static_cast(i), *session_logger_)); } diff --git a/onnxruntime/test/optimizer/graph_transform_test_builder.cc b/onnxruntime/test/optimizer/graph_transform_test_builder.cc index 52f4aa372d..28a5e01af9 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_builder.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_builder.cc @@ -25,7 +25,8 @@ void TransformerTester(const std::function& buil double per_sample_tolerance, double relative_per_sample_tolerance, std::unique_ptr transformer, - const std::function& add_session_options) { + const std::function& add_session_options, + const InlinedHashSet& disabled_optimizers) { // Build the model for this test. std::unordered_map domain_to_version; domain_to_version[kOnnxDomain] = opset_version; @@ -58,6 +59,8 @@ void TransformerTester(const std::function& buil ASSERT_STATUS_OK(session.Load(model_data.data(), static_cast(model_data.size()))); if (transformer) { ASSERT_STATUS_OK(session.RegisterGraphTransformer(std::move(transformer), level)); + } else if (!disabled_optimizers.empty()) { + ASSERT_STATUS_OK(session.FilterEnabledOptimizers(InlinedHashSet{disabled_optimizers})); } ASSERT_STATUS_OK(session.Initialize()); diff --git a/onnxruntime/test/optimizer/graph_transform_test_builder.h b/onnxruntime/test/optimizer/graph_transform_test_builder.h index 41aee08ae6..a721a940d6 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_builder.h +++ b/onnxruntime/test/optimizer/graph_transform_test_builder.h @@ -282,7 +282,8 @@ void TransformerTester(const std::function& buil double per_sample_tolerance = 0.0, double relative_per_sample_tolerance = 0.0, std::unique_ptr transformer = nullptr, - const std::function& add_session_options = {}); + const std::function& add_session_options = {}, + const InlinedHashSet& disabled_optimizers = {}); } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc index 6e878ee88d..787cbc0d70 100644 --- a/onnxruntime/test/optimizer/qdq_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc @@ -1622,7 +1622,9 @@ TEST(QDQTransformerTests, ConvTranspose_DQForward) { TransformerTester(build_test_case, check_graph, TransformerLevel::Level1, - TransformerLevel::Level2); + TransformerLevel::Level2, + 12, 0.0, 0.0, nullptr, {}, // defaults that we're not overriding + {"TransposeOptimizer"}); // disable TransposeOptimizer for simplicity }; test_case({1, 13, 13, 23}, {30, 23, 3, 3}, {0, 3, 1, 2}); @@ -1695,7 +1697,9 @@ TEST(QDQTransformerTests, DQForward_MutilpleSteps) { check_graph, TransformerLevel::Level1, TransformerLevel::Level2, - 13 /*opset_version*/); + 13 /*opset_version*/, + 0.0, 0.0, nullptr, {}, // defaults that we're not overriding + {"TransposeOptimizer"}); // disable TransposeOptimizer for simplicity }; test_case({1, 13, 13, 23}, {30, 23, 3, 3}, {0, 3, 1, 2}); @@ -1957,7 +1961,9 @@ TEST(QDQTransformerTests, QDQPropagation_DQForward) { TransformerTester(build_test_case, check_graph, TransformerLevel::Default, - TransformerLevel::Level1); + TransformerLevel::Level1, + 12, 0.0, 0.0, nullptr, {}, // defaults that we're not overriding + {"TransposeOptimizer"}); // disable TransposeOptimizer for simplicity }; test_case({1, 13, 13, 23}, 4, {0, 3, 1, 2}, false, false); @@ -2103,7 +2109,9 @@ TEST(QDQTransformerTests, QDQPropagation_Per_Layer_No_Propagation) { TransformerTester(build_test_case, check_graph, TransformerLevel::Default, - TransformerLevel::Level1); + TransformerLevel::Level1, + 12, 0.0, 0.0, nullptr, {}, // defaults that we're not overriding + {"TransposeOptimizer"}); // disable TransposeOptimizer for simplicity }; test_case({1, 13, 13, 23}, {0, 2, 3, 1}); @@ -2266,6 +2274,63 @@ TEST(QDQTransformerTests, QDQ_Selector_Test) { } } +// regression test to validate TransposeOptimizer and QDQ Propagation don't loop +// see https://github.com/microsoft/onnxruntime/issues/11605 +TEST(QDQTransformerTests, QDQPropagation_GH11605) { + auto test_case = [&]() { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput({1, 4, 4}, + std::numeric_limits::min(), + std::numeric_limits::max()); + // add DQ + auto* dq_output = builder.MakeIntermediate(); + builder.AddDequantizeLinearNode(input_arg, 0.123f, uint8_t(0), dq_output); + + // add Transpose 0, 2, 1 + const std::vector& perms{0, 2, 1}; + auto* transpose_output = builder.MakeIntermediate(); + Node& transpose_node = builder.AddNode("Transpose", {dq_output}, {transpose_output}); + transpose_node.AddAttribute("perm", perms); + + // add Softmax with axis=2 (to block the Transpose moving past it due to the transpose perms) + auto* softmax_output = builder.MakeIntermediate(); + Node& softmax_node = builder.AddNode("Softmax", {transpose_output}, {softmax_output}); + softmax_node.AddAttribute("axis", int64_t(2)); + + // add second Transpose. this is so the check in TransposeOptimizer::ProcessTranspose for outputs leading to + // a Transpose is satisfied, allowing the first Transpose to move past the Q/DQ inserted by QDQ Propagation + Node& transpose_node2 = builder.AddNode("Transpose", {softmax_output}, {builder.MakeOutput()}); + transpose_node2.AddAttribute("perm", perms); + }; + + // check that an edge case where transpose optimization gets blocked is handled gracefully. + // Original: DQ -> Tr -> SoftM -> Tr + // QDQ Prop inserts a Q/DQ pair to create a QDQ node group for the Transpose: DQ -> Tr -> Q -> DQ -> SoftM -> Tr + // Transpose opt phase 1 moves the Tr down until it blocks on the SoftMax: DQ -> Q -> DQ -> Tr -> SoftM -> Tr + // Transpose opt phase 2 flips the Tr to prior to the DQ as it's not part of a QDQ node group at that point, as + // running the transpose on 8-bit data should be cheaper: DQ -> Q -> Tr -> DQ -> SoftM -> Tr + // QDQ cleanup in Level2 removes the unnecessary DQ/Q pair at the start: Tr -> DQ -> SoftM -> Tr + // this is the optimal result as the Transpose is using 8-bit data and we have no surplus Q/DQ pairs + auto check_graph = [&](InferenceSessionWrapper& session) { + std::vector expected_op_types_in_order{ + "Transpose", + "DequantizeLinear", + "Softmax", + "Transpose"}; + + const auto op_types_in_order = GetNodeOpTypesInTopologicalOrder(session.GetGraph()); + EXPECT_EQ(op_types_in_order, expected_op_types_in_order); + }; + + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Default, + TransformerLevel::Level2); + }; + + test_case(); +} + // test removal of Q->DQ pairs by QDQFinalCleanupTransformer TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicQDQCleanup) { auto test_case = [&](const std::vector>& input_shapes,