saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-06 04:28:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix clash between QDQ propagation and TransposeOptimizer (#11636)

Browse source 
* Initial changes with comments on potential unit test changes.

* Update tests to disable TransposeOptimizer as that's simpler.
Add some extra comments.
Cleanup.

* Update comments in TransformGraph

* Add regression test.
Add limitation that transpose optimizer will ignore assigned nodes that do not match the context EP if that is set.

* Fix test. I removed a trailing Transpose after initial validation to simplify but that changed things so that the transpose optimizer didn't kick in, and the DQ -> Transpose -> Q was actually converted to a single Transpose by the CPU EP QDQ handling. Same end result in most builds so the subtle difference wasn't noticed, but in a build without contrib ops the CPU EP QDQ handling is disabled so the end result was different.

Update the test to re-instate the trailing Transpose so transpose optimizer alters the graph as desired.

* Don't run level 1 optimizers after partitioning as they don't guarantee to handle EP assignment for new nodes they create.
This commit is contained in:
Scott McKay 2022-06-04 09:16:35 +10:00 committed by GitHub
parent 95a16c1ffe
commit ef64b2ee52
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 136 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/core/optimizer/graph_transformer_utils.cc
View file

@ -185,13 +185,15 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
transformers.emplace_back(std::make_unique<ReshapeFusion>());
transformers.emplace_back(std::make_unique<FreeDimensionOverrideTransformer>(
session_options.free_dimension_overrides));
auto cpu_allocator = cpu_execution_provider.GetAllocator(0, OrtMemTypeDefault);
transformers.emplace_back(std::make_unique<TransposeOptimizer>(std::move(cpu_allocator)));
if (!disable_quant_qdq) {
transformers.emplace_back(std::make_unique<QDQPropagationTransformer>());
}
// 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<TransposeOptimizer>(std::move(cpu_allocator)));
} break;
case TransformerLevel::Level2: {

2
onnxruntime/core/optimizer/nhwc_transformer.cc
View file

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

7
onnxruntime/core/optimizer/transpose_optimizer/ort_transpose_optimizer.h
View file

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

51
onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc
View file

@ -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<api::NodeRef>& 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<api::NodeRef>& 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

20
onnxruntime/core/session/inference_session.cc
View file

@ -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<int>(TransformerLevel::Level1); i <= static_cast<int>(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<int>(TransformerLevel::Level2); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
ORT_RETURN_IF_ERROR_SESSIONID_(
graph_transformer_mgr.ApplyTransformers(graph, static_cast<TransformerLevel>(i), *session_logger_));
}

5
onnxruntime/test/optimizer/graph_transform_test_builder.cc
View file

@ -25,7 +25,8 @@ void TransformerTester(const std::function<void(ModelTestBuilder& helper)>& buil
double per_sample_tolerance,
double relative_per_sample_tolerance,
std::unique_ptr<GraphTransformer> transformer,
const std::function<void(SessionOptions&)>& add_session_options) {
const std::function<void(SessionOptions&)>& add_session_options,
const InlinedHashSet<std::string>& disabled_optimizers) {
// Build the model for this test.
std::unordered_map<std::string, int> domain_to_version;
domain_to_version[kOnnxDomain] = opset_version;
@ -58,6 +59,8 @@ void TransformerTester(const std::function<void(ModelTestBuilder& helper)>& buil
ASSERT_STATUS_OK(session.Load(model_data.data(), static_cast<int>(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<std::string>{disabled_optimizers}));
}
ASSERT_STATUS_OK(session.Initialize());

3
onnxruntime/test/optimizer/graph_transform_test_builder.h
View file

@ -282,7 +282,8 @@ void TransformerTester(const std::function<void(ModelTestBuilder& helper)>& buil
double per_sample_tolerance = 0.0,
double relative_per_sample_tolerance = 0.0,
std::unique_ptr<GraphTransformer> transformer = nullptr,
const std::function<void(SessionOptions&)>& add_session_options = {});
const std::function<void(SessionOptions&)>& add_session_options = {},
const InlinedHashSet<std::string>& disabled_optimizers = {});
} // namespace test
} // namespace onnxruntime

73
onnxruntime/test/optimizer/qdq_transformer_test.cc
View file

@ -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<uint8_t>({1, 4, 4},
std::numeric_limits<uint8_t>::min(),
std::numeric_limits<uint8_t>::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<int64_t>& 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<std::string> 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<std::vector<int64_t>>& input_shapes,