From 0404763f23e4d2e1b06128b3545ad7c1042a7a66 Mon Sep 17 00:00:00 2001 From: Ashwini Khade Date: Tue, 30 Jun 2020 14:30:59 -0700 Subject: [PATCH] Update function body initialization for ONNX functions (#4332) * Update function body initialization * minor fix * changes per review comments * minor fix * format fix * add function initialization in mixed precision transformer * more updates * more fixes --- include/onnxruntime/core/graph/graph.h | 21 ++++++- .../core/framework/graph_partitioner.cc | 7 ++- onnxruntime/core/graph/graph.cc | 58 ++++++++++++------- .../core/graph/mixed_precision_transformer.cc | 9 ++- 4 files changed, 68 insertions(+), 27 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 5f41d7f819..69632df596 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -101,7 +101,21 @@ class Node { /** Gets the Node's Node::Type. */ Node::Type NodeType() const noexcept; - /** Gets the function body if the #NodeType is fused, or nullptr if not. */ + /** + Gets the function body if applicable otherwise nullptr + @param try_init_func_body If not already intialized, initialize the function body + (only applicable to operators which are defined as function in ONNX spec). + Function body can be initialized in 2 cases : + 1. For nodes of type "Fused" + 2. For nodes which are defined as functions in ONNX spec (example: DynamicQuantizeLinear) + For all other cases this will always return nullptr. + Nodes of type "Fused" are created during partitioning and the function body + initialization for such nodes also happens during node creation. Therefore, + initialization of function body will happen via this method only in case 2 mentioned above. + */ + const Function* GetFunctionBody(bool try_int_func_body = true) noexcept; + + /** Gets the function body if applicable otherwise nullptr. */ const Function* GetFunctionBody() const noexcept; /** Gets the node description. */ @@ -779,7 +793,10 @@ class Graph { @param node Node with Node::Type of Node::Type::Fused @returns Status indicating success or providing an error message. */ - Status InlineFunction(Node& node); + Status InlineFunction(Node& node); + + /** Initialize function body for the given node */ + void InitFunctionBodyForNode(Node& node); /** Mark a NodeArg name as coming from the outer scope when programmatically constructing a Graph that will be used as a GraphProto attribute in another Node.. diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 3466ee8a47..1a604048ba 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -197,12 +197,13 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f if (nullptr == node_func) { continue; } - nodes_need_inline.push_back(&node); + nodes_need_inline.push_back(&node); } - } + } + for (auto* node : nodes_need_inline) { // If the node has a functionbody with no kernel and cannot be inlined - // it is a invalid function + // it is an invalid function ORT_RETURN_IF_ERROR(graph.InlineFunction(*node)); } diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 8d6fc7457a..f6cca85ff5 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -178,8 +178,8 @@ bool NodeArg::HasTensorOrScalarShape() const { const auto type_case = type->value_case(); switch (type_case) { case TypeProto::kTensorType: - case TypeProto::kSparseTensorType: - // Standard tensor has a valid shape field while + case TypeProto::kSparseTensorType: + // Standard tensor has a valid shape field while // scalar's shape is empty. Thus, we don't need to // check shape here. return true; @@ -433,6 +433,19 @@ void Node::SetNodeType(Node::Type node_type) noexcept { node_type_ = node_type; } +const Function* Node::GetFunctionBody(bool try_init_func_body) noexcept { + if (nullptr != func_body_) { + return func_body_; + } + + // Initialize function body + if (try_init_func_body) { + graph_->InitFunctionBodyForNode(*this); + } + + return func_body_; +} + const Function* Node::GetFunctionBody() const noexcept { return func_body_; } @@ -1314,7 +1327,6 @@ void Graph::ReverseDFSFrom(const std::vector& from, const std::function& enter, const std::function& leave, const std::function& comp) const { - ReverseDFSFrom(from, enter, leave, comp, {}); } @@ -2037,22 +2049,6 @@ Status Graph::VerifyNodeAndOpMatch(const ResolveOptions& options) { node.op_ = nullptr; } - if (node.op_ && (node.op_->HasFunction() || node.op_->HasContextDependentFunction())) { - onnx::FunctionProto onnx_function_proto; - onnx::FunctionBodyBuildContextImpl function_body_ctx(node_proto); - if (node.op_->HasContextDependentFunction()) { - node.op_->BuildContextDependentFunction(function_body_ctx, onnx_function_proto); - } else { - onnx_function_proto = *(node.op_->GetFunction()); - } - - auto func_ptr = onnxruntime::make_unique(*this, node.Index(), onnx_function_proto, - logger_); - - function_container_.emplace_back(std::move(func_ptr)); - node.SetFunctionBody(*function_container_.back()); - } - if (!node.op_) { return Status(ONNXRUNTIME, FAIL, "Fatal error: " + node.OpType() + " is not a registered function/op"); } @@ -2100,6 +2096,26 @@ Status Graph::VerifyNodeAndOpMatch(const ResolveOptions& options) { return Status::OK(); } +void Graph::InitFunctionBodyForNode(Node& node) { + if (node.op_ && (node.op_->HasFunction() || node.op_->HasContextDependentFunction())) { + onnx::FunctionProto onnx_function_proto; + if (node.op_->HasContextDependentFunction()) { + NodeProto node_proto; + node.ToProto(node_proto); + onnx::FunctionBodyBuildContextImpl function_body_ctx(node_proto); + node.op_->BuildContextDependentFunction(function_body_ctx, onnx_function_proto); + } else { + onnx_function_proto = *(node.op_->GetFunction()); + } + + auto func_ptr = onnxruntime::make_unique(*this, node.Index(), onnx_function_proto, + logger_); + + function_container_.emplace_back(std::move(func_ptr)); + node.SetFunctionBody(*function_container_.back()); + } +} + void Graph::FindAllSubgraphs(std::vector& subgraphs) { for (auto& node : Nodes()) { for (auto& subgraph : node.MutableSubgraphs()) { @@ -2398,8 +2414,8 @@ const std::vector& Graph::GetValueInfo() const noexcept { return value_info_; } -void Graph::AddValueInfo(const NodeArg* new_value_info){ - for(const auto* info : value_info_){ +void Graph::AddValueInfo(const NodeArg* new_value_info) { + for (const auto* info : value_info_) { ORT_ENFORCE(info->Name() != new_value_info->Name(), "Error: trying to add an existing value info."); } value_info_.push_back(new_value_info); diff --git a/orttraining/orttraining/core/graph/mixed_precision_transformer.cc b/orttraining/orttraining/core/graph/mixed_precision_transformer.cc index 82b66eb7a4..4b00f4c0a6 100644 --- a/orttraining/orttraining/core/graph/mixed_precision_transformer.cc +++ b/orttraining/orttraining/core/graph/mixed_precision_transformer.cc @@ -376,7 +376,14 @@ Status TransformGraphForMixedPrecision(Graph& graph, const std::unordered_set& weights_to_train, bool use_fp16_initializer, std::unordered_map& fp32_weight_name_to_fp16_node_arg) { - // Stag 1: Convert whole graph including forward and backward to FP16 + // Stage 1: Convert whole graph including forward and backward to FP16 + // Initialize function body for all function nodes + // This is required to make sure after converting inputs\weights to FP16 + // the new NodeArg updates are correctly propagated to the function body nodes as well. + for (auto& node : graph.Nodes()) { + graph.InitFunctionBodyForNode(node); + } + // Insert Cast node to convert inputs from FP32 to FP16 for (const NodeArg* input : graph.GetInputs()) { if (input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {