From eba730500fd7d10fafe057fad364f9543346ad7a Mon Sep 17 00:00:00 2001 From: satyajandhyala Date: Thu, 10 Feb 2022 13:31:12 -0800 Subject: [PATCH] Remove file-scope non-constant static variables to support multiple inference sessions (#10481) * Changed file-scope static variables to automatic variables or function-scope static const. * Reduce load time overhead by using constexpr. * Use node indices instead of node names to track inserted, deleted and changed nodes. Co-authored-by: Satya Jandhyala --- .../core/optimizer/propagate_cast_ops.cc | 1059 +++++++++-------- .../core/optimizer/propagate_cast_ops.h | 9 + 2 files changed, 552 insertions(+), 516 deletions(-) diff --git a/onnxruntime/core/optimizer/propagate_cast_ops.cc b/onnxruntime/core/optimizer/propagate_cast_ops.cc index 740e915f0a..0e344dc003 100644 --- a/onnxruntime/core/optimizer/propagate_cast_ops.cc +++ b/onnxruntime/core/optimizer/propagate_cast_ops.cc @@ -9,47 +9,47 @@ using namespace ONNX_NAMESPACE; using namespace onnxruntime::common; /********************************************************************************************************************************* -* PropagateCastOps transformation tries to expand float16 computation regions in the graph by converting float operations to float16. -* In order to perform the the transformation, certain operations are considered FP16 safe, i.e. the computation may be performed -* without effecting the numerical result, ex. transpose, shape, etc. The transformation supports three levels of optimization, 0, 1 -* and 2. Level 2 being the most agressive, may consider moving float operations to float16 which may result in different numerical results -* due to loss of precision. The user may choose level 0, whereby the user chooses the opcodes which are "FP16 Safe" instead of a list -* predetermined opcodes as in levels 1 and 2. -* Currently three strategies are available, None, InsertAndReduce and FloodFill. -* None: -* Although no new cast operations are inserted or propagated using this strategy some optimizations are performed -* 1. Remove back-to-back casts -* 2. Fuse subgraphs -* 3. Remove unnecessary casts -* InsertAndReduce : -* This transformation converts all FP16 operations to float16. The transformation first -* 1. Inserts float16 cast operation on all the float inputs -* 2. Changes all float outputs to float16 -* 3. Inserts float cast operations on all float outputs as expected -* After inserting the FP16 and FP32 cast operation nodes around all the nodes with FP16 Safe opcodes, the transformation reduces the -* cast operations, using the following operations on the graph, iteratively until no more reduction is possible. -* 1. Remove back-to-back casts -* 2. Remove redundant casts -* 3. Fuse sibling subgraphs with same cast operation -* FloodFill: -* Operations are converted from float to float16 by propagating float16 cast operations up the graph or float cast operations down the -* graph. Using this strategy, for each pre-existing float/float16 cast operations the transformation first finds the possible expansion of -* float16 region up/down the graph using DFS/ReverseDFS and (TODO) identifies loss/gain by performing such expansion, considering -* the gain by reducing float operations lower precision and loss due to newly inserted cast operations (/TODO). -* In addition to propagating cast operations up/down the graph, in this strategy, the above mentioned cast op reduction functions -* are also used. -* InsertAndReduce exhaustively inserts cast operations before and after all the nodes with the allowed opcodes whereas FloodFill only -* propagates existing casts up or down the graph, inserting new casts as it expands the float16 regions. -*********************************************************************************************************************************/ + * PropagateCastOps transformation tries to expand float16 computation regions in the graph by converting float operations to float16. + * In order to perform the transformation, certain operations are considered FP16 safe, i.e. the computation may be performed + * without effecting the numerical result, ex. transpose, shape, etc. The transformation supports three levels of optimization, 0, 1 + * and 2. Level 2 being the most aggressive, may consider moving float operations to float16 which may result in different numerical results + * due to loss of precision. The user may choose level 0, whereby the user chooses the opcodes which are "FP16 Safe" instead of a list + * predetermined opcodes as in levels 1 and 2. + * Currently three strategies are available, None, InsertAndReduce and FloodFill. + * None: + * Although no new cast operations are inserted or propagated using this strategy some optimizations are performed + * 1. Remove back-to-back casts + * 2. Fuse subgraphs + * 3. Remove unnecessary casts + * InsertAndReduce : + * This transformation converts all FP16 operations to float16. The transformation first + * 1. Inserts float16 cast operation on all the float inputs + * 2. Changes all float outputs to float16 + * 3. Inserts float cast operations on all float outputs as expected + * After inserting the FP16 and FP32 cast operation nodes around all the nodes with FP16 Safe opcodes, the transformation reduces the + * cast operations, using the following operations on the graph, iteratively until no more reduction is possible. + * 1. Remove back-to-back casts + * 2. Remove redundant casts + * 3. Fuse sibling subgraphs with same cast operation + * FloodFill: + * Operations are converted from float to float16 by propagating float16 cast operations up the graph or float cast operations down the + * graph. Using this strategy, for each pre-existing float/float16 cast operations the transformation first finds the possible expansion of + * float16 region up/down the graph using DFS/ReverseDFS and (TODO) identifies loss/gain by performing such expansion, considering + * the gain by reducing float operations lower precision and loss due to newly inserted cast operations (/TODO). + * In addition to propagating cast operations up/down the graph, in this strategy, the above mentioned cast op reduction functions + * are also used. + * InsertAndReduce exhaustively inserts cast operations before and after all the nodes with the allowed opcodes whereas FloodFill only + * propagates existing casts up or down the graph, inserting new casts as it expands the float16 regions. + *********************************************************************************************************************************/ namespace onnxruntime { // NodeArg to Select consumer node map. typedef std::unordered_map> NodeArgToConsumerMap; -/* -* ConcatNames -* Collects all the names from the pointers of the objects stores in the container class C -* the class should have a member functions returning a string (or a ref). -*/ +/* + * ConcatNames + * Collects all the names from the pointers of the objects stores in the container class C + * the class should have a member functions returning a string (or a ref). + */ template static std::string ConcatNames( C const& items, std::string (*f)(const T& n) = [](const T& n) { return n->Name(); }) { @@ -63,53 +63,36 @@ static std::string ConcatNames( static std::string GetName(const std::pair>& p) { return p.first->Name() + " feeding " + ConcatNames(p.second) + "; "; }; -/* -* The collection fp16_allow_ops, specifies for a given propagate_cast_ops level, a vector of node op_types that -* the code is allowed to propage Cast operations across. The user may specify a custom list of optypes using level 0. -* The opcodes are split into multiple levels. Cast propagation is done based on the level. Level 2 op code -* list includes Level 1 list also. -*/ -static std::vector> fp16_allow_ops = { - /* Level 0 */ {}, - /* Level 1 */ {"Expand", "Transpose", "Relu", "Reshape", "Split", "Tanh", "Squeeze", "Unsqueeze"}, - /* Level 2 */ {"Add", "BiasGelu", "Dropout", "FastGelu", "Gather", "Gelu", "LayerNormalization", "Where"}}; + +using NodeIndices = std::unordered_set; +using FP16AllowOps = PropagateCastOps::FP16AllowOps; + +static constexpr std::array level1_fp16_allow_ops = {"Expand", "Transpose", "Relu", "Reshape", "Split", "Tanh", "Squeeze", "Unsqueeze"}; +static constexpr std::array level2_fp16_allow_ops = {"Add", "BiasGelu", "Dropout", "FastGelu", "Gather", "Gelu", "LayerNormalization", "Where"}; /* -* The following two maps specify the opcode to input and opcode to output mappings to list the inputs/outputs to consider while propagating -* cast operations. All other inputs/outputs not listed in this table are not relevant for deciding whether an operation -* performed in float or float16. If an opcode is not listed in these tables, the code will look at all the inputs and outputs to validate -* transformation. -*/ -static std::unordered_map> opcode_to_input_map = { - {"Gather", {0}}, - {"Reshape", {0}}, - {"Dropout", {0}}, - {"Expand", {0}}, - {"LayerNormalization", {0, 1, 2}}, - {"Squeeze", {0}}, - {"Unsqueeze", {0}} -}; - -static std::unordered_map> opcode_to_output_map = { - {"Gather", {0}}, - {"Reshape", {0}}, - {"Dropout", {0}}, - {"Expand", {0}}, - {"LayerNormalization", {0}}, - {"Squeeze", {0}}, - {"Unsqueeze", {0}} -}; - -static std::unordered_set inserted_node_names; // Names of the nodes inserted -static std::unordered_set converted_node_names; // Names of the nodes converted to FP16 - -/* -* Check if the input is relevant to consider for cast propagation for the given node. -* Return true if the opcode is not found in the opcode_to_input map. -*/ + * Check if the input is relevant to consider for cast propagation for the given node. + * Return true if the opcode is not found in the opcode_to_input map. + */ static bool IsRelevantInput(const Node* node, const NodeArg* input) { - if (opcode_to_input_map.find(node->OpType()) != opcode_to_input_map.end()) { - const std::vector& selected_inputs = opcode_to_input_map[node->OpType()]; + /* + * The following map specifies the opcode to input mapping to list the inputs to consider while propagating + * cast operations. All other inputs not listed in this table are not relevant for deciding whether an operation + * performed in float or float16. If an opcode is not listed in these tables, the code will look at all the inputs to validate + * transformation. + */ + static const std::unordered_map> opcode_to_input_map = { + {"Gather", {0}}, + {"Reshape", {0}}, + {"Dropout", {0}}, + {"Expand", {0}}, + {"LayerNormalization", {0, 1, 2}}, + {"Squeeze", {0}}, + {"Unsqueeze", {0}}}; + + auto it = opcode_to_input_map.find(node->OpType()); + if (it != opcode_to_input_map.cend()) { + const std::vector& selected_inputs = it->second; int input_index = optimizer_utils::IndexOfNodeInput(*node, *input); return std::find(selected_inputs.begin(), selected_inputs.end(), input_index) != selected_inputs.end(); } @@ -117,12 +100,28 @@ static bool IsRelevantInput(const Node* node, const NodeArg* input) { } /* -* Check if the output is relevant to consider for cast propagation for the given node. -* Return true if the opcode is not found in the opcode_to_output map. -*/ + * Check if the output is relevant to consider for cast propagation for the given node. + * Return true if the opcode is not found in the opcode_to_output map. + */ static bool IsRelevantOutput(const Node* node, const NodeArg* output) { - if (opcode_to_output_map.find(node->OpType()) != opcode_to_output_map.end()) { - const std::vector& selected_outputs = opcode_to_output_map[node->OpType()]; + /* + * The following map specifies the opcode to output mapping to list the outputs to consider while propagating + * cast operations. All other outputs not listed in this table are not relevant for deciding whether an operation + * performed in float or float16. If an opcode is not listed in these tables, the code will look at all the outputs to validate + * transformation. + */ + static const std::unordered_map> opcode_to_output_map = { + {"Gather", {0}}, + {"Reshape", {0}}, + {"Dropout", {0}}, + {"Expand", {0}}, + {"LayerNormalization", {0}}, + {"Squeeze", {0}}, + {"Unsqueeze", {0}}}; + + auto it = opcode_to_output_map.find(node->OpType()); + if (it != opcode_to_output_map.cend()) { + const std::vector& selected_outputs = it->second; int input_index = optimizer_utils::IndexOfNodeOutput(*node, *output); return std::find(selected_outputs.begin(), selected_outputs.end(), input_index) != selected_outputs.end(); } @@ -130,10 +129,10 @@ static bool IsRelevantOutput(const Node* node, const NodeArg* output) { } // Check whether the given opcode is fp16 allowed for the given level of optimization. -static bool IsFP16Allow(const std::string& op_type, size_t level) { +static bool IsFP16Allow(const std::string& op_type, size_t level, const FP16AllowOps& fp16_allow_ops) { bool fp16_allow = false; for (size_t i = 0; i <= level && i < fp16_allow_ops.size() && !fp16_allow; ++i) { - fp16_allow = std::find(fp16_allow_ops[i].begin(), fp16_allow_ops[i].end(), op_type) != fp16_allow_ops[i].end(); + fp16_allow = fp16_allow_ops[i].find(op_type) != fp16_allow_ops[i].end(); } return fp16_allow; } @@ -158,36 +157,35 @@ static bool IsType(const NodeArg& node_arg, TensorProto_DataType data_type) { } /* InsertCastNodes -* Insert a new Cast node after each NodeArg in the require_cast map, feeding the nodes (consumer) in the vector mapped to -* the NodeArg. The other consumers of the NodeArg will not be changed. The cast node is FLOAT16 if is_fp16 is True -* and FLOAT otherwise. This function fixes the graph edges in addition to inserting the cast nodes. -* -* In the following example only the first two consumers, Opcode0 and Opcode1 get casted and the third consumer Opcode2 does not. -* -* Input0/NodeArg Input/NodeArg -* ___ ____|________________ |____________________ -* | | | | | -* | | | _____V______ | -* | | | | CastFP16 | | -* | | _____V_____ |_or_FP32__| ______V___ -* | | | Opcode2 | | | Opcode2| -* | | |_________| _______|___________ |________| -* | | | | -* | | ---\ _____V______ _____V_____ -* | | ---/ | Opcode0 | | Opcode1 | -* _____V_____ _____V______ |__________| |_________| -* | Opcode0 | | Opcode1 | | | -* |_________| |__________| -* | | -*/ + * Insert a new Cast node after each NodeArg in the require_cast map, feeding the nodes (consumer) in the vector mapped to + * the NodeArg. The other consumers of the NodeArg will not be changed. The cast node is FLOAT16 if is_fp16 is True + * and FLOAT otherwise. This function fixes the graph edges in addition to inserting the cast nodes. + * + * In the following example only the first two consumers, Opcode0 and Opcode1 get casted and the third consumer Opcode2 does not. + * + * Input0/NodeArg Input/NodeArg + * ___ ____|________________ |____________________ + * | | | | | + * | | | _____V______ | + * | | | | CastFP16 | | + * | | _____V_____ |_or_FP32__| ______V___ + * | | | Opcode2 | | | Opcode2| + * | | |_________| _______|___________ |________| + * | | | | + * | | ---\ _____V______ _____V_____ + * | | ---/ | Opcode0 | | Opcode1 | + * _____V_____ _____V______ |__________| |_________| + * | Opcode0 | | Opcode1 | | | + * |_________| |__________| + * | | + */ static Status InsertCastNodes(Graph& graph, const NodeArgToConsumerMap& require_cast, bool is_fp16, - std::deque& removed_nodes) { - //Create required new Cast nodes. - for (std::pair> element : require_cast) { - NodeArg* node_arg = element.first; - std::vector nodes = element.second; + const NodeIndices& removed_nodes, + NodeIndices& inserted_nodes) { + // Create required new Cast nodes. + for (const auto& [node_arg, nodes] : require_cast) { if (!node_arg->Exists()) { continue; } @@ -220,7 +218,7 @@ static Status InsertCastNodes(Graph& graph, cast_inputs, cast_outputs, &attributes); - inserted_node_names.insert(cast.Name()); + inserted_nodes.insert(cast.Index()); Node* producer = graph.GetMutableProducerNode(node_arg->Name()); std::vector consumers = graph.GetMutableConsumerNodes(node_arg->Name()); std::vector other_nodes; @@ -245,7 +243,7 @@ static Status InsertCastNodes(Graph& graph, // Update consumers of node_arg to use the output of the cast node int cast_output_index = optimizer_utils::IndexOfNodeOutput(cast, cast_output); for (Node* consumer : consumers) { - if (std::find(removed_nodes.begin(), removed_nodes.end(), consumer->Index()) == removed_nodes.end()) { + if (removed_nodes.find(consumer->Index()) == removed_nodes.end()) { if (std::find(nodes.begin(), nodes.end(), consumer) == nodes.end()) { // Consumers not getting casted need to replace input-def if the producer's output-def is changed if (is_node_arg_cast_output) { @@ -281,48 +279,48 @@ static Status InsertCastNodes(Graph& graph, } /* RemoveCastNodesChain -* Remove the cast nodes specified in casts vector and fix the graph edges accordingly. If the output of the cast -* is also a graph output then insert an Identity node if the input of the cast node feeds other nodes. In the -* trivial case the chain has only one cast node. The caller is responsible for the validity of removing casts. -* -* _____|______ -* | Opcode1 | -* |__________| -* | -* _____V______ ____________ -* | Cast | | Opcode 1 | -* |__________| |__________| -* | | -* . | -* . ---\ _____V______ -* . ---/ | Opcode2 | -* _____V______ |__________| -* | Cast | | -* |__________| V -* | -* _____V______ -* | Opcode2 | -* |__________| -* | -* V -* -* OR -* _____|______ _____|_____ -* | Opcode1 | | Opcode1 | -* |__________| |_________| -* | | -* _____V______ ---\ _____V______ -* | Cast | ---/ | Opcode 2 | -* |__________| |__________| -* | | -* _____V______ V -* | Opcode2 | -* |__________| -* | -* V -*/ + * Remove the cast nodes specified in casts vector and fix the graph edges accordingly. If the output of the cast + * is also a graph output then insert an Identity node if the input of the cast node feeds other nodes. In the + * trivial case the chain has only one cast node. The caller is responsible for the validity of removing casts. + * + * _____|______ + * | Opcode1 | + * |__________| + * | + * _____V______ ____________ + * | Cast | | Opcode 1 | + * |__________| |__________| + * | | + * . | + * . ---\ _____V______ + * . ---/ | Opcode2 | + * _____V______ |__________| + * | Cast | | + * |__________| V + * | + * _____V______ + * | Opcode2 | + * |__________| + * | + * V + * + * OR + * _____|______ _____|_____ + * | Opcode1 | | Opcode1 | + * |__________| |_________| + * | | + * _____V______ ---\ _____V______ + * | Cast | ---/ | Opcode 2 | + * |__________| |__________| + * | | + * _____V______ V + * | Opcode2 | + * |__________| + * | + * V + */ -static Status RemoveCastNodesChain(Graph& graph, std::vector casts, std::deque& removed_nodes) { +static Status RemoveCastNodesChain(Graph& graph, const std::vector& casts, NodeIndices& removed_nodes) { ORT_ENFORCE(casts.size() > 0); Node* lead_cast = casts.front(); Node* trail_cast = casts.back(); @@ -360,7 +358,7 @@ static Status RemoveCastNodesChain(Graph& graph, std::vector casts, std:: int cast_output_index = optimizer_utils::IndexOfNodeOutput(*trail_cast, *cast_output); for (Node* consumer : consumers) { if (nullptr != consumer && - std::find(removed_nodes.begin(), removed_nodes.end(), consumer->Index()) == removed_nodes.end()) { + removed_nodes.find(consumer->Index()) == removed_nodes.end()) { auto& consumer_inputs = consumer->MutableInputDefs(); int input_index = optimizer_utils::IndexOfNodeInput(*consumer, *cast_output); graph.RemoveEdge(trail_cast->Index(), consumer->Index(), cast_output_index, input_index); @@ -374,69 +372,69 @@ static Status RemoveCastNodesChain(Graph& graph, std::vector casts, std:: } for (auto cast : casts) { graph_utils::RemoveNodeOutputEdges(graph, *cast); - removed_nodes.push_front(cast->Index()); + removed_nodes.insert(cast->Index()); } return Status::OK(); } /* -* RemoveBackToBackCasts -* Remove FLOAT and FLOAT16 casts back-to-back, only if the parent cast is from FLOAT16 to FLOAT -* and the child cast is from FLOAT to FLOAT16 or from FLOAT to FLOAT. -* The trivial case is parent cast has only one output. A non-trivial case handled by this function -* is, parent has multiple children and one or more child node output is also a graph output. -* In the non-trivial case, when a child-cast nullifies the parent cast, the consumer nodes are moved -* from the child-cast to the producer of the parent-cast input. The code handles cornercases such as -* the child-cast output is also a graph output, in addition to possibly other nodes, and the -* parent-cast has other output nodes. The first check CanRemoveNode rules-out the possibility of -* parent-cast output also being the graph output. -* The inputs is Cast to FLOAT16 -* With the possibility of child/parent cast feeding other nodes/graph-output, the transformation is either -* _____|______ -* | Opcode1 | -* |__________| -* | -* _____V______ ____________ -* | Cast | | Opcode 1 | -* |FP16->FP32| |__________| -* |__________| | -* | | -* | ---\ _____V______ -* | ---/ | Opcode2 | -* _____V______ |__________| -* | Cast | - |FP32->Fp16| -* |__________| -* | -* _____V______ -* | Opcode2 | -* |__________| -* -* or -* -* _____|______ -* | Opcode1 | -* |__________| -* | -* _____V______ ____________ -* | Cast | | Opcode 1 | -* |FP16->FP32| |__________| -* |__________| | -* | | -* | ---\ _____V______ -* | ---/ | Cast FP32| -* _____V______ |__________| -* | Cast | | -* |FP32->FP32| | -* |__________| _____V______ -* | | Opcode2 | -* _____V______ |__________| -* | Opcode2 | | -* |__________| -* | -*/ + * RemoveBackToBackCasts + * Remove FLOAT and FLOAT16 casts back-to-back, only if the parent cast is from FLOAT16 to FLOAT + * and the child cast is from FLOAT to FLOAT16 or from FLOAT to FLOAT. + * The trivial case is parent cast has only one output. A non-trivial case handled by this function + * is, parent has multiple children and one or more child node output is also a graph output. + * In the non-trivial case, when a child-cast nullifies the parent cast, the consumer nodes are moved + * from the child-cast to the producer of the parent-cast input. The code handles cornercases such as + * the child-cast output is also a graph output, in addition to possibly other nodes, and the + * parent-cast has other output nodes. The first check CanRemoveNode rules-out the possibility of + * parent-cast output also being the graph output. + * The inputs is Cast to FLOAT16 + * With the possibility of child/parent cast feeding other nodes/graph-output, the transformation is either + * _____|______ + * | Opcode1 | + * |__________| + * | + * _____V______ ____________ + * | Cast | | Opcode 1 | + * |FP16->FP32| |__________| + * |__________| | + * | | + * | ---\ _____V______ + * | ---/ | Opcode2 | + * _____V______ |__________| + * | Cast | + * |FP32->Fp16| + * |__________| + * | + * _____V______ + * | Opcode2 | + * |__________| + * + * or + * + * _____|______ + * | Opcode1 | + * |__________| + * | + * _____V______ ____________ + * | Cast | | Opcode 1 | + * |FP16->FP32| |__________| + * |__________| | + * | | + * | ---\ _____V______ + * | ---/ | Cast FP32| + * _____V______ |__________| + * | Cast | | + * |FP32->FP32| | + * |__________| _____V______ + * | | Opcode2 | + * _____V______ |__________| + * | Opcode2 | | + * |__________| + * | + */ static bool RemoveBackToBackCasts(Graph& graph, Node* parent, - std::deque& removed_nodes, + NodeIndices& removed_nodes, const logging::Logger& logger) { ORT_ENFORCE(IsCastTo(parent, TensorProto::FLOAT)); bool modified = false; @@ -445,7 +443,7 @@ static bool RemoveBackToBackCasts(Graph& graph, Node* parent, std::vector children = graph.GetMutableConsumerNodes(cast_output->Name()); if (children.size() == 1) { Node* child = children[0]; - if (std::find(removed_nodes.begin(), removed_nodes.end(), child->Index()) == removed_nodes.end()) { + if (removed_nodes.find(child->Index()) == removed_nodes.end()) { if (IsCastTo(child, TensorProto::FLOAT16)) { // The parent and child cancell out LOGS(logger, VERBOSE) << "RemoveBackToBackCasts: Removed Cast nodes " << parent->Name() << " and " << child->Name(); @@ -465,7 +463,7 @@ static bool RemoveBackToBackCasts(Graph& graph, Node* parent, std::vector new_consumers; size_t children_count = children.size(); for (Node* child : children) { - if (std::find(removed_nodes.begin(), removed_nodes.end(), child->Index()) == removed_nodes.end()) { + if (removed_nodes.find(child->Index()) == removed_nodes.end()) { if (IsCastTo(child, TensorProto::FLOAT16)) { // The parent and child cancell out // Remove the child node without effecting the other nodes. @@ -495,7 +493,7 @@ static bool RemoveBackToBackCasts(Graph& graph, Node* parent, } modified = true; graph.RemoveEdge(parent->Index(), child->Index(), 0, 0); - removed_nodes.push_front(child->Index()); + removed_nodes.insert(child->Index()); children_count--; } else if (IsCastTo(child, TensorProto::FLOAT)) { // Child is a duplicate of parent @@ -521,20 +519,21 @@ static bool RemoveBackToBackCasts(Graph& graph, Node* parent, } /* -* SearchUpstream: -* ReverseDFS, traverse bottom-up, the graph upstream collecting all the NodeArgs that require a cast -* in order to move an FP16 Cast operation up the graph. -* Visited float NodeArgs are either in require_cast or require_type_change so that the same -* nodearg is traversed not more than once. -* If the level is 2, the functions traverses up the graph identifying required FP32 casts even if -* multiple consumers for the node outputs are found while level 0 or 1 quit traversing up. -*/ + * SearchUpstream: + * ReverseDFS, traverse bottom-up, the graph upstream collecting all the NodeArgs that require a cast + * in order to move an FP16 Cast operation up the graph. + * Visited float NodeArgs are either in require_cast or require_type_change so that the same + * nodearg is traversed not more than once. + * If the level is 2, the functions traverses up the graph identifying required FP32 casts even if + * multiple consumers for the node outputs are found while level 0 or 1 quit traversing up. + */ static void SearchUpstream(Graph& graph, NodeArg* node_arg, Node* dst_node, NodeArgToConsumerMap& require_cast, NodeArgToConsumerMap& require_cast_fp32, std::unordered_set& require_type_change, - std::deque& removed_nodes, - size_t level) { + const NodeIndices& removed_nodes, + size_t level, + const FP16AllowOps& fp16_allow_ops) { Node* node = graph.GetMutableProducerNode(node_arg->Name()); // If the Cast input feeds more than one node or the cast node feeds a graph output and at least one // node then it cannot propagate. @@ -551,13 +550,13 @@ static void SearchUpstream(Graph& graph, NodeArg* node_arg, Node* dst_node, if (IsType(*node_arg, TensorProto_DataType_FLOAT)) { require_cast[node_arg].push_back(dst_node); } - } else if (std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end()) { + } else if (removed_nodes.find(node->Index()) == removed_nodes.end()) { if (IsCastTo(node, TensorProto_DataType_FLOAT)) { // This Cast node and the Cast node that will be created later will cancel out require_cast[node_arg].push_back(dst_node); } else { std::string op_type = node->OpType(); - if (!IsFP16Allow(op_type, level)) { + if (!IsFP16Allow(op_type, level, fp16_allow_ops)) { // Cannot traverse-up beyond this point if (node_arg->Exists() && IsType(*node_arg, TensorProto_DataType_FLOAT)) { require_cast[node_arg].push_back(dst_node); @@ -579,7 +578,7 @@ static void SearchUpstream(Graph& graph, NodeArg* node_arg, Node* dst_node, if (level >= 2) { for (Node* consumer : graph.GetMutableConsumerNodes(node_arg->Name())) { if (nullptr != consumer && consumer != dst_node && consumer->OpType() != "Cast" && - std::find(removed_nodes.begin(), removed_nodes.end(), consumer->Index()) == removed_nodes.end()) { + removed_nodes.find(consumer->Index()) == removed_nodes.end()) { require_cast_fp32[node_arg].push_back(consumer); } } @@ -591,7 +590,7 @@ static void SearchUpstream(Graph& graph, NodeArg* node_arg, Node* dst_node, if (IsRelevantInput(node, node_input) && IsType(*node_input, TensorProto_DataType_FLOAT) && require_cast.find(node_input) == require_cast.end() && require_type_change.find(node_input) == require_type_change.end()) { - SearchUpstream(graph, node_input, node, require_cast, require_cast_fp32, require_type_change, removed_nodes, level); + SearchUpstream(graph, node_input, node, require_cast, require_cast_fp32, require_type_change, removed_nodes, level, fp16_allow_ops); if (require_cast.find(node_input) == require_cast.end() && require_cast_fp32.find(node_input) == require_cast_fp32.end()) { require_type_change.insert(node_input); @@ -604,17 +603,17 @@ static void SearchUpstream(Graph& graph, NodeArg* node_arg, Node* dst_node, } /* -* SearchDownstream: -* Recursively DFS traverse the graph downstream collecting all the NodeArgs that require a cast -* in order to remove an FP32 Cast operation up the graph. Also collect the NodeArgs that need to -* be converted from float to float16 along the way. -*/ + * SearchDownstream: + * Recursively DFS traverse the graph downstream collecting all the NodeArgs that require a cast + * in order to remove an FP32 Cast operation up the graph. Also collect the NodeArgs that need to + * be converted from float to float16 along the way. + */ static void SearchDownstream(Graph& graph, NodeArg* node_arg, NodeArgToConsumerMap& require_cast, NodeArgToConsumerMap& require_cast_fp16, std::unordered_set& require_type_change, - std::deque& removed_nodes, - size_t level) { + size_t level, + const FP16AllowOps& fp16_allow_ops) { for (Node* node : graph.GetMutableConsumerNodes(node_arg->Name())) { if (node) { std::string op_type = node->OpType(); @@ -622,7 +621,7 @@ static void SearchDownstream(Graph& graph, NodeArg* node_arg, // This Cast node and the Cast node that will be created later will cancel out require_cast[node_arg].push_back(node); } else { - if (!IsFP16Allow(op_type, level)) { + if (!IsFP16Allow(op_type, level, fp16_allow_ops)) { if (node_arg->Exists() && IsType(*node_arg, TensorProto_DataType_FLOAT)) { require_cast[node_arg].push_back(node); @@ -646,7 +645,7 @@ static void SearchDownstream(Graph& graph, NodeArg* node_arg, if (IsRelevantOutput(node, node_output) && IsType(*node_output, TensorProto_DataType_FLOAT) && require_cast.find(node_output) == require_cast.end() && require_type_change.find(node_output) == require_type_change.end()) { - SearchDownstream(graph, node_output, require_cast, require_cast_fp16, require_type_change, removed_nodes, level); + SearchDownstream(graph, node_output, require_cast, require_cast_fp16, require_type_change, level, fp16_allow_ops); if (require_cast.find(node_output) == require_cast.end()) { require_type_change.insert(node_output); } @@ -662,7 +661,10 @@ static void SearchDownstream(Graph& graph, NodeArg* node_arg, } // Change the elem_type of the given NodeArgs from FLOAT to FLOAT16. -static void ChangeTypeToFP16(Graph& graph, std::unordered_set& require_type_change, bool is_forward, const logging::Logger& logger) { +static void ChangeTypeToFP16(Graph& graph, std::unordered_set& require_type_change, bool is_forward, + NodeIndices& converted_nodes, + const NodeIndices& inserted_nodes, + const logging::Logger& logger) { ONNX_NAMESPACE::TypeProto type_proto; type_proto.mutable_tensor_type()->set_elem_type(TensorProto::FLOAT16); for (NodeArg* node_arg : require_type_change) { @@ -672,13 +674,15 @@ static void ChangeTypeToFP16(Graph& graph, std::unordered_set& require if (is_forward) { // Propagating forwards. Count consumers. for (const Node* node : graph.GetConsumerNodes(node_arg->Name())) { - converted_node_names.insert(node->Name()); + if (inserted_nodes.find(node->Index()) == inserted_nodes.end()) { + converted_nodes.insert(node->Index()); + } } } else { // Propagating backwards. Count producers. const Node* producer = graph.GetProducerNode(node_arg->Name()); - if (nullptr != producer) { - converted_node_names.insert(producer->Name()); + if (nullptr != producer && inserted_nodes.find(producer->Index()) == inserted_nodes.end()) { + converted_nodes.insert(producer->Index()); } } } @@ -686,19 +690,22 @@ static void ChangeTypeToFP16(Graph& graph, std::unordered_set& require } /* -* PropagateForwards -* Propagate FP32 Cast operations forwards (downstream) -* Using SearchDownStream search the graph for Cast FP16 safe/allowed operations to expand -* the float16 computation region. -* The required_cast vector is the collection of nodes that require float cast. -* All nodeargs on a path down to any of the -* frontier nodes require type change from FLOAT to FLOAT16. -* require_type_change consists of such nodes. All the frontier nodes require fp32 cast -* The input node is expected to be non-nullptr -*/ + * PropagateForwards + * Propagate FP32 Cast operations forwards (downstream) + * Using SearchDownStream search the graph for Cast FP16 safe/allowed operations to expand + * the float16 computation region. + * The required_cast vector is the collection of nodes that require float cast. + * All nodeargs on a path down to any of the + * frontier nodes require type change from FLOAT to FLOAT16. + * require_type_change consists of such nodes. All the frontier nodes require fp32 cast + * The input node is expected to be non-nullptr + */ static bool PropagateForwards(Graph& graph, Node* node, - std::deque& removed_nodes, + NodeIndices& removed_nodes, size_t level, + const FP16AllowOps& fp16_allow_ops, + NodeIndices& converted_nodes, + NodeIndices& inserted_nodes, const logging::Logger& logger) { ORT_ENFORCE(nullptr != node); bool modified = false; @@ -706,39 +713,42 @@ static bool PropagateForwards(Graph& graph, Node* node, NodeArgToConsumerMap require_cast_fp16; std::unordered_set require_type_change; NodeArg* cast_output = node->MutableOutputDefs()[0]; - SearchDownstream(graph, cast_output, require_cast, require_cast_fp16, require_type_change, removed_nodes, level); + SearchDownstream(graph, cast_output, require_cast, require_cast_fp16, require_type_change, level, fp16_allow_ops); if (require_cast.size() > 0 && require_cast.find(cast_output) == require_cast.end()) { if (require_cast_fp16.size() > 0) { - ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast_fp16, true, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast_fp16, true, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateForwards: Inserted FP16 Cast nodes " << ConcatNames(require_cast_fp16, GetName); } // Remove Cast operation LOGS(logger, VERBOSE) << "PropagateForwards: Removed Cast node " << node->Name(); ORT_THROW_IF_ERROR(RemoveCastNodesChain(graph, {node}, removed_nodes)); - ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast, false, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast, false, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateForwards: Inserted Cast FP32 nodes " << ConcatNames(require_cast, GetName); - ChangeTypeToFP16(graph, require_type_change, true, logger); + ChangeTypeToFP16(graph, require_type_change, true, converted_nodes, inserted_nodes, logger); modified = true; } return modified; } /* -* PropagateBackwards -* Propagate FP16 Cast operations backwards (upstream) -* Using SearchUpstream search the graph for Cast FP16 safe/allowed operations and expand -* float16 computation regsion and -* find the frontiers of the float16 computation region. -* The required_cast or require_cast_fp32 vector is a collection of -* FP16-cast-frontiers of the cast node. All node-args on the path from any of the -* frontier nodes to the cast node require type change from FLOAT to FLOAT16. -* Each of the frontier nodes requires an fp16 cast or fp32 cast. -* The input node is expected be non-nullptr. -*/ + * PropagateBackwards + * Propagate FP16 Cast operations backwards (upstream) + * Using SearchUpstream search the graph for Cast FP16 safe/allowed operations and expand + * float16 computation regsion and + * find the frontiers of the float16 computation region. + * The required_cast or require_cast_fp32 vector is a collection of + * FP16-cast-frontiers of the cast node. All node-args on the path from any of the + * frontier nodes to the cast node require type change from FLOAT to FLOAT16. + * Each of the frontier nodes requires an fp16 cast or fp32 cast. + * The input node is expected be non-nullptr. + */ static bool PropagateBackwards(Graph& graph, Node* node, - std::deque& removed_nodes, + NodeIndices& removed_nodes, size_t level, + const FP16AllowOps& fp16_allow_ops, + NodeIndices& converted_nodes, + NodeIndices& inserted_nodes, const logging::Logger& logger) { bool modified = false; ORT_ENFORCE(nullptr != node); @@ -746,24 +756,24 @@ static bool PropagateBackwards(Graph& graph, Node* node, NodeArgToConsumerMap require_cast_fp32; NodeArg* cast_input = node->MutableInputDefs()[0]; std::unordered_set require_type_change; - SearchUpstream(graph, cast_input, node, require_cast, require_cast_fp32, require_type_change, removed_nodes, level); + SearchUpstream(graph, cast_input, node, require_cast, require_cast_fp32, require_type_change, removed_nodes, level, fp16_allow_ops); if (require_cast_fp32.empty()) { require_type_change.insert(cast_input); } // TODO need a heuristic when to insert FP32 Cast if (require_cast.size() > 0 && require_cast.find(cast_input) == require_cast.end() /* && require_cast.size() >= require_cast_fp32.size() */) { if (require_cast_fp32.size() > 0) { - ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast_fp32, false, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast_fp32, false, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateBackwards: Inserted FP32 Cast nodes " << ConcatNames(require_cast_fp32, GetName); } // Remove Cast operations ORT_THROW_IF_ERROR(RemoveCastNodesChain(graph, {node}, removed_nodes)); LOGS(logger, VERBOSE) << "PropagateBackwards: Removed Cast node " << node->Name(); - ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast, true, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, require_cast, true, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateBackwards: Inserted Cast nodes " << ConcatNames(require_cast, GetName); - ChangeTypeToFP16(graph, require_type_change, false, logger); + ChangeTypeToFP16(graph, require_type_change, false, converted_nodes, inserted_nodes, logger); LOGS(logger, VERBOSE) << "PropagateBackwards: Changed the type from float to float16 : " << ConcatNames>(require_type_change); modified = true; @@ -772,29 +782,30 @@ static bool PropagateBackwards(Graph& graph, Node* node, } /* -* FuseNodes -* Fuse all (cast) nodes, and replace with a single (cast) node. -* Assumptions: -* 1. all nodes are Cast ops and are of the same Cast type -* 2. all the nodes have the same input -* Input0/NodeArg -* ___ ____|______ Input0/NodeArg -* | | | -* _____V____ _____V______ _____V______ -* |Cast FP16| | Cast FP16| | CastFP16 | -* |or_FP32__| |_or_FP32 _| |_or_FP32__| -* | | | -* | | _______|___________ -* | | | | -* | | ---\ _____V______ _____V_____ -* | | ---/ | Opcode0 | | Opcode1 | -* _____V_____ _____V______ |__________| |_________| -* | Opcode0 | | Opcode1 | | | -* |_________| |__________| -* | | -*/ + * FuseNodes + * Fuse all (cast) nodes, and replace with a single (cast) node. + * Assumptions: + * 1. all nodes are Cast ops and are of the same Cast type + * 2. all the nodes have the same input + * Input0/NodeArg + * ___ ____|______ Input0/NodeArg + * | | | + * _____V____ _____V______ _____V______ + * |Cast FP16| | Cast FP16| | CastFP16 | + * |or_FP32__| |_or_FP32 _| |_or_FP32__| + * | | | + * | | _______|___________ + * | | | | + * | | ---\ _____V______ _____V_____ + * | | ---/ | Opcode0 | | Opcode1 | + * _____V_____ _____V______ |__________| |_________| + * | Opcode0 | | Opcode1 | | | + * |_________| |__________| + * | | + */ static void FuseNodes(Graph& graph, const NodeArg* input, std::vector nodes, - std::deque& removed_nodes) { + NodeIndices& removed_nodes, + NodeIndices& inserted_nodes) { ORT_ENFORCE(nodes.size() > 0); Node* node = nodes[0]; const Node* producer = graph.GetProducerNode(input->Name()); @@ -811,7 +822,7 @@ static void FuseNodes(Graph& graph, const NodeArg* input, std::vector nod if (nullptr != producer) { graph.AddEdge(producer->Index(), new_cast.Index(), output_index, 0); } - inserted_node_names.insert(new_cast.Name()); + inserted_nodes.insert(new_cast.Index()); std::vector consumers; for (Node* cast : nodes) { for (NodeArg* output : cast->MutableOutputDefs()) { @@ -833,14 +844,15 @@ static void FuseNodes(Graph& graph, const NodeArg* input, std::vector nod graph.AddConsumerNode(input->Name(), &new_cast); graph.UpdateConsumerNodes(new_output.Name(), consumers); for (Node* n : nodes) { - removed_nodes.push_front(n->Index()); + removed_nodes.insert(n->Index()); graph_utils::RemoveNodeOutputEdges(graph, *n); } } // Traverse the graph recursively searching/collecting sibling Cast op nodes to fuse and call FuseNodes. static bool FuseSiblingCasts(Graph& graph, const NodeArg* node_arg, - std::deque& removed_nodes, + NodeIndices& removed_nodes, + NodeIndices& inserted_nodes, const logging::Logger& logger) { bool modified = false; std::vector cast_fp16_siblings; @@ -848,7 +860,7 @@ static bool FuseSiblingCasts(Graph& graph, const NodeArg* node_arg, for (Node* node : graph.GetMutableConsumerNodes(node_arg->Name())) { // If a cast node feeds a graph output then it is not a candidate for fusion. if (nullptr == node || node->OpType() != "Cast" || - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) != removed_nodes.end() || + removed_nodes.find(node->Index()) != removed_nodes.end() || graph.IsOutput(node->OutputDefs()[0])) { continue; } @@ -860,12 +872,12 @@ static bool FuseSiblingCasts(Graph& graph, const NodeArg* node_arg, } if (cast_fp16_siblings.size() > 1) { modified = true; - FuseNodes(graph, node_arg, cast_fp16_siblings, removed_nodes); + FuseNodes(graph, node_arg, cast_fp16_siblings, removed_nodes, inserted_nodes); LOGS(logger, VERBOSE) << "FusedSubgraphs: Fused Cast nodes : " << ConcatNames>(cast_fp16_siblings); } if (cast_fp32_siblings.size() > 1) { modified = true; - FuseNodes(graph, node_arg, cast_fp32_siblings, removed_nodes); + FuseNodes(graph, node_arg, cast_fp32_siblings, removed_nodes, inserted_nodes); LOGS(logger, VERBOSE) << "FusedSubgraphs: Fused Cast nodes : " << ConcatNames>(cast_fp32_siblings); } return modified; @@ -873,11 +885,12 @@ static bool FuseSiblingCasts(Graph& graph, const NodeArg* node_arg, // Overloaded function which goes through all the output args of a node and calls FuseSiblingCasts static bool FuseSiblingCasts(Graph& graph, const Node* parent, - std::deque& removed_nodes, + NodeIndices& removed_nodes, + NodeIndices& inserted_nodes, const logging::Logger& logger) { bool modified = false; for (const NodeArg* output : parent->OutputDefs()) { - modified |= FuseSiblingCasts(graph, output, removed_nodes, logger); + modified |= FuseSiblingCasts(graph, output, removed_nodes, inserted_nodes, logger); } return modified; } @@ -885,7 +898,7 @@ static bool FuseSiblingCasts(Graph& graph, const Node* parent, // RemoveUnnecessaryCasts // Remove a cast if the input elem_type is same the required cast type. static bool RemoveUnnecessaryCasts(Graph& graph, Node* node, - std::deque& removed_nodes, + NodeIndices& removed_nodes, const logging::Logger& logger) { bool modified = false; if (node->InputDefs().size() == 1) { @@ -902,49 +915,52 @@ static bool RemoveUnnecessaryCasts(Graph& graph, Node* node, } /* -* PropagateFP32CastsFromInputsToOutputs -* This non-recursive fusion, checks whether the given node is fp16 safe op and -* whether all floatingpoint inputs are cast to fp32 -* and propagates cast op to the floatingpoint outputs. -* Convert the following graph -* -* Input0/NodeArg Input1/NodeArg -* | | -* _____V____ _____V______ -* |Cast FP32| | Cast FP32| -* |_________| |__________| -* | | -* __V______________V___ -* | Opcode | -* |(operation performed | -* | in float32) | -* |_____________________| -* | | -* | | -* -* and produce the following output -* -* Input0/NodeArg Input1/NodeArg -* | | -* __V_______________V___ -* | Opcode | -* |(operation performed | -* | in float16) | -* |_____________________| -* | | -* _____V____ _____V______ -* |Cast FP32| | Cast FP32| -* |_________| |__________| -* | | -* -* -*/ + * PropagateFP32CastsFromInputsToOutputs + * This non-recursive fusion, checks whether the given node is fp16 safe op and + * whether all floatingpoint inputs are cast to fp32 + * and propagates cast op to the floatingpoint outputs. + * Convert the following graph + * + * Input0/NodeArg Input1/NodeArg + * | | + * _____V____ _____V______ + * |Cast FP32| | Cast FP32| + * |_________| |__________| + * | | + * __V______________V___ + * | Opcode | + * |(operation performed | + * | in float32) | + * |_____________________| + * | | + * | | + * + * and produce the following output + * + * Input0/NodeArg Input1/NodeArg + * | | + * __V_______________V___ + * | Opcode | + * |(operation performed | + * | in float16) | + * |_____________________| + * | | + * _____V____ _____V______ + * |Cast FP32| | Cast FP32| + * |_________| |__________| + * | | + * + * + */ static bool PropagateFP32CastsFromInputsToOutputs(Graph& graph, Node* node, - std::deque& removed_nodes, + NodeIndices& removed_nodes, size_t level, + const FP16AllowOps& fp16_allow_ops, + NodeIndices& converted_nodes, + NodeIndices& inserted_nodes, const logging::Logger& logger) { bool modified = false; - if (IsFP16Allow(node->OpType(), level)) { + if (IsFP16Allow(node->OpType(), level, fp16_allow_ops)) { bool has_float_inputs = false; bool all_float_inputs_have_casts = true; std::vector casts; @@ -961,7 +977,7 @@ static bool PropagateFP32CastsFromInputsToOutputs(Graph& graph, Node* node, has_float_inputs = true; Node* producer = graph.GetMutableProducerNode(input->Name()); if (nullptr != producer) { - if (std::find(removed_nodes.begin(), removed_nodes.end(), producer->Index()) == removed_nodes.end()) { + if (removed_nodes.find(producer->Index()) == removed_nodes.end()) { if (IsCastTo(producer, TensorProto::FLOAT) && producer->GetOutputEdgesCount() == 1 && !graph.IsOutput(input)) { @@ -981,7 +997,7 @@ static bool PropagateFP32CastsFromInputsToOutputs(Graph& graph, Node* node, } if (has_float_inputs && (level >= 2 || all_float_inputs_have_casts) && casts.size() > 1) { if (non_cast_producers_map.size() > 0) { - ORT_THROW_IF_ERROR(InsertCastNodes(graph, non_cast_producers_map, true, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, non_cast_producers_map, true, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateFP32CastsFromInputsToOutputs: Inserted FP16 Cast node to " << ConcatNames(non_cast_producers_map, GetName); } @@ -997,59 +1013,62 @@ static bool PropagateFP32CastsFromInputsToOutputs(Graph& graph, Node* node, node_args_map.insert(std::make_pair(output, graph.GetMutableConsumerNodes(output->Name()))); } } - ORT_THROW_IF_ERROR(InsertCastNodes(graph, node_args_map, false, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, node_args_map, false, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateFP32CastsFromInputsToOutputs: Inserted FP32 Cast node to " << ConcatNames(node_args_map, GetName); - ChangeTypeToFP16(graph, require_type_change, true, logger); + ChangeTypeToFP16(graph, require_type_change, true, converted_nodes, inserted_nodes, logger); modified = true; } } return modified; } /* -* PropagateFP16CastsFromOutputsToInputs -* This non-recursive fusion, checks whether the given node is fp16 safe op and -* whether all floatingpoint outputs are cast to fp16 -* and propagates cast op to the floatingpoint inputs. -* Convert the following graph -* -* Input0/NodeArg Input1/NodeArg -* | | -* __V______________V___ -* | Opcode | -* |(operation performed | -* | in float32) | -* |_____________________| -* | | -* _____V____ _____V______ -* |Cast FP16| | Cast FP16| -* |_________| |__________| -* | | -* | | -* V V -* -* and produce the following output -* -* Input0/NodeArg Input1/NodeArg -* | | -* _____V____ _____V______ -* |Cast FP16| | Cast FP16| -* |_________| |__________| -* | | -* __V______________V___ -* | Opcode | -* |(operation performed | -* | in float16) | -* |_____________________| -* | | -* V V -*/ + * PropagateFP16CastsFromOutputsToInputs + * This non-recursive fusion, checks whether the given node is fp16 safe op and + * whether all floatingpoint outputs are cast to fp16 + * and propagates cast op to the floatingpoint inputs. + * Convert the following graph + * + * Input0/NodeArg Input1/NodeArg + * | | + * __V______________V___ + * | Opcode | + * |(operation performed | + * | in float32) | + * |_____________________| + * | | + * _____V____ _____V______ + * |Cast FP16| | Cast FP16| + * |_________| |__________| + * | | + * | | + * V V + * + * and produce the following output + * + * Input0/NodeArg Input1/NodeArg + * | | + * _____V____ _____V______ + * |Cast FP16| | Cast FP16| + * |_________| |__________| + * | | + * __V______________V___ + * | Opcode | + * |(operation performed | + * | in float16) | + * |_____________________| + * | | + * V V + */ static bool PropagateFP16CastsFromOutputsToInputs(Graph& graph, Node* node, - std::deque& removed_nodes, + NodeIndices& removed_nodes, size_t level, + const FP16AllowOps& fp16_allow_ops, + NodeIndices& converted_nodes, + NodeIndices& inserted_nodes, const logging::Logger& logger) { bool modified = false; - if (IsFP16Allow(node->OpType(), level)) { + if (IsFP16Allow(node->OpType(), level, fp16_allow_ops)) { bool has_float_outputs = false; bool all_float_outputs_have_casts = true; std::vector casts; // Cast nodes to propagate. @@ -1066,7 +1085,7 @@ static bool PropagateFP16CastsFromOutputsToInputs(Graph& graph, Node* node, for (auto node_iter = consumers.begin(); node_iter != consumers.end() && (level >= 2 || all_float_outputs_have_casts); ++node_iter) { Node* consumer = *node_iter; if (nullptr != consumer && - std::find(removed_nodes.begin(), removed_nodes.end(), consumer->Index()) == removed_nodes.end()) { + removed_nodes.find(consumer->Index()) == removed_nodes.end()) { if (IsCastTo(consumer, TensorProto::FLOAT16)) { casts.push_back(consumer); } else { @@ -1084,7 +1103,7 @@ static bool PropagateFP16CastsFromOutputsToInputs(Graph& graph, Node* node, } if (has_float_outputs && (level >= 2 || all_float_outputs_have_casts) && casts.size() > 1) { if (non_cast_consumers_map.size() > 0) { - ORT_THROW_IF_ERROR(InsertCastNodes(graph, non_cast_consumers_map, false, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, non_cast_consumers_map, false, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateFP16CastsFromOutputsToInputs: Inserted FP32 Cast node to " << ConcatNames(non_cast_consumers_map, GetName); } @@ -1100,10 +1119,10 @@ static bool PropagateFP16CastsFromOutputsToInputs(Graph& graph, Node* node, node_args_map.insert(std::make_pair(input, std::vector({node}))); } } - ORT_THROW_IF_ERROR(InsertCastNodes(graph, node_args_map, true, removed_nodes)); + ORT_THROW_IF_ERROR(InsertCastNodes(graph, node_args_map, true, removed_nodes, inserted_nodes)); LOGS(logger, VERBOSE) << "PropagateFP16CastsFromOutputsToInputs: Inserted FP16 Cast node to " << ConcatNames(node_args_map, GetName); - ChangeTypeToFP16(graph, require_type_change, false, logger); + ChangeTypeToFP16(graph, require_type_change, false, converted_nodes, inserted_nodes, logger); modified = true; } } @@ -1111,12 +1130,13 @@ static bool PropagateFP16CastsFromOutputsToInputs(Graph& graph, Node* node, } /* -* CreateCast -* Create a cast node based on the node_arg for the given data type. If the node_arg is a graph output is_graph_outut is set. -* If the node_arg is not a graph output then the node_arg is the input of the new cast node. Otherwise the node_arg is the output -* of the new cast node. This function is used by InsertFP16Cast or InsertFP32Casts. -*/ -static Node& CreateCast(Graph& graph, NodeArg* node_arg, TensorProto_DataType data_type, bool is_graph_output = false) { + * CreateCast + * Create a cast node based on the node_arg for the given data type. If the node_arg is a graph output is_graph_outut is set. + * If the node_arg is not a graph output then the node_arg is the input of the new cast node. Otherwise the node_arg is the output + * of the new cast node. This function is used by InsertFP16Cast or InsertFP32Casts. + */ +static Node& CreateCast(Graph& graph, NodeArg* node_arg, TensorProto_DataType data_type, + NodeIndices& inserted_nodes, bool is_graph_output = false) { TypeProto type_proto; type_proto.mutable_tensor_type()->set_elem_type(is_graph_output ? (data_type == TensorProto::FLOAT ? TensorProto::FLOAT16 : TensorProto::FLOAT) : data_type); NodeArg& new_node_arg = graph.GetOrCreateNodeArg(graph.GenerateNodeArgName(node_arg->Name()), &type_proto); @@ -1135,42 +1155,43 @@ static Node& CreateCast(Graph& graph, NodeArg* node_arg, TensorProto_DataType da inputs, outputs, &attributes); - inserted_node_names.insert(node.Name()); + inserted_nodes.insert(node.Index()); return node; } /* InsertFP16Cast -* Insert a new cast input for the given float input for the give node. This function should be used with InsertFP32Casts -* in order to compute FP16 allowed operations in 16 bit precision instead of 32 bit precision. -* -* Input0/NodeArg Input1/NodeArg -* | | -* __V______________V___ -* | Opcode | -* |(operation performed | -* | in float32) | -* |_____________________| -* | | -* V V -* -* Is eventually translated to the following, after all inputs are casted -* -* Input0/NodeArg Input1/NodeArg -* | | -* _____V____ _____V______ -* |Cast FP16| | Cast FP16| -* |_________| |__________| -* | | -* __V______________V___ -* | Opcode | -* |(operation performed | -* | in float16) | -* |_____________________| -* | | -* V V -* -*/ -static void InsertFP16Cast(Graph& graph, NodeArg* input_arg, Node* node, const logging::Logger& logger) { - Node& cast = CreateCast(graph, input_arg, TensorProto::FLOAT16); + * Insert a new cast input for the given float input for the give node. This function should be used with InsertFP32Casts + * in order to compute FP16 allowed operations in 16 bit precision instead of 32 bit precision. + * + * Input0/NodeArg Input1/NodeArg + * | | + * __V______________V___ + * | Opcode | + * |(operation performed | + * | in float32) | + * |_____________________| + * | | + * V V + * + * Is eventually translated to the following, after all inputs are casted + * + * Input0/NodeArg Input1/NodeArg + * | | + * _____V____ _____V______ + * |Cast FP16| | Cast FP16| + * |_________| |__________| + * | | + * __V______________V___ + * | Opcode | + * |(operation performed | + * | in float16) | + * |_____________________| + * | | + * V V + * + */ +static void InsertFP16Cast(Graph& graph, NodeArg* input_arg, Node* node, + NodeIndices& inserted_nodes, const logging::Logger& logger) { + Node& cast = CreateCast(graph, input_arg, TensorProto::FLOAT16, inserted_nodes); NodeArg* new_input_arg = cast.MutableOutputDefs()[0]; std::vector& inputs = node->MutableInputDefs(); Node* producer = graph.GetMutableProducerNode(input_arg->Name()); @@ -1190,42 +1211,43 @@ static void InsertFP16Cast(Graph& graph, NodeArg* input_arg, Node* node, const l } /* InsertFP32Casts -* Insert float casts on the given output for each consumer. This function should be used with InsertFP16Casts -* in order to compute FP16 allowed operations in 16 bit precision instead of 32 bit precision. -* -* | | -* __V_______________V___ -* | Opcode | -* |(operation performed | -* | in float16) | -* |_____________________| -* _______|_______ -* | | -* _____V____ _____V______ -* |Consumer0| |Consumer1 | -* |_________| |__________| -* | | -* -* will be translated to the following -* -* | | -* __V_______________V___ -* | Opcode | -* |(operation performed | -* | in float16) | -* |_____________________| -* _______|_______ -* | | -* _____V____ _____V______ -* |Cast FP32| | Cast FP32| -* |_________| |__________| -* | | -* _____V____ _____V______ -* |Consumer0| |Consumer1 | -* |_________| |__________| -* | | -*/ -static void InsertFP32Casts(Graph& graph, NodeArg* output_arg, const logging::Logger& logger) { + * Insert float casts on the given output for each consumer. This function should be used with InsertFP16Casts + * in order to compute FP16 allowed operations in 16 bit precision instead of 32 bit precision. + * + * | | + * __V_______________V___ + * | Opcode | + * |(operation performed | + * | in float16) | + * |_____________________| + * _______|_______ + * | | + * _____V____ _____V______ + * |Consumer0| |Consumer1 | + * |_________| |__________| + * | | + * + * will be translated to the following + * + * | | + * __V_______________V___ + * | Opcode | + * |(operation performed | + * | in float16) | + * |_____________________| + * _______|_______ + * | | + * _____V____ _____V______ + * |Cast FP32| | Cast FP32| + * |_________| |__________| + * | | + * _____V____ _____V______ + * |Consumer0| |Consumer1 | + * |_________| |__________| + * | | + */ +static void InsertFP32Casts(Graph& graph, NodeArg* output_arg, + NodeIndices& inserted_nodes, const logging::Logger& logger) { NodeArg* orig_output_arg = output_arg; Node* producer = graph.GetMutableProducerNode(output_arg->Name()); std::vector consumers = graph.GetMutableConsumerNodes(output_arg->Name()); @@ -1238,7 +1260,7 @@ static void InsertFP32Casts(Graph& graph, NodeArg* output_arg, const logging::Lo } // Create a new output_arg on the producer if the output_arg is a graph output if (graph.IsOutput(output_arg)) { - Node& cast = CreateCast(graph, output_arg, TensorProto::FLOAT, true); + Node& cast = CreateCast(graph, output_arg, TensorProto::FLOAT, inserted_nodes, true); graph.UpdateProducerNode(output_arg->Name(), cast.Index()); NodeArg* new_output_arg = cast.MutableInputDefs()[0]; graph.UpdateProducerNode(new_output_arg->Name(), producer->Index()); @@ -1254,7 +1276,7 @@ static void InsertFP32Casts(Graph& graph, NodeArg* output_arg, const logging::Lo } // Create a new cast node for each consumer for (Node* consumer : consumers) { - Node& cast = CreateCast(graph, output_arg, TensorProto::FLOAT); + Node& cast = CreateCast(graph, output_arg, TensorProto::FLOAT, inserted_nodes); NodeArg* new_output_arg = cast.MutableOutputDefs()[0]; std::vector& inputs = consumer->MutableInputDefs(); int output_index = optimizer_utils::IndexOfNodeOutput(cast, *new_output_arg); @@ -1269,20 +1291,22 @@ static void InsertFP32Casts(Graph& graph, NodeArg* output_arg, const logging::Lo graph.UpdateConsumerNodes(output_arg->Name(), new_consumers); } /* -* Expand FP16 compute regions on the graph by example float16 compute nodes, -* propagating float32 Cast operation down the graph and propagating float16 -* Cast operations up the graph. The following functions are performed -* 1. Fuse subgraphs -* 2. Propagate fp32 casts forwards -* 3. Propagate fp16 casts back -* 4. Insert casts before and after allowed operations (InsertAndReduce strategy) -* 5. Remove back to back casts -* 6. Remove redundant casts -* 7. Move FP32 casts from inputs to outputs -* 8. Move FP16 casts from outputs to inputs -*/ + * Expand FP16 compute regions on the graph by example float16 compute nodes, + * propagating float32 Cast operation down the graph and propagating float16 + * Cast operations up the graph. The following functions are performed + * 1. Fuse subgraphs + * 2. Propagate fp32 casts forwards + * 3. Propagate fp16 casts back + * 4. Insert casts before and after allowed operations (InsertAndReduce strategy) + * 5. Remove back to back casts + * 6. Remove redundant casts + * 7. Move FP32 casts from inputs to outputs + * 8. Move FP16 casts from outputs to inputs + */ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { bool local_modified = false; + NodeIndices inserted_nodes; // Names of the nodes inserted + NodeIndices converted_nodes; // Names of the nodes converted to FP16 // First apply the transformation to the subgraphs. { GraphViewer graph_viewer(graph); @@ -1300,19 +1324,19 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level // Using InsertFP16Cast and InsertFP32Casts insert float16 casts on all inputs and float casts on all outputs. // Each consumer of each output gets a separate float cast inserted. Doing so will convert the computation of // curren node from 32 bit float to 16 bit float operation. These cast operations will be eventually reduced. - if (IsFP16Allow(node.OpType(), level_)) { + if (IsFP16Allow(node.OpType(), level_, fp16_allow_ops_)) { // Insert FP16 Cast on all float inputs - converted_node_names.insert(node.Name()); + converted_nodes.insert(node.Index()); for (NodeArg* input_arg : node.MutableInputDefs()) { if (IsRelevantInput(&node, input_arg) && IsType(*input_arg, TensorProto::FLOAT)) { - InsertFP16Cast(graph, input_arg, node_ptr, logger); + InsertFP16Cast(graph, input_arg, node_ptr, inserted_nodes, logger); local_modified = true; } } // Convert all output args to FP16 and insert FP32 cast for all consumers for (NodeArg* output_arg : node.MutableOutputDefs()) { if (IsRelevantOutput(&node, output_arg) && IsType(*output_arg, TensorProto::FLOAT)) { - InsertFP32Casts(graph, output_arg, logger); + InsertFP32Casts(graph, output_arg, inserted_nodes, logger); local_modified = true; } } @@ -1320,11 +1344,11 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level } } } - std::unordered_set removed_node_names; + std::vector removed_node_names; int pass = 0; do { LOGS(logger, VERBOSE) << "Propagate Cast Operations Pass " << pass << ":"; - std::deque removed_nodes; + NodeIndices removed_nodes; if (local_modified) { ORT_RETURN_IF_ERROR(graph.Resolve()); @@ -1338,7 +1362,7 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end() && + removed_nodes.find(node->Index()) == removed_nodes.end() && node->OpType() == "Cast") { local_modified |= RemoveUnnecessaryCasts(graph, node, removed_nodes, logger); } @@ -1348,19 +1372,19 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end()) { - local_modified |= FuseSiblingCasts(graph, node, removed_nodes, logger); + removed_nodes.find(node->Index()) == removed_nodes.end()) { + local_modified |= FuseSiblingCasts(graph, node, removed_nodes, inserted_nodes, logger); } } for (const NodeArg* node_arg : graph.GetInputs()) { - local_modified |= FuseSiblingCasts(graph, node_arg, removed_nodes, logger); + local_modified |= FuseSiblingCasts(graph, node_arg, removed_nodes, inserted_nodes, logger); } // Remove back to back Casts, with FLOAT->FLOAT16 followed by FLOAT16->FLOAT, but not the other way. for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end() && + removed_nodes.find(node->Index()) == removed_nodes.end() && IsCastTo(node, TensorProto::FLOAT)) { local_modified |= RemoveBackToBackCasts(graph, node, removed_nodes, logger); } @@ -1372,8 +1396,8 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end()) { - local_modified |= PropagateFP16CastsFromOutputsToInputs(graph, node, removed_nodes, level_, logger); + removed_nodes.find(node->Index()) == removed_nodes.end()) { + local_modified |= PropagateFP16CastsFromOutputsToInputs(graph, node, removed_nodes, level_, fp16_allow_ops_, converted_nodes, inserted_nodes, logger); } } @@ -1381,8 +1405,8 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end()) { - local_modified |= PropagateFP32CastsFromInputsToOutputs(graph, node, removed_nodes, level_, logger); + removed_nodes.find(node->Index()) == removed_nodes.end()) { + local_modified |= PropagateFP32CastsFromInputsToOutputs(graph, node, removed_nodes, level_, fp16_allow_ops_, converted_nodes, inserted_nodes, logger); } } @@ -1390,9 +1414,9 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end() && + removed_nodes.find(node->Index()) == removed_nodes.end() && IsCastTo(node, TensorProto::FLOAT)) { - local_modified |= PropagateForwards(graph, node, removed_nodes, level_, logger); + local_modified |= PropagateForwards(graph, node, removed_nodes, level_, fp16_allow_ops_, converted_nodes, inserted_nodes, logger); } } @@ -1400,14 +1424,22 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level for (auto node_index : node_topology_list) { Node* node = graph.GetNode(node_index); if (nullptr != node && - std::find(removed_nodes.begin(), removed_nodes.end(), node->Index()) == removed_nodes.end() && + removed_nodes.find(node->Index()) == removed_nodes.end() && IsCastTo(node, TensorProto::FLOAT16)) { - local_modified |= PropagateBackwards(graph, node, removed_nodes, level_, logger); + local_modified |= PropagateBackwards(graph, node, removed_nodes, level_, fp16_allow_ops_, converted_nodes, inserted_nodes, logger); } } } + // In order to generate summary collect only removed node names found in the input graph + // and remove all removed nodes from inserted_nodes and converted nodes collections. for (NodeIndex removed_node : removed_nodes) { - removed_node_names.insert(graph.GetNode(removed_node)->Name()); + auto it = inserted_nodes.find(removed_node); + if (it == inserted_nodes.end()) { + removed_node_names.push_back(graph.GetNode(removed_node)->Name()); + } else { + inserted_nodes.erase(it); + } + converted_nodes.erase(removed_node); graph.RemoveNode(removed_node); } modified |= local_modified; @@ -1419,29 +1451,24 @@ Status PropagateCastOps::ApplyImpl(Graph& graph, bool& modified, int graph_level LOGS(logger, INFO) << "Propagate Cast operations summary:"; LOGS(logger, INFO) << "Number of passes = " << pass; LOGS(logger, INFO) << "Nodes Inserted:"; - std::for_each(inserted_node_names.begin(), inserted_node_names.end(), [removed_node_names, logger](std::string name) { - if (removed_node_names.find(name) == removed_node_names.end()) { LOGS(logger, INFO) << name; } }); + std::for_each(inserted_nodes.begin(), inserted_nodes.end(), [&](NodeIndex idx) { LOGS(logger, INFO) << graph.GetNode(idx)->Name(); }); LOGS(logger, INFO) << "Nodes Removed:"; - std::for_each(removed_node_names.begin(), removed_node_names.end(), [logger](std::string name) { - if (inserted_node_names.find(name) == inserted_node_names.end()) { LOGS(logger, INFO) << name; } }); + std::for_each(removed_node_names.begin(), removed_node_names.end(), [&](const std::string& name) { LOGS(logger, INFO) << name; }); LOGS(logger, INFO) << "Nodes Converted to FP16:"; - std::for_each(converted_node_names.begin(), converted_node_names.end(), [removed_node_names, logger](std::string name) { - if (removed_node_names.find(name) == removed_node_names.end() && inserted_node_names.find(name) == inserted_node_names.end()) { - LOGS(logger, INFO) << name; } }); + std::for_each(converted_nodes.begin(), converted_nodes.end(), [&](NodeIndex idx) { LOGS(logger, INFO) << graph.GetNode(idx)->Name(); }); } - inserted_node_names.clear(); - converted_node_names.clear(); return Status::OK(); } PropagateCastOps::PropagateCastOps(GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy strategy, - size_t level, const std::vector& _allow_list, + size_t level, const std::vector& allow_list, const std::unordered_set& compatible_execution_providers) noexcept - : GraphTransformer("PropagateCastOps", compatible_execution_providers), level_(level), strategy_(strategy) { - fp16_allow_ops[0].clear(); // Remove previously added op types if any. - std::copy(_allow_list.begin(), _allow_list.end(), std::inserter(fp16_allow_ops[0], fp16_allow_ops[0].begin())); + : GraphTransformer("PropagateCastOps", compatible_execution_providers), level_(level), fp16_allow_ops_(3), strategy_(strategy) { + std::copy(allow_list.begin(), allow_list.end(), std::inserter(fp16_allow_ops_[0], fp16_allow_ops_[0].begin())); + std::transform(level1_fp16_allow_ops.begin(), level1_fp16_allow_ops.end(), std::inserter(fp16_allow_ops_[1], fp16_allow_ops_[1].end()), [](const std::string_view& sv) { return std::string(sv); }); + std::transform(level2_fp16_allow_ops.begin(), level2_fp16_allow_ops.end(), std::inserter(fp16_allow_ops_[2], fp16_allow_ops_[2].end()), [](const std::string_view& sv) { return std::string(sv); }); } } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/propagate_cast_ops.h b/onnxruntime/core/optimizer/propagate_cast_ops.h index 86f4b63b5a..d4963a9817 100644 --- a/onnxruntime/core/optimizer/propagate_cast_ops.h +++ b/onnxruntime/core/optimizer/propagate_cast_ops.h @@ -15,6 +15,14 @@ Propagate FP16 Cast operations up the graph and FP32 Cast operations down the gr */ class PropagateCastOps : public GraphTransformer { public: + /* + * The collection FP16AllowOps, specifies for a given propagate_cast_ops level, a collection of node op_types that + * the code is allowed to propage Cast operations across. The user may specify a custom list of optypes using level 0. + * The opcodes are split into multiple levels. Cast propagation is done based on the level. Level 2 op code + * list includes Level 1 list also. + */ + typedef std::vector> FP16AllowOps; + PropagateCastOps(GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy strategy = GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy::FloodFill, size_t level = 0, const std::vector& allow_list = {}, @@ -24,6 +32,7 @@ class PropagateCastOps : public GraphTransformer { private: size_t level_; + FP16AllowOps fp16_allow_ops_; GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy strategy_; };