saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make transpose optimizer able to look past DQ node for const initializer (#17618)

Browse source 
### Description
<!-- Describe your changes. -->
Add ability for transpose optimizer to look past a DQ node if it has a
constant initializer as input. This allows UnsqueezeInput/TransposeInput
to modify the initializer in-place in the same way it would for a
non-QDQ format model.

Shared initializers are also handled, and any additional
Squeeze/Transpose added to the other usages of the initializer should
cancel out when we push the same Transpose though them.

The in-place modification means we don't need to run QDQ fixup and
constant folding after layout transformation. This means we do not need
to enable those optimizers in a minimal build to get an optimal model
post-layout transformation.


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
Ensure layout transformation produces optimal model in full and minimal
builds.

---------

Co-authored-by: Edward Chen <18449977+edgchen1@users.noreply.github.com>
This commit is contained in:
Scott McKay 2023-09-27 21:17:16 +10:00 committed by GitHub
parent 33295ed883
commit a99c965d05
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 608 additions and 94 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

452
onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc
View file

@ -4,9 +4,11 @@
#include "onnx_transpose_optimization.h"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "core/common/gsl.h"
@ -93,6 +95,55 @@ static std::unique_ptr<api::NodeRef> MakeSqueezeOrUnsqueeze(int64_t opset, api::
return graph.AddNode(op_type, inputs, /*num_outputs*/ 1);
}
/// <summary>
/// Return a DequantizeLinear node if it's input is a constant initializer with known consumers.
/// In this case the initializer can be updated in-place by UnsqueezeInput or TransposeInput.
/// </summary>
/// <param name="graph">Current graph</param>
/// <param name="dq_output_name">Value to check if produced by a DQ node who's input is a constant initializer</param>
/// <returns>NodeRef for DQ node if it meets the requirements.</returns>
static std::unique_ptr<api::NodeRef> GetDQWithConstInitializerInput(const api::GraphRef& graph,
std::string_view dq_output_name) {
std::unique_ptr<api::NodeRef> dq_node;
auto maybe_dq_node = graph.GetNodeProducingOutput(dq_output_name);
if (maybe_dq_node && maybe_dq_node->OpType() == "DequantizeLinear") {
do {
auto dq_input = maybe_dq_node->Inputs()[0];
auto dq_constant = graph.GetConstant(dq_input);
// input to DQ must be a constant initializer
if (!dq_constant) {
break;
}
// For now keep it simple and don't support per-axis quantization as that would require updating the
// scale and zero point values in the DQ node to re-order if transposing, or reshape if unsqueezing.
// the rank of the `scale` and `zero point` inputs must match so we only need to check `scale`.
auto dq_scale = graph.GetConstant(maybe_dq_node->Inputs()[1]);
if (!dq_scale || dq_scale->NumElements() != 1) {
break;
}
// need to know all the initializer consumers as we're potentially going to modify it directly
auto initializer_consumers = graph.GetValueConsumers(dq_input);
if (!initializer_consumers->comprehensive) {
break;
}
// DQ output is only used by the node we're modifying.
auto dq_consumers = graph.GetValueConsumers(dq_output_name);
if (!dq_consumers->comprehensive || dq_consumers->nodes.size() != 1) {
break;
}
dq_node = std::move(maybe_dq_node);
} while (false);
}
return dq_node;
}
// Returns whether perm is a valid permutation (contains each value from 0 to perm.size() - 1 exactly once)
static bool IsValidPerm(const std::vector<int64_t>& perm) {
size_t rank = perm.size();
@ -345,6 +396,12 @@ static std::vector<int64_t> SortedAxesForTransposedInput(const std::vector<int64
return new_axes;
}
static void UpdateDQNodeInputAndShape(api::GraphRef& graph, api::NodeRef& dq, std::string_view new_input) {
dq.SetInput(0, new_input);
auto new_shape = *graph.GetValueInfo(new_input)->Shape();
graph.GetValueInfo(dq.Outputs()[0])->SetShape(&new_shape);
}
/////// </Helper Utils> ///////
/////// <Core Helpers> ///////
@ -357,51 +414,126 @@ static std::string_view HelpHandleUnsqueeze(HandlerArgs& args, const std::vector
// broadcasting.
static void UnsqueezeInput(OptimizerCtx& ctx, api::NodeRef& node, size_t i, const std::vector<int64_t>& axes) {
std::string_view input = node.Inputs()[i];
// Remove this node as a consumer
node.SetInput(i, "");
std::unique_ptr<api::TensorRef> constant = ctx.graph.GetLocalConstant(input);
auto consumers = ctx.graph.GetValueConsumers(input);
// allow a constant initializer coming via a DQ node with a single consumer
std::unique_ptr<api::NodeRef> dq_node;
std::string_view constant_dq_input;
if (!constant) {
// look past a DQ node for a constant initializer. essentially we pretend the DQ node doesn't exist
// to enable directly making changes to the initializer. any nodes added for other consumers of the initializer
// in 'Case 1' are prior to the DQ so we don't break up any QDQ node units.
dq_node = GetDQWithConstInitializerInput(ctx.graph, input);
if (dq_node) {
// underlying string for the input name is in the Node so it's safe to store in string_view constant_dq_input
constant_dq_input = dq_node->Inputs()[0];
constant = ctx.graph.GetLocalConstant(constant_dq_input);
// remove the DQ node as a consumer of the initializer while we modify things
dq_node->SetInput(0, "");
}
}
// Clear the input, which also removes this node's input as a consumer of the value.
// NOTE: the node may have multiple inputs consuming the value.
node.SetInput(i, "");
auto value_to_modify = dq_node ? constant_dq_input : input;
auto consumers = ctx.graph.GetValueConsumers(value_to_modify);
// Case 1: input is a constant with a known list of consumer nodes
if (constant != nullptr && consumers->comprehensive) {
// We will reshape the initializer. If there are existing consumers, still reshape it but add Squeeze nodes
// We will reshape the initializer. If there are existing consumers, reshape it and add Squeeze nodes
// to counteract its effect. If they later Unsqueeze the same input, the Squeeze nodes will simply be deleted
// (see Case 2).
if (consumers->nodes.size() > 0) {
auto squeeze_ptr = MakeSqueezeOrUnsqueeze(ctx.opset, ctx.graph, "Squeeze", input, axes);
// record the consumer node input as being special cased for use in Case 2 if a DQ node, and IsConstant
for (auto& consumer : consumers->nodes) {
auto& consumer_node_inputs = ctx.nodes_using_updated_shared_initializer[consumer->Id()];
// find input id/s for consumer
auto consumer_inputs = consumer->Inputs();
for (size_t input_idx = 0; input_idx < consumer_inputs.size(); ++input_idx) {
if (consumer_inputs[input_idx] == value_to_modify) {
consumer_node_inputs.push_back(input_idx);
}
}
}
auto squeeze_ptr = MakeSqueezeOrUnsqueeze(ctx.opset, ctx.graph, "Squeeze", value_to_modify, axes);
api::NodeRef& squeeze = *squeeze_ptr;
std::string_view sq_out = squeeze.Outputs()[0];
ctx.graph.CopyValueInfo(input, sq_out);
ReplaceValueReferences(consumers->nodes, input, sq_out);
ctx.graph.CopyValueInfo(value_to_modify, sq_out);
ReplaceValueReferences(consumers->nodes, value_to_modify, sq_out);
}
auto new_shape = UnsqueezeShape(constant->Shape(), axes);
ctx.graph.ReshapeInitializer(input, new_shape);
node.SetInput(i, input);
ctx.graph.ReshapeInitializer(value_to_modify, new_shape);
if (dq_node) {
UpdateDQNodeInputAndShape(ctx.graph, *dq_node, constant_dq_input);
}
node.SetInput(i, input); // restore the original connection
return;
}
// Case 2: input is a Squeeze node with matching axes
std::unique_ptr<api::NodeRef> inp_node = ctx.graph.GetNodeProducingOutput(input);
// check if this is a special-cased DQ node where we put the Squeeze on input 0 of the DQ in 'Case 1' above
if (inp_node && inp_node->OpType() == "DequantizeLinear" &&
std::find_if(ctx.nodes_using_updated_shared_initializer.begin(),
ctx.nodes_using_updated_shared_initializer.end(),
[&inp_node](const auto& entry) {
const auto id = entry.first;
const auto& input_idxs = entry.second;
// check Id matches and the entry was for input 0 of the DQ node
return id == inp_node->Id() &&
std::find(input_idxs.begin(), input_idxs.end(), size_t(0)) != input_idxs.end();
}) != ctx.nodes_using_updated_shared_initializer.end()) {
// set things up so we can look past the DQ node to the Squeeze that was inserted in front of the reshaped
// constant initializer that was shared with this node.
dq_node = std::move(inp_node);
auto dq_input = dq_node->Inputs()[0];
inp_node = ctx.graph.GetNodeProducingOutput(dq_input);
consumers = ctx.graph.GetValueConsumers(dq_input);
}
if (inp_node != nullptr && inp_node->IsOp("Squeeze")) {
const std::vector<std::string_view>& inp_node_inputs = inp_node->Inputs();
std::optional<std::vector<int64_t>> squeeze_axes = std::nullopt;
squeeze_axes = ReadFromAttrOrInput(ctx, *inp_node, "axes", /*inp_index*/ 1, /*opset*/ 13);
if (squeeze_axes != std::nullopt && *squeeze_axes == axes) {
if (dq_node) {
UpdateDQNodeInputAndShape(ctx.graph, *dq_node, inp_node_inputs[0]);
node.SetInput(i, dq_node->Outputs()[0]);
} else {
node.SetInput(i, inp_node_inputs[0]);
}
// Remove the Squeeze node if possible
if (consumers->comprehensive && consumers->nodes.size() == 0) {
// if there's a DQ node the `consumers` list still includes it so allow for that.
// in that case UpdateDQNodeInputAndShape already updated the input of the DQ node so it's safe to remove it.
if (consumers->comprehensive && consumers->nodes.size() == size_t(dq_node ? 1 : 0)) {
ctx.graph.RemoveNode(*inp_node);
if (ctx.opset >= 13 && !ctx.graph.HasValueConsumers(inp_node_inputs[1])) {
ctx.graph.RemoveInitializer(inp_node_inputs[1]);
}
}
node.SetInput(i, inp_node_inputs[0]);
return;
}
// Axes don't match. Fall through to Case 3.
}
// any DQ node special casing doesn't apply anymore, so go back to the original inp_node
if (dq_node) {
inp_node = std::move(dq_node);
}
// Case 3: Add an Unsqueeze node.
auto unsqueeze_ptr = MakeSqueezeOrUnsqueeze(ctx.opset, ctx.graph, "Unsqueeze", input, axes);
api::NodeRef& unsqueeze = *unsqueeze_ptr;
@ -453,83 +585,197 @@ static void Permute1DConstant(api::GraphRef& graph, api::NodeRef& node, api::Ten
// Replaces ith input to node with transposed value. Might create a new Transpose node, find an existing one,
// or transpose an initializer.
void TransposeInput(api::GraphRef& graph, api::NodeRef& node, size_t i,
const std::vector<int64_t>& perm, const std::vector<int64_t>& perm_inv) {
static void TransposeInputImpl(api::GraphRef& graph,
NodeIdToInputIdxsMap* nodes_using_updated_shared_initializer,
api::NodeRef& node, size_t i, const std::vector<int64_t>& perm,
const std::vector<int64_t>& perm_inv) {
std::string_view input = node.Inputs()[i];
// Remove this node as a consumer
node.SetInput(i, "");
// Only local constants are editable
std::unique_ptr<api::TensorRef> constant = graph.GetLocalConstant(input);
auto consumers = graph.GetValueConsumers(input);
// allow a constant initializer coming via a DQ node with a single consumer
std::unique_ptr<api::NodeRef> dq_node;
std::string_view constant_dq_input;
if (!constant) {
// look past a DQ node for a constant initializer. essentially we pretend the DQ node doesn't exist
// to enable directly making changes to the initializer. any nodes added for other consumers of the initializer
// in 'Case 1' are prior to the DQ so we don't break up any QDQ node units.
dq_node = GetDQWithConstInitializerInput(graph, input);
if (dq_node) {
// underlying string for the input name is in the Node so it's safe to store in string_view constant_dq_input
constant_dq_input = dq_node->Inputs()[0];
constant = graph.GetLocalConstant(constant_dq_input);
// remove the DQ node as a consumer of the initializer while we modify things
dq_node->SetInput(0, "");
}
}
// Clear the input, which also removes this node's input as a consumer of the value.
// NOTE: the node may have multiple inputs consuming the value.
node.SetInput(i, "");
auto constant_to_modify = dq_node ? constant_dq_input : input;
auto consumers = graph.GetValueConsumers(constant_to_modify);
// Case 1: input is a constant with a known list of consumer nodes
if (constant != nullptr && consumers->comprehensive) {
// Input is scalar, return early.
if (constant->Shape().size() == 1 && constant->Shape()[0] == 0) {
// we modify the initializer in-place and need to reconnect things up when we're done. this helper will
// do that when it goes out of scope. if we have manually reconnected, input or constant_dq_input is
// set to an empty string.
auto reconnect_nodes = gsl::finally([i, &node, &dq_node, &input, &constant_dq_input] {
if (!input.empty()) {
node.SetInput(i, input);
}
if (!constant_dq_input.empty()) {
dq_node->SetInput(0, constant_dq_input);
}
});
// If there is only one element return early as the transpose won't change the data
if (constant->NumElements() == 1) {
return;
}
// This is a special case where the constant is 1D with length == perm.
// TODO: TransposeInitializer should be updated to handle this case.
// e.g. it provides a set of values that are relative to the input axes like the `sizes` input for Resize
// Permute1DConstant permutes the constant and adds a new initializer. The old initializer is removed only if
// there are no other consumers.
if (constant->Shape().size() == 1 && constant->Shape()[0] == gsl::narrow_cast<int64_t>(perm.size())) {
Permute1DConstant(graph, node, *constant, i, input, perm);
auto& node_to_update = dq_node ? *dq_node : node;
Permute1DConstant(graph, node_to_update, *constant, i, constant_to_modify, perm);
// unset updated input so reconnect_nodes doesn't change it back
if (dq_node) {
constant_dq_input = "";
} else {
input = "";
}
return;
}
if (consumers->nodes.size() > 0) {
// Transpose the initializer. If there are existing consumers, add Transpose nodes to them using perm_inv
// to counteract the effect. These Transposes will hopefully be optimized out later.
auto transpose_inv_ptr = MakeTranspose(graph, input, perm_inv);
// record the consumer node's input as being special cased for use in Case 2 if a DQ node, and IsConstant
if (nodes_using_updated_shared_initializer) {
for (auto& consumer : consumers->nodes) {
auto& consumer_node_inputs = (*nodes_using_updated_shared_initializer)[consumer->Id()];
// find input id/s for consumer
auto consumer_inputs = consumer->Inputs();
for (size_t input_idx = 0; input_idx < consumer_inputs.size(); ++input_idx) {
if (consumer_inputs[input_idx] == constant_to_modify) {
consumer_node_inputs.push_back(input_idx);
}
}
}
}
auto transpose_inv_ptr = MakeTranspose(graph, constant_to_modify, perm_inv);
api::NodeRef& transpose_inv = *transpose_inv_ptr;
std::string_view transpose_out = transpose_inv.Outputs()[0];
graph.CopyValueInfo(input, transpose_out);
ReplaceValueReferences(consumers->nodes, input, transpose_out);
graph.CopyValueInfo(constant_to_modify, transpose_out);
ReplaceValueReferences(consumers->nodes, constant_to_modify, transpose_out);
}
graph.TransposeInitializer(input, perm);
node.SetInput(i, input);
graph.TransposeInitializer(constant_to_modify, perm);
if (dq_node) {
UpdateDQNodeInputAndShape(graph, *dq_node, constant_to_modify);
constant_dq_input = ""; // DQ input was already updated so we don't need reconnect_nodes to handle it
}
return;
}
// Case 2: input is a Transpose node
std::unique_ptr<api::NodeRef> inp_node = graph.GetNodeProducingOutput(input);
// check if this is a special-cased DQ node where we put the Transpose on input 0 of the DQ in 'Case 1' above
if (inp_node && inp_node->OpType() == "DequantizeLinear" &&
nodes_using_updated_shared_initializer &&
std::find_if(nodes_using_updated_shared_initializer->begin(), nodes_using_updated_shared_initializer->end(),
[&inp_node](const auto entry) {
const auto id = entry.first;
const auto& input_idxs = entry.second;
// id matches and the entry is for input 0 of the DQ node
return id == inp_node->Id() &&
std::find(input_idxs.begin(), input_idxs.end(), size_t(0)) != input_idxs.end();
}) != nodes_using_updated_shared_initializer->end()) {
// set things up so we can look past the DQ node to the Transpose that was inserted in front of the reshaped
// constant initializer that was shared with this node.
dq_node = std::move(inp_node);
auto dq_input = dq_node->Inputs()[0];
inp_node = graph.GetNodeProducingOutput(dq_input);
consumers = graph.GetValueConsumers(dq_input);
}
if (inp_node != nullptr && inp_node->IsOp("Transpose")) {
std::optional<std::vector<int64_t>> perm2 = GetPermAttrIfValid(*inp_node);
if (perm2 != std::nullopt && perm2->size() == perm.size()) {
// If they cancel, use pre_transpose_value and remove Transpose if possible.
if (*perm2 == perm_inv) {
std::string_view pre_transpose_value = inp_node->Inputs()[0];
if (dq_node) {
UpdateDQNodeInputAndShape(graph, *dq_node, pre_transpose_value);
node.SetInput(i, dq_node->Outputs()[0]);
} else {
node.SetInput(i, pre_transpose_value);
}
// Remove the Transpose node if possible
// if there's a DQ node the `consumers` list still includes it so allow for that.
// in that case UpdateDQNodeInputAndShape already updated the input of the DQ node so it's safe to remove it.
if (consumers->comprehensive && consumers->nodes.size() == size_t(dq_node ? 1 : 0)) {
graph.RemoveNode(*inp_node);
}
return;
}
// NOTE: We expect the Transpose to cancel out when handling a special-cased DQ node that was originally
// connected to a shared constant initializer, so we don't expect to get here if dq_node is not nullptr.
// If there was a dq_node where the Transpose didn't cancel out we fall through to the next case
// so we retain the potential to cancel out for any other usages of the shared initializer.
assert(!dq_node); // assert in debug build to investigate. fall through to next case in release build to be safe.
if (!dq_node) {
// Otherwise, compose the perm and Transpose pre_transpose_value. Cost is the same and we may be able to remove
// the other Transpose.
const std::vector<int64_t>& perm_combined = ComposePerm(*perm2, perm);
auto transpose_ptr = MakeTranspose(graph, inp_node->Inputs()[0], perm_combined);
api::NodeRef& transpose = *transpose_ptr;
std::string_view transpose_out = transpose.Outputs()[0];
graph.CopyValueInfo(input, transpose_out);
graph.GetValueInfo(transpose_out)->PermuteDims(perm);
if (consumers->comprehensive && consumers->nodes.size() == 0) {
graph.RemoveNode(*inp_node);
}
node.SetInput(i, pre_transpose_value);
return;
} else if (*perm2 == perm) {
// we are trying to add a duplicate transpose.
// do nothing and return
return;
}
// Otherwise, compose the perm and Transpose pre_transpose_value. Cost is the same and we may be able to remove
// the other Transpose.
const std::vector<int64_t>& perm_combined = ComposePerm(*perm2, perm);
auto transpose_ptr = MakeTranspose(graph, inp_node->Inputs()[0], perm_combined);
api::NodeRef& transpose = *transpose_ptr;
std::string_view transpose_out = transpose.Outputs()[0];
graph.CopyValueInfo(input, transpose_out);
graph.GetValueInfo(transpose_out)->PermuteDims(perm);
if (consumers->comprehensive && consumers->nodes.size() == 0) {
graph.RemoveNode(*inp_node);
node.SetInput(i, transpose_out);
return;
}
node.SetInput(i, transpose_out);
return;
}
}
// Case 3: A Transpose op might already exist
for (size_t j = 0; j < consumers->nodes.size(); ++j) {
api::NodeRef& consumer = *consumers->nodes[j];
if (consumer.IsOp("Transpose") && GetPermAttrIfValid(consumer) == perm) {
node.SetInput(i, consumer.Outputs()[0]);
// any DQ node special casing doesn't apply anymore, so go back to the original inp_node
if (dq_node) {
inp_node = std::move(dq_node);
consumers = graph.GetValueConsumers(input);
}
// Case 3: A Transpose op with the same perms might already exist
for (auto& consumer : consumers->nodes) {
if (consumer->IsOp("Transpose") && GetPermAttrIfValid(*consumer) == perm) {
node.SetInput(i, consumer->Outputs()[0]);
return;
}
}
@ -540,11 +786,25 @@ void TransposeInput(api::GraphRef& graph, api::NodeRef& node, size_t i,
std::string_view transpose_out = transpose.Outputs()[0];
graph.CopyValueInfo(input, transpose_out);
graph.GetValueInfo(transpose_out)->PermuteDims(perm);
node.SetInput(i, transpose_out);
}
void TransposeInput(api::GraphRef& graph, api::NodeRef& node, size_t i,
const std::vector<int64_t>& perm,
const std::vector<int64_t>& perm_inv) {
// this TransposeInput is used by the layout transformer to wrap a node in Transpose ops. there's no OptimizerCtx
// in that scenario and we're not tracking special-cased DQ nodes as we only do that when pushing Transpose nodes.
TransposeInputImpl(graph, /* nodes_using_updated_shared_initializer */ nullptr, node, i, perm, perm_inv);
}
static void TransposeInput(OptimizerCtx& ctx, api::NodeRef& node, size_t i, const std::vector<int64_t>& perm,
const std::vector<int64_t>& perm_inv) {
TransposeInputImpl(ctx.graph, &ctx.nodes_using_updated_shared_initializer, node, i, perm, perm_inv);
}
// Unsqueezes inputs of node to have uniform rank. Returns false if input ranks are unknown or exceed the target rank.
static bool NormalizeInputRanks(OptimizerCtx ctx, api::NodeRef& node, size_t target_rank,
static bool NormalizeInputRanks(OptimizerCtx& ctx, api::NodeRef& node, size_t target_rank,
const std::vector<size_t>& input_indices) {
auto inputs = node.Inputs();
@ -579,7 +839,7 @@ void TransposeInputs(OptimizerCtx& ctx, api::NodeRef& node, const std::vector<in
const std::vector<size_t>& input_indices) {
auto perm_inv = InvertPerm(perm);
for (size_t j : input_indices) {
TransposeInput(ctx.graph, node, j, perm, perm_inv);
TransposeInput(ctx, node, j, perm, perm_inv);
}
}
@ -670,23 +930,74 @@ static bool CanLikelyRemoveTranspose(const api::GraphRef& graph, api::NodeRef& t
return true;
}
// return true if
// - the value is a constant initializer
// - the value is the output of a DQ node who's input is a constant initializer
// - UnsqueezeInput/TranposeInput can look past the DQ to update the constant initializer directly
// - DQ node is currently ignored if it uses per-channel quantization
// - supporting per-channel quantization requires modifying the scales and zero point data, which can be done
// if/when there's a use-case to justify the development cost.
// - the input was originally connected to a shared constant initializer that was updated in place by UnsqueezeInput
// or TransposeInput, and usage by this node had Squeeze/Transpose nodes inserted to counteract the effect of the
// in-place update. if we push the same transpose through this node it should cancel out that Squeeze/Transpose
//
// in all these cases we expect pushing the transpose through to not require a runtime Transpose node
static bool IsConstant(const api::GraphRef& graph, const api::NodeRef& node,
size_t input_id,
std::string_view value_name,
const NodeIdToInputIdxsMap& nodes_using_updated_shared_initializer) {
std::unique_ptr<api::NodeRef> producer_node = graph.GetNodeProducingOutput(value_name);
if (!producer_node) {
// initializer. may or may not be constant depending on whether it has a matching graph input
std::unique_ptr<api::TensorRef> constant = graph.GetConstant(value_name);
return constant != nullptr;
}
auto node_id_to_check = node.Id();
// handle potentially looking past a DQ node
if (producer_node->OpType() == "DequantizeLinear") {
std::unique_ptr<api::NodeRef> dq_node = GetDQWithConstInitializerInput(graph, value_name);
if (dq_node != nullptr) {
// DQ node pointing to an initializer that has not been updated in-place yet
return true;
}
// could also be a DQ that was connected to a shared initializer that was updated in-place.
// update the info on the node/input index to check and fall through
node_id_to_check = producer_node->Id();
input_id = 0; // can only be input 0 of a DQ node
}
auto entry = nodes_using_updated_shared_initializer.find(node_id_to_check);
if (entry != nodes_using_updated_shared_initializer.end()) {
if (std::find(entry->second.begin(), entry->second.end(), input_id) != entry->second.end()) {
return true;
}
}
return false;
}
// Estimates the cost of transposing an input. Currently uses rank heuristic. Negative if transpose is removed.
// Feel free to improve as needed.
static int EstimateTransposeValueCost(const api::GraphRef& graph, std::string_view input,
static int EstimateTransposeValueCost(const api::GraphRef& graph, const api::NodeRef& node,
size_t input_id, std::string_view input,
const std::vector<int64_t>& perm_inv,
const HandlerMap& extended_handlers) {
const HandlerMap& extended_handlers,
const NodeIdToInputIdxsMap& nodes_using_updated_shared_initializer) {
// Case 1: Transposing constants probably costs nothing.
std::unique_ptr<api::TensorRef> constant = graph.GetConstant(input);
if (constant != nullptr) {
if (IsConstant(graph, node, input_id, input, nodes_using_updated_shared_initializer)) {
return 0;
}
// Case 2: Transposing a transpose either cancels it or composes the permutations.
std::unique_ptr<api::NodeRef> node = graph.GetNodeProducingOutput(input);
if (node != nullptr && node->IsOp("Transpose")) {
std::optional<std::vector<int64_t>> perm2 = GetPermAttrIfValid(*node);
std::unique_ptr<api::NodeRef> producer_node = graph.GetNodeProducingOutput(input);
if (producer_node != nullptr && producer_node->IsOp("Transpose")) {
std::optional<std::vector<int64_t>> perm2 = GetPermAttrIfValid(*producer_node);
if (perm2 != std::nullopt) {
if (*perm2 == perm_inv && CanLikelyRemoveTranspose(graph, *node, extended_handlers)) {
if (*perm2 == perm_inv && CanLikelyRemoveTranspose(graph, *producer_node, extended_handlers)) {
return -EstimateValueRank(graph, input);
} else {
return 0;
@ -702,11 +1013,13 @@ static int EstimateTransposeValueCost(const api::GraphRef& graph, std::string_vi
static int EstimateTransposeInputsCost(const api::GraphRef& graph, const api::NodeRef& node,
const std::vector<int64_t>& perm_inv,
const std::vector<size_t>& input_indices,
const HandlerMap& extended_handlers) {
const HandlerMap& extended_handlers,
const NodeIdToInputIdxsMap& nodes_using_updated_shared_initializer) {
auto inputs = node.Inputs();
int cost = 0;
for (size_t j : input_indices) {
cost += EstimateTransposeValueCost(graph, inputs[j], perm_inv, extended_handlers);
cost += EstimateTransposeValueCost(graph, node, j, inputs[j], perm_inv, extended_handlers,
nodes_using_updated_shared_initializer);
}
return cost;
}
@ -734,8 +1047,10 @@ static bool HandleSimpleNodeBase(HandlerArgs& args, bool broadcast_inputs) {
if (broadcast_inputs && !NormalizeInputRanks(args.ctx, args.node, rank, args.transposible_inputs)) {
return false;
}
TransposeInputs(args.ctx, args.node, args.perm_inv, args.transposible_inputs);
TransposeOutputs(args.ctx, args.node, args.perm);
return true;
}
@ -1787,12 +2102,14 @@ static int CalculateCost(const api::GraphRef& graph, const api::NodeRef& node,
const std::unordered_set<std::string>& outputs_leading_to_transpose,
const HandlerInfo& info,
const std::vector<size_t>& input_indices,
const HandlerMap& extended_handlers) {
const HandlerMap& extended_handlers,
const NodeIdToInputIdxsMap& nodes_using_updated_shared_initializer) {
// We require the input cost (number of transposes before the op) and the total cost to strictly decrease.
// Strict decrease of the input cost ensures the optimization is stable, since the total cost decrease is just an
// estimate (the transpose after the op may or may not cancel with a subsequent transpose). We don't want
// repeated runs of the optimizer to have a transpose toggle between two inputs of a binary op.
int cost = EstimateTransposeInputsCost(graph, node, perm, input_indices, extended_handlers);
int cost = EstimateTransposeInputsCost(graph, node, perm, input_indices, extended_handlers,
nodes_using_updated_shared_initializer);
if (cost < 0 && info.transposes_outputs) {
// If the output will be transposed and won't ultimately cancel, factor in that cost.
@ -1822,13 +2139,14 @@ static bool ShouldPushTranspose(const api::GraphRef& graph, const api::NodeRef&
const std::unordered_set<std::string>& outputs_leading_to_transpose,
const HandlerInfo& info,
const std::vector<size_t> transposable_input_indices,
const HandlerMap& extended_handlers) {
const HandlerMap& extended_handlers,
const NodeIdToInputIdxsMap& nodes_using_updated_shared_initializer) {
if (node.IsOp("Transpose")) {
return true;
}
int cost = CalculateCost(graph, node, perm, outputs_leading_to_transpose, info, transposable_input_indices,
extended_handlers);
extended_handlers, nodes_using_updated_shared_initializer);
return cost < 0;
}
@ -1855,7 +2173,7 @@ bool ProcessTranspose(OptimizerCtx& ctx, api::NodeRef& transpose, api::NodeRef&
if (cost == CostCheckResult::kFallThrough) {
cost = ShouldPushTranspose(ctx.graph, node, perm, outputs_leading_to_transpose, *info, input_indices,
ctx.extended_handlers)
ctx.extended_handlers, ctx.nodes_using_updated_shared_initializer)
? CostCheckResult::kPushTranspose
: CostCheckResult::kStop;
}
@ -1889,7 +2207,7 @@ std::optional<OptimizerCtx> MakeOptimizerContext(api::GraphRef& graph,
return std::nullopt;
}
OptimizerCtx ctx{*opset, graph, provider_type, cost_check_fn, extended_handlers};
OptimizerCtx ctx{*opset, graph, provider_type, cost_check_fn, extended_handlers, {}};
return ctx;
}

29
onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.h
View file

@ -3,6 +3,7 @@
#pragma once
#include <unordered_map>
#include <vector>
// implementation details of the transpose optimizer API defined in optimizer_api.h.
@ -38,6 +39,8 @@ struct HandlerInfo {
bool transposes_outputs = true;
};
using NodeIdToInputIdxsMap = std::unordered_map<int64_t, std::vector<size_t>>;
struct OptimizerCtx {
int64_t opset;
api::GraphRef& graph;
@ -48,6 +51,32 @@ struct OptimizerCtx {
// Handlers for ops that are not in the ONNX opset, or for ONNX ops where special handling is required.
// If a handler is not found in this map, the default handlers will be used.
const HandlerMap& extended_handlers;
// When we update a shared constant initializer as part of pushing a transpose through a node we update the
// initializer in-place and insert Squeeze (in UnsqueezeInput if the initializer is broadcast) or
// Transpose (in TransposeInput) nodes between the updated initializer and the other usages.
// This map contains the set of nodes that had a Squeeze or Transpose added between them and the initializer.
// The entry contains the node id (key) and original input index/es (value) that were connected to the initializer
// prior to the insertion of the Squeeze/Transpose.
//
// Assuming we also transpose the other usages of the initializer in the same way (which would be expected) the
// Squeeze and Transpose nodes would be cancelled out, and the other usages will end up using the original
// initializer that was updated in-place.
//
// We use this information in two ways.
//
// 1. In the IsConstant calculation that determines the cost of pushing a transpose through a node.
// - as we expect the transpose to be making the same modification to all shared usages of the initializer we
// expect the Squeeze/Transpose nodes to be cancelled out, resulting in no runtime cost to push the transpose
// through that input.
//
// 2. To enable and track a special case in a QDQ format model where there is the added complexity of a DQ node
// between the initializer and each usage.
// - we look past a DQ node in UnsqueezeInput and TransposeInput to determine if there is a constant initializer
// that can be updated in-place as the DQ node is not sensitive to any rank or layout changes
// - NOTE we currently ignore DQ nodes with per-channel quantization as they are sensitive to changes
// - we also look past DQ nodes when processing the other usages in order to cancel out the Squeeze/Transpose
NodeIdToInputIdxsMap nodes_using_updated_shared_initializer;
};
/// <summary>

6
onnxruntime/core/optimizer/transpose_optimization/optimizer_api.h
View file

@ -242,6 +242,12 @@ class NodeRef {
/// <returns>since version or default value -1</returns>
virtual int SinceVersion() const = 0;
/// <summary>
/// Get the unique id of the node.
/// </summary>
/// <returns>Id</returns>
virtual int64_t Id() const = 0;
virtual ~NodeRef(){};
};

7
onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc
View file

@ -95,7 +95,8 @@ class ApiNode final : public api::NodeRef {
void ClearAttribute(std::string_view name) override;
void SetInput(size_t i, std::string_view name) override;
std::string_view GetExecutionProviderType() const override;
virtual int SinceVersion() const override;
int SinceVersion() const override;
int64_t Id() const override;
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ApiNode);
@ -417,6 +418,10 @@ int ApiNode::SinceVersion() const {
return node_.SinceVersion();
}
int64_t ApiNode::Id() const {
return node_.Index();
}
// </ApiNode>
std::optional<int64_t> ApiGraph::Opset(std::string_view domain) const {

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

@ -1005,18 +1005,22 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
layout_transformation::TransformLayoutForEP(graph_to_transform, modified, execution_provider,
std::move(cpu_allocator), debug_graph_fn));
if (modified) {
ORT_RETURN_IF_ERROR_SESSIONID_(
graph_transformer_mgr_.ApplyTransformers(graph_to_transform, TransformerLevel::Level1, *session_logger_));
// debug the graph after the L1 transformers have run against any layout transformation changes.
// this is prior to GraphPartitioner::GetCapabilityForEP calling IExecutionProvider::GetCapability the second
// time to validate the EP that requested the layout transformation can take all nodes using the new layout.
// if that fails, this allows debugging the graph used in that GetCapability call.
if (debug_graph_fn) {
debug_graph_fn(graph_to_transform);
}
}
// Previously we ran the L1 transformers to handle constant folding of any initializers that were transposed in
// a QDQ format model. The transpose optimizer can now look past DQ nodes to directly update initializers which
// takes care of most models without needing this.
//
// if (modified) {
// ORT_RETURN_IF_ERROR_SESSIONID_(
// graph_transformer_mgr_.ApplyTransformers(graph_to_transform, TransformerLevel::Level1, *session_logger_));
//
// debug the graph after the L1 transformers have run against any layout transformation changes.
// this is prior to GraphPartitioner::GetCapabilityForEP calling IExecutionProvider::GetCapability the second
// time to validate the EP that requested the layout transformation can take all nodes using the new layout.
// if that fails, this allows debugging the graph used in that GetCapability call.
// if (debug_graph_fn) {
// debug_graph_fn(graph_to_transform);
//}
//}
return Status::OK();
};

120
onnxruntime/test/optimizer/transpose_optimizer_test.cc
View file

@ -12,6 +12,9 @@
#include "core/graph/node_attr_utils.h"
#include "core/framework/op_node_proto_helper.h"
#include "core/framework/utils.h"
#include "core/optimizer/transpose_optimization/onnx_transpose_optimization.h"
#include "core/optimizer/transpose_optimization/optimizer_api.h"
#include "core/optimizer/transpose_optimization/ort_optimizer_utils.h"
#include "core/session/onnxruntime_session_options_config_keys.h"
#include "test/test_environment.h"
@ -19,6 +22,7 @@
#include "test/providers/internal_testing/internal_testing_execution_provider.h"
#include "test/util/include/asserts.h"
#include "test/util/include/inference_session_wrapper.h"
#include "test/util/include/test_utils.h"
namespace onnxruntime {
namespace test {
@ -4395,9 +4399,9 @@ TEST(TransposeOptimizerTests, RegressionTest_GitHubIssue9671) {
SessionOptions so;
so.session_logid = "TransposeOptimizerTests.RegressionTest_GitHubIssue9671";
InferenceSession session_object{so, GetEnvironment()};
ASSERT_STATUS_OK(session_object.Load(model_uri));
ASSERT_STATUS_OK(session_object.Initialize()); // optimizers run during initialization
InferenceSession session{so, GetEnvironment()};
ASSERT_STATUS_OK(session.Load(model_uri));
ASSERT_STATUS_OK(session.Initialize()); // optimizers run during initialization
}
// regression test for a model where the transpose optimizations incorrectly removed a node providing an implicit
@ -4409,9 +4413,9 @@ TEST(TransposeOptimizerTests, RegressionTest_GitHubIssue10305) {
SessionOptions so;
so.session_logid = "TransposeOptimizerTests.RegressionTest_GitHubIssue10305";
InferenceSession session_object{so, GetEnvironment()};
ASSERT_STATUS_OK(session_object.Load(model_uri));
ASSERT_STATUS_OK(session_object.Initialize()); // optimizers run during initialization
InferenceSession session{so, GetEnvironment()};
ASSERT_STATUS_OK(session.Load(model_uri));
ASSERT_STATUS_OK(session.Initialize()); // optimizers run during initialization
}
// regression test for a model with DQ node with per-axis dequantization followed by a Transpose.
@ -4432,18 +4436,18 @@ TEST(TransposeOptimizerTests, RegressionTest_GitHubIssue12151) {
{
so.graph_optimization_level = TransformerLevel::Default; // off
InferenceSession session_object{so, GetEnvironment()};
ASSERT_STATUS_OK(session_object.Load(model_uri));
ASSERT_STATUS_OK(session_object.Initialize());
ASSERT_STATUS_OK(session_object.Run(feeds, output_names, &fetches_orig));
InferenceSession session{so, GetEnvironment()};
ASSERT_STATUS_OK(session.Load(model_uri));
ASSERT_STATUS_OK(session.Initialize());
ASSERT_STATUS_OK(session.Run(feeds, output_names, &fetches_orig));
}
{
so.graph_optimization_level = TransformerLevel::Level1; // enable transpose optimizer
InferenceSession session_object{so, GetEnvironment()};
ASSERT_STATUS_OK(session_object.Load(model_uri));
ASSERT_STATUS_OK(session_object.Initialize());
ASSERT_STATUS_OK(session_object.Run(feeds, output_names, &fetches));
InferenceSession session{so, GetEnvironment()};
ASSERT_STATUS_OK(session.Load(model_uri));
ASSERT_STATUS_OK(session.Initialize());
ASSERT_STATUS_OK(session.Run(feeds, output_names, &fetches));
}
ASSERT_THAT(fetches_orig[0].Get<Tensor>().DataAsSpan<float>(),
@ -4497,6 +4501,13 @@ TEST(TransposeOptimizerTests, QnnTransposeReshape) {
for (const auto& node : graph.Nodes()) {
EXPECT_TRUE(node.GetExecutionProviderType() == expected_ep) << node.OpType() << " node named '" << node.Name()
<< "' was not assigned to the internal testing EP.";
if (node.Name() == "Mul_212" || node.Name() == "Add_213") {
// check that the special case in TransposeInputs for a single element input reconnects things back up correctly
const auto& inputs = node.InputDefs();
EXPECT_EQ(inputs.size(), size_t(2));
EXPECT_TRUE(inputs[1]->Exists());
}
}
#endif
}
@ -4543,5 +4554,86 @@ TEST(TransposeOptimizerTests, QnnTransposeReshapeQDQ) {
}
#endif
}
static void CheckSharedInitializerHandling(bool broadcast) {
auto model_uri = broadcast ? ORT_TSTR("testdata/transpose_optimizer_shared_initializers_broadcast.onnx")
: ORT_TSTR("testdata/transpose_optimizer_shared_initializers.onnx");
RandomValueGenerator random{123};
std::vector<int64_t> input_dims{1, 2, 2, 3};
std::vector<float> input_data = random.Gaussian<float>(input_dims, 0.0f, 1.0f);
OrtValue input;
CreateMLValue<float>(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], input_dims, input_data, &input);
NameMLValMap feeds{{"input0", input}};
std::vector<std::string> output_names{"output0"};
std::vector<OrtValue> fetches_orig;
std::vector<OrtValue> fetches;
SessionOptions so;
ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kDebugLayoutTransformation, "1"));
ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsDisableQuantQDQ, "1"));
// get results with no modifications to the model
{
so.graph_optimization_level = TransformerLevel::Default; // off
InferenceSessionWrapper session{so, GetEnvironment()};
ASSERT_STATUS_OK(session.Load(model_uri));
ASSERT_STATUS_OK(session.Initialize());
ASSERT_STATUS_OK(session.Run(feeds, output_names, &fetches_orig));
}
{
InferenceSessionWrapper session{so, GetEnvironment()};
ASSERT_STATUS_OK(session.Load(model_uri));
// we call the ONNX transpose optimizer directly to simplify the model required to exercise the shared initializer
// handling. this means we don't need to disable optimizers that might alter the graph before the
// transpose optimizer runs (at a minimum ConstantFolding, CommonSubexpressionElimination and ConstantSharing).
Graph& graph = session.GetMutableGraph();
CPUAllocator allocator;
using namespace onnx_transpose_optimization;
auto api_graph = MakeApiGraph(graph, TestCPUExecutionProvider()->CreatePreferredAllocators()[0],
/*new_node_ep*/ nullptr);
// default optimization cost check
OptimizeResult result = Optimize(*api_graph);
ASSERT_EQ(result.error_msg, std::nullopt);
ASSERT_TRUE(result.graph_modified);
ASSERT_TRUE(graph.GraphResolveNeeded());
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
EXPECT_EQ(op_to_count["Transpose"], 0) << "The Transpose nodes should have been pushed through and canceled out.";
ASSERT_STATUS_OK(graph.Resolve());
ASSERT_STATUS_OK(session.Initialize());
ASSERT_STATUS_OK(session.Run(feeds, output_names, &fetches));
}
ASSERT_THAT(fetches_orig[0].Get<Tensor>().DataAsSpan<float>(),
testing::ContainerEq(fetches[0].Get<Tensor>().DataAsSpan<float>()));
}
// test we re-use a modified shared initializer wherever possible. model has one initializer that is used by 3 DQ nodes
// and one initializer that is used by 2 Add nodes. both cases should be handled with the initializer being
// modified in-place for the first usage, and the Transpose added to the second usage being cancelled out when the
// original Transpose at the start of the model is pushed down.
TEST(TransposeOptimizerTests, SharedInitializerHandling) {
CheckSharedInitializerHandling(/*broadcast*/ false);
}
// same setup as the above test, however the initializer is broadcast to bring UnsqueezeInput into play.
// the in-place modification of the initializer for the first usage results in
// <initializer> -> Transpose -> Squeeze -> {DQ | Add}
// the later usages of the initializer should attempt to cancel out the Squeeze in UnsqueezeInput,
// followed by canceling out the Transpose in TransposeInput.
TEST(TransposeOptimizerTests, SharedInitializerHandlingBroadcast) {
CheckSharedInitializerHandling(/*broadcast*/ true);
}
} // namespace test
} // namespace onnxruntime

BIN
onnxruntime/test/testdata/transpose_optimizer_shared_initializers.onnx vendored Normal file
View file

Binary file not shown.

60
onnxruntime/test/testdata/transpose_optimizer_shared_initializers.py vendored Normal file
View file

@ -0,0 +1,60 @@
import numpy as np
import onnx
from onnx import TensorProto, helper
# Create a model with shared initializers that can be updated in-place by the transpose optimizer,
# including ones behind a DQ node. The transpose optimizer updates the first usage and inserts
# Transpose/Unsqueeze ops on the others (see UnsqueezeInput and TransposeInput).
# When we push the Transpose past other usages we should be able to cancel out those Transpose/Unsqueeze ops.
# We need 3 DQ nodes to ensure the Transpose or Unsqueeze added by the transpose optimizer is not
# removed prematurely.
def create_model(broadcast_weights: bool):
if broadcast_weights:
bias_shape = [2, 2]
bias_values = np.random.randn(2, 2)
else:
bias_shape = [1, 3, 2, 2]
bias_values = np.random.randn(1, 3, 2, 2)
graph = helper.make_graph(
name="graph",
inputs=[
helper.make_tensor_value_info("input0", TensorProto.FLOAT, [1, 2, 2, 3]),
],
initializer=[
helper.make_tensor("bias_quant", TensorProto.UINT8, bias_shape, bias_values.astype(np.uint8)),
helper.make_tensor("bias_fp32", TensorProto.FLOAT, bias_shape, bias_values.astype(np.float32)),
helper.make_tensor("dq_scale0", TensorProto.FLOAT, [], [1.5]),
helper.make_tensor("dq_zp0", TensorProto.UINT8, [], [5]),
helper.make_tensor("dq_scale1", TensorProto.FLOAT, [], [0.5]),
],
nodes=[
# Transpose input from channels last to channels first
helper.make_node("Transpose", ["input0"], ["input_T"], perm=[0, 3, 1, 2]),
helper.make_node("DequantizeLinear", ["bias_quant", "dq_scale0", "dq_zp0"], ["DQ0"], "DQ0"),
helper.make_node("Add", ["input_T", "DQ0"], ["A0"], "A0"),
helper.make_node("DequantizeLinear", ["bias_quant", "dq_scale1"], ["DQ1"], "DQ1"),
helper.make_node("Add", ["A0", "DQ1"], ["A1"], "A1"),
helper.make_node("DequantizeLinear", ["bias_quant", "dq_scale0"], ["DQ2"], "DQ2"),
helper.make_node("Add", ["A1", "DQ2"], ["A2"], "A2"),
helper.make_node("Add", ["A2", "bias_fp32"], ["A3"], "A3"),
helper.make_node("Add", ["A3", "bias_fp32"], ["A4"], "A4"),
# NCHW to NHWC
helper.make_node("Transpose", ["A4"], ["output0"], perm=[0, 2, 3, 1]),
],
outputs=[
helper.make_tensor_value_info("output0", TensorProto.FLOAT, [1, 2, 2, 3]),
],
)
model = helper.make_model(graph)
onnx.checker.check_model(model, full_check=True)
return model
if __name__ == "__main__":
model = create_model(broadcast_weights=False)
onnx.save(model, "transpose_optimizer_shared_initializers.onnx")
model = create_model(broadcast_weights=True)
onnx.save(model, "transpose_optimizer_shared_initializers_broadcast.onnx")

BIN
onnxruntime/test/testdata/transpose_optimizer_shared_initializers_broadcast.onnx vendored Normal file
View file

Binary file not shown.