diff --git a/onnxruntime/core/framework/compute_capability.h b/onnxruntime/core/framework/compute_capability.h index 5f21ba2f01..de479ed029 100644 --- a/onnxruntime/core/framework/compute_capability.h +++ b/onnxruntime/core/framework/compute_capability.h @@ -2,8 +2,11 @@ // Licensed under the MIT License. #pragma once +#include #include "core/common/common.h" #include "core/graph/indexed_sub_graph.h" +#include "core/graph/graph.h" + namespace onnxruntime { // A structure encodes a subgraph and the method to run it. @@ -21,5 +24,19 @@ struct ComputeCapability { ComputeCapability(std::unique_ptr t_sub_graph) : sub_graph(std::move(t_sub_graph)) {} + + // optional function to optimize this ComputeCapability + // this will be called by ORT once the ComputeCapability is assigned to the EP + // Optimization: std::function + std::function optimization_func; + + // optional ComputeCapability instances for sets of nodes within this ComputeCapability that should be optimized. + // when an optimization is applied, ORT will update this ComputeCapability to reflect the changes made. + // IndexedSubGraph.nodes: + // - update based on RemovedNode/AddNode calls + // IndexedSubGraph.MetaDef (if present): + // - inputs and outputs will be unchanged + // - constant_initializers MAY change if we constant fold an initializer during optimization + std::vector> nodes_to_optimize; }; } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/constant_folding.h b/onnxruntime/core/optimizer/constant_folding.h index 14eb2a9c5f..4fe2658e89 100644 --- a/onnxruntime/core/optimizer/constant_folding.h +++ b/onnxruntime/core/optimizer/constant_folding.h @@ -28,6 +28,8 @@ class ConstantFolding : public GraphTransformer { const InlinedHashSet& compatible_execution_providers = {}, const InlinedHashSet& excluded_initializers = {}) noexcept; + virtual bool AllowConstantFolding(const Node&) const { return true; } + private: Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; diff --git a/onnxruntime/core/optimizer/qdq_transformer/constant_folding_dq_node.cc b/onnxruntime/core/optimizer/qdq_transformer/constant_folding_dq_node.cc new file mode 100644 index 0000000000..da31176f7a --- /dev/null +++ b/onnxruntime/core/optimizer/qdq_transformer/constant_folding_dq_node.cc @@ -0,0 +1,359 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "core/optimizer/qdq_transformer/constant_folding_dq_node.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/utils.h" +#include "core/graph/graph_utils.h" +#include "core/optimizer/optimizer_execution_frame.h" +#include "core/optimizer/utils.h" +#include "core/framework/op_kernel.h" +#include "core/framework/tensorprotoutils.h" + +using namespace onnxruntime::common; + +namespace onnxruntime { + +ConstantFoldingDQ::ConstantFoldingDQ(const IExecutionProvider& execution_provider, + bool skip_dequantize_linear, + const ConfigOptions& config_options, + const InlinedHashSet& compatible_execution_providers, + const InlinedHashSet& excluded_initializers, + const InlinedHashSet& node_index_in_compute_capability) noexcept + : ConstantFolding(execution_provider,skip_dequantize_linear, config_options, compatible_execution_providers, excluded_initializers), + node_index_in_compute_capability_(node_index_in_compute_capability), + skip_dequantize_linear_(skip_dequantize_linear), + config_options_(config_options), + excluded_initializers_(excluded_initializers), + execution_provider_(execution_provider) { +} + +// We need to handle a Shape node separately as the input doesn't need to be a constant initializer for +// Shape to be able to be constant folded. +static bool ConstantFoldShapeNode(Graph& graph, Node& node) { + // Opset-15 Shape supports slicing using a 'start' and 'end' attribute + const auto& shape_attributes = node.GetAttributes(); + + int64_t start = 0; + int64_t end = std::numeric_limits::max(); + + for (const auto& attr : shape_attributes) { + if (attr.first == "start") { + start = attr.second.i(); + } else if (attr.first == "end") { + end = attr.second.i(); + } + } + + auto shape = node.MutableInputDefs()[0]->Shape(); + bool is_concrete_shape = true; + std::vector dim_values; + if (shape != nullptr) { + for (int dim_index = 0; dim_index < shape->dim_size(); dim_index++) { + auto dim = shape->dim(dim_index); + if (!utils::HasDimValue(dim)) { + is_concrete_shape = false; + break; + } + dim_values.push_back(dim.dim_value()); + } + } else { + is_concrete_shape = false; + } + + if (is_concrete_shape) { + int64_t rank = static_cast(dim_values.size()); + + // We ascertain the "true" starts/ends (if they were provided) + // Opset-15 Shape op supports slicing shape values + + // Deal with negatives and clamp + start = start < 0 ? start + rank : start; + start = start < 0 ? 0 : ((start > rank) ? rank : start); + + end = end < 0 ? end + rank : end; + end = end < 0 ? 0 : ((end > rank) ? rank : end); + + int64_t slice_length = end - start; + size_t clamped_slice_length = slice_length < 0 ? 0 : static_cast(slice_length); + + ONNX_NAMESPACE::TensorProto shape_constant; + auto* constant_arg_out = node.MutableOutputDefs()[0]; + shape_constant.set_name(constant_arg_out->Name()); + shape_constant.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + shape_constant.add_dims(clamped_slice_length); + utils::SetRawDataInTensorProto(shape_constant, dim_values.data() + start, clamped_slice_length * sizeof(int64_t)); + ONNX_NAMESPACE::TensorShapeProto result_shape; + result_shape.add_dim()->set_dim_value(clamped_slice_length); + constant_arg_out->SetShape(result_shape); + graph.AddInitializedTensor(shape_constant); + } + + return is_concrete_shape; // convert to constant if this is true +} + +// This function inlines the appropriate subgraph. It does not literally fold it. +static Status ConstantFoldIfNode(Graph& graph, Node& if_node, const logging::Logger& logger, bool& folded) { + folded = false; + // First, find out which subgraph to inline + // We need to fetch the constant argument. + assert(if_node.InputDefs().size() == 1); + const auto* condition_def = if_node.InputDefs()[0]; + + // We need to check if the condition is a constant. + constexpr bool check_outer_scope_true = true; + const ONNX_NAMESPACE::TensorProto* initializer = + graph.GetConstantInitializer(condition_def->Name(), check_outer_scope_true); + if (initializer == nullptr) { + return Status::OK(); + } + + // This is a boolean initializer with a single element. + Initializer condition{*initializer}; + ORT_RETURN_IF_NOT(condition.size() == 1, "If node condition initializer: `", condition_def->Name(), + "' is expected to have a single boolean element"); + + const bool condition_value = *condition.data(); + + auto status = graph.InlineIfSubgraph(condition_value, if_node, logger); + + if (!status.IsOK()) { + LOGS(logger, WARNING) << "Unable to constant fold. InlineIfSubgraph failed " + << " node '" << if_node.Name() << "': " + << status.ErrorMessage(); + return status; + } + + graph_utils::RemoveNodeOutputEdges(graph, if_node); + graph.RemoveNode(if_node.Index()); + + folded = true; + return status; +} + +bool ConstantFoldingDQ::AllowConstantFolding(const Node& node) const { + if (node_index_in_compute_capability_.find(node.Index()) != node_index_in_compute_capability_.end()) { + return true; + } + return false; +} + +Status ConstantFoldingDQ::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { + bool have_updated_nodes = false; + GraphViewer graph_viewer(graph); + auto& order = graph_viewer.GetNodesInTopologicalOrder(); + +#if !defined(DISABLE_SPARSE_TENSORS) + std::function is_sparse_initializer_check = [&graph](const std::string& name) -> bool { + return graph.IsSparseInitializer(name); + }; +#endif + + for (NodeIndex i : order) { + auto* node = graph.GetNode(i); + if (!node) { + continue; + } + + if (!AllowConstantFolding(*node)) { + continue; + } + + ORT_RETURN_IF_ERROR(Recurse(*node, modified, graph_level, logger)); + + // Updating a node may allow shape inferencing to infer output shapes of following nodes, + // so re-run the shape inferencing. use have_updated_nodes as that only applies to this Graph + // (vs. 'modified' which is passed into subgraphs and applies to the main graph and all subgraphs) + // Ignore any control flow node containing subgraphs as UpdateShapeInference is not intended to be used on it. + if (have_updated_nodes && !node->ContainsSubgraph()) { + ORT_RETURN_IF_ERROR(graph.UpdateShapeInference(*node)); + } + + bool converted_to_constant = false; + if (node->OpType().compare("If") == 0) { + // This process constant folds the If node only, + // but inlines the nodes of the corresponding branch graph. + // It does not convert the node to a constant in a common sense. + // We call it constant folding because the `If` node constant condition + // may enable us to inline the corresponding branch graph. + bool folded = false; + ORT_RETURN_IF_ERROR(ConstantFoldIfNode(graph, *node, logger, folded)); + if (folded) { + // Node removal is done within ConstantFoldIfNode() + modified = true; + have_updated_nodes = true; + } + } else if (node->OpType().compare("Shape") == 0) { + converted_to_constant = ConstantFoldShapeNode(graph, *node); + } else { + InitializedTensorSet constant_inputs; + + // Check if constant folding can be applied on this node. + const auto can_constant_fold_node = [&](const Node& n, bool skip_inputs_constant_check = false) { + return graph_utils::IsSupportedProvider(n, GetCompatibleExecutionProviders()) && + optimizer_utils::IsOperationDeterministic(n.Domain(), n.OpType()) && + // constant folding does not support executing a node that includes subgraphs (control flow operators, + // such as If/Loop/Scan, fall into this category). individual nodes in the subgraph will be processed + // by the Recurse call above + !n.ContainsSubgraph() && + (skip_inputs_constant_check || + graph_utils::AllNodeInputsAreConstant(graph, n, constant_inputs, excluded_initializers_)); + }; + + if (!can_constant_fold_node(*node)) { + continue; + } + + // if skip_dequantize_linear is true we want to maintain QDQ node units so avoid constant folding + // DequantizeLinear unless we can fold the whole QDQ node unit + if (skip_dequantize_linear_ && node->OpType() == "DequantizeLinear") { + bool can_constant_fold_qdq_node_unit = false; + + // Simplest scenario where the whole QDQ node unit of (DQ -> X -> Q) can be constant folded is if: + // - the DQ node does not produce a graph output, and its output is only consumed by X + // - X is a deterministic node with a single input and single output + // - the output from X is not a graph output and is only consumed by a Q node + if (optimizer_utils::CheckOutputEdges(graph, *node, 1)) { // DQ does not produce graph output, single consumer + const Node& node_x = *node->OutputNodesBegin(); + if (node_x.InputDefs().size() == 1 && + node_x.OutputDefs().size() == 1 && + optimizer_utils::CheckOutputEdges(graph, node_x, 1)) { + const Node& probably_q = *node_x.OutputNodesBegin(); + + if (probably_q.OpType() == "QuantizeLinear") { + // the inputs to these nodes are not const yet, but will be if we constant fold, + // so set skip_const_check to simulate that having happened + constexpr bool skip_const_check = true; + can_constant_fold_qdq_node_unit = can_constant_fold_node(node_x, skip_const_check) && + can_constant_fold_node(probably_q, skip_const_check); + } + } + } + + if (!can_constant_fold_qdq_node_unit) { + continue; + } + } + +#if !defined(DISABLE_SPARSE_TENSORS) + // Create execution frame for executing constant nodes. + OptimizerExecutionFrame::Info info({node}, constant_inputs, graph.ModelPath(), execution_provider_, + is_sparse_initializer_check, logger); +#else + // Create execution frame for executing constant nodes. + OptimizerExecutionFrame::Info info( + {node}, constant_inputs, graph.ModelPath(), execution_provider_, [](const std::string&) { return false; }, + logger); +#endif + + std::vector fetch_mlvalue_idxs; + for (const auto* node_out : node->OutputDefs()) { + fetch_mlvalue_idxs.push_back(info.GetMLValueIndex(node_out->Name())); + } + + const bool node_on_cpu_ep = node->GetExecutionProviderType() == kCpuExecutionProvider; + + std::unique_ptr kernel; + + if (!node_on_cpu_ep) { + // We need to copy the string here instead of taking a reference to it since node->SetExecutionProviderType + // will change the value of the reference + auto ep_type = node->GetExecutionProviderType(); + + // override the EP assigned to the node so that it will use the CPU kernel for Compute. + node->SetExecutionProviderType(kCpuExecutionProvider); + + kernel = info.CreateKernel(node, config_options_); + + // undo the EP change to the value that was assigned at graph partitioning time + node->SetExecutionProviderType(ep_type); + } else { + kernel = info.CreateKernel(node, config_options_); + } + + // We currently constant fold using the CPU EP only. + // If we can't find a CPU kernel for this node, then we can't proceed with constant folding. + // + // TODO(adrianlizarraga): Support constant folding with other execution providers. For example, we may be able + // to use a CUDA kernel to constant fold operators with data types not supported by the CPU EP kernel. + if (kernel == nullptr) { + LOGS(logger, WARNING) << "Could not find a CPU kernel and hence " + << "can't constant fold " << node->OpType() << " node '" << node->Name() << "'"; + + // Move on to the next candidate node + continue; + } + + OptimizerExecutionFrame frame(info, fetch_mlvalue_idxs); +#ifdef _WIN32 +#pragma warning(push) +#pragma warning(disable : 6387) +#endif + OpKernelContext op_kernel_context(&frame, kernel.get(), /*stream*/ nullptr, nullptr, logger); + ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context)); +#ifdef _WIN32 +#pragma warning(pop) +#endif + + std::vector fetches; + ORT_RETURN_IF_ERROR(frame.GetOutputs(fetches)); + + // Go over all output node args and substitute them with the newly computed tensors, which will be + // added to the graph as initializers. + ORT_ENFORCE(fetches.size() == node->OutputDefs().size()); + converted_to_constant = true; + for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { + const auto& constant_arg_out = *node->OutputDefs()[fetch_idx]; + // XXX: Add support for SparseTensors outputs when we have sparse outputs + if (!utils::HasTensorType(*constant_arg_out.TypeAsProto())) { + LOGS(logger, INFO) << "Unsupported output type of " << constant_arg_out.Type() + << ". Can't constant fold " << node->OpType() << " node '" << node->Name() << "'"; + converted_to_constant = false; + break; + } + } + + if (converted_to_constant) { + for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { + OrtValue& ort_value = fetches[fetch_idx]; + // Build the TensorProto that corresponds to the computed OrtValue and add it as initializer to the graph. + auto* constant_arg_out = node->MutableOutputDefs()[fetch_idx]; + const Tensor& out_tensor = ort_value.Get(); + ONNX_NAMESPACE::TensorProto out_tensorproto = utils::TensorToTensorProto(out_tensor, constant_arg_out->Name()); + + ONNX_NAMESPACE::TensorShapeProto result_shape; + for (auto& dim : out_tensor.Shape().GetDims()) { + result_shape.add_dim()->set_dim_value(dim); + } + + constant_arg_out->SetShape(result_shape); + graph.AddInitializedTensor(out_tensorproto); + } + } + } + + if (converted_to_constant) { + // Remove single-output node chain for inputs of the node + auto p_ip_node = node->InputNodesBegin(); + const auto p_ip_node_end = node->InputNodesEnd(); + while (p_ip_node != p_ip_node_end) { + const auto& input_node = *p_ip_node; + // Update the node iterator before removing the corresponding node because removing + // the node will invalidate the node iterator + ++p_ip_node; + graph_utils::RemoveNodesWithOneOutputBottomUp(graph, input_node); + } + + // Remove the output edges of the constant node and then remove the node itself. + graph_utils::RemoveNodeOutputEdges(graph, *node); + graph.RemoveNode(node->Index()); + modified = true; + have_updated_nodes = true; + } + } + + return Status::OK(); +} +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/qdq_transformer/constant_folding_dq_node.h b/onnxruntime/core/optimizer/qdq_transformer/constant_folding_dq_node.h new file mode 100644 index 0000000000..3a9d8006ae --- /dev/null +++ b/onnxruntime/core/optimizer/qdq_transformer/constant_folding_dq_node.h @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" +#include "core/optimizer/constant_folding.h" +#include "core/framework/ort_value.h" +#include +#include "core/framework/execution_provider.h" + +namespace onnxruntime { + +/** +@class ConstantFolding + +Transformer that traverses the graph top-down and performs constant folding, i.e., +it statically computes parts of the graph that rely only on constant initializers. +*/ +class ConstantFoldingDQ : public ConstantFolding { + public: + /*! Constant folding will not be applied to nodes that have one of initializers from excluded_initializers as input. + For pre-training, the trainable weights are those initializers to be excluded. + \param execution_provider Execution provider instance to execute constant folding. + */ + ConstantFoldingDQ(const IExecutionProvider& execution_provider, + bool skip_dequantize_linear, + const ConfigOptions& config_options, + const InlinedHashSet& compatible_execution_providers = {}, + const InlinedHashSet& excluded_initializers = {}, + const InlinedHashSet& node_index_in_compute_capability = {}) noexcept; + + bool AllowConstantFolding(const Node&) const; + + private: + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; + + const InlinedHashSet& node_index_in_compute_capability_; + bool skip_dequantize_linear_; + const ConfigOptions& config_options_; + const InlinedHashSet excluded_initializers_; + const IExecutionProvider& execution_provider_; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index 5a179ec622..9e8116177c 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -1221,6 +1221,8 @@ struct ProviderHost { virtual Status LoadDynamicLibrary(onnxruntime::PathString library_name) = 0; #endif + virtual Status GetEPOptimizerByName(const std::string& optimizer_name, std::function>(const GraphViewer&)>& selection_func) = 0; + // ModelMetadefIdGenerator virtual std::unique_ptr ModelMetadefIdGenerator__construct() = 0; virtual void ModelMetadefIdGenerator__operator_delete(ModelMetadefIdGenerator* p) = 0; diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index af39edae20..682f361d96 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -203,6 +203,29 @@ common::Status LoadDynamicLibraryFromProvider(onnxruntime::PathString library_na } #endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) +Status ApplyConstantFoldingOnDQ(const Graph&, const ComputeCapability& this_optimization, ComputeCapability& cc_to_update) { + + return Status::OK(); +} + +std::vector> dq_nodes_to_constant_fold(const GraphViewer& graph_viewer) { + std::vector> result; + std::unique_ptr sub_graph = std::make_unique(); + const std::vector& node_index = graph_viewer.GetNodesInTopologicalOrder(ExecutionOrder::PRIORITY_BASED /*priority-based topological sort*/); + for (const auto& index : node_index) { + const auto& node = graph_viewer.GetNode(index); + if (node->OpType() != "DequantizeLinear") { + continue; + } + sub_graph->nodes.push_back(index); + std::cout << node->Name() << ", op type: " << node->OpType() << std::endl; + } + + result.push_back(std::make_unique(std::move(sub_graph))); + result.back()->optimization_func = ApplyConstantFoldingOnDQ; + return result; +} + #if defined(_MSC_VER) && !defined(__clang__) #pragma warning(push) #pragma warning(disable : 26436) @@ -1485,7 +1508,14 @@ struct ProviderHostImpl : ProviderHost { #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) Status LoadDynamicLibrary(onnxruntime::PathString library_name) override { return LoadDynamicLibraryFromProvider(library_name); }; #endif + + Status GetEPOptimizerByName(const std::string& optimizer_name, std::function>(const GraphViewer&)>& selection_func) override { + selection_func = dq_nodes_to_constant_fold; + return Status::OK(); + }; + } provider_host_; + #if defined(_MSC_VER) && !defined(__clang__) #pragma warning(pop) #endif